ilgeco
568 lines
26 KiB
Python
568 lines
26 KiB
Python
#!/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())
|