From 7a3a808ae861362aec5812ba4debf9ead5b82e36 Mon Sep 17 00:00:00 2001 From: ilgeco Date: Thu, 6 Aug 2026 14:34:22 +0200 Subject: [PATCH] Some tool drawio and sequence diagram --- .../tools/conv_lowering_tiles_drawio.py | 567 ++++++++++++++++++ validation/tools/pim_sequence_diagram.py | 438 ++++++++++++++ .../tools/pimcomp/compare_raptor_pimcomp.py | 48 +- .../pimcomp/run_pimcomp_paper_latency.py | 262 ++++++-- 4 files changed, 1246 insertions(+), 69 deletions(-) create mode 100644 validation/tools/conv_lowering_tiles_drawio.py create mode 100644 validation/tools/pim_sequence_diagram.py diff --git a/validation/tools/conv_lowering_tiles_drawio.py b/validation/tools/conv_lowering_tiles_drawio.py new file mode 100644 index 0000000..bdfe437 --- /dev/null +++ b/validation/tools/conv_lowering_tiles_drawio.py @@ -0,0 +1,567 @@ +#!/usr/bin/env python3 +"""Generate an editable diagrams.net library for Conv lowering comparisons.""" + +from __future__ import annotations + +import argparse +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from pathlib import Path + + +TILE = 760 +GAP = 40 +COLS = 4 +INK = "#172033" +MUTED = "#667085" +GRID = "#d7deea" +PALE = "#f8fafc" +REFERENCE = "#5f6b7a" +PIMCOMP = "#ef8354" +RAPTOR = "#3b82f6" +INPUT = ("#f4a261", "#52b788", "#4895ef") +WEIGHT = ("#c86418", "#237a57", "#2768b2") +OUTPUT = ("#8b5cf6", "#ec4899", "#06b6d4", "#eab308") +SPATIAL = ("00", "01", "02", "10", "11", "12", "20", "21", "22") + + +@dataclass(frozen=True) +class Tile: + slug: str + owner: str + title: str + subtitle: str + scene: str + formula: str + algorithm: str + + +TILES = ( + Tile("classic-reference", "REFERENCE", "Classic Conv + exact weight unfolding", + "Original OIHW weights; the two implementations choose different K orders.", + "reference", "Y[p,o] = Σc,kh,kw Xpatch[p,c,kh,kw] · W[o,c,kh,kw]", + "At every output position, multiply the patch by one filter and add every product."), + Tile("pimcomp-element", "PIMCOMP", "Element pipeline", + "One patch vector per input cycle; mapped weights stay fixed.", + "pimcomp_element", "patchPIM[1×K] · WflatPIM[K×O] → Yp[1×O]", + "Keep Wflat in the arrays; stream one patch each cycle to produce all O outputs."), + Tile("pimcomp-batch", "PIMCOMP", "Batch / replicated pipeline", + "Complete Wflat copies divide patches or input samples.", + "pimcomp_batch", "for replica r: Yr = patchr[1×K] · WflatPIM[K×O]", + "Copy all weights R times and send different patches to the copies in parallel."), + Tile("raptor-legacy-im2col", "RAPTOR", "Legacy explicit im2col", + "Every patch becomes one row of a global P×K matrix.", + "legacy", "Y[P×O] = im2col(X)[P×K] · WflatR[K×O]", + "Write every image patch as one matrix row, then multiply the two large matrices."), + Tile("raptor-packed-im2col", "RAPTOR", "Packed im2col", + "Pack q patch rows and repeat Wflat on a block diagonal.", + "packed", "packedY[1×qO] = [patch0|…|patchq−1] · diag(WflatR,…,WflatR)", + "Join q patches and use diagonal weight copies so one multiply computes q independent outputs."), + Tile("raptor-streamed-patch", "RAPTOR", "Streamed patch", + "Gather one patch into bounded scratch; avoid global im2col.", + "streamed_patch", "Yp[1×O] = scratchPatchp[1×K] · WflatR[K×O]", + "Gather one patch, multiply it, write its output, and reuse scratch for the next patch."), + Tile("raptor-streamed-packed", "RAPTOR", "Streamed packed", + "Gather q patch rows in bounded scratch, then block-diagonal pack them.", + "streamed_packed", "packedY = packedScratch[1×qK] · diag(WflatR×q)[qK×qO]", + "Gather q patches in small scratch, join them, multiply by diagonal weights, then unpack q outputs."), + Tile("raptor-depthwise", "RAPTOR", "Depthwise special case", + "Each channel owns one row-major 3×3 kernel; channels never reduce together.", + "depthwise", "Y[p,c] = Σkh,kw Xpatch[p,c,kh,kw] · W[c,kh,kw]", + "For each channel separately, multiply its nine patch values by its nine weights and add."), + Tile("raptor-output-channel-tiled", "RAPTOR", "Output-channel tiled", + "Every O tile retains all channel-major K rows and selects output columns.", + "c_tiled", "Y[:,Oj] = patch[1×K] · Wflat[:,Oj][K×|Oj|]; concat j", + "Reuse the full patch for each output-filter group, then join the output groups."), + Tile("raptor-input-k-tiled", "RAPTOR", "Input-K tiled", + "Split matching K ranges; add their partial output vectors.", + "k_tiled", "Y[1×O] = Σi patch[Ki] · Wflat[Ki,:]", + "Multiply matching K slices independently, then add their partial output vectors."), + Tile("raptor-tiled-2d", "RAPTOR", "Two-dimensional tiled", + "Partition both K rows and output-filter columns.", + "tiled_2d", "Y[:,Oj] = Σi patch[Ki] · Wflat[Ki,Oj]; concat j", + "Split both directions: add results down K and join results across output groups."), + Tile("raptor-row-strip", "RAPTOR", "Pixel-major row-strip", + "A lane forms patches across one output row and slices K.", + "row_strip", "for x: Y[r,x,:] = Σi patch[r,x,Ki] · Wflat[Ki,:]", + "Move across one output row; at each x form a patch, multiply its K slices, and add."), + Tile("raptor-row-strip-c-tiled", "RAPTOR", "Row-strip + output tiling", + "Each row lane is duplicated across disjoint output-column tiles.", + "row_strip_c", "for x,j: Y[r,x,Oj] = patch[r,x,:] · Wflat[:,Oj]", + "Give each output-filter group a copy of the row lane, then join their output columns."), +) + + +class Drawio: + def __init__(self) -> None: + rows = (len(TILES) + COLS - 1) // COLS + self.mxfile = ET.Element("mxfile", host="app.diagrams.net", compressed="false") + diagram = ET.SubElement(self.mxfile, "diagram", id="conv-lowering-library", + name="Conv lowering tile library") + model = ET.SubElement( + diagram, "mxGraphModel", dx="1200", dy="900", grid="1", gridSize="10", + guides="1", tooltips="1", connect="1", arrows="1", fold="1", + page="0", pageScale="1", pageWidth=str(COLS * (TILE + GAP)), + pageHeight=str(rows * (TILE + GAP)), math="0", shadow="0", + ) + self.root = ET.SubElement(model, "root") + ET.SubElement(self.root, "mxCell", id="0") + ET.SubElement(self.root, "mxCell", id="1", parent="0") + self.counter = 2 + + def _id(self, prefix: str = "c") -> str: + value = f"{prefix}-{self.counter}" + self.counter += 1 + return value + + def vertex(self, parent: str, x: float, y: float, w: float, h: float, + value: str = "", style: str = "", *, cell_id: str | None = None) -> str: + cell_id = cell_id or self._id() + cell = ET.SubElement(self.root, "mxCell", id=cell_id, value=value, + style=style, vertex="1", parent=parent) + ET.SubElement(cell, "mxGeometry", x=str(x), y=str(y), width=str(w), + height=str(h), **{"as": "geometry"}) + return cell_id + + def group(self, x: float, y: float, name: str) -> str: + return self.vertex("1", x, y, TILE, TILE, name, + "group;connectable=0;", cell_id=f"tile-{name}") + + def edge(self, parent: str, source: str, target: str, value: str = "") -> str: + style = ( + "edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;" + f"html=1;endArrow=block;endFill=1;strokeWidth=2;strokeColor={INK};" + f"fontSize=10;fontColor={INK};labelBackgroundColor=#ffffff;" + ) + cell_id = self._id("e") + cell = ET.SubElement(self.root, "mxCell", id=cell_id, value=value, + style=style, edge="1", parent=parent, + source=source, target=target) + ET.SubElement(cell, "mxGeometry", relative="1", **{"as": "geometry"}) + return cell_id + + def write(self, path: Path) -> None: + ET.indent(self.mxfile, space=" ") + path.parent.mkdir(parents=True, exist_ok=True) + ET.ElementTree(self.mxfile).write(path, encoding="utf-8", xml_declaration=True) + + +def rect_style(fill: str = "#ffffff", stroke: str = GRID, *, rounded: bool = True, + font_size: int = 11, font_color: str = INK, align: str = "center", + stroke_width: int = 1) -> str: + return ( + f"rounded={int(rounded)};whiteSpace=wrap;html=1;fillColor={fill};" + f"strokeColor={stroke};strokeWidth={stroke_width};fontColor={font_color};" + f"fontSize={font_size};fontFamily=Helvetica;align={align};verticalAlign=middle;" + ) + + +def text_style(size: int, *, color: str = INK, align: str = "left", bold: bool = False) -> str: + return ( + "text;html=1;strokeColor=none;fillColor=none;whiteSpace=wrap;" + f"fontSize={size};fontColor={color};fontFamily=Helvetica;align={align};" + f"verticalAlign=middle;fontStyle={1 if bold else 0};" + ) + + +def add_text(d: Drawio, parent: str, x: float, y: float, w: float, h: float, + value: str, size: int = 11, *, color: str = INK, + align: str = "left", bold: bool = False) -> str: + return d.vertex(parent, x, y, w, h, value, + text_style(size, color=color, align=align, bold=bold)) + + +def add_box(d: Drawio, parent: str, x: float, y: float, w: float, h: float, + value: str, *, fill: str = "#ffffff", stroke: str = GRID, + size: int = 11, rounded: bool = True, stroke_width: int = 1) -> str: + return d.vertex(parent, x, y, w, h, value, + rect_style(fill, stroke, rounded=rounded, font_size=size, + stroke_width=stroke_width)) + + +def matrix3(d: Drawio, parent: str, x: float, y: float, size: float, + color: str, *, labels: bool = True) -> None: + cell = size / 3 + for row in range(3): + for col in range(3): + value = SPATIAL[row * 3 + col] if labels else "" + d.vertex(parent, x + col * cell, y + row * cell, cell, cell, value, + rect_style(color, "#ffffff", rounded=False, + font_size=7, stroke_width=1)) + d.vertex(parent, x, y, size, size, "", + "rounded=0;whiteSpace=wrap;html=1;fillColor=none;" + f"strokeColor={INK};strokeWidth=1;") + + +def source_panel(d: Drawio, parent: str) -> None: + add_text(d, parent, 34, 132, 230, 24, "Input patch X[:,C,3,3]", 12, bold=True) + for channel, color in enumerate(INPUT): + x = 34 + channel * 70 + matrix3(d, parent, x, 166, 48, color) + add_text(d, parent, x, 216, 48, 18, f"C{channel}", 9, + color=color, align="center", bold=True) + + add_text(d, parent, 340, 132, 380, 24, "Original W[O,C,3,3] (OIHW)", + 12, bold=True, align="center") + for output, outline in enumerate(OUTPUT): + x = 340 + output * 96 + add_box(d, parent, x, 158, 88, 78, "", fill="#ffffff", + stroke=outline, stroke_width=2) + for channel, color in enumerate(WEIGHT): + gx = x + 7 + channel * 25 + matrix3(d, parent, gx, 181, 19, color, labels=False) + add_text(d, parent, gx, 163, 19, 16, f"C{channel}", 7, + color=color, align="center", bold=True) + add_text(d, parent, x, 216, 88, 18, f"filter O{output}", 9, + color=outline, align="center", bold=True) + + +def order_sequence(order: str) -> list[tuple[int, str]]: + if order == "raptor": + return [(channel, spatial) for channel in range(3) for spatial in SPATIAL] + if order == "pimcomp": + return [(channel, spatial) for spatial in SPATIAL for channel in range(3)] + raise ValueError(order) + + +def order_vector(d: Drawio, parent: str, y: float, order: str, *, + weight: bool, label: str) -> None: + x, width, height = 125, 595, 28 + sequence = order_sequence(order) + cell_w = width / len(sequence) + colors = WEIGHT if weight else INPUT + for index, (channel, spatial) in enumerate(sequence): + value = f"w{spatial}" if weight else (spatial if order == "raptor" else f"C{channel}") + font_color = "#ffffff" if weight else INK + d.vertex(parent, x + index * cell_w, y, cell_w, height, value, + rect_style(colors[channel], "#ffffff", rounded=False, + font_size=7, font_color=font_color)) + d.vertex(parent, x, y, width, height, "", + f"rounded=0;fillColor=none;strokeColor={INK};strokeWidth=1;") + add_text(d, parent, 35, y, 82, height, label, 8, color=MUTED, + align="right", bold=True) + if order == "raptor": + for channel in range(3): + add_text(d, parent, x + channel * width / 3, y - 18, width / 3, 16, + f"C{channel}: row-major 00→01→02→10→…→22", 8, + color=colors[channel], align="center", bold=True) + else: + for spatial_index, spatial in enumerate(SPATIAL): + add_text(d, parent, x + spatial_index * width / 9, y - 18, + width / 9, 16, f"({spatial}) C0,C1,C2", 7, + color=MUTED, align="center", bold=True) + + +def node(d: Drawio, parent: str, x: float, y: float, w: float, h: float, + label: str, *, fill: str = PALE, stroke: str = RAPTOR) -> str: + return add_box(d, parent, x, y, w, h, label, fill=fill, stroke=stroke, + size=10, stroke_width=2) + + +def layout_strip(d: Drawio, parent: str, x: float, y: float, w: float, h: float, + order: str, colors: tuple[str, ...], *, channel: int | None = None, + repeat: int = 1) -> None: + sequence = ([channel] * 9 if channel is not None + else [item[0] for item in order_sequence(order)]) + row_h = h / repeat + for copy in range(repeat): + cell_w = w / len(sequence) + for index, color_index in enumerate(sequence): + d.vertex(parent, x + index * cell_w, y + copy * row_h, + cell_w, row_h, "", + rect_style(colors[color_index], "#ffffff", rounded=False, + font_size=1)) + d.vertex(parent, x, y, w, h, "", + f"rounded=0;fillColor=none;strokeColor={INK};strokeWidth=1;") + + +def layout_node(d: Drawio, parent: str, x: float, y: float, w: float, h: float, + label: str, *, order: str | None = None, + colors: tuple[str, ...] = INPUT, channel: int | None = None, + repeat: int = 1, stroke: str = RAPTOR) -> str: + result = node(d, parent, x, y, w, h, "", fill="#ffffff", stroke=stroke) + label_height = h - (28 if order is not None else 10) + add_text(d, parent, x + 6, y + 5, w - 12, label_height, label, 10, + align="center", bold=True) + if order is not None: + layout_strip(d, parent, x + 8, y + h - 19, w - 16, 12, order, colors, + channel=channel, repeat=repeat) + return result + + +def scene_linear(d: Drawio, parent: str, + items: tuple[tuple[str, dict | None], ...], *, + y: float = 320) -> None: + margin, gap = 42, 34 + width = (TILE - 2 * margin - gap * (len(items) - 1)) / len(items) + ids = [] + for index, (label, layout) in enumerate(items): + x = margin + index * (width + gap) + ids.append(layout_node(d, parent, x, y, width, 92, label, + **(layout or {}))) + for left, right in zip(ids, ids[1:]): + d.edge(parent, left, right) + + +def scene_batch(d: Drawio, parent: str) -> None: + for row in range(3): + y = 260 + row * 110 + patch = layout_node(d, parent, 42, y, 210, 72, + f"patches p{row}, p{row + 3}, … [1×K]", + order="pimcomp", colors=INPUT, stroke=PIMCOMP) + weights = layout_node(d, parent, 302, y, 220, 72, + f"replica R{row}: Wflat [K×O]", + order="pimcomp", colors=WEIGHT, stroke=PIMCOMP) + result = node(d, parent, 610, y, 106, 72, f"Yp\n[1×O]", + fill=OUTPUT[row], stroke=PIMCOMP) + d.edge(parent, patch, weights, "×") + d.edge(parent, weights, result) + + +def scene_depthwise(d: Drawio, parent: str) -> None: + for channel in range(3): + y = 260 + channel * 110 + patch = layout_node(d, parent, 46, y, 190, 72, + f"patch C{channel} [1×9]", order="raptor", + colors=INPUT, channel=channel) + kernel = layout_node(d, parent, 300, y, 210, 72, + f"W[C{channel},0,:,:] [9×1]", order="raptor", + colors=WEIGHT, channel=channel) + result = node(d, parent, 578, y, 136, 72, + f"Y channel {channel}", fill=OUTPUT[channel]) + d.edge(parent, patch, kernel, "×") + d.edge(parent, kernel, result) + + +def scene_c_tiled(d: Drawio, parent: str) -> None: + for row, outputs in enumerate(((0, 1), (2, 3))): + y = 280 + row * 130 + patch = layout_node(d, parent, 42, y, 190, 84, + "same full patch [1×K]", order="raptor", + colors=INPUT) + weights = layout_node( + d, parent, 292, y, 244, 84, + f"Wflat[:,O{outputs[0]}:O{outputs[-1] + 1}]\n[K×2], all K rows", + order="raptor", colors=WEIGHT, stroke=OUTPUT[outputs[0]]) + result = node(d, parent, 610, y, 106, 84, + f"Y tile {row}\n[1×2]", fill=OUTPUT[outputs[0]]) + d.edge(parent, patch, weights, "×") + d.edge(parent, weights, result) + add_text(d, parent, 250, 552, 260, 22, "concatenate tile 0 | tile 1 along O", + 10, color=MUTED, align="center", bold=True) + + +def scene_k_tiled(d: Drawio, parent: str) -> None: + for row in range(3): + y = 250 + row * 100 + patch = layout_node(d, parent, 42, y, 188, 70, + f"patch Ki{row}: C{row} [1×9]", order="raptor", + colors=INPUT, channel=row) + weights = layout_node(d, parent, 298, y, 224, 70, + f"Wflat[Ki{row},:] [9×O]", order="raptor", + colors=WEIGHT, channel=row) + partial = node(d, parent, 596, y, 120, 70, + f"partial {row}\n[1×O]", fill=PALE) + d.edge(parent, patch, weights, "×") + d.edge(parent, weights, partial) + node(d, parent, 300, 570, 160, 46, "VADD Σ → Y [1×O]", fill="#ffffff") + add_text(d, parent, 470, 578, 238, 28, + "partial 0 + partial 1 + partial 2", 9, + color=MUTED, align="center", bold=True) + + +def scene_2d(d: Drawio, parent: str) -> None: + for row in range(3): + y = 250 + row * 100 + patch = layout_node(d, parent, 38, y, 160, 70, + f"patch Ki{row} [1×9]", order="raptor", + colors=INPUT, channel=row) + for col in range(2): + layout_node(d, parent, 270 + col * 230, y, 180, 70, + f"× Wflat[Ki{row},Oj{col}] [9×2]", + order="raptor", colors=WEIGHT, channel=row, + stroke=OUTPUT[col * 2]) + add_text(d, parent, 270, 566, 440, 24, + "Σ tile rows along K; concatenate tile columns along O", 10, + color=MUTED, align="center", bold=True) + + +def scene_row_strip_c(d: Drawio, parent: str) -> None: + patch = layout_node(d, parent, 42, 330, 206, 88, + "row-window patch for x [1×K]", order="raptor", + colors=INPUT) + for col, outputs in enumerate(((0, 1), (2, 3))): + y = 270 + col * 150 + weights = layout_node( + d, parent, 310, y, 238, 88, + f"lane r × O{outputs[0]}:O{outputs[-1] + 1}\nWflat[:,Oj] [K×2]", + order="raptor", colors=WEIGHT, stroke=OUTPUT[outputs[0]]) + result = node(d, parent, 614, y, 104, 88, + f"Yj\n[1×2]", fill=OUTPUT[outputs[0]]) + d.edge(parent, patch, weights, "reuse ×") + d.edge(parent, weights, result) + node(d, parent, 500, 570, 190, 46, "concat O → output row [1×4]") + + +def operation_scene(d: Drawio, parent: str, scene: str) -> None: + add_text(d, parent, 30, 204, 700, 22, "LOWERED COMPUTE", 11, + color=MUTED, align="center", bold=True) + if scene == "pimcomp_element": + scene_linear(d, parent, ( + ("input cycle p\npatchPIM [1×K]", + {"order": "pimcomp", "colors": INPUT, "stroke": PIMCOMP}), + ("mapped Array Group\nWflatPIM [K×O]", + {"order": "pimcomp", "colors": WEIGHT, "stroke": PIMCOMP}), + ("Yp [1×O]", {"stroke": PIMCOMP}), + )) + elif scene == "pimcomp_batch": + scene_batch(d, parent) + elif scene == "legacy": + scene_linear(d, parent, ( + ("global im2col\n[P×K]", {"order": "raptor", "colors": INPUT}), + ("WflatR\n[K×O]", {"order": "raptor", "colors": WEIGHT}), + ("Y rows\n[P×O]", None), + )) + elif scene == "packed": + scene_linear(d, parent, ( + ("q patch rows\n[q×K]", + {"order": "raptor", "colors": INPUT, "repeat": 2}), + ("pack → [1×qK]", + {"order": "raptor", "colors": INPUT, "repeat": 2}), + ("diag(WflatR×q)\n[qK×qO]", + {"order": "raptor", "colors": WEIGHT, "repeat": 2}), + ("packed Y\n[1×qO]", None), + )) + elif scene == "streamed_patch": + scene_linear(d, parent, ( + ("gather one patch", None), + ("bounded scratch\n[1×K]", {"order": "raptor", "colors": INPUT}), + ("WflatR\n[K×O]", {"order": "raptor", "colors": WEIGHT}), + ("Yp\n[1×O]", None), + )) + elif scene == "streamed_packed": + scene_linear(d, parent, ( + ("q patches\n[q×K]", + {"order": "raptor", "colors": INPUT, "repeat": 2}), + ("packedScratch\n[1×qK]", + {"order": "raptor", "colors": INPUT, "repeat": 2}), + ("diag(WflatR×q)\n[qK×qO]", + {"order": "raptor", "colors": WEIGHT, "repeat": 2}), + ("packed Y\n[1×qO]", None), + )) + elif scene == "depthwise": + scene_depthwise(d, parent) + elif scene == "c_tiled": + scene_c_tiled(d, parent) + elif scene == "k_tiled": + scene_k_tiled(d, parent) + elif scene == "tiled_2d": + scene_2d(d, parent) + elif scene == "row_strip": + scene_linear(d, parent, ( + ("lane r row windows", None), + ("for x: patch\n[1×K]", {"order": "raptor", "colors": INPUT}), + ("K-sliced Wflat\n[Ki×O]", {"order": "raptor", "colors": WEIGHT}), + ("Σ partials →\noutput row", None), + )) + elif scene == "row_strip_c": + scene_row_strip_c(d, parent) + else: + raise ValueError(scene) + + +def reference_body(d: Drawio, parent: str) -> None: + add_box(d, parent, 24, 252, 712, 158, "", fill="#ffffff", stroke=PIMCOMP) + add_text(d, parent, 40, 260, 680, 22, + "PIMCOMP: spatial-major, row-wise positions; C interleaved", 11, + color=PIMCOMP, align="center", bold=True) + order_vector(d, parent, 306, "pimcomp", weight=False, label="Input patch") + order_vector(d, parent, 366, "pimcomp", weight=True, label="Matching W") + add_text(d, parent, 40, 394, 680, 16, + "k=((kh·Kw)+kw)·Cin+c — c changes fastest", 9, + color=PIMCOMP, align="center", bold=True) + + add_box(d, parent, 24, 424, 712, 158, "", fill="#ffffff", stroke=RAPTOR) + add_text(d, parent, 40, 432, 680, 22, + "RAPTOR: channel-major; each 3×3 plane is row-major", 11, + color=RAPTOR, align="center", bold=True) + order_vector(d, parent, 478, "raptor", weight=False, label="Input patch") + order_vector(d, parent, 538, "raptor", weight=True, label="Matching W") + add_text(d, parent, 40, 566, 680, 16, + "k=((c·Kh)+kh)·Kw+kw — kw changes fastest", 9, + color=RAPTOR, align="center", bold=True) + add_box(d, parent, 120, 590, 520, 42, "", fill=PALE, stroke=GRID) + add_text(d, parent, 132, 594, 496, 34, + "Shade key: light = activation; dark = matching weight row. " + "Depthwise: independent C0/C1/C2 row-major Kc=9 vectors.", + 9, color=MUTED, align="center", bold=True) + + +def layout_reference(d: Drawio, parent: str, tile: Tile, accent: str) -> None: + if tile.scene == "depthwise": + layout = "independent row-major Kc=9 per channel" + elif tile.owner == "PIMCOMP": + layout = "PIMCOMP spatial-major K order" + else: + layout = "RAPTOR channel-major K order" + add_box(d, parent, 24, 140, 712, 42, "", fill=PALE, stroke=accent) + add_text(d, parent, 38, 146, 684, 30, + f"LAYOUT → see REFERENCE tile: {layout}", 10, + color=accent, align="center", bold=True) + + +def algorithm_card(d: Drawio, parent: str, tile: Tile, accent: str) -> None: + add_box(d, parent, 24, 638, 712, 102, "", fill="#ffffff", stroke=accent) + add_text(d, parent, 40, 646, 90, 34, "ALGORITHM", 9, + color=accent, bold=True) + add_text(d, parent, 132, 644, 588, 38, tile.algorithm, 10) + add_text(d, parent, 40, 690, 90, 34, "MATH", 9, + color=accent, bold=True) + add_text(d, parent, 132, 686, 588, 42, tile.formula, 10, + color=accent, bold=True) + + +def render_tile(d: Drawio, tile: Tile, index: int) -> None: + col, row = index % COLS, index // COLS + parent = d.group(col * (TILE + GAP), row * (TILE + GAP), tile.slug) + accent = {"REFERENCE": REFERENCE, "PIMCOMP": PIMCOMP, "RAPTOR": RAPTOR}[tile.owner] + add_box(d, parent, 0, 0, TILE, TILE, "", fill="#fbfcff", stroke=accent, + stroke_width=3) + add_box(d, parent, 24, 20, 106, 28, tile.owner, fill=accent, stroke=accent, + size=10) + add_text(d, parent, 24, 56, 712, 36, tile.title, 22, bold=True) + add_text(d, parent, 24, 92, 712, 30, tile.subtitle, 11, color=MUTED) + if tile.scene == "reference": + source_panel(d, parent) + reference_body(d, parent) + else: + layout_reference(d, parent, tile, accent) + operation_scene(d, parent, tile.scene) + algorithm_card(d, parent, tile, accent) + + +def validate(root: ET.Element) -> None: + ids = [cell.get("id") for cell in root.findall(".//mxCell")] + assert len(ids) == len(set(ids)), "draw.io cell IDs must be unique" + assert {"0", "1"}.issubset(ids), "draw.io root cells are required" + groups = [cell for cell in root.findall(".//mxCell") + if cell.get("style") == "group;connectable=0;"] + assert len(groups) == len(TILES), "one editable group is required per tile" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("-o", "--output", type=Path, + default=Path("conv_lowering_tiles.drawio")) + args = parser.parse_args() + diagram = Drawio() + for index, tile in enumerate(TILES): + render_tile(diagram, tile, index) + validate(diagram.mxfile) + diagram.write(args.output) + print(args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/tools/pim_sequence_diagram.py b/validation/tools/pim_sequence_diagram.py new file mode 100644 index 0000000..033cd78 --- /dev/null +++ b/validation/tools/pim_sequence_diagram.py @@ -0,0 +1,438 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import re +import struct +import sys +from collections import Counter, namedtuple +from pathlib import Path + + +Event = namedtuple("Event", "op peer size ordinal instruction") +Program = namedtuple("Program", "events operations starts_inactive ends_with_send") +Transfer = namedtuple( + "Transfer", + "sender receiver size sender_ordinal receiver_ordinal sender_instruction receiver_instruction " + "sender_last_instruction receiver_last_instruction count", +) +CORE_FILE = re.compile(r"core_(\d+)\.(json|pim)$") +HEADER = struct.Struct("<4sII") +RECORD = struct.Struct(" str: + counts = Counter(op for op in operations if op not in IGNORED_OPS and op not in ("send", "recv")) + ordered = sorted(counts, key=lambda op: (DISPLAY_RANK.get(op, len(DISPLAY_RANK)), op)) + groups = ([op for op in ordered if op in MEMORY_OPS], [op for op in ordered if op not in MEMORY_OPS]) + rows = [ + [(counts[op], op) for op in group[index : index + 2]] + for group in groups + for index in range(0, len(group), 2) + ] + if not rows: + return "" + count_width = max(len(str(count)) for row in rows for count, _ in row) + op_width = max(len(op) for row in rows for _, op in row) + return '""' + "\\n".join( + " ".join(f"{count:>{count_width}}x {op:<{op_width}}" for count, op in row).rstrip() + for row in rows + ) + '""' + + +def read_json(path: Path) -> Program: + instructions = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(instructions, list): + raise ValueError(f"{path}: expected a JSON instruction array") + + operations = tuple( + str(instruction.get("op", "unknown")) if isinstance(instruction, dict) else "unknown" + for instruction in instructions + ) + events = [] + for index, instruction in enumerate(instructions): + if not isinstance(instruction, dict): + continue + op = instruction.get("op") + if op not in ("send", "recv"): + continue + try: + peer = int(instruction["core"]) + size = int(instruction["size"]) + except (KeyError, TypeError, ValueError) as error: + raise ValueError(f"{path}: invalid communication instruction {instruction!r}") from error + events.append(Event(op, peer, size, len(events), index)) + starts_inactive = bool( + events + and events[0].op == "recv" + and all( + isinstance(instruction, dict) and instruction.get("op") == "sldi" + for instruction in instructions[:events[0].instruction] + ) + ) + last_op = instructions[-1].get("op") if instructions and isinstance(instructions[-1], dict) else None + return Program(events, operations, starts_inactive, last_op == "send") + + +def read_binary(path: Path) -> Program: + data = path.read_bytes() + if len(data) < HEADER.size: + raise ValueError(f"{path}: binary core file is too small") + magic, version, count = HEADER.unpack_from(data) + expected_size = HEADER.size + count * RECORD.size + if magic != b"PIMB": + raise ValueError(f"{path}: invalid PIM binary magic") + if version != 1: + raise ValueError(f"{path}: unsupported PIM binary version {version}") + if len(data) != expected_size: + raise ValueError(f"{path}: expected {expected_size} bytes, found {len(data)}") + + events = [] + operations = [] + last_opcode = None + for index in range(count): + opcode, _, _, _, peer, _, _, size = RECORD.unpack_from(data, HEADER.size + index * RECORD.size) + last_opcode = opcode + operations.append(OPCODE_NAMES[opcode] if opcode < len(OPCODE_NAMES) else f"opcode_{opcode}") + if opcode in (29, 30): + events.append(Event("send" if opcode == 29 else "recv", peer, size, len(events), index)) + starts_inactive = bool( + events + and events[0].op == "recv" + and all(data[HEADER.size + index * RECORD.size] == 1 for index in range(events[0].instruction)) + ) + return Program(events, tuple(operations), starts_inactive, last_opcode == 29) + + +def read_programs(directory: Path, artifact_format: str) -> dict[int, Program]: + candidates: dict[str, dict[int, Path]] = {"json": {}, "pim": {}} + for path in directory.iterdir(): + match = CORE_FILE.fullmatch(path.name) + if match: + candidates[match.group(2)][int(match.group(1))] = path + + if artifact_format == "auto": + artifact_format = "json" if candidates["json"] else "pim" + files = candidates[artifact_format] + if not files: + raise ValueError(f"{directory}: no core_*.{artifact_format} files found") + + reader = read_json if artifact_format == "json" else read_binary + return {core: reader(path) for core, path in files.items()} + + +def match_transfers(programs: dict[int, Program]) -> list[Transfer]: + positions = {core: 0 for core in programs} + transfers = [] + + while True: + made_progress = False + for core in sorted(programs): + events = programs[core].events + position = positions[core] + if position == len(events): + continue + event = events[position] + + if event.peer not in programs: + sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core) + transfers.append( + Transfer( + sender, + receiver, + event.size, + event.ordinal if event.op == "send" else None, + event.ordinal if event.op == "recv" else None, + event.instruction if event.op == "send" else None, + event.instruction if event.op == "recv" else None, + event.instruction if event.op == "send" else None, + event.instruction if event.op == "recv" else None, + 1, + ) + ) + positions[core] += 1 + made_progress = True + continue + + peer_events = programs[event.peer].events + peer_position = positions[event.peer] + if peer_position == len(peer_events): + continue + peer_event = peer_events[peer_position] + if ( + event.peer == core + or peer_event.peer != core + or peer_event.op == event.op + or peer_event.size != event.size + ): + continue + + send = event if event.op == "send" else peer_event + receive = peer_event if event.op == "send" else event + sender, receiver = (core, event.peer) if event.op == "send" else (event.peer, core) + transfers.append( + Transfer( + sender, + receiver, + event.size, + send.ordinal, + receive.ordinal, + send.instruction, + receive.instruction, + send.instruction, + receive.instruction, + 1, + ) + ) + positions[core] += 1 + positions[event.peer] += 1 + made_progress = True + + if made_progress: + continue + remaining = { + core: programs[core].events[position] + for core, position in positions.items() + if position < len(programs[core].events) + } + if not remaining: + return transfers + details = ", ".join( + f"core {core}: {event.op} {event.peer} ({event.size} B)" + for core, event in sorted(remaining.items()) + ) + raise ValueError(f"communication streams cannot be matched at {details}") + + +def visible_transfers( + transfers: list[Transfer], + selected: set[int], + programs: dict[int, Program], +) -> list[Transfer]: + visible = [transfer for transfer in transfers if transfer.sender in selected or transfer.receiver in selected] + grouped = [] + for transfer in visible: + previous = grouped[-1] if grouped else None + can_group = ( + previous is not None + and transfer.sender_ordinal is not None + and previous.sender == transfer.sender + and previous.receiver == transfer.receiver + and previous.sender_ordinal + previous.count == transfer.sender_ordinal + and ( + transfer.receiver_ordinal is None + or previous.receiver_ordinal + previous.count == transfer.receiver_ordinal + ) + ) + if can_group: + sender_gap = programs[transfer.sender].operations[ + previous.sender_last_instruction + 1 : transfer.sender_instruction + ] + receiver_gap = ( + programs[transfer.receiver].operations[ + previous.receiver_last_instruction + 1 : transfer.receiver_instruction + ] + if transfer.receiver in programs + else () + ) + can_group = not summarize_operations(sender_gap) and not summarize_operations(receiver_gap) + + if can_group: + grouped[-1] = previous._replace( + size=previous.size + transfer.size, + sender_last_instruction=transfer.sender_last_instruction, + receiver_last_instruction=transfer.receiver_last_instruction, + count=previous.count + transfer.count, + ) + else: + grouped.append(transfer) + return grouped + + +def collect_operation_notes( + cores: list[int], + transfers: list[Transfer], + programs: dict[int, Program], +) -> dict[int, list[tuple[int, str]]]: + notes: dict[int, list[tuple[int, str]]] = {} + last_instructions = {core: -1 for core in cores} + anchors = {core: -1 for core in cores} + + for transfer_index, transfer in enumerate(transfers): + endpoints = { + transfer.sender: (transfer.sender_instruction, transfer.sender_last_instruction), + transfer.receiver: (transfer.receiver_instruction, transfer.receiver_last_instruction), + } + for core, (instruction, last_instruction) in endpoints.items(): + if core not in last_instructions or instruction is None: + continue + summary = summarize_operations( + programs[core].operations[last_instructions[core] + 1 : instruction] + ) + if summary: + notes.setdefault(anchors[core], []).append((core, summary)) + last_instructions[core] = last_instruction + anchors[core] = transfer_index + + for core in cores: + summary = summarize_operations(programs[core].operations[last_instructions[core] + 1 :]) + if summary: + notes.setdefault(anchors[core], []).append((core, summary)) + return notes + + +def note_height(summary: str) -> float: + return 2.5 + 1.6 * (summary.count("\\n") + 1) + + +def parallel_note_order( + notes: list[tuple[int, str]], + cores: list[int], +) -> list[tuple[int, str]]: + positions = {core: position for position, core in enumerate(cores)} + packed: list[list[tuple[int, str]]] = [] + for note in sorted(notes, key=lambda item: positions[item[0]]): + row = next( + ( + row + for row in packed + if positions[note[0]] - positions[row[-1][0]] > 1 + ), + None, + ) + if row is None: + row = [] + packed.append(row) + row.append(note) + return [note for row in packed for note in row] + + +def render_text(cores: list[int], transfers: list[Transfer], programs: dict[int, Program]) -> str: + aliases = {core: f"C{core}" for core in cores} + positions = {core: position for position, core in enumerate(cores)} + lines = [f'participant "Core {core}" as {aliases[core]}' for core in cores] + first_receives = { + core: next((event.ordinal for event in programs[core].events if event.op == "recv"), None) + for core in cores + } + active = {core for core in cores if not programs[core].starts_inactive} + lines.extend(f"activate {aliases[core]}" for core in cores if core in active) + notes = collect_operation_notes(cores, transfers, programs) + availability = [0.0] * (2 * len(cores) + 1) + cursor = 0.0 + + def emit(line: str, left: int, right: int, height: float) -> None: + nonlocal cursor + start = max(availability[left : right + 1]) + gap = start - cursor + if abs(gap) >= 0.05: + value = f"{gap:.1f}".rstrip("0").rstrip(".") + lines.append(f"space {value}") + lines.append(line) + cursor = start + height + availability[left : right + 1] = [cursor] * (right - left + 1) + + def emit_notes(anchor: int) -> None: + for core, summary in parallel_note_order(notes.get(anchor, []), cores): + center = 2 * positions[core] + 1 + emit(f"note over {aliases[core]}:{summary}", center - 1, center + 1, note_height(summary)) + + emit_notes(-1) + + for transfer_index, transfer in enumerate(transfers): + label = f"{transfer.size} B" + if transfer.count > 1: + label = f"{transfer.count} sends, {label}" + if transfer.sender in aliases and transfer.receiver in aliases: + sender = 2 * positions[transfer.sender] + 1 + receiver = 2 * positions[transfer.receiver] + 1 + line = f"{aliases[transfer.sender]}->(1){aliases[transfer.receiver]}:{label}" + left, right = sorted((sender, receiver)) + elif transfer.receiver in aliases: + line = f"[->(1){aliases[transfer.receiver]}:{label}" + left, right = 0, 2 * positions[transfer.receiver] + 1 + else: + line = f"{aliases[transfer.sender]}->(1)]:{label}" + left, right = 2 * positions[transfer.sender] + 1, len(availability) - 1 + emit(line, left, right, ARROW_HEIGHT) + + receiver = transfer.receiver + if ( + receiver in aliases + and receiver not in active + and transfer.receiver_ordinal == first_receives[receiver] + ): + lines.append(f"activate {aliases[receiver]}") + active.add(receiver) + + sender = transfer.sender + if ( + sender in active + and programs[sender].ends_with_send + and transfer.sender_ordinal + transfer.count == len(programs[sender].events) + ): + lines.append(f"deactivate {aliases[sender]}") + active.remove(sender) + + emit_notes(transfer_index) + + lines.extend(f"deactivateafter {aliases[core]}" for core in cores if core in active) + return "\n".join(lines) + "\n" + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Generate SequenceDiagram.org text for PIM communication and intervening work." + ) + parser.add_argument("pim_dir", type=Path, help="Directory containing core_.json or core_.pim files") + selection = parser.add_mutually_exclusive_group(required=True) + selection.add_argument("--cores", nargs="+", type=int, help="Core lifelines to display, in column order") + selection.add_argument("--all-cores", action="store_true", help="Display every used core, ordered by core ID") + parser.add_argument("--format", choices=("auto", "json", "pim"), default="auto") + parser.add_argument("-o", "--output", type=Path, help="Text output path (default: stdout)") + args = parser.parse_args() + + try: + if not args.pim_dir.is_dir(): + raise ValueError(f"{args.pim_dir}: not a directory") + if args.cores is not None and len(set(args.cores)) != len(args.cores): + raise ValueError("--cores contains duplicates") + programs = read_programs(args.pim_dir, args.format) + cores = ( + [core for core in sorted(programs) if programs[core].operations] + if args.all_cores + else args.cores + ) + missing = [core for core in cores if core not in programs] + if missing: + raise ValueError(f"missing artifact for selected core(s): {', '.join(map(str, missing))}") + transfers = visible_transfers(match_transfers(programs), set(cores), programs) + diagram = render_text(cores, transfers, programs) + if args.output: + args.output.write_text(diagram, encoding="utf-8") + else: + sys.stdout.write(diagram) + except (OSError, ValueError) as error: + parser.error(str(error)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/tools/pimcomp/compare_raptor_pimcomp.py b/validation/tools/pimcomp/compare_raptor_pimcomp.py index 4dc4f1f..97e8b9a 100755 --- a/validation/tools/pimcomp/compare_raptor_pimcomp.py +++ b/validation/tools/pimcomp/compare_raptor_pimcomp.py @@ -285,6 +285,7 @@ def load_effective_hardware(args: argparse.Namespace) -> dict[str, int]: def select_pimsim_config(args: argparse.Namespace, hardware: dict[str, int]) -> Path: + fallback: Path | None = None for path in sorted(PIMSIM_CONFIG_DIR.glob(f"*/{args.pimsim_mode}_config.json")): with open(path, encoding="utf-8") as f: config = json.load(f) @@ -297,15 +298,44 @@ def select_pimsim_config(args: argparse.Namespace, hardware: dict[str, int]) -> and matrix["xbar_size"] == [hardware["crossbar_size"]] * 2 and network["layout"] == [hardware["mesh_rows"], hardware["mesh_cols"]] and config["sim_config"]["sim_mode"] == (1 if args.pimsim_mode == "latency" else 0) - and config["sim_config"]["sim_time"] == args.pimsim_time_ms ): - return path + if config["sim_config"]["sim_time"] == args.pimsim_time_ms: + return path + fallback = fallback or path + if fallback is not None: + return fallback raise ValueError( - f"No pre-generated {args.pimsim_mode} pimsim-nn config matches " - f"{hardware} with sim_time={args.pimsim_time_ms}" + f"No pre-generated {args.pimsim_mode} pimsim-nn config matches {hardware}" ) +def prepare_pimsim_config( + args: argparse.Namespace, + hardware: dict[str, int], + out_dir: Path, +) -> Path: + source = select_pimsim_config(args, hardware) + with open(source, encoding="utf-8") as f: + config = json.load(f) + if config["sim_config"]["sim_time"] == args.pimsim_time_ms: + return source + + config["sim_config"]["sim_time"] = args.pimsim_time_ms + target = out_dir / "pimsim_config.json" + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8") + + network_path = Path(config["chip_config"]["network_config"]["net_config_file_path"]) + if not network_path.is_absolute(): + network_path = source.parent / network_path + target_network = target.parent / Path( + config["chip_config"]["network_config"]["net_config_file_path"] + ) + target_network.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(network_path, target_network) + return target + + def compile_reference( args: argparse.Namespace, model_path: Path, @@ -1221,6 +1251,8 @@ def main(): help="Return a non-zero status if a stage or semantic validation fails.", ) args = parser.parse_args() + if args.pimsim_time_ms <= 0: + parser.error("--pimsim-time-ms must be positive") if args.pimcomp_pipeline is None: args.pimcomp_pipeline = "element" if args.pimsim_mode == "latency" else "batch" @@ -1485,10 +1517,11 @@ def main(): if not args.skip_pimsim_nn and hardware["core_count"] > 0: written_config = try_stage( failures, - "Select pimsim-nn config", - select_pimsim_config, + "Prepare pimsim-nn config", + prepare_pimsim_config, args, hardware, + out_dir, ) if written_config is not None: pimsim_config = written_config @@ -1593,8 +1626,11 @@ def main(): "model": str(model_path), "hardware": hardware, "pimsim_mode": args.pimsim_mode, + "pimsim_time_ms": args.pimsim_time_ms, "pimcomp_pipeline": args.pimcomp_pipeline, "pimcomp_replication": args.pimcomp_replication, + "pimcomp_config": str(args.pimcomp_config), + "raptor_extra_args": args.raptor_extra_arg, "reused_raptor_report": optional_path(args.reuse_raptor_report.resolve()) if reuse_raptor else None, "failures": failures, "steps": [asdict(step) for step in steps], diff --git a/validation/tools/pimcomp/run_pimcomp_paper_latency.py b/validation/tools/pimcomp/run_pimcomp_paper_latency.py index ab1b331..5f29d1e 100755 --- a/validation/tools/pimcomp/run_pimcomp_paper_latency.py +++ b/validation/tools/pimcomp/run_pimcomp_paper_latency.py @@ -20,57 +20,87 @@ from raptor_validation.pimsim_nn import parse_pimsim_nn_metrics # noqa: E402 from raptor_validation.validate_one import STAGE_COLORS # noqa: E402 PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN" -PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json" +PIMCOMP_CONFIGS = REPO / "validation/pimsim_configs/pimcomp" COMPARE = Path(__file__).resolve().with_name("compare_raptor_pimcomp.py") +ARCHES = tuple(sorted(path.name for path in PIMCOMP_CONFIGS.iterdir() if path.is_dir())) MODELS = { "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx", "resnet18": SUITE / "resnet18/resnet18-v1-7.onnx", "resnet34": SUITE / "resnet34/resnet34-v1-7.onnx", "googlenet": SUITE / "googlenet/googlenet-12-latency.onnx", } +COMPARISONS = ( + ("latency", 1, "element"), + ("throughput", 2, "batch"), + ("throughput", 4, "batch"), + ("throughput", 8, "batch"), +) -def result_dir(root: Path | None, name: str) -> Path: - return root / name if root is not None else MODELS[name].parent +def result_dir(root: Path | None, name: str, mode: str, pipeline: int) -> Path: + base = root / name if root is not None else MODELS[name].parent + suffix = "latency" if mode == "latency" else f"throughput/pipeline{pipeline}" + return base / suffix -def write_results_csv(root: Path | None) -> Path: +def write_results_csv(root: Path | None, arch: str, models: list[str]) -> Path: output = (root or SUITE) / "results.csv" fields = ( "model", + "arch", + "mode", + "raptor_pipeline", + "pimcomp_pipeline", + "status", + "raptor_throughput_samples_s", + "pimcomp_throughput_samples_s", "raptor_latency_ms", "pimcomp_latency_ms", + "raptor_power_mw", + "pimcomp_power_mw", "raptor_energy_pj", "pimcomp_energy_pj", - "faster_compiler", + "better_compiler", "speedup", ) rows = [] - for name in MODELS: - report_path = result_dir(root, name) / "pimcomp/comparison_report.json" - if not report_path.exists(): - continue - report = json.loads(report_path.read_text(encoding="utf-8")) - raptor = report.get("raptor_performance") or {} - pimcomp = report.get("pimcomp_performance") or {} - raptor_latency = raptor.get("latency_ms") - pimcomp_latency = pimcomp.get("latency_ms") - if raptor_latency is None or pimcomp_latency is None: - continue - raptor_energy = (raptor.get("average_energy_pj") - or parse_pimsim_nn_metrics(raptor.get("raw_output", "")).get("average_energy_pj")) - pimcomp_energy = (pimcomp.get("average_energy_pj") - or parse_pimsim_nn_metrics(pimcomp.get("raw_output", "")).get("average_energy_pj")) - faster = "raptor" if raptor_latency < pimcomp_latency else "pimcomp" - rows.append({ - "model": name, - "raptor_latency_ms": f"{raptor_latency:.6f}", - "pimcomp_latency_ms": f"{pimcomp_latency:.6f}", - "raptor_energy_pj": "" if raptor_energy is None else f"{raptor_energy:.6f}", - "pimcomp_energy_pj": "" if pimcomp_energy is None else f"{pimcomp_energy:.6f}", - "faster_compiler": faster, - "speedup": f"{max(raptor_latency, pimcomp_latency) / min(raptor_latency, pimcomp_latency):.2f}", - }) + for name in models: + for mode, pipeline, pimcomp_pipeline in COMPARISONS: + report_path = result_dir(root, name, mode, pipeline) / "pimcomp/comparison_report.json" + if not report_path.exists(): + continue + report = json.loads(report_path.read_text(encoding="utf-8")) + raptor = report.get("raptor_performance") or {} + pimcomp = report.get("pimcomp_performance") or {} + raptor_values = performance_values(raptor) + pimcomp_values = performance_values(pimcomp) + raptor_metric = raptor_values["throughput"] if mode == "throughput" else raptor_values["latency"] + pimcomp_metric = pimcomp_values["throughput"] if mode == "throughput" else pimcomp_values["latency"] + status = "PASS" if comparison_passed(report) else "FAIL" + if raptor_metric is None or pimcomp_metric is None: + better = "" + speedup = "" + else: + better = comparison_winner(mode, raptor_metric, pimcomp_metric) + speedup = f"{max(raptor_metric, pimcomp_metric) / min(raptor_metric, pimcomp_metric):.2f}" + rows.append({ + "model": name, + "arch": arch, + "mode": mode, + "raptor_pipeline": pipeline, + "pimcomp_pipeline": pimcomp_pipeline, + "status": status, + "raptor_throughput_samples_s": format_value(raptor_values["throughput"]), + "pimcomp_throughput_samples_s": format_value(pimcomp_values["throughput"]), + "raptor_latency_ms": format_value(raptor_values["latency"]), + "pimcomp_latency_ms": format_value(pimcomp_values["latency"]), + "raptor_power_mw": format_value(raptor_values["power"]), + "pimcomp_power_mw": format_value(pimcomp_values["power"]), + "raptor_energy_pj": format_value(raptor_values["energy"]), + "pimcomp_energy_pj": format_value(pimcomp_values["energy"]), + "better_compiler": better, + "speedup": speedup, + }) with open(output, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields, lineterminator="\n") writer.writeheader() @@ -78,6 +108,47 @@ def write_results_csv(root: Path | None) -> Path: return output +def performance_values(performance: dict) -> dict[str, float | None]: + parsed = parse_pimsim_nn_metrics(performance.get("raw_output", "")) + return { + "throughput": performance.get("throughput") or parsed.get("throughput"), + "latency": ( + performance.get("latency_ms") + or performance.get("average_latency_ms") + or parsed.get("latency_ms") + or parsed.get("average_latency_ms") + ), + "power": performance.get("average_power_mw") or parsed.get("average_power_mw"), + "energy": performance.get("average_energy_pj") or parsed.get("average_energy_pj"), + } + + +def comparison_passed(report: dict) -> bool: + if report.get("failures"): + return False + for key in ("raptor_validation", "pimcomp_validation"): + result = report.get(key) or {} + if result.get("status") != "done" or not result.get("passed"): + return False + for key in ("raptor_performance", "pimcomp_performance"): + performance = report.get(key) or {} + if performance.get("error") or performance.get("skipped"): + return False + return True + + +def comparison_winner(mode: str, raptor: float, pimcomp: float) -> str: + if raptor == pimcomp: + return "tie" + if mode == "throughput": + return "raptor" if raptor > pimcomp else "pimcomp" + return "raptor" if raptor < pimcomp else "pimcomp" + + +def format_value(value: float | None) -> str: + return "" if value is None else f"{value:.6f}" + + def print_stage(title: str, color: str) -> None: print("\n" + Style.BRIGHT + color + f"[{title}]" + Style.RESET_ALL, flush=True) @@ -98,7 +169,17 @@ def validate_pimcomp_source() -> None: raise RuntimeError(f"PIMCOMP paper setting is missing: {setting}") -def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[str]: +def comparison_command( + model: Path, + result_dir: Path, + config: Path, + mode: str, + pipeline: int, + pimcomp_pipeline: str, + pimsim_time_ms: int, + timeout: float, +) -> list[str]: + time_args = ["--pimsim-time-ms", str(pimsim_time_ms)] if mode == "throughput" else [] return [ sys.executable, str(COMPARE), @@ -109,32 +190,49 @@ def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[st "--pimcomp-dir", str(PIMCOMP_SOURCE), "--pimcomp-config", - str(PIMCOMP_CONFIG), - "--core-count", - "168", - "--crossbar-count", - "96", - "--crossbar-size", - "128", - "--mesh-rows", - "12", - "--mesh-cols", - "14", + str(config), "--pimsim-mode", - "latency", + mode, + *time_args, "--pimcomp-pipeline", - "element", + pimcomp_pipeline, "--pimcomp-replication", "GA", + f"--raptor-extra-arg=--pipeline={pipeline}", "--timeout-seconds", str(timeout), "--fail-on-error", ] +def config_path(arch: str, mode: str) -> Path: + path = PIMCOMP_CONFIGS / arch / f"{mode}_config.json" + if not path.exists(): + raise ValueError(f"{arch} has no {mode} config: {path}") + return path + + +def core_count(config: Path) -> int: + with open(config, encoding="utf-8") as f: + return int(json.load(f)["chip_config"]["core_cnt"]) + + +def completed_report(path: Path, mode: str, pipeline: int, config: Path, pimsim_time_ms: int) -> bool: + if not path.exists(): + return False + report = json.loads(path.read_text(encoding="utf-8")) + return ( + report.get("pimsim_mode") == mode + and report.get("pimcomp_pipeline") == ("element" if mode == "latency" else "batch") + and report.get("pimsim_time_ms") == pimsim_time_ms + and report.get("pimcomp_config") == str(config.resolve()) + and f"--pipeline={pipeline}" in report.get("raptor_extra_args", []) + ) + + def main() -> int: parser = argparse.ArgumentParser( - description="Reproduce the serial Arch-A latency comparison from the PIMCOMP paper." + description="Compare supported PIMCOMP models with Raptor latency and throughput schedules." ) parser.add_argument( "--out-dir", @@ -142,6 +240,15 @@ def main() -> int: help="Result root (default: artifacts beside each model under validation/).", ) parser.add_argument("--models", nargs="+", choices=MODELS, default=list(MODELS)) + parser.add_argument( + "--arch", choices=ARCHES, default="arch-a", help="PIM architecture (default: arch-a)." + ) + parser.add_argument( + "--pimsim-time-ms", + type=int, + default=100, + help="throughput pimsim-nn horizon in ms (default: 100).", + ) parser.add_argument("--timeout-seconds", type=float, default=3600.0) parser.add_argument( "--resume", @@ -151,6 +258,16 @@ def main() -> int: parser.add_argument("--dry-run", action="store_true", help="Print commands without modifying files.") args = parser.parse_args() + if args.pimsim_time_ms <= 0: + parser.error("--pimsim-time-ms must be positive") + configs = {mode: config_path(args.arch, mode) for mode, _, _ in COMPARISONS} + unsupported = [pipeline for mode, pipeline, _ in COMPARISONS if core_count(configs[mode]) % pipeline] + if unsupported: + parser.error( + f"{args.arch} has {core_count(configs['throughput'])} cores; " + f"throughput pipelines must divide that count (invalid: {unsupported})" + ) + out_dir = args.out_dir.resolve() if args.out_dir is not None else None missing = [str(MODELS[name]) for name in args.models if not MODELS[name].exists()] @@ -162,6 +279,8 @@ def main() -> int: out_dir.mkdir(parents=True, exist_ok=True) print(Style.BRIGHT + f"Found {len(args.models)} PIMCOMP model(s) to compare." + Style.RESET_ALL) + print(f"Architecture: {args.arch}") + print(f"Throughput pimsim time: {args.pimsim_time_ms} ms") print(f"Results root: {out_dir or SUITE}") print("=" * 72) @@ -176,34 +295,51 @@ def main() -> int: failed = [] for index, name in enumerate(args.models, start=1): - model_result_dir = result_dir(out_dir, name) - print( - "\n" + Fore.CYAN + f"[{index}/{len(args.models)}]" + Style.RESET_ALL - + f" {Style.BRIGHT}Comparing {name}{Style.RESET_ALL}", - flush=True, - ) - if args.resume and (model_result_dir / "pimcomp/comparison_report.json").exists(): + for mode, pipeline, pimcomp_pipeline in COMPARISONS: + model_result_dir = result_dir(out_dir, name, mode, pipeline) print( - Fore.YELLOW + " Completed report exists; skipping" + Style.RESET_ALL, + "\n" + Fore.CYAN + f"[{index}/{len(args.models)}]" + Style.RESET_ALL + + f" {Style.BRIGHT}Comparing {name} ({mode}, pipeline={pipeline}){Style.RESET_ALL}", flush=True, ) - continue - returncode = run( - comparison_command(MODELS[name], model_result_dir, args.timeout_seconds), - dry_run=args.dry_run, - check=False, - ) - if returncode: - failed.append(name) + if args.resume and completed_report( + model_result_dir / "pimcomp/comparison_report.json", + mode, + pipeline, + configs[mode], + args.pimsim_time_ms, + ): + print( + Fore.YELLOW + " Completed report exists; skipping" + Style.RESET_ALL, + flush=True, + ) + continue + returncode = run( + comparison_command( + MODELS[name], + model_result_dir, + configs[mode], + mode, + pipeline, + pimcomp_pipeline, + args.pimsim_time_ms, + args.timeout_seconds, + ), + dry_run=args.dry_run, + check=False, + ) + if returncode: + failed.append(f"{name}/{mode}/pipeline{pipeline}") if args.dry_run: return 1 if failed else 0 - results_path = write_results_csv(out_dir) + results_path = write_results_csv(out_dir, args.arch, args.models) print_stage("Results", STAGE_COLORS["Compare Outputs"]) print(results_path.read_text(encoding="utf-8"), end="") print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL) - print(Style.BRIGHT + f"Passed: {len(args.models) - len(failed)}" + Style.RESET_ALL) + total_jobs = len(args.models) * len(COMPARISONS) + print(Style.BRIGHT + f"Passed: {total_jobs - len(failed)}" + Style.RESET_ALL) print(Style.BRIGHT + f"Failed: {len(failed)}" + Style.RESET_ALL) print(Style.BRIGHT + f"Results: {results_path}" + Style.RESET_ALL) if failed: