[TOGSim] ONNXim frontend and ACS for tutorial

TOGSim/acs/__main__.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import argparse
+import os
+import re
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
+    os.path.abspath(__file__)))))
+
+from TOGSim.acs import cycles as acs_cycles
+from TOGSim.acs import node_mapping
+from TOGSim.onnx_frontend.config import read_hardware
+
+TILE_NAMES = {
+    "gemm": ("TM", "TN"),
+    "conv": ("TILE_M", "TILE_N"),
+    "softmax": ("TM", "TN"),
+    "layernorm": ("TM", "TN"),
+    "embed_layernorm": ("TSEQ", "TDIM"),
+    "global_avgpool": ("TC", "THW"),
+    "bias_act": ("TM", "TN"),
+    "bias_gelu": ("TM", "TN"),
+    "maxpool": ("TROW", "TCOL"),
+    "adaptive_avgpool": ("TROW", "TCOL"),
+    "attention": ("Q_LEN", "DHEAD"),
+    "kvcache_concat": ("TTOK", "HIDDEN"),
+    "concat": ("TROW", "TD"),
+    "flatten": (),
+}
+
+
+def read_constant(source: str, name: str) -> int | None:
+    m = re.search(rf"static\s+const\s+\w+(?:_t)?\s+{re.escape(name)}\s*=\s*"
+                  r"([^;]+);", source)
+    if not m:
+        return None
+    text = m.group(1).strip()
+    if re.fullmatch(r"-?\d+", text):
+        return int(text)
+    return None
+
+
+def tile_from_source(stem: str, source: str) -> dict[str, int]:
+    names = TILE_NAMES.get(stem)
+    if names is None:
+        raise KeyError(
+            f"no tile constants recorded for kernel {stem!r}; add it to "
+            f"TILE_NAMES in acs/__main__.py")
+    if not names:
+        return {"rows": 1, "elems": 1}
+
+    values = []
+    for n in names:
+        v = read_constant(source, n)
+        if v is None:
+            raise KeyError(
+                f"{stem}: constant {n!r} is not declared, or is not a plain "
+                f"integer. Pass --rows and --elems instead.")
+        values.append(v)
+    rows = values[0]
+    elems = rows * values[1] if len(values) > 1 else rows
+    return {"rows": rows, "elems": elems}
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(
+        description="Derive a TOGSim cycle table for one kernel, without gem5.")
+    ap.add_argument("kernel", help="path to the kernel .cpp")
+    ap.add_argument("--config", required=True)
+    ap.add_argument("-o", "--out", default=None,
+                    help="output .tsv (default: cycles.tsv beside the kernel)")
+    ap.add_argument("--sa", type=int, default=None,
+                    help="systolic array dimension; inferred from the config name")
+    ap.add_argument("--rows", type=int, default=None,
+                    help="override the tile row count")
+    ap.add_argument("--elems", type=int, default=None,
+                    help="override the tile element count")
+    args = ap.parse_args(argv)
+
+    hw = read_hardware(args.config, args.sa)
+    stem = os.path.splitext(os.path.basename(args.kernel))[0]
+    source = open(args.kernel).read()
+
+    desc = node_mapping.lookup(stem)
+    if args.rows is not None or args.elems is not None:
+        tile = {"rows": args.rows or 1, "elems": args.elems or 1}
+    else:
+        tile = tile_from_source(stem, source)
+
+    rows = acs_cycles.build_table(desc, tile, hw)
+    out = args.out or os.path.join(os.path.dirname(args.kernel), "cycles.tsv")
+    acs_cycles.write_table(rows, out)
+
+    print(f"kernel   : {stem} ({len(rows)} compute nodes)")
+    print(f"hardware : SA={hw.sa} lanes={hw.vlane} vlen={hw.vlen_bits}b "
+          f"({hw.throughput} elem/cycle)")
+    print(f"tile     : rows={tile['rows']} elems={tile['elems']}")
+    print(f"table    : {out}")
+    for (c, o), node in zip(rows, desc.nodes):
+        print(f"  {node.name:<14} {node.compute_type:<8} {c:>6}\t{o:>6}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

TOGSim/acs/cycles.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import math
+
+from ..onnx_frontend.tiling import Hardware
+from .node_mapping import ELEMENTWISE, REDUCTION, NodeDesc
+
+#: vslidedown at 2,4,6,8...; each vfmax consumes the last -- one chain, so linear.
+FOLD_WIDTH = 8
+
+
+def preload_cycles(hw: Hardware) -> tuple[int, int]:
+    return 2 * hw.sa - 1, 0
+
+
+def matmul_cycles(rows: int, hw: Hardware) -> tuple[int, int]:
+    cycles = 2 * hw.sa - 2 + rows
+    return cycles, min(hw.sa, rows)
+
+
+def vector_cycles(node: NodeDesc, tile_elems: int, hw: Hardware) -> tuple[int, int]:
+    total = 0
+    for op in node.ops:
+        if op.kind == ELEMENTWISE:
+            total += math.ceil(tile_elems / hw.throughput) * op.n
+        elif op.kind == REDUCTION:
+            total += (FOLD_WIDTH - 1) * (1 + op.n)
+        else:
+            raise ValueError(f"unknown op kind {op.kind!r} in node {node.name!r}")
+    # nothing in front of a VPU node to hide behind: fully exposed
+    return total, total
+
+
+def default_cost(node: NodeDesc, tile: dict, hw: Hardware) -> tuple[int, int]:
+    if node.compute_type == "preload":
+        return preload_cycles(hw)
+    if node.compute_type == "matmul":
+        return matmul_cycles(int(tile.get("rows", hw.sa)), hw)
+    if node.compute_type == "vector":
+        return vector_cycles(node, int(tile.get("elems", 1)), hw)
+    raise ValueError(f"unknown compute type {node.compute_type!r}")
+
+_cost_function = default_cost
+
+
+def set_cost_function(fn) -> None:
+    global _cost_function
+    _cost_function = fn
+
+
+def build_table(kernel_desc, tile: dict, hw: Hardware) -> list[tuple[int, int]]:
+    rows = []
+    for node in kernel_desc.nodes:
+        cycles, interval = _cost_function(node, tile, hw)
+        # column 2 is overlapping, not the interval: cycles - overlapping is
+        # what the core sustains. Writing the interval there inverts it.
+        rows.append((cycles, max(0, cycles - interval)))
+    return rows
+
+
+def write_table(rows: list[tuple[int, int]], path: str, origin: str = "acs") -> None:
+    with open(path, "w") as fh:
+        for cycles, overlapping in rows:
+            fh.write(f"{cycles}\t{overlapping}\n")
+        fh.write(f"# origins: {origin}\n")

TOGSim/acs/node_mapping.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+ELEMENTWISE = "elementwise"
+REDUCTION = "reduction"
+
+#: max/min: compare, then select.
+MAXMIN = 2
+
+
+@dataclass(frozen=True)
+
+
+class Op:
+    kind: str
+    n: int          # instructions per element (elementwise) or per fold step
+    comment: str
+
+
+@dataclass(frozen=True)
+
+
+class NodeDesc:
+    name: str
+    compute_type: str       # "vector" | "matmul" | "preload"
+    ops: tuple = ()
+
+
+@dataclass(frozen=True)
+
+
+class KernelDesc:
+    name: str
+    nodes: tuple
+
+
+def _vec(name, *ops):
+    return NodeDesc(name, "vector", tuple(ops))
+
+#: Kernel name -> its compute nodes, in the order the .cpp emits them.
+KERNELS: dict[str, KernelDesc] = {
+
+    "gemm": KernelDesc("gemm", (
+        _vec("acc_init", Op(ELEMENTWISE, 1, "zero the accumulator")),
+        NodeDesc("preload", "preload"),
+        NodeDesc("matmul", "matmul"),
+    )),
+
+    "conv": KernelDesc("conv", (
+        _vec("acc_init", Op(ELEMENTWISE, 1, "zero the accumulator")),
+        NodeDesc("preload", "preload"),
+        NodeDesc("matmul", "matmul"),
+    )),
+
+    "softmax": KernelDesc("softmax", (
+        _vec("max_reduce", Op(ELEMENTWISE, MAXMIN, "running max")),
+        _vec("max_write", Op(REDUCTION, MAXMIN, "fold the max")),
+        _vec("sum_reduce",
+             Op(ELEMENTWISE, 1, "x - max"),
+             Op(ELEMENTWISE, 1, "exp"),
+             Op(ELEMENTWISE, 1, "accumulate")),
+        _vec("sum_write", Op(REDUCTION, 1, "fold the sum")),
+        _vec("softmax",
+             Op(ELEMENTWISE, 1, "x - max"),
+             Op(ELEMENTWISE, 1, "exp"),
+             Op(ELEMENTWISE, 1, "/ sum")),
+    )),
+
+    "layernorm": KernelDesc("layernorm", (
+        _vec("stats_reduce",
+             Op(ELEMENTWISE, 1, "accumulate x"),
+             Op(ELEMENTWISE, 2, "accumulate x*x")),
+        _vec("write_mean", Op(REDUCTION, 1, "fold the sum")),
+        _vec("write_var", Op(REDUCTION, 1, "fold the sum of squares")),
+        _vec("normalize",
+             Op(ELEMENTWISE, 1, "x - mean"),
+             Op(ELEMENTWISE, 1, "* rstd"),
+             Op(ELEMENTWISE, 2, "scale and shift")),
+    )),
+
+    "embed_layernorm": KernelDesc("embed_layernorm", (
+        _vec("gather",
+             Op(ELEMENTWISE, 1, "gather the embedding"),
+             Op(ELEMENTWISE, 1, "x - mean"),
+             Op(ELEMENTWISE, 1, "* rstd"),
+             Op(ELEMENTWISE, 2, "scale and shift")),
+    )),
+
+    "global_avgpool": KernelDesc("global_avgpool", (
+        _vec("sum_reduce", Op(ELEMENTWISE, 1, "accumulate")),
+        _vec("sum_write", Op(REDUCTION, 1, "fold the sum")),
+        _vec("scale", Op(ELEMENTWISE, 1, "* 1/HW")),
+    )),
+
+    "bias_act": KernelDesc("bias_act", (
+        _vec("bias_act",
+             Op(ELEMENTWISE, 1, "+ bias"),
+             Op(ELEMENTWISE, 1, "max(x, 0)")),
+    )),
+
+    "bias_gelu": KernelDesc("bias_gelu", (
+        _vec("bias_gelu",
+             Op(ELEMENTWISE, 1, "+ bias"),
+             Op(ELEMENTWISE, 1, "x / sqrt(2)"),
+             Op(ELEMENTWISE, 2, "erf"),
+             Op(ELEMENTWISE, 1, "1 + erf"),
+             Op(ELEMENTWISE, 1, "* 0.5x")),
+    )),
+
+    "maxpool": KernelDesc("maxpool", (
+        _vec("maxpool", Op(ELEMENTWISE, MAXMIN, "running max over the window")),
+    )),
+
+    "adaptive_avgpool": KernelDesc("adaptive_avgpool", (
+        _vec("avgpool",
+             Op(ELEMENTWISE, 1, "accumulate the window"),
+             Op(ELEMENTWISE, 1, "* 1/n")),
+    )),
+
+    # pure DMA: no compute node
+    "flatten": KernelDesc("flatten", ()),
+
+    "attention": KernelDesc("attention", (
+        NodeDesc("qk_preload", "preload"),
+        NodeDesc("qk_matmul", "matmul"),
+        _vec("rowmax", Op(REDUCTION, MAXMIN, "running row max")),
+        _vec("sub", Op(ELEMENTWISE, 1, "s - max")),
+        _vec("exp", Op(ELEMENTWISE, 1, "exp")),
+        _vec("rowsum", Op(REDUCTION, 1, "fold the row sum")),
+        _vec("mac", Op(ELEMENTWISE, 2, "running l")),
+        _vec("rescale", Op(ELEMENTWISE, 2, "rescale the accumulator")),
+        NodeDesc("sv_preload", "preload"),
+        NodeDesc("sv_matmul", "matmul"),
+        _vec("normalize",
+             Op(ELEMENTWISE, 1, "/ l"),
+             Op(ELEMENTWISE, 1, "scale")),
+    )),
+
+    "concat": KernelDesc("concat", ()),
+
+    "kvcache_concat": KernelDesc("kvcache_concat", (
+        _vec("split", Op(ELEMENTWISE, 1, "split QKV into query, key, value")),
+    )),
+}
+
+
+def lookup(kernel_name: str) -> KernelDesc:
+    try:
+        return KERNELS[kernel_name]
+    except KeyError:
+        raise KeyError(
+            f"no compute-node table for kernel {kernel_name!r}. Add one to "
+            f"acs/node_mapping.py: without it the cost of every node in this "
+            f"kernel would be invented.") from None