Upload twill/visualization.py with huggingface_hub

twill/visualization.py

@@ -0,0 +1,253 @@
+"""
+Visualization: Generate schedule diagrams and warp assignment views.
+
+Based on the figures in the paper (Figures 1, 3, 7, 9).
+"""
+
+import numpy as np
+from typing import Dict, List, Optional
+from twill.graph import DependenceGraph
+from twill.smt_joint import JointSWPWSResult
+from twill.modulo_scheduler import ModuloScheduleResult
+
+
+# Color palette for different instruction types / warps
+WARP_COLORS = [
+    '#4CAF50',  # Green - variable latency / producer
+    '#E91E63',  # Pink - compute / TC
+    '#2196F3',  # Blue - compute / EXP  
+    '#FF9800',  # Orange - compute / misc
+    '#9C27B0',  # Purple
+    '#00BCD4',  # Cyan
+    '#795548',  # Brown
+    '#607D8B',  # Blue-grey
+]
+
+FU_COLORS = {
+    'TC': '#E91E63',   # Pink for Tensor Core
+    'EXP': '#2196F3',  # Blue for Exponential
+    'TMA': '#4CAF50',  # Green for TMA loads
+    'TMEM': '#FF9800', # Orange for Tensor Memory
+}
+
+
+def visualize_schedule(
+    graph: DependenceGraph,
+    result: JointSWPWSResult,
+    output_path: Optional[str] = None,
+    title: str = "Twill Schedule",
+) -> str:
+    """Generate a text-based visualization of the schedule.
+    
+    Shows a timeline with instructions placed at their scheduled cycles,
+    colored by warp assignment (in the text representation, shown as markers).
+    
+    Args:
+        graph: The dependence graph
+        result: The joint SWP+WS result
+        output_path: If provided, also generate a matplotlib figure
+        title: Title for the visualization
+    
+    Returns:
+        String representation of the schedule
+    """
+    I = result.I
+    L = result.length
+    n_copies = result.num_copies
+    M = result.schedule
+    wa = result.warp_assignment
+
+    lines = []
+    lines.append(f"╔{'═' * 60}╗")
+    lines.append(f"║ {title:^58s} ║")
+    lines.append(f"║ I={I}, L={L}, copies={n_copies}{' ' * (58 - len(f'I={I}, L={L}, copies={n_copies}'))}║")
+    lines.append(f"╠{'═' * 60}╣")
+
+    # Header: functional units
+    fu_names = graph.machine.functional_units
+    header = "Cycle │"
+    for fu in fu_names:
+        header += f" {fu:^8s} │"
+    header += " Warp "
+    lines.append(f"║ {header:<58s} ║")
+    lines.append(f"║ {'─' * (len(header)):^58s} ║")
+
+    # Build timeline
+    for t in range(L):
+        # Find what's scheduled at this cycle
+        active_ops = []
+        for v in graph.V:
+            for i in range(n_copies):
+                abs_time = M[v.name] + i * I
+                if abs_time == t:
+                    active_ops.append((v, i))
+
+        if active_ops:
+            for v, i in active_ops:
+                warp = wa.warp_of(v.name)
+                # Build functional unit usage string
+                fu_str = f" {t:3d}  │"
+                for f_idx in range(len(fu_names)):
+                    usage = int(v.rrt[:, f_idx].sum())
+                    if usage > 0:
+                        fu_str += f" {v.name:^8s} │"
+                    else:
+                        fu_str += f" {'·':^8s} │"
+                fu_str += f"  W{warp} "
+                if i > 0:
+                    fu_str += f"(i+{i})"
+                lines.append(f"║ {fu_str:<58s} ║")
+        else:
+            fu_str = f" {t:3d}  │"
+            for _ in fu_names:
+                fu_str += f" {'·':^8s} │"
+            fu_str += "     "
+            lines.append(f"║ {fu_str:<58s} ║")
+
+    lines.append(f"╠{'═' * 60}╣")
+
+    # Warp assignment summary
+    lines.append(f"║ {'Warp Assignments:':^58s} ║")
+    for w in range(graph.machine.num_warps):
+        instrs = wa.instructions_on_warp(w)
+        if instrs:
+            label = wa.warp_names.get(w, f"Warp {w}")
+            instr_str = f"  {label}: {', '.join(instrs)}"
+            lines.append(f"║ {instr_str:<58s} ║")
+
+    # Cross-warp barriers
+    barriers = []
+    for edge in graph.E:
+        src_warp = wa.warp_of(edge.src)
+        dst_warp = wa.warp_of(edge.dst)
+        if src_warp != dst_warp:
+            barriers.append(f"  {edge.src}(W{src_warp}) → {edge.dst}(W{dst_warp})")
+    
+    if barriers:
+        lines.append(f"║ {'':^58s} ║")
+        lines.append(f"║ {'Cross-Warp Barriers:':^58s} ║")
+        for b in barriers:
+            lines.append(f"║ {b:<58s} ║")
+
+    lines.append(f"╚{'═' * 60}╝")
+
+    text_viz = "\n".join(lines)
+
+    # Optionally generate matplotlib figure
+    if output_path:
+        _generate_matplotlib_figure(graph, result, output_path, title)
+
+    return text_viz
+
+
+def _generate_matplotlib_figure(
+    graph: DependenceGraph,
+    result: JointSWPWSResult,
+    output_path: str,
+    title: str,
+):
+    """Generate a matplotlib figure of the schedule (Gantt chart style)."""
+    try:
+        import matplotlib.pyplot as plt
+        import matplotlib.patches as mpatches
+    except ImportError:
+        print("matplotlib not available for figure generation")
+        return
+
+    I = result.I
+    L = result.length
+    n_copies = result.num_copies
+    M = result.schedule
+    wa = result.warp_assignment
+    machine = graph.machine
+
+    fig, ax = plt.subplots(1, 1, figsize=(14, max(6, L * 0.4)))
+
+    # Y-axis: time (cycles), X-axis: functional units
+    fu_names = machine.functional_units
+    n_fus = len(fu_names)
+    bar_width = 0.8
+
+    for v in graph.V:
+        warp = wa.warp_of(v.name)
+        color = WARP_COLORS[warp % len(WARP_COLORS)]
+
+        for i in range(n_copies):
+            abs_time = M[v.name] + i * I
+            if abs_time >= L:
+                continue
+
+            for c in range(v.cycles):
+                for f_idx in range(n_fus):
+                    if v.rrt[c, f_idx] > 0:
+                        rect = mpatches.FancyBboxPatch(
+                            (f_idx - bar_width / 2, abs_time + c),
+                            bar_width, 1,
+                            boxstyle="round,pad=0.05",
+                            facecolor=color,
+                            edgecolor='black',
+                            linewidth=0.5,
+                            alpha=0.8,
+                        )
+                        ax.add_patch(rect)
+                        label = f"{v.name}" if i == 0 else f"{v.name}+{i}"
+                        ax.text(f_idx, abs_time + c + 0.5, label,
+                               ha='center', va='center', fontsize=8,
+                               fontweight='bold', color='white')
+
+    # Formatting
+    ax.set_xlim(-0.5, n_fus - 0.5)
+    ax.set_ylim(-0.5, L + 0.5)
+    ax.set_xticks(range(n_fus))
+    ax.set_xticklabels(fu_names, fontsize=10)
+    ax.set_yticks(range(L))
+    ax.set_ylabel("Clock Cycle", fontsize=12)
+    ax.set_xlabel("Functional Unit", fontsize=12)
+    ax.set_title(f"{title}\nI={I}, L={L}, copies={n_copies}", fontsize=14)
+    ax.invert_yaxis()
+    ax.grid(True, alpha=0.3)
+
+    # Legend for warps
+    legend_patches = []
+    for w in range(machine.num_warps):
+        instrs = wa.instructions_on_warp(w)
+        if instrs:
+            label = wa.warp_names.get(w, f"Warp {w}")
+            legend_patches.append(
+                mpatches.Patch(color=WARP_COLORS[w % len(WARP_COLORS)],
+                             label=f"{label}: {', '.join(instrs)}")
+            )
+    ax.legend(handles=legend_patches, loc='upper right', fontsize=8)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"Schedule figure saved to {output_path}")
+
+
+def print_modular_rrt(
+    graph: DependenceGraph,
+    schedule: ModuloScheduleResult,
+) -> str:
+    """Print the modular RRT as a table."""
+    from twill.modulo_scheduler import compute_modular_rrt
+    
+    mod_rrt = compute_modular_rrt(graph, schedule)
+    I = schedule.I
+    fu_names = graph.machine.functional_units
+
+    lines = [f"Modular RRT (I={I}):"]
+    header = "  t  │ " + " │ ".join(f"{fu:^8s}" for fu in fu_names) + " │"
+    lines.append(header)
+    lines.append("─" * len(header))
+    
+    for t in range(I):
+        row = f" {t:2d}  │ "
+        for f_idx in range(len(fu_names)):
+            val = mod_rrt[t, f_idx]
+            cap = graph.machine.capacity_vector[f_idx]
+            marker = "!" if val > cap else " "
+            row += f" {val:^6d}{marker} │ "
+        lines.append(row)
+
+    return "\n".join(lines)