render_html.py

"""Render parsed CUTLASS pipeline structure as a self-contained HTML/SVG
dataflow diagram. Producer column on the left, SMEM buffers in the middle,
consumer column on the right, loopback arc at the bottom."""

from __future__ import annotations

import html
import json
from pathlib import Path


OP_STYLE = {
    "producer_acquire": ("#c98bff", "acq"),
    "producer_commit": ("#c98bff", "cmt"),
    "tma_load": ("#7cc7ff", "TMA"),
    "consumer_wait": ("#7ee0a6", "wait"),
    "wgmma": ("#ffd27c", "WGMMA"),
    "wgmma_arrive": ("#ffd27c", "wg.arr"),
    "wgmma_wait": ("#ffd27c", "wg.wait"),
    "consumer_release": ("#7ee0a6", "rel"),
}

BUF_COLOR = "#ff9ec7"
BG = "#0e1016"
PANEL = "#161922"
INK = "#e7eaf0"
MUTED = "#93a0b4"
LINE = "#272d3a"

EDGE_COLOR = {
    "acquire→tma": "#c98bff",
    "tma→smem": "#7cc7ff",
    "tma→commit": "#c98bff",
    "tma→mbarrier": "#ff9ec7",
    "smem→wait": "#7ee0a6",
    "wait→wgmma": "#ffd27c",
    "wgmma→release": "#7ee0a6",
    "release→acquire": "#ff7b7b",
}


def _node_id(s: str) -> str:
    return s.replace(":", "_").replace("#", "_").replace("→", "_to_")


def _op_label(op: dict) -> str:
    style_color, short = OP_STYLE.get(op["op"], ("#93a0b4", op["op"]))
    detail = op.get("detail") or {}
    extra = ""
    if op["op"] == "tma_load":
        extra = f" {detail.get('tma_desc','?')} → {detail.get('dst_buffer','?')}"
    elif op["op"] == "wgmma":
        extra = f" {detail.get('src_a','?')}·{detail.get('src_b','?')} → {detail.get('accum','?')}"
    return f"{short}{extra}"


def _lay_out_column(ops: list[dict], x: float, y0: float, dy: float) -> dict[str, tuple[float, float]]:
    positions = {}
    for i, op in enumerate(ops):
        positions[f"L{op['line']}"] = (x, y0 + i * dy)
    return positions


def render(parsed: dict) -> str:
    warp_groups = parsed["warp_groups"]
    buffers = parsed["buffers"]
    edges = parsed["edges"]

    prod = next((w for w in warp_groups if w["role"] == "producer"), None)
    cons = next((w for w in warp_groups if w["role"] == "consumer"), None)

    col_w = 280
    x_prod = 40
    x_buf = x_prod + col_w + 100
    x_cons = x_buf + 220
    y0 = 120
    dy = 62

    box_w = col_w
    box_h = 44

    prod_ops = prod["ops"] if prod else []
    cons_ops = cons["ops"] if cons else []

    prod_pos = _lay_out_column(prod_ops, x_prod, y0, dy) if prod else {}
    cons_pos = _lay_out_column(cons_ops, x_cons, y0, dy) if cons else {}

    buf_h = 60
    buf_gap = 40
    buf_pos = {b["name"]: (x_buf, y0 + i * (buf_h + buf_gap)) for i, b in enumerate(buffers)}

    max_rows = max(len(prod_ops), len(cons_ops), len(buffers) * 2)
    height = y0 + max(max_rows * dy, len(buffers) * (buf_h + buf_gap)) + 200
    width = x_cons + col_w + 40

    def node_center(node_id: str) -> tuple[float, float] | None:
        if node_id.startswith("buffer:"):
            name = node_id.split(":", 1)[1]
            if name in buf_pos:
                bx, by = buf_pos[name]
                return (bx + col_w / 2, by + buf_h / 2)
        if "#L" in node_id:
            side, tag = node_id.rsplit("#", 1)
            if side.startswith("producer") and tag in prod_pos:
                px, py = prod_pos[tag]
                return (px + box_w, py + box_h / 2)
            if side.startswith("consumer") and tag in cons_pos:
                px, py = cons_pos[tag]
                return (px, py + box_h / 2)
        if node_id.startswith("mbarrier:"):
            return (x_buf + col_w / 2, y0 - 60)
        return None

    svg_parts: list[str] = []
    svg_parts.append(
        f'<svg id="graph" width="{width}" height="{height}" '
        f'viewBox="0 0 {width} {height}" xmlns="http://www.w3.org/2000/svg" '
        f'font-family="SFMono-Regular,Consolas,monospace">'
    )
    svg_parts.append(
        '<defs>'
        '<marker id="arr" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="6" markerHeight="6" orient="auto-start-reverse">'
        '<path d="M0,0 L10,5 L0,10 z" fill="#93a0b4"/></marker>'
        + "".join(
            f'<marker id="arr_{i}" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="6" markerHeight="6" orient="auto-start-reverse">'
            f'<path d="M0,0 L10,5 L0,10 z" fill="{c}"/></marker>'
            for i, c in enumerate(EDGE_COLOR.values())
        )
        + "</defs>"
    )

    svg_parts.append(
        f'<text x="{x_prod}" y="{y0 - 60}" fill="{MUTED}" font-size="12">PRODUCER WARP-GROUP</text>'
        f'<text x="{x_prod}" y="{y0 - 42}" fill="{INK}" font-size="16" font-weight="600">'
        f'{html.escape(prod["function"]) if prod else "(none)"}</text>'
    )
    svg_parts.append(
        f'<text x="{x_buf}" y="{y0 - 60}" fill="{MUTED}" font-size="12">SHARED MEMORY BUFFERS</text>'
        f'<text x="{x_buf}" y="{y0 - 42}" fill="{INK}" font-size="16" font-weight="600">'
        f'{len(buffers)} × {html.escape(str(parsed.get("stages_expr","?")))} stages</text>'
    )
    svg_parts.append(
        f'<text x="{x_cons}" y="{y0 - 60}" fill="{MUTED}" font-size="12">CONSUMER WARP-GROUP</text>'
        f'<text x="{x_cons}" y="{y0 - 42}" fill="{INK}" font-size="16" font-weight="600">'
        f'{html.escape(cons["function"]) if cons else "(none)"}</text>'
    )

    for op in prod_ops:
        tag = f"L{op['line']}"
        x, y = prod_pos[tag]
        color, _ = OP_STYLE.get(op["op"], (MUTED, op["op"]))
        node_id = _node_id(f"{prod['id']}#{tag}")
        svg_parts.append(
            f'<g class="node" data-id="{prod["id"]}#{tag}" data-line="{op["line"]}" '
            f'data-detail="{html.escape(json.dumps(op["detail"]))}">'
            f'<rect x="{x}" y="{y}" width="{box_w}" height="{box_h}" rx="8" ry="8" '
            f'fill="{PANEL}" stroke="{color}" stroke-width="1.5"/>'
            f'<text x="{x + 12}" y="{y + 18}" fill="{INK}" font-size="12">{html.escape(_op_label(op))}</text>'
            f'<text x="{x + 12}" y="{y + 34}" fill="{MUTED}" font-size="10">L{op["line"]}</text>'
            f'</g>'
        )

    for op in cons_ops:
        tag = f"L{op['line']}"
        x, y = cons_pos[tag]
        color, _ = OP_STYLE.get(op["op"], (MUTED, op["op"]))
        svg_parts.append(
            f'<g class="node" data-id="{cons["id"]}#{tag}" data-line="{op["line"]}" '
            f'data-detail="{html.escape(json.dumps(op["detail"]))}">'
            f'<rect x="{x}" y="{y}" width="{box_w}" height="{box_h}" rx="8" ry="8" '
            f'fill="{PANEL}" stroke="{color}" stroke-width="1.5"/>'
            f'<text x="{x + 12}" y="{y + 18}" fill="{INK}" font-size="12">{html.escape(_op_label(op))}</text>'
            f'<text x="{x + 12}" y="{y + 34}" fill="{MUTED}" font-size="10">L{op["line"]}</text>'
            f'</g>'
        )

    for b in buffers:
        bx, by = buf_pos[b["name"]]
        stages_label = html.escape(str(b.get("stages", "?")))
        svg_parts.append(
            f'<g class="node buffer" data-id="buffer:{b["name"]}" '
            f'data-detail="{html.escape(json.dumps(b))}">'
            f'<rect x="{bx}" y="{by}" width="{col_w}" height="{buf_h}" rx="8" ry="8" '
            f'fill="{PANEL}" stroke="{BUF_COLOR}" stroke-width="1.5" stroke-dasharray="6 3"/>'
            f'<text x="{bx + 12}" y="{by + 22}" fill="{INK}" font-size="13" font-weight="600">{html.escape(b["name"])}</text>'
            f'<text x="{bx + 12}" y="{by + 40}" fill="{MUTED}" font-size="10">'
            f'{html.escape(b["layout"])} · {stages_label} stages</text>'
        )
        stage_bar_x = bx + col_w - 90
        stage_bar_y = by + 32
        for si in range(3):
            svg_parts.append(
                f'<rect x="{stage_bar_x + si * 24}" y="{stage_bar_y}" width="20" height="12" '
                f'rx="2" ry="2" fill="{BG}" stroke="{BUF_COLOR}" stroke-width="0.5"/>'
            )
        svg_parts.append(
            f'<text x="{stage_bar_x + 72}" y="{stage_bar_y + 10}" fill="{MUTED}" font-size="9">...</text>'
        )
        svg_parts.append("</g>")

    edge_colors = list(EDGE_COLOR.items())
    color_idx = {k: i for i, (k, _) in enumerate(edge_colors)}
    for e in edges:
        src = node_center(e["src"])
        dst = node_center(e["dst"])
        if src is None or dst is None:
            continue
        color = EDGE_COLOR.get(e["kind"], "#93a0b4")
        marker_id = f"arr_{color_idx.get(e['kind'], 0)}"
        x1, y1 = src
        x2, y2 = dst
        if e["kind"] == "release→acquire":
            path_d = (
                f"M{x1},{y1} "
                f"C{x1 + 120},{y1 + 180} {x2 - 120},{y2 + 180} {x2},{y2}"
            )
            stroke_dash = 'stroke-dasharray="6 4"'
        elif e["kind"] in ("acquire→tma", "tma→commit", "wait→wgmma", "wgmma→release"):
            mid = (x1 + x2) / 2
            path_d = f"M{x1},{y1} C{mid},{y1} {mid},{y2} {x2},{y2}"
            stroke_dash = ""
        else:
            mid_x = (x1 + x2) / 2
            path_d = f"M{x1},{y1} C{mid_x},{y1} {mid_x},{y2} {x2},{y2}"
            stroke_dash = 'stroke-dasharray="4 3"' if e["kind"] == "smem→wait" else ""
        svg_parts.append(
            f'<path class="edge" data-src="{e["src"]}" data-dst="{e["dst"]}" '
            f'd="{path_d}" fill="none" stroke="{color}" stroke-width="1.3" '
            f'marker-end="url(#{marker_id})" {stroke_dash} opacity="0.55"/>'
        )
        if e.get("label"):
            lx, ly = (x1 + x2) / 2, (y1 + y2) / 2 - 6
            svg_parts.append(
                f'<text x="{lx}" y="{ly}" fill="{MUTED}" font-size="9" text-anchor="middle">{html.escape(e["label"])}</text>'
            )

    legend_y = y0 + max(len(prod_ops), len(cons_ops), len(buffers) * 2) * dy + 60
    lx = x_prod
    for i, (kind, color) in enumerate(edge_colors):
        rx = lx + (i % 4) * 220
        ry = legend_y + (i // 4) * 22
        svg_parts.append(
            f'<line x1="{rx}" y1="{ry}" x2="{rx + 30}" y2="{ry}" stroke="{color}" stroke-width="2"/>'
            f'<text x="{rx + 38}" y="{ry + 4}" fill="{MUTED}" font-size="11">{html.escape(kind)}</text>'
        )

    svg_parts.append("</svg>")

    svg = "".join(svg_parts)

    parsed_json = json.dumps(parsed, indent=2)

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>CUTLASS Pipeline Dataflow · {html.escape(Path(parsed["source_file"]).name)}</title>
<style>
  body {{ margin: 0; background: {BG}; color: {INK};
    font: 14px/1.5 -apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif; }}
  header {{ padding: 16px 24px 12px; border-bottom: 1px solid {LINE};
    background: linear-gradient(180deg,#192033,{BG}); }}
  h1 {{ margin: 0 0 4px; font-size: 18px; }}
  header p {{ margin: 0; color: {MUTED}; font-size: 12px;
    font-family: "SFMono-Regular",Consolas,monospace; }}
  main {{ display: grid; grid-template-columns: 1fr 340px; }}
  .diag {{ padding: 24px; overflow: auto; }}
  .side {{ background: {PANEL}; border-left: 1px solid {LINE}; padding: 18px;
    height: calc(100vh - 60px); overflow: auto; }}
  .side h2 {{ font-size: 12px; text-transform: uppercase; color: {MUTED};
    letter-spacing: .06em; margin: 0 0 8px; }}
  .side pre {{ font-family: "SFMono-Regular",Consolas,monospace; font-size: 11px;
    background: {BG}; color: {INK}; padding: 10px; border-radius: 6px;
    border: 1px solid {LINE}; overflow: auto; white-space: pre-wrap; }}
  .side .kv {{ display: grid; grid-template-columns: 100px 1fr; gap: 4px 10px;
    font-size: 12px; margin-bottom: 12px; }}
  .side .kv b {{ color: {MUTED}; font-weight: normal; }}
  .node {{ cursor: pointer; }}
  .node.dim {{ opacity: 0.25; }}
  .edge {{ transition: opacity 0.15s, stroke-width 0.15s; }}
  .edge.dim {{ opacity: 0.08; }}
  .edge.hi {{ opacity: 1; stroke-width: 2.2; }}
</style>
</head>
<body>
<header>
  <h1>CUTLASS Pipeline Dataflow</h1>
  <p>{html.escape(parsed["source_file"])}</p>
</header>
<main>
  <div class="diag">{svg}</div>
  <aside class="side">
    <h2>Selection</h2>
    <div id="sel">Click a node to inspect it.</div>
    <h2 style="margin-top:22px">Pipeline meta</h2>
    <div class="kv">
      <b>type</b><span>{html.escape(parsed["pipelines"][0]["type"] if parsed["pipelines"] else "?")}</span>
      <b>stages</b><span>{html.escape(str(parsed.get("stages_expr","?")))}</span>
      <b>buffers</b><span>{len(parsed["buffers"])}</span>
      <b>producer</b><span>{html.escape(prod["function"] if prod else "-")}</span>
      <b>consumer</b><span>{html.escape(cons["function"] if cons else "-")}</span>
    </div>
    <h2 style="margin-top:22px">Parsed JSON</h2>
    <pre id="raw">{html.escape(parsed_json)}</pre>
  </aside>
</main>
<script>
const nodes = document.querySelectorAll('.node');
const edges = document.querySelectorAll('.edge');
const sel = document.getElementById('sel');

function reset() {{
  nodes.forEach(n => n.classList.remove('dim'));
  edges.forEach(e => {{ e.classList.remove('dim'); e.classList.remove('hi'); }});
}}

nodes.forEach(n => {{
  n.addEventListener('click', ev => {{
    ev.stopPropagation();
    const id = n.dataset.id;
    const line = n.dataset.line;
    const detail = n.dataset.detail ? JSON.parse(n.dataset.detail) : {{}};
    reset();
    nodes.forEach(m => {{ if (m !== n) m.classList.add('dim'); }});
    edges.forEach(e => {{
      if (e.dataset.src === id || e.dataset.dst === id) {{
        e.classList.add('hi');
        const other = e.dataset.src === id ? e.dataset.dst : e.dataset.src;
        document.querySelectorAll('.node').forEach(m => {{
          if (m.dataset.id === other) m.classList.remove('dim');
        }});
      }} else {{
        e.classList.add('dim');
      }}
    }});
    sel.innerHTML = '<div class="kv">' +
      '<b>id</b><span>' + id + '</span>' +
      (line ? '<b>line</b><span>' + line + '</span>' : '') +
      '</div><pre style="background:#0e1016;padding:8px;border-radius:6px;border:1px solid #272d3a;font-size:11px">' +
      JSON.stringify(detail, null, 2) + '</pre>';
  }});
}});

document.addEventListener('click', () => {{ reset(); sel.textContent = 'Click a node to inspect it.'; }});
</script>
</body>
</html>
"""