diff --git a/TOGSim/acs/__main__.py b/TOGSim/acs/__main__.py new file mode 100644 index 000000000..14e2a58a0 --- /dev/null +++ b/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()) diff --git a/TOGSim/acs/cycles.py b/TOGSim/acs/cycles.py new file mode 100644 index 000000000..f96ba8599 --- /dev/null +++ b/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") diff --git a/TOGSim/acs/node_mapping.py b/TOGSim/acs/node_mapping.py new file mode 100644 index 000000000..b6b9c9c11 --- /dev/null +++ b/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 diff --git a/TOGSim/example/adaptive_avgpool/adaptive_avgpool.cpp b/TOGSim/example/adaptive_avgpool/adaptive_avgpool.cpp new file mode 100644 index 000000000..6a7ca9979 --- /dev/null +++ b/TOGSim/example/adaptive_avgpool/adaptive_avgpool.cpp @@ -0,0 +1,92 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t ROWS = 1024; +static const int64_t COLS = 16; + +static const int64_t TROW = 256; +static const int64_t TCOL = 16; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_ROW = (ROWS + TROW - 1) / TROW; +static const int64_t TILES_COL = (COLS + TCOL - 1) / TCOL; + +static const int64_t IN_ROW_STRIDE = 64; +static const int64_t IN_W = 32; + +static const int32_t ARG_X = 0; +static const int32_t ARG_OUT = 1; + +static const int64_t BUF_W3[1] = {0}; +static const int64_t BUF_W0[1] = {1}; +static const int64_t BUF_W1[1] = {2}; +static const int64_t BUF_W2[1] = {3}; +static const int64_t BUF_OUT[1] = {4}; + +static const int64_t VEC_READ[4] = {0, 1, 2, 3}; + +static const int64_t TILE[2] = {TROW, TCOL}; + +static const int64_t STRIDE_IN[2] = {IN_ROW_STRIDE, 2}; +static const int64_t STRIDE_OUT[2] = {TCOL, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_W0 = 0; +static const int32_t TAG_W1 = 1; +static const int32_t TAG_W2 = 2; +static const int32_t TAG_W3 = 3; +static const int32_t TAG_OUT = 4; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_AVG = 0; + +static void avgpool_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t r0 = iv[0], c0 = iv[1]; + const int64_t base = r0 * IN_ROW_STRIDE + c0 * 2; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)base, + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W0, 0, nullptr, 0, BUF_W0, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + 1), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W1, 0, nullptr, 0, BUF_W1, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + IN_W), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W2, 0, nullptr, 0, BUF_W2, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + IN_W + 1), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W3, 0, nullptr, 0, BUF_W3, 1); + // COMPUTE + togsim_compute(ctx, TID_AVG, CT_VECTOR, 0, nullptr, + VEC_READ, 4, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(r0 * TCOL + c0), + 2, TILE, STRIDE_OUT, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t ri = 0; ri < TILES_ROW; ++ri) { + for (int64_t ci = 0; ci < TILES_COL; ++ci) { + int64_t iv[2] = {ri * TROW, ci * TCOL}; + togsim_dispatch(ctx, avgpool_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/attention/attention.cpp b/TOGSim/example/attention/attention.cpp new file mode 100644 index 000000000..285c4dfdd --- /dev/null +++ b/TOGSim/example/attention/attention.cpp @@ -0,0 +1,181 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// SA: 8x8 x2 per core +// VLANE: 8 +// mapping: autotune +static const int64_t HEADS = 8; +static const int64_t KV_HEADS = 8; +static const int64_t SEQ = 128; +static const int64_t DHEAD = 64; + +static const int64_t Q_LEN = 8; +static const int64_t KV_BLOCK = 64; + +static const int64_t SA = 8; +static const int64_t HEADS_PER_KV = HEADS / KV_HEADS; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_Q = (SEQ + Q_LEN - 1) / Q_LEN; +static const int64_t TILES_KV = (SEQ + KV_BLOCK - 1) / KV_BLOCK; + +static const int64_t QK_S_STEPS = (KV_BLOCK + SA - 1) / SA; +static const int64_t QK_K_STEPS = (DHEAD + SA - 1) / SA; +static const int64_t SV_K_STEPS = (DHEAD + SA - 1) / SA; +static const int64_t SV_S_STEPS = (KV_BLOCK + SA - 1) / SA; + +static const int32_t ARG_Q = 0; +static const int32_t ARG_K = 1; +static const int32_t ARG_V = 2; +static const int32_t ARG_OUT = 3; + +static const int64_t BUF_K[1] = {0}; +static const int64_t BUF_Q[1] = {1}; +static const int64_t BUF_V[1] = {2}; +static const int64_t BUF_SA_QK[1] = {3}; +static const int64_t BUF_LOGIT[1] = {4}; +static const int64_t BUF_STAT[1] = {5}; +static const int64_t BUF_ACC[1] = {6}; +static const int64_t BUF_SA_SV[1] = {7}; + +static const int64_t QK_PRELOAD_READ[2] = {0, 1}; +static const int64_t QK_MM_READ[2] = {3, 4}; +static const int64_t SV_PRELOAD_READ[2] = {2, 4}; +static const int64_t SV_MM_READ[2] = {6, 7}; + +static const int64_t STAT_READ[2] = {4, 5}; +static const int64_t LOGIT_READ[2] = {4, 5}; +static const int64_t ACC_READ[2] = {5, 6}; + +static const int64_t TILE_Q[2] = {Q_LEN, DHEAD}; +static const int64_t TILE_KV[2] = {KV_BLOCK, DHEAD}; +static const int64_t TILE_OUT[2] = {Q_LEN, DHEAD}; + +static const int64_t STRIDE_Q[2] = {DHEAD, 1}; +static const int64_t STRIDE_KV[2] = {DHEAD, 1}; +static const int64_t STRIDE_OUT[2] = {DHEAD, 1}; + +static const int32_t SYNC = 0; +static const int32_t ASYNC = 1; + +static const int32_t TAG_K = 0; +static const int32_t TAG_Q = 1; +static const int32_t TAG_V = 2; +static const int32_t TAG_O = 3; + +static const int32_t CT_VECTOR = 0; +static const int32_t CT_MATMUL = 1; +static const int32_t CT_PRELOAD = 2; + +static const uint64_t TID_QK_PRELOAD = 0; +static const uint64_t TID_QK_MATMUL = 1; +static const uint64_t TID_ROWMAX = 2; +static const uint64_t TID_SUB = 3; +static const uint64_t TID_EXP = 4; +static const uint64_t TID_ROWSUM = 5; +static const uint64_t TID_MAC = 6; +static const uint64_t TID_RESCALE = 7; +static const uint64_t TID_SV_PRELOAD = 8; +static const uint64_t TID_SV_MATMUL = 9; +static const uint64_t TID_NORM = 10; + +static void attention_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t h = iv[0], q0 = iv[1]; + + for (int64_t kv0 = 0; kv0 < SEQ; kv0 += KV_BLOCK) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_K, + (uint64_t)(h * SEQ * DHEAD + kv0 * DHEAD), + 2, TILE_KV, STRIDE_KV, ELEM_BITS, + ASYNC, TAG_K, 0, nullptr, 0, BUF_K, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_Q, + (uint64_t)(h * SEQ * DHEAD + q0 * DHEAD), + 2, TILE_Q, STRIDE_Q, ELEM_BITS, + ASYNC, TAG_Q, 0, nullptr, 0, BUF_Q, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_V, + (uint64_t)(h * SEQ * DHEAD + kv0 * DHEAD), + 2, TILE_KV, STRIDE_KV, ELEM_BITS, + ASYNC, TAG_V, 0, nullptr, 0, BUF_V, 1); + + for (int64_t s = 0; s < QK_S_STEPS; ++s) { + for (int64_t k = 0; k < QK_K_STEPS; ++k) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_K, 0, BUF_K, 1); + // COMPUTE + togsim_compute(ctx, TID_QK_PRELOAD, CT_PRELOAD, 0, nullptr, + QK_PRELOAD_READ, 2, BUF_SA_QK, 1); + for (int64_t hh = 0; hh < HEADS_PER_KV; ++hh) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_Q, 0, BUF_Q, 1); + // COMPUTE + togsim_compute(ctx, TID_QK_MATMUL, CT_MATMUL, 0, nullptr, + QK_MM_READ, 2, BUF_LOGIT, 1); + } + } + } + + // COMPUTE + togsim_compute(ctx, TID_ROWMAX, CT_VECTOR, 0, nullptr, + BUF_LOGIT, 1, BUF_STAT, 1); + // COMPUTE + togsim_compute(ctx, TID_SUB, CT_VECTOR, 0, nullptr, + STAT_READ, 2, BUF_LOGIT, 1); + // COMPUTE + togsim_compute(ctx, TID_EXP, CT_VECTOR, 0, nullptr, + BUF_LOGIT, 1, BUF_LOGIT, 1); + // COMPUTE + togsim_compute(ctx, TID_ROWSUM, CT_VECTOR, 0, nullptr, + BUF_LOGIT, 1, BUF_STAT, 1); + // COMPUTE + togsim_compute(ctx, TID_MAC, CT_VECTOR, 0, nullptr, + LOGIT_READ, 2, BUF_STAT, 1); + // COMPUTE + togsim_compute(ctx, TID_RESCALE, CT_VECTOR, 0, nullptr, + ACC_READ, 2, BUF_ACC, 1); + + for (int64_t k = 0; k < SV_K_STEPS; ++k) { + for (int64_t s = 0; s < SV_S_STEPS; ++s) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_V, 0, BUF_V, 1); + // COMPUTE + togsim_compute(ctx, TID_SV_PRELOAD, CT_PRELOAD, 0, nullptr, + SV_PRELOAD_READ, 2, BUF_SA_SV, 1); + for (int64_t hh = 0; hh < HEADS_PER_KV; ++hh) { + // COMPUTE + togsim_compute(ctx, TID_SV_MATMUL, CT_MATMUL, 0, nullptr, + SV_MM_READ, 2, BUF_ACC, 1); + } + } + } + } + + // COMPUTE + togsim_compute(ctx, TID_NORM, CT_VECTOR, 0, nullptr, + ACC_READ, 2, BUF_ACC, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, + (uint64_t)(h * SEQ * DHEAD + q0 * DHEAD), + 2, TILE_OUT, STRIDE_OUT, ELEM_BITS, + SYNC, TAG_O, 0, BUF_ACC, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t h = 0; h < KV_HEADS; ++h) { + for (int64_t qi = 0; qi < TILES_Q; ++qi) { + int64_t iv[2] = {h, qi * Q_LEN}; + togsim_dispatch(ctx, attention_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/bias_act/bias_act.cpp b/TOGSim/example/bias_act/bias_act.cpp new file mode 100644 index 000000000..f032b5196 --- /dev/null +++ b/TOGSim/example/bias_act/bias_act.cpp @@ -0,0 +1,78 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t M = 128; +static const int64_t N = 256; + +static const int64_t TM = 128; +static const int64_t TN = 64; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_M = (M + TM - 1) / TM; +static const int64_t TILES_N = (N + TN - 1) / TN; + +static const int32_t ARG_X = 0; +static const int32_t ARG_BIAS = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_BIAS[1] = {0}; +static const int64_t BUF_X[1] = {1}; +static const int64_t BUF_OUT[1] = {2}; + +static const int64_t VEC_READ[2] = {0, 1}; + +static const int64_t TILE[2] = {TM, TN}; + +static const int64_t STRIDE_2D[2] = {N, 1}; +static const int64_t STRIDE_BCAST[2] = {0, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_X = 0; +static const int32_t TAG_BIAS = 1; +static const int32_t TAG_OUT = 2; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_BIAS_ACT = 0; + +static void bias_gelu_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0], n0 = iv[1]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_BIAS, (uint64_t)n0, + 2, TILE, STRIDE_BCAST, ELEM_BITS, + SYNC, TAG_BIAS, 0, nullptr, 0, BUF_BIAS, 1); + + // COMPUTE + togsim_compute(ctx, TID_BIAS_ACT, CT_VECTOR, 0, nullptr, + VEC_READ, 2, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t mi = 0; mi < TILES_M; ++mi) { + for (int64_t ni = 0; ni < TILES_N; ++ni) { + int64_t iv[2] = {mi * TM, ni * TN}; + togsim_dispatch(ctx, bias_gelu_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/bias_gelu/bias_gelu.cpp b/TOGSim/example/bias_gelu/bias_gelu.cpp new file mode 100644 index 000000000..53d62bc41 --- /dev/null +++ b/TOGSim/example/bias_gelu/bias_gelu.cpp @@ -0,0 +1,78 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t M = 128; +static const int64_t N = 256; + +static const int64_t TM = 128; +static const int64_t TN = 64; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_M = (M + TM - 1) / TM; +static const int64_t TILES_N = (N + TN - 1) / TN; + +static const int32_t ARG_X = 0; +static const int32_t ARG_BIAS = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_BIAS[1] = {0}; +static const int64_t BUF_X[1] = {1}; +static const int64_t BUF_OUT[1] = {2}; + +static const int64_t VEC_READ[2] = {0, 1}; + +static const int64_t TILE[2] = {TM, TN}; + +static const int64_t STRIDE_2D[2] = {N, 1}; +static const int64_t STRIDE_BCAST[2] = {0, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_X = 0; +static const int32_t TAG_BIAS = 1; +static const int32_t TAG_OUT = 2; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_BIAS_GELU = 0; + +static void bias_gelu_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0], n0 = iv[1]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_BIAS, (uint64_t)n0, + 2, TILE, STRIDE_BCAST, ELEM_BITS, + SYNC, TAG_BIAS, 0, nullptr, 0, BUF_BIAS, 1); + + // COMPUTE + togsim_compute(ctx, TID_BIAS_GELU, CT_VECTOR, 0, nullptr, + VEC_READ, 2, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t mi = 0; mi < TILES_M; ++mi) { + for (int64_t ni = 0; ni < TILES_N; ++ni) { + int64_t iv[2] = {mi * TM, ni * TN}; + togsim_dispatch(ctx, bias_gelu_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/concat/concat.cpp b/TOGSim/example/concat/concat.cpp new file mode 100644 index 000000000..40cfeae7a --- /dev/null +++ b/TOGSim/example/concat/concat.cpp @@ -0,0 +1,74 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t A_ROWS = 128; +static const int64_t B_ROWS = 128; +static const int64_t D = 256; + +static const int64_t TROW = 128; +static const int64_t TD = 16; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_D = (D + TD - 1) / TD; + +static const int32_t ARG_A = 0; +static const int32_t ARG_B = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_A[1] = {0}; +static const int64_t BUF_B[1] = {1}; + +static const int64_t TILE[2] = {TROW, TD}; +static const int64_t STRIDE[2] = {D, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_A_IN = 0; +static const int32_t TAG_A_OUT = 1; +static const int32_t TAG_B_IN = 2; +static const int32_t TAG_B_OUT = 3; + +static void concat_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t d0 = iv[1]; + + for (int64_t r = 0; r < A_ROWS; r += TROW) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_A, (uint64_t)(r * D + d0), + 2, TILE, STRIDE, ELEM_BITS, + SYNC, TAG_A_IN, 0, nullptr, 0, BUF_A, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(r * D + d0), + 2, TILE, STRIDE, ELEM_BITS, + SYNC, TAG_A_OUT, 0, BUF_A, 1, nullptr, 0); + } + + for (int64_t r = 0; r < B_ROWS; r += TROW) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_B, (uint64_t)(r * D + d0), + 2, TILE, STRIDE, ELEM_BITS, + SYNC, TAG_B_IN, 0, nullptr, 0, BUF_B, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)((r + A_ROWS) * D + d0), + 2, TILE, STRIDE, ELEM_BITS, + SYNC, TAG_B_OUT, 0, BUF_B, 1, nullptr, 0); + } +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t di = 0; di < TILES_D; ++di) { + int64_t iv[2] = {0, di * TD}; + togsim_dispatch(ctx, concat_tile, iv, 2); + } +} diff --git a/TOGSim/example/conv/conv.cpp b/TOGSim/example/conv/conv.cpp new file mode 100644 index 000000000..36d15cd17 --- /dev/null +++ b/TOGSim/example/conv/conv.cpp @@ -0,0 +1,149 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// SA: 8x8 x2 per core +// VLANE: 8 +// mapping: autotune +static const int64_t I_C = 64; +static const int64_t I_H = 16; +static const int64_t I_W = 16; +static const int64_t O_C = 64; +static const int64_t O_H = 16; +static const int64_t O_W = 16; +static const int64_t K_H = 3; +static const int64_t K_W = 3; + +static const int64_t TILE_M = 16; +static const int64_t TILE_N = 64; +static const int64_t TILE_K = 64; +static const int64_t TILE_I_H = 8; + +static const int64_t SA = 8; +static const int64_t VLANE = 8; +static const int32_t ELEM_BITS = 32; + +static const int64_t PRE_ROWS = 8; +static const int64_t PRE_COLS = 8; +static const int64_t MM_ROWS = 8; +static const int64_t MM_COLS = 2; + +static const int64_t W_SUBTILES = 8; +static const int64_t X_SUBTILES = 16; +static const int64_t ROW_PAIR = (I_W + 2) * I_C; +static const int64_t ROW_HALF = TILE_I_H * I_C; + +static const int32_t ARG_X = 0; +static const int32_t ARG_W = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_ACC[1] = {0}; +static const int64_t BUF_W[1] = {1}; +static const int64_t BUF_X[1] = {2}; +static const int64_t BUF_SA[1] = {3}; + +static const int64_t PRELOAD_READ[2] = {1, 2}; +static const int64_t MM_READ[2] = {0, 3}; + +static const int64_t TILE_X[4] = {1, 1, TILE_I_H, I_C}; +static const int64_t TILE_W[4] = {1, 1, O_C, TILE_I_H}; +static const int64_t TILE_OUT[4] = {1, O_C, TILE_I_H, O_W}; + +static const int64_t STRIDE_X[4] = {(I_H + 2) * (I_W + 2) * I_C, (I_W + 2) * I_C, I_C, 1}; +static const int64_t STRIDE_W[4] = {I_C * O_C * K_W, I_C * O_C, I_C, 1}; +static const int64_t STRIDE_OUT[4] = {0, O_H * O_W, O_W, 1}; + +static const int32_t SYNC = 0; +static const int32_t ASYNC = 1; + +static const int32_t TAG_X = 0; +static const int32_t TAG_W = 1; +static const int32_t TAG_OUT = 2; + +static const int32_t CT_VECTOR = 0; +static const int32_t CT_MATMUL = 1; +static const int32_t CT_PRELOAD = 2; + +static const uint64_t TID_ACC_INIT = 0; +static const uint64_t TID_PRELOAD = 1; +static const uint64_t TID_MATMUL = 2; + +static void conv_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t oc0 = iv[0], oh0 = iv[1], ow0 = iv[2]; + + // COMPUTE + togsim_compute(ctx, TID_ACC_INIT, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_ACC, 1); + + for (int64_t kh = 0; kh < K_H; ++kh) { + for (int64_t kw = 0; kw < K_W; ++kw) { + const int64_t pos = kh * K_W + kw; + const int64_t x_off = (oh0 + kh) * (I_W + 2) * I_C + (ow0 + kw) * I_C; + const int64_t w_off = oc0 * K_H * K_W * I_C; + + for (int64_t s = 0; s < W_SUBTILES; ++s) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_W, (uint64_t)(w_off + s * SA), + 4, TILE_W, STRIDE_W, ELEM_BITS, + ASYNC, TAG_W, (uint64_t)s, nullptr, 0, BUF_W, 1); + } + for (int64_t s = 0; s < X_SUBTILES / 2; ++s) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, + (uint64_t)(x_off + s * ROW_PAIR), + 4, TILE_X, STRIDE_X, ELEM_BITS, + ASYNC, TAG_X, (uint64_t)(2 * s), nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, + (uint64_t)(x_off + s * ROW_PAIR + ROW_HALF), + 4, TILE_X, STRIDE_X, ELEM_BITS, + ASYNC, TAG_X, (uint64_t)(2 * s + 1), nullptr, 0, BUF_X, 1); + } + + for (int64_t pr = 0; pr < PRE_ROWS; ++pr) { + for (int64_t pc = 0; pc < PRE_COLS; ++pc) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_W, (uint64_t)pc, BUF_W, 1); + // COMPUTE + togsim_compute(ctx, TID_PRELOAD, CT_PRELOAD, 0, nullptr, + PRELOAD_READ, 2, BUF_SA, 1); + + for (int64_t mr = 0; mr < MM_ROWS; ++mr) { + for (int64_t mc = 0; mc < MM_COLS; ++mc) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_X, (uint64_t)(mr * MM_COLS + mc), BUF_X, 1); + // COMPUTE + togsim_compute(ctx, TID_MATMUL, CT_MATMUL, 0, nullptr, + MM_READ, 2, BUF_ACC, 1); + } + } + } + } + } + } + + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, + (uint64_t)(oc0 * O_H * O_W + oh0 * O_W + ow0), + 4, TILE_OUT, STRIDE_OUT, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_ACC, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t oc = 0; oc < O_C; oc += TILE_N) { + for (int64_t oh = 0; oh < O_H; oh += TILE_I_H) { + for (int64_t ow = 0; ow < O_W; ow += TILE_M) { + int64_t iv[3] = {oc, oh, ow}; + togsim_dispatch(ctx, conv_tile, iv, 3); + } + } + } +} diff --git a/TOGSim/example/embed_layernorm/embed_layernorm.cpp b/TOGSim/example/embed_layernorm/embed_layernorm.cpp new file mode 100644 index 000000000..012e36fb1 --- /dev/null +++ b/TOGSim/example/embed_layernorm/embed_layernorm.cpp @@ -0,0 +1,78 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t SEQ = 128; +static const int64_t DIM = 256; + +static const int64_t TSEQ = 128; +static const int64_t TDIM = 64; + +static const int32_t ELEM_BITS = 32; +static const int32_t IDX_BITS = 64; + +static const int64_t TILES_DIM = (DIM + TDIM - 1) / TDIM; + +static const int32_t ARG_IDS = 0; +static const int32_t ARG_TABLE = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_IDS[1] = {0}; +static const int64_t BUF_ROW[1] = {1}; + +static const int64_t VEC_READ[2] = {0, 1}; + +static const int64_t TILE[2] = {TSEQ, TDIM}; + +static const int64_t STRIDE_IDS[2] = {1, 0}; +static const int64_t STRIDE_ROW[2] = {0, 1}; +static const int64_t STRIDE_OUT[2] = {DIM, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_IDS = 0; +static const int32_t TAG_TABLE = 1; +static const int32_t TAG_OUT = 2; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_GATHER = 0; + +static void embed_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t s0 = iv[0], d0 = iv[1]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_IDS, (uint64_t)s0, + 2, TILE, STRIDE_IDS, IDX_BITS, + SYNC, TAG_IDS, 0, nullptr, 0, BUF_IDS, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_TABLE, (uint64_t)d0, + 2, TILE, STRIDE_ROW, ELEM_BITS, + SYNC, TAG_TABLE, 0, BUF_IDS, 1, BUF_ROW, 1); + + // COMPUTE + togsim_compute(ctx, TID_GATHER, CT_VECTOR, 0, nullptr, + VEC_READ, 2, nullptr, 0); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(s0 * DIM + d0), + 2, TILE, STRIDE_OUT, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_ROW, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t s0 = 0; s0 < SEQ; s0 += TSEQ) { + for (int64_t di = 0; di < TILES_DIM; ++di) { + int64_t iv[2] = {s0, di * TDIM}; + togsim_dispatch(ctx, embed_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/flatten/flatten.cpp b/TOGSim/example/flatten/flatten.cpp new file mode 100644 index 000000000..43b7fd88e --- /dev/null +++ b/TOGSim/example/flatten/flatten.cpp @@ -0,0 +1,16 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)ctx; (void)shape_args; (void)n; +} diff --git a/TOGSim/example/gemm/gemm.cpp b/TOGSim/example/gemm/gemm.cpp new file mode 100644 index 000000000..0cb8bd5e5 --- /dev/null +++ b/TOGSim/example/gemm/gemm.cpp @@ -0,0 +1,114 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// SA: 8x8 x2 per core +// VLANE: 8 +// mapping: autotune +static const int64_t M = 128; +static const int64_t K = 256; +static const int64_t N = 256; + +static const int64_t TM = 32; +static const int64_t TN = 256; +static const int64_t TK = 64; + +static const int64_t SA = 8; +static const int64_t VLANE = 8; +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_M = (M + TM - 1) / TM; +static const int64_t TILES_N = (N + TN - 1) / TN; + +static const int64_t STEPS_M = TM / SA; +static const int64_t STEPS_N = TN / (SA * VLANE); + +static const int32_t ARG_A = 0; +static const int32_t ARG_B = 1; +static const int32_t ARG_C = 2; + +static const int64_t BUF_C[1] = {0}; +static const int64_t BUF_B[1] = {1}; +static const int64_t BUF_A[1] = {2}; +static const int64_t BUF_SA[1] = {3}; + +static const int64_t PRELOAD_READ[2] = {1, 2}; +static const int64_t MM_READ[2] = {0, 3}; + +static const int64_t TILE_A[2] = {TM, TK}; +static const int64_t TILE_B[2] = {TK, TN}; +static const int64_t TILE_C[2] = {TM, TN}; +static const int64_t STRIDE[2] = {N, 1}; + +static const int32_t SYNC = 0; +static const int32_t ASYNC = 1; + +static const int32_t TAG_A = 0; +static const int32_t TAG_B = 1; +static const int32_t TAG_C = 2; + +static const int32_t CT_VECTOR = 0; +static const int32_t CT_MATMUL = 1; +static const int32_t CT_PRELOAD = 2; + +static const uint64_t TID_ACC_INIT = 0; +static const uint64_t TID_PRELOAD = 1; +static const uint64_t TID_MATMUL = 2; + +static void gemm_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0], n0 = iv[1]; + // COMPUTE + togsim_compute(ctx, TID_ACC_INIT, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_C, 1); + + for (int64_t k0 = 0; k0 < K; k0 += TK) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_A, (uint64_t)(m0 * K + k0), + 2, TILE_A, STRIDE, ELEM_BITS, + ASYNC, TAG_A, 0, nullptr, 0, BUF_A, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_B, (uint64_t)(k0 * N + n0), + 2, TILE_B, STRIDE, ELEM_BITS, + ASYNC, TAG_B, 0, nullptr, 0, BUF_B, 1); + + for (int64_t k = 0; k < TK; ++k) { + for (int64_t n = 0; n < STEPS_N; ++n) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_B, 0, BUF_B, 1); + // COMPUTE + togsim_compute(ctx, TID_PRELOAD, CT_PRELOAD, 0, nullptr, + PRELOAD_READ, 2, BUF_SA, 1); + + for (int64_t m = 0; m < STEPS_M; ++m) { + // MEMORY_BAR + togsim_memory_barrier(ctx, TAG_A, 0, BUF_A, 1); + // COMPUTE + togsim_compute(ctx, TID_MATMUL, CT_MATMUL, 0, nullptr, + MM_READ, 2, BUF_C, 1); + } + } + } + } + + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_C, (uint64_t)(m0 * N + n0), + 2, TILE_C, STRIDE, ELEM_BITS, + SYNC, TAG_C, 0, BUF_C, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t mi = 0; mi < TILES_M; ++mi) { + for (int64_t ni = 0; ni < TILES_N; ++ni) { + int64_t iv[2] = {mi * TM, ni * TN}; + togsim_dispatch(ctx, gemm_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/global_avgpool/global_avgpool.cpp b/TOGSim/example/global_avgpool/global_avgpool.cpp new file mode 100644 index 000000000..70804bf0b --- /dev/null +++ b/TOGSim/example/global_avgpool/global_avgpool.cpp @@ -0,0 +1,103 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t C = 256; +static const int64_t HW = 64; + +static const int64_t TC = 256; +static const int64_t THW = 64; +static const int64_t TC_SCALE = 128; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_C = (C + TC - 1) / TC; +static const int64_t TILES_C_SCALE = (C + TC_SCALE - 1) / TC_SCALE; + +static const int32_t ARG_X = 0; +static const int32_t ARG_SUM = 1; +static const int32_t ARG_OUT = 2; + +static const int64_t BUF_X[1] = {0}; +static const int64_t BUF_SUM[1] = {1}; +static const int64_t BUF_OUT[1] = {2}; + +static const int64_t TILE[2] = {TC, THW}; +static const int64_t TILE_ROW[1] = {TC}; +static const int64_t TILE_SCALE[1] = {TC_SCALE}; + +static const int64_t STRIDE_2D[2] = {HW, 1}; +static const int64_t STRIDE_1D[1] = {1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_X = 0; +static const int32_t TAG_SUM = 1; +static const int32_t TAG_OUT = 2; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_REDUCE = 0; +static const uint64_t TID_WRITE = 1; +static const uint64_t TID_SCALE = 2; + +static void sum_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t c0 = iv[0]; + + for (int64_t hw0 = 0; hw0 < HW; hw0 += THW) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(c0 * HW + hw0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // COMPUTE + togsim_compute(ctx, TID_REDUCE, CT_VECTOR, 0, nullptr, + BUF_X, 1, nullptr, 0); + } + + // COMPUTE + togsim_compute(ctx, TID_WRITE, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_SUM, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_SUM, (uint64_t)c0, + 1, TILE_ROW, STRIDE_1D, ELEM_BITS, + SYNC, TAG_SUM, 0, BUF_SUM, 1, nullptr, 0); +} + +static void scale_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t c0 = iv[0]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_SUM, (uint64_t)c0, + 1, TILE_SCALE, STRIDE_1D, ELEM_BITS, + SYNC, TAG_SUM, 0, nullptr, 0, BUF_SUM, 1); + + // COMPUTE + togsim_compute(ctx, TID_SCALE, CT_VECTOR, 0, nullptr, + BUF_SUM, 1, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)c0, + 1, TILE_SCALE, STRIDE_1D, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + + for (int64_t ci = 0; ci < TILES_C; ++ci) { + int64_t iv[1] = {ci * TC}; + togsim_dispatch(ctx, sum_tile, iv, 1); + } + for (int64_t ci = 0; ci < TILES_C_SCALE; ++ci) { + int64_t iv[1] = {ci * TC_SCALE}; + togsim_dispatch(ctx, scale_tile, iv, 1); + } +} diff --git a/TOGSim/example/kvcache_concat/kvcache_concat.cpp b/TOGSim/example/kvcache_concat/kvcache_concat.cpp new file mode 100644 index 000000000..a17f39990 --- /dev/null +++ b/TOGSim/example/kvcache_concat/kvcache_concat.cpp @@ -0,0 +1,94 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t TOKENS = 128; +static const int64_t HIDDEN = 256; +static const int64_t NUM_HEADS = 8; +static const int64_t NUM_KV_HEADS = 8; +static const int64_t PAST = 128; + +static const int64_t CACHE_DIM = HIDDEN / NUM_HEADS * NUM_KV_HEADS; +static const int64_t QKV_WIDTH = HIDDEN + 2 * CACHE_DIM; + +static const int64_t TTOK = 32; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_TOK = (TOKENS + TTOK - 1) / TTOK; + +static const int32_t ARG_QKV = 0; +static const int32_t ARG_QUERY = 1; +static const int32_t ARG_KEY = 2; +static const int32_t ARG_VALUE = 3; + +static const int64_t BUF_QKV[1] = {0}; +static const int64_t BUF_QUERY[1] = {1}; +static const int64_t BUF_KEY[1] = {2}; +static const int64_t BUF_VALUE[1] = {3}; + +static const int64_t SPLIT_WRITE[3] = {1, 2, 3}; + +static const int64_t TILE_QKV[2] = {TTOK, QKV_WIDTH}; +static const int64_t TILE_QUERY[2] = {TTOK, HIDDEN}; +static const int64_t TILE_CACHE[2] = {TTOK, CACHE_DIM}; + +static const int64_t STRIDE_QKV[2] = {QKV_WIDTH, 1}; +static const int64_t STRIDE_QUERY[2] = {HIDDEN, 1}; +static const int64_t STRIDE_CACHE[2] = {CACHE_DIM, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_QKV = 0; +static const int32_t TAG_QUERY = 1; +static const int32_t TAG_KEY = 2; +static const int32_t TAG_VALUE = 3; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_SPLIT = 0; + +static void split_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t t0 = iv[0]; + + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_QKV, (uint64_t)(t0 * QKV_WIDTH), + 2, TILE_QKV, STRIDE_QKV, ELEM_BITS, + SYNC, TAG_QKV, 0, nullptr, 0, BUF_QKV, 1); + + // COMPUTE + togsim_compute(ctx, TID_SPLIT, CT_VECTOR, 0, nullptr, + BUF_QKV, 1, SPLIT_WRITE, 3); + + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_QUERY, (uint64_t)(t0 * HIDDEN), + 2, TILE_QUERY, STRIDE_QUERY, ELEM_BITS, + SYNC, TAG_QUERY, 0, BUF_QUERY, 1, nullptr, 0); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_KEY, + (uint64_t)((PAST + t0) * CACHE_DIM), + 2, TILE_CACHE, STRIDE_CACHE, ELEM_BITS, + SYNC, TAG_KEY, 0, BUF_KEY, 1, nullptr, 0); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_VALUE, + (uint64_t)((PAST + t0) * CACHE_DIM), + 2, TILE_CACHE, STRIDE_CACHE, ELEM_BITS, + SYNC, TAG_VALUE, 0, BUF_VALUE, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t ti = 0; ti < TILES_TOK; ++ti) { + int64_t iv[1] = {ti * TTOK}; + togsim_dispatch(ctx, split_tile, iv, 1); + } +} diff --git a/TOGSim/example/layernorm/layernorm.cpp b/TOGSim/example/layernorm/layernorm.cpp new file mode 100644 index 000000000..6bfd64caf --- /dev/null +++ b/TOGSim/example/layernorm/layernorm.cpp @@ -0,0 +1,155 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t M = 128; +static const int64_t N = 256; + +static const int64_t TM = 64; +static const int64_t TN = 64; +static const int64_t TM_STATS = 128; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_M = (M + TM - 1) / TM; +static const int64_t TILES_N = (N + TN - 1) / TN; +static const int64_t TILES_M_STATS = (M + TM_STATS - 1) / TM_STATS; + +static const int32_t ARG_X = 0; +static const int32_t ARG_SKIP = 1; +static const int32_t ARG_MEAN = 2; +static const int32_t ARG_VAR = 3; +static const int32_t ARG_W = 4; +static const int32_t ARG_B = 5; +static const int32_t ARG_OUT = 6; + +static const int64_t BUF_X[1] = {0}; +static const int64_t BUF_W[1] = {1}; +static const int64_t BUF_VAR[1] = {2}; +static const int64_t BUF_MEAN[1] = {3}; +static const int64_t BUF_SKIP[1] = {4}; +static const int64_t BUF_B[1] = {5}; +static const int64_t BUF_OUT[1] = {6}; + +static const int64_t READ_X_SKIP[2] = {0, 4}; +static const int64_t READ_ALL[6] = {0, 1, 2, 3, 4, 5}; + +static const int64_t TILE[2] = {TM, TN}; +static const int64_t TILE_STATS[2] = {TM_STATS, TN}; +static const int64_t TILE_ROW[1] = {TM_STATS}; + +static const int64_t STRIDE_2D[2] = {N, 1}; +static const int64_t STRIDE_1D[1] = {1}; +static const int64_t STRIDE_ROW[2] = {1, 0}; +static const int64_t STRIDE_COL[2] = {0, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_X = 0; +static const int32_t TAG_SKIP = 1; +static const int32_t TAG_MEAN = 2; +static const int32_t TAG_VAR = 3; +static const int32_t TAG_W = 4; +static const int32_t TAG_B = 5; +static const int32_t TAG_OUT = 6; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_REDUCE = 0; +static const uint64_t TID_WRITE_MEAN = 1; +static const uint64_t TID_WRITE_VAR = 2; +static const uint64_t TID_NORM = 3; + +static void stats_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0]; + + for (int64_t n0 = 0; n0 < N; n0 += TN) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE_STATS, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_SKIP, (uint64_t)(m0 * N + n0), + 2, TILE_STATS, STRIDE_2D, ELEM_BITS, + SYNC, TAG_SKIP, 0, nullptr, 0, BUF_SKIP, 1); + // COMPUTE + togsim_compute(ctx, TID_REDUCE, CT_VECTOR, 0, nullptr, + READ_X_SKIP, 2, nullptr, 0); + } + + // COMPUTE + togsim_compute(ctx, TID_WRITE_MEAN, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_MEAN, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_MEAN, (uint64_t)m0, + 1, TILE_ROW, STRIDE_1D, ELEM_BITS, + SYNC, TAG_MEAN, 0, BUF_MEAN, 1, nullptr, 0); + // COMPUTE + togsim_compute(ctx, TID_WRITE_VAR, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_VAR, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_VAR, (uint64_t)m0, + 1, TILE_ROW, STRIDE_1D, ELEM_BITS, + SYNC, TAG_VAR, 0, BUF_VAR, 1, nullptr, 0); +} + +static void norm_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0], n0 = iv[1]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_SKIP, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_SKIP, 0, nullptr, 0, BUF_SKIP, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_MEAN, (uint64_t)m0, + 2, TILE, STRIDE_ROW, ELEM_BITS, + SYNC, TAG_MEAN, 0, nullptr, 0, BUF_MEAN, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_VAR, (uint64_t)m0, + 2, TILE, STRIDE_ROW, ELEM_BITS, + SYNC, TAG_VAR, 0, nullptr, 0, BUF_VAR, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_W, (uint64_t)n0, + 2, TILE, STRIDE_COL, ELEM_BITS, + SYNC, TAG_W, 0, nullptr, 0, BUF_W, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_B, (uint64_t)n0, + 2, TILE, STRIDE_COL, ELEM_BITS, + SYNC, TAG_B, 0, nullptr, 0, BUF_B, 1); + + // COMPUTE + togsim_compute(ctx, TID_NORM, CT_VECTOR, 0, nullptr, + READ_ALL, 6, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + + for (int64_t mi = 0; mi < TILES_M_STATS; ++mi) { + int64_t iv[1] = {mi * TM_STATS}; + togsim_dispatch(ctx, stats_tile, iv, 1); + } + for (int64_t mi = 0; mi < TILES_M; ++mi) { + for (int64_t ni = 0; ni < TILES_N; ++ni) { + int64_t iv[2] = {mi * TM, ni * TN}; + togsim_dispatch(ctx, norm_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/maxpool/maxpool.cpp b/TOGSim/example/maxpool/maxpool.cpp new file mode 100644 index 000000000..55a318d53 --- /dev/null +++ b/TOGSim/example/maxpool/maxpool.cpp @@ -0,0 +1,92 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t ROWS = 1024; +static const int64_t COLS = 16; + +static const int64_t TROW = 256; +static const int64_t TCOL = 16; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_ROW = (ROWS + TROW - 1) / TROW; +static const int64_t TILES_COL = (COLS + TCOL - 1) / TCOL; + +static const int64_t IN_ROW_STRIDE = 64; +static const int64_t IN_W = 32; + +static const int32_t ARG_X = 0; +static const int32_t ARG_OUT = 1; + +static const int64_t BUF_W3[1] = {0}; +static const int64_t BUF_W0[1] = {1}; +static const int64_t BUF_W1[1] = {2}; +static const int64_t BUF_W2[1] = {3}; +static const int64_t BUF_OUT[1] = {4}; + +static const int64_t VEC_READ[4] = {0, 1, 2, 3}; + +static const int64_t TILE[2] = {TROW, TCOL}; + +static const int64_t STRIDE_IN[2] = {IN_ROW_STRIDE, 2}; +static const int64_t STRIDE_OUT[2] = {TCOL, 1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_W0 = 0; +static const int32_t TAG_W1 = 1; +static const int32_t TAG_W2 = 2; +static const int32_t TAG_W3 = 3; +static const int32_t TAG_OUT = 4; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_MAX = 0; + +static void maxpool_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t r0 = iv[0], c0 = iv[1]; + const int64_t base = r0 * IN_ROW_STRIDE + c0 * 2; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)base, + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W0, 0, nullptr, 0, BUF_W0, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + 1), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W1, 0, nullptr, 0, BUF_W1, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + IN_W), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W2, 0, nullptr, 0, BUF_W2, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(base + IN_W + 1), + 2, TILE, STRIDE_IN, ELEM_BITS, + SYNC, TAG_W3, 0, nullptr, 0, BUF_W3, 1); + // COMPUTE + togsim_compute(ctx, TID_MAX, CT_VECTOR, 0, nullptr, + VEC_READ, 4, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(r0 * TCOL + c0), + 2, TILE, STRIDE_OUT, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + for (int64_t ri = 0; ri < TILES_ROW; ++ri) { + for (int64_t ci = 0; ci < TILES_COL; ++ci) { + int64_t iv[2] = {ri * TROW, ci * TCOL}; + togsim_dispatch(ctx, maxpool_tile, iv, 2); + } + } +} diff --git a/TOGSim/example/softmax/softmax.cpp b/TOGSim/example/softmax/softmax.cpp new file mode 100644 index 000000000..c8106e7f2 --- /dev/null +++ b/TOGSim/example/softmax/softmax.cpp @@ -0,0 +1,151 @@ +#include +#include +using std::size_t; + +#include "togsim_runtime.h" + +int32_t togsim_abi_version(void) { return TOGSIM_ABI_VERSION; } + +// config: systolic_ws_8x8_c1_simple_noc_tpuv3.yml +// VLANE: 8 +// mapping: autotune +static const int64_t M = 128; +static const int64_t N = 256; + +static const int64_t TM = 128; +static const int64_t TN = 64; + +static const int32_t ELEM_BITS = 32; + +static const int64_t TILES_M = (M + TM - 1) / TM; +static const int64_t TILES_N = (N + TN - 1) / TN; + +static const int32_t ARG_X = 0; +static const int32_t ARG_MAX = 1; +static const int32_t ARG_SUM = 2; +static const int32_t ARG_OUT = 3; + +static const int64_t BUF_SUM[1] = {0}; +static const int64_t BUF_X[1] = {1}; +static const int64_t BUF_MAX[1] = {2}; +static const int64_t BUF_OUT[1] = {3}; + +static const int64_t READ_X_MAX[2] = {1, 2}; +static const int64_t READ_X_MAX_SUM[3] = {0, 1, 2}; + +static const int64_t TILE[2] = {TM, TN}; +static const int64_t TILE_ROW[1] = {TM}; + +static const int64_t STRIDE_2D[2] = {N, 1}; +static const int64_t STRIDE_BCAST[2] = {1, 0}; +static const int64_t STRIDE_1D[1] = {1}; + +static const int32_t SYNC = 0; + +static const int32_t TAG_X = 0; +static const int32_t TAG_MAX = 1; +static const int32_t TAG_SUM = 2; +static const int32_t TAG_OUT = 3; + +static const int32_t CT_VECTOR = 0; + +static const uint64_t TID_MAX_REDUCE = 0; +static const uint64_t TID_MAX_WRITE = 1; +static const uint64_t TID_SUM_REDUCE = 2; +static const uint64_t TID_SUM_WRITE = 3; +static const uint64_t TID_SOFTMAX = 4; + +static void max_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0]; + + for (int64_t n0 = 0; n0 < N; n0 += TN) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // COMPUTE + togsim_compute(ctx, TID_MAX_REDUCE, CT_VECTOR, 0, nullptr, + BUF_X, 1, nullptr, 0); + } + + // COMPUTE + togsim_compute(ctx, TID_MAX_WRITE, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_MAX, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_MAX, (uint64_t)m0, + 1, TILE_ROW, STRIDE_1D, ELEM_BITS, + SYNC, TAG_MAX, 0, BUF_MAX, 1, nullptr, 0); +} + +static void sum_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0]; + + for (int64_t n0 = 0; n0 < N; n0 += TN) { + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_MAX, (uint64_t)m0, + 2, TILE, STRIDE_BCAST, ELEM_BITS, + SYNC, TAG_MAX, 0, nullptr, 0, BUF_MAX, 1); + // COMPUTE + togsim_compute(ctx, TID_SUM_REDUCE, CT_VECTOR, 0, nullptr, + READ_X_MAX, 2, nullptr, 0); + } + + // COMPUTE + togsim_compute(ctx, TID_SUM_WRITE, CT_VECTOR, 0, nullptr, + nullptr, 0, BUF_SUM, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_SUM, (uint64_t)m0, + 1, TILE_ROW, STRIDE_1D, ELEM_BITS, + SYNC, TAG_SUM, 0, BUF_SUM, 1, nullptr, 0); +} + +static void softmax_tile(EmitCtx* ctx, int64_t* iv, int32_t n_iv) { + (void)n_iv; + const int64_t m0 = iv[0], n0 = iv[1]; + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_X, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_X, 0, nullptr, 0, BUF_X, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_MAX, (uint64_t)m0, + 2, TILE, STRIDE_BCAST, ELEM_BITS, + SYNC, TAG_MAX, 0, nullptr, 0, BUF_MAX, 1); + // MOVIN + togsim_dma(ctx, TOGSIM_DMA_LOAD, ARG_SUM, (uint64_t)m0, + 2, TILE, STRIDE_BCAST, ELEM_BITS, + SYNC, TAG_SUM, 0, nullptr, 0, BUF_SUM, 1); + + // COMPUTE + togsim_compute(ctx, TID_SOFTMAX, CT_VECTOR, 0, nullptr, + READ_X_MAX_SUM, 3, BUF_OUT, 1); + // MOVOUT + togsim_dma(ctx, TOGSIM_DMA_STORE, ARG_OUT, (uint64_t)(m0 * N + n0), + 2, TILE, STRIDE_2D, ELEM_BITS, + SYNC, TAG_OUT, 0, BUF_OUT, 1, nullptr, 0); +} + +// DISPATCH +extern "C" void togsim_kernel(EmitCtx* ctx, int64_t* shape_args, int32_t n) { + (void)shape_args; (void)n; + + for (int64_t mi = 0; mi < TILES_M; ++mi) { + int64_t iv[1] = {mi * TM}; + togsim_dispatch(ctx, max_tile, iv, 1); + } + for (int64_t mi = 0; mi < TILES_M; ++mi) { + int64_t iv[1] = {mi * TM}; + togsim_dispatch(ctx, sum_tile, iv, 1); + } + for (int64_t mi = 0; mi < TILES_M; ++mi) { + for (int64_t ni = 0; ni < TILES_N; ++ni) { + int64_t iv[2] = {mi * TM, ni * TN}; + togsim_dispatch(ctx, softmax_tile, iv, 2); + } + } +} diff --git a/TOGSim/onnx_frontend/config.py b/TOGSim/onnx_frontend/config.py new file mode 100644 index 000000000..2f233873a --- /dev/null +++ b/TOGSim/onnx_frontend/config.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import os +import re + +import yaml + +from .tiling import Hardware + +#: systolic_ws_x_... -- the only place SA is written down. +_SA_IN_NAME = re.compile(r"_(\d+)x(\d+)_") + + +def array_dim_from_name(config_path: str) -> int | None: + m = _SA_IN_NAME.search(os.path.basename(config_path)) + if not m or m.group(1) != m.group(2): + return None + return int(m.group(1)) + + +def read_hardware(config_path: str, sa: int | None = None, + elem_bits: int = 32) -> Hardware: + with open(config_path) as fh: + cfg = yaml.safe_load(fh) or {} + + resolved = sa if sa is not None else array_dim_from_name(config_path) + if resolved is None: + raise ValueError( + f"cannot tell the systolic array dimension from " + f"{os.path.basename(config_path)!r}: it is not a TOGSim config key and " + f"the name does not carry it. Pass --sa.") + + try: + vlane = int(cfg["vpu_num_lanes"]) + vlen = int(cfg["vpu_vector_length_bits"]) + except KeyError as exc: + raise ValueError(f"{config_path}: missing required key {exc}") from None + + return Hardware(sa=resolved, vlane=vlane, vlen_bits=vlen, elem_bits=elem_bits) diff --git a/TOGSim/onnx_frontend/emit.py b/TOGSim/onnx_frontend/emit.py new file mode 100644 index 000000000..c9bfa0b28 --- /dev/null +++ b/TOGSim/onnx_frontend/emit.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import os +import re +import subprocess + + +def _assignment(name: str) -> re.Pattern: + return re.compile( + rf"^(\s*static\s+const\s+\w+(?:_t)?\s+{re.escape(name)}\s*=\s*)([^;]+)(;)", + re.MULTILINE) + + +def retarget(source: str, values: dict[str, int]) -> str: + text = source + for name, value in values.items(): + pattern = _assignment(name) + text, n = pattern.subn(rf"\g<1>{int(value)}\g<3>", text, count=1) + if n == 0: + raise KeyError( + f"constant {name!r} is not declared in this kernel; the tile " + f"names in tiling.py no longer match the .cpp") + return text + + +def emit(kernel_path: str, values: dict[str, int], out_cpp: str) -> str: + with open(kernel_path) as fh: + source = fh.read() + os.makedirs(os.path.dirname(out_cpp) or ".", exist_ok=True) + with open(out_cpp, "w") as fh: + fh.write(retarget(source, values)) + return out_cpp + + +def compile_so(cpp_path: str, include_dir: str, out_so: str) -> str: + if os.path.exists(out_so) and \ + os.path.getmtime(out_so) >= os.path.getmtime(cpp_path): + return out_so + subprocess.run( + ["g++", "-shared", "-fPIC", "-std=gnu++17", "-O0", + "-I", include_dir, cpp_path, "-o", out_so], + check=True, capture_output=True, text=True) + return out_so diff --git a/TOGSim/onnx_frontend/graph.py b/TOGSim/onnx_frontend/graph.py new file mode 100644 index 000000000..86f4b4667 --- /dev/null +++ b/TOGSim/onnx_frontend/graph.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import os +import tempfile + +import onnx + + +def optimize(model_path: str) -> str: + try: + import onnxruntime as ort + except ImportError: + return model_path + + out = os.path.join(tempfile.mkdtemp(prefix="onnx_opt_"), "optimized.onnx") + opts = ort.SessionOptions() + opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED + opts.optimized_model_filepath = out + # one thread: ORT otherwise logs an affinity failure per core + opts.intra_op_num_threads = 1 + opts.inter_op_num_threads = 1 + ort.set_default_logger_severity(3) + try: + ort.InferenceSession(model_path, opts, + providers=["CPUExecutionProvider"]) + except Exception: + return model_path # not fatal: the unfused graph still maps + return out if os.path.exists(out) else model_path + + +def infer_shapes(model) -> dict[str, list[int]]: + model = onnx.shape_inference.infer_shapes(model) + shapes: dict[str, list[int]] = {} + for init in model.graph.initializer: + shapes[init.name] = list(init.dims) + for group in (model.graph.input, model.graph.value_info, model.graph.output): + for vi in group: + dims = [d.dim_value if d.dim_value > 0 else 1 + for d in vi.type.tensor_type.shape.dim] + if dims: + shapes[vi.name] = dims + return shapes + + +def load(model_path: str, optimize_graph: bool = True): + original = onnx.load(model_path) + if not optimize_graph: + return original, infer_shapes(original) + + path = optimize(model_path) + if path == model_path: + return original, infer_shapes(original) + + # ORT drops value_info, and shape inference cannot rebuild it because the + # fused nodes it introduces are com.microsoft ops with no standard schema. + # The unfused graph still names the tensors that survive, so its shapes + # fill the gap; the optimized graph wins wherever both have an entry. + model = onnx.load(path) + shapes = {**infer_shapes(original), **infer_shapes(model)} + return model, shapes diff --git a/TOGSim/onnx_frontend/onnx_ops.py b/TOGSim/onnx_frontend/onnx_ops.py new file mode 100644 index 000000000..85ac04136 --- /dev/null +++ b/TOGSim/onnx_frontend/onnx_ops.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from dataclasses import dataclass + +SYSTOLIC = "systolic" +REDUCING = "reducing" +POINTWISE = "pointwise" + + +@dataclass(frozen=True) + + +class KernelRef: + directory: str + stem: str + tile_class: str + size_names: tuple = () + tile_names: tuple = () + hw_names: tuple = () + +#: op_type -> its kernel. +OPS: dict[str, list[KernelRef]] = { + "Gemm": [KernelRef("gemm", "gemm", SYSTOLIC, ("M", "K", "N"), + ("TM", "TN", "TK"), ("SA", "VLANE"))], + "MatMul": [KernelRef("gemm", "gemm", SYSTOLIC, ("M", "K", "N"), + ("TM", "TN", "TK"), ("SA", "VLANE"))], + + "Conv": [KernelRef("conv", "conv", SYSTOLIC, + ("I_C", "I_H", "I_W", "O_C", "O_H", "O_W", "K_H", "K_W"), + ("TILE_M", "TILE_N", "TILE_K"), ("SA", "VLANE"))], + + "Attention": [KernelRef("attention", "attention", SYSTOLIC, + ("HEADS", "KV_HEADS", "SEQ", "DHEAD"), (), ("SA",))], + "MultiHeadAttention": [KernelRef("attention", "attention", SYSTOLIC, + ("HEADS", "KV_HEADS", "SEQ", "DHEAD"), (), + ("SA",))], + + "Softmax": [ + KernelRef("softmax", "softmax", REDUCING, ("M", "N"), ("TM", "TN"))], + + "LayerNormalization": [ + KernelRef("layernorm", "layernorm", REDUCING, ("M", "N"), ("TM", "TN"))], + + "EmbedLayerNormalization": [ + KernelRef("embed_layernorm", "embed_layernorm", REDUCING, + ("SEQ", "DIM"), ("TSEQ", "TDIM"))], + + "GlobalAveragePool": [ + KernelRef("global_avgpool", "global_avgpool", REDUCING, + ("C", "HW"), ("TC", "THW"))], + + "Relu": [KernelRef("bias_act", "bias_act", POINTWISE, ("M", "N"), ("TM", "TN"))], + "Add": [KernelRef("bias_act", "bias_act", POINTWISE, ("M", "N"), ("TM", "TN"))], + "Gelu": [KernelRef("bias_gelu", "bias_gelu", POINTWISE, ("M", "N"), ("TM", "TN"))], + + "MaxPool": [KernelRef("maxpool", "maxpool", POINTWISE, + ("ROWS", "COLS"), ("TROW", "TCOL"))], + "AveragePool": [KernelRef("adaptive_avgpool", "adaptive_avgpool", POINTWISE, + ("ROWS", "COLS"), ("TROW", "TCOL"))], + + "Concat": [KernelRef("concat", "concat", POINTWISE, ("A_ROWS",), ("TROW",))], + "Flatten": [KernelRef("flatten", "flatten", POINTWISE, (), ())], +} + +#: Fused nodes with no kernel, and the operators they decompose to. +DECOMPOSE: dict[str, list[str]] = { + "FusedMatMul": ["MatMul"], + "FusedConv": ["Conv", "Relu"], + "BiasGelu": ["Add", "Gelu"], + "FastGelu": ["Add", "Gelu"], + "SkipLayerNormalization": ["Add", "LayerNormalization"], +} + +#: Metadata-only operators: no compute node to charge. +NO_COMPUTE = frozenset({ + "Reshape", "Squeeze", "Unsqueeze", "Identity", "Constant", "Shape", + "Dropout", "Cast", "Transpose", "Gather", "Slice", "ConstantOfShape", +}) + + +def resolve(op_type: str, _depth: int = 0) -> list[KernelRef] | None: + if op_type in NO_COMPUTE: + return [] + if op_type in OPS: + return list(OPS[op_type]) + if op_type in DECOMPOSE and _depth < 4: + out: list[KernelRef] = [] + for part in DECOMPOSE[op_type]: + got = resolve(part, _depth + 1) + if got is None: + return None + out.extend(got) + return out + return None diff --git a/TOGSim/onnx_frontend/run.py b/TOGSim/onnx_frontend/run.py new file mode 100644 index 000000000..38fcd1c8f --- /dev/null +++ b/TOGSim/onnx_frontend/run.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys + +from ..acs import cycles as acs_cycles +from ..acs import node_mapping +from . import emit, graph, shapes +from .config import read_hardware +from .onnx_ops import resolve + +HERE = os.path.dirname(os.path.abspath(__file__)) +TOGSIM = os.path.dirname(HERE) +EXAMPLE_DIR = os.path.join(TOGSIM, "example") +INCLUDE_DIR = os.path.join(TOGSIM, "include") +SIMULATOR = os.path.join(TOGSIM, "build", "bin", "Simulator") + +CYCLE_RE = re.compile(r"Total execution cycles:\s*(\d+)") + + +def simulate(so: str, table: str, config: str) -> int: + cmd = [SIMULATOR, "--config", config, "--trace_so", so, "--cycle_table", table] + res = subprocess.run(cmd, capture_output=True, text=True) + m = CYCLE_RE.search(res.stdout + res.stderr) + if not m: + tail = (res.stderr or res.stdout)[-400:] + raise RuntimeError(f"TOGSim reported no cycle count: {tail}") + return int(m.group(1)) + + +def run_node(node, kernel, shp, hw, config, out_dir): + values = shapes.constants(node, kernel, shp, hw) + tag = "_".join(str(values[n]) for n in kernel.size_names if n in values) + stem = f"{kernel.stem}_{tag}" if tag else kernel.stem + + cpp = os.path.join(out_dir, stem + ".cpp") + emit.emit(os.path.join(EXAMPLE_DIR, kernel.directory, kernel.stem + ".cpp"), + values, cpp) + so = emit.compile_so(cpp, INCLUDE_DIR, os.path.join(out_dir, stem + ".so")) + + desc = node_mapping.lookup(kernel.stem) + tile = shapes.tile_for_cost(kernel, values) + table = os.path.join(out_dir, stem + ".tsv") + acs_cycles.write_table(acs_cycles.build_table(desc, tile, hw), table) + + return simulate(so, table, config) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser( + description="Run an ONNX graph on TOGSim from the kernel library.", + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("model") + ap.add_argument("--config", required=True) + ap.add_argument("--out", default=os.path.join(HERE, "result")) + ap.add_argument("--sa", type=int, default=None, + help="systolic array dimension; inferred from the config name") + ap.add_argument("--no-optimize", action="store_true", + help="skip the ORT fusion pass") + args = ap.parse_args(argv) + + hw = read_hardware(args.config, args.sa) + os.makedirs(args.out, exist_ok=True) + model, shp = graph.load(args.model, optimize_graph=not args.no_optimize) + + print(f"config : {os.path.basename(args.config)}") + print(f"hardware : SA={hw.sa} lanes={hw.vlane} vlen={hw.vlen_bits}b " + f"({hw.throughput} elem/cycle)") + print() + + total = 0 + ran = no_compute = 0 + unmapped: dict[str, int] = {} + rows = [] + + for node in model.graph.node: + kernels = resolve(node.op_type) + if kernels is None: + unmapped[node.op_type] = unmapped.get(node.op_type, 0) + 1 + continue + if not kernels: + no_compute += 1 + continue + node_cycles = 0 + try: + for kernel in kernels: + node_cycles += run_node(node, kernel, shp, hw, args.config, + args.out) + except (KeyError, ValueError, RuntimeError, + subprocess.CalledProcessError) as exc: + reason = f"{node.op_type} ({type(exc).__name__})" + unmapped[reason] = unmapped.get(reason, 0) + 1 + print(f" !! {node.op_type:<18} {exc}", file=sys.stderr) + continue + + total += node_cycles + ran += 1 + rows.append({"op": node.op_type, "name": node.name or f"#{ran}", + "cycles": node_cycles}) + print(f" {node.op_type:<20} {node_cycles:>12,} cycles") + + print() + print(f"nodes run : {ran}") + print(f"nodes no-op : {no_compute}") + print(f"total cycles : {total:,}") + if unmapped: + print("\nUNMAPPED (not counted in the total):") + for op, n in sorted(unmapped.items(), key=lambda kv: -kv[1]): + print(f" {op}: {n}") + + report = {"model": args.model, "config": args.config, + "hardware": {"sa": hw.sa, "vlane": hw.vlane, + "vlen_bits": hw.vlen_bits}, + "total_cycles": total, "nodes_run": ran, + "nodes_no_compute": no_compute, "unmapped": unmapped, + "nodes": rows} + with open(os.path.join(args.out, "report.json"), "w") as fh: + json.dump(report, fh, indent=2) + print(f"\nreport : {os.path.join(args.out, 'report.json')}") + return 0 + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/TOGSim/onnx_frontend/shapes.py b/TOGSim/onnx_frontend/shapes.py new file mode 100644 index 000000000..b6fbe97f7 --- /dev/null +++ b/TOGSim/onnx_frontend/shapes.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +from .onnx_ops import KernelRef, POINTWISE, REDUCING, SYSTOLIC +from .tiling import (Hardware, check_systolic, pointwise_tile, reducing_tile, + systolic_tile) + + +def rows_cols(shape: list[int]) -> tuple[int, int]: + if not shape: + return 1, 1 + if len(shape) == 1: + return 1, shape[0] + rows = 1 + for d in shape[:-1]: + rows *= d + return rows, shape[-1] + + +def attr(node, name, default): + for a in node.attribute: + if a.name == name: + return list(a.ints) if a.ints else a.i + return default + + +def _matmul_dims(node, shapes) -> tuple[int, int, int]: + a = shapes[node.input[0]] + b = shapes[node.input[1]] + m, k = rows_cols(a) + if len(b) == 1: + n = b[0] + else: + n = b[-1] + if attr(node, "transB", 0): + n = b[-2] + return m, k, n + + +def _conv_dims(node, shapes) -> dict[str, int]: + x = shapes[node.input[0]] + w = shapes[node.input[1]] + if attr(node, "group", 1) != 1: + raise ValueError("grouped convolution: no library kernel takes a group count") + o_c, _, k_h, k_w = w + strides = attr(node, "strides", [1, 1]) + s = strides[0] if isinstance(strides, list) else strides + pads = attr(node, "pads", [0, 0, 0, 0]) + p = pads[0] if isinstance(pads, list) else pads + i_c, i_h, i_w = x[1], x[2], x[3] + o_h = (i_h + 2 * p - k_h) // s + 1 + o_w = (i_w + 2 * p - k_w) // s + 1 + return dict(I_C=i_c, I_H=i_h, I_W=i_w, O_C=o_c, + O_H=max(1, o_h), O_W=max(1, o_w), K_H=k_h, K_W=k_w) + + +def _attention_heads(node, q_shape) -> int: + n = attr(node, "num_heads", 0) + return int(n) if n else 1 + + +def _pool_dims(node, shapes) -> tuple[int, int]: + x = shapes[node.input[0]] + kernel = attr(node, "kernel_shape", [1, 1]) + strides = attr(node, "strides", [1, 1]) + s = strides[0] if isinstance(strides, list) else strides + if len(x) == 4: + return x[1] * max(1, x[2] // max(1, s)), max(1, x[3] // max(1, s)) + return rows_cols(x) + + +def constants(node, kernel: KernelRef, shapes: dict, hw: Hardware) -> dict[str, int]: + out: dict[str, int] = {} + + if kernel.directory == "conv": + dims = _conv_dims(node, shapes) + out.update(dims) + m = dims["O_H"] * dims["O_W"] + n = dims["O_C"] + k = dims["I_C"] * dims["K_H"] * dims["K_W"] + check_systolic(n, hw) + tm, tn, tk = systolic_tile(m, n, k, hw) + out.update(TILE_M=tm, TILE_N=tn, TILE_K=tk) + + elif kernel.stem == "attention": + q = shapes[node.input[0]] + seq = q[1] if len(q) >= 3 else q[0] + heads = _attention_heads(node, q) + kv_heads = int(attr(node, "num_kv_heads", 0)) or heads + dhead = max(1, q[-1] // max(1, heads)) + out.update(HEADS=heads, KV_HEADS=kv_heads, SEQ=seq, DHEAD=dhead) + + elif kernel.tile_class == SYSTOLIC: + m, k, n = _matmul_dims(node, shapes) + check_systolic(n, hw) + tm, tn, tk = systolic_tile(m, n, k, hw) + out.update(M=m, K=k, N=n, TM=tm, TN=tn, TK=tk) + + elif kernel.stem == "global_avgpool": + x = shapes[node.input[0]] + c = x[1] if len(x) >= 2 else x[0] + hw_elems = (x[2] * x[3]) if len(x) == 4 else 1 + tc, thw = reducing_tile(c, hw_elems, hw) + out.update(C=c, HW=hw_elems, TC=tc, THW=thw) + + elif not kernel.size_names: + pass + + else: + if kernel.directory in ("maxpool", "adaptive_avgpool"): + rows, cols = _pool_dims(node, shapes) + else: + rows, cols = rows_cols(shapes[node.input[0]]) + if kernel.tile_class == REDUCING: + tm, tn = reducing_tile(rows, cols, hw) + else: + tm, tn = pointwise_tile(rows, cols, hw) + # each kernel names these its own way + size = kernel.size_names + tiles = kernel.tile_names + out[size[0]] = rows + if len(size) > 1: + out[size[1]] = cols + if tiles: + out[tiles[0]] = tm + if len(tiles) > 1: + out[tiles[1]] = tn + + for name in kernel.hw_names: + out[name] = hw.sa if name == "SA" else hw.vlane + return out + + +def tile_for_cost(kernel: KernelRef, values: dict[str, int]) -> dict[str, int]: + if kernel.directory == "conv": + return {"rows": values["TILE_M"], + "elems": values["TILE_M"] * values["TILE_N"]} + names = kernel.tile_names + if not names: # flatten, attention: no tile constants + return {"rows": values.get("SEQ", 1), + "elems": values.get("SEQ", 1) * values.get("DHEAD", 1)} + rows = values[names[0]] + elems = rows * values[names[1]] if len(names) > 1 else rows + return {"rows": rows, "elems": elems} diff --git a/TOGSim/onnx_frontend/tests/test_frontend.py b/TOGSim/onnx_frontend/tests/test_frontend.py new file mode 100644 index 000000000..28e5cf2c2 --- /dev/null +++ b/TOGSim/onnx_frontend/tests/test_frontend.py @@ -0,0 +1,169 @@ +import os +import re +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname( + os.path.dirname(os.path.abspath(__file__)))))) + +from TOGSim.acs import cycles as acs_cycles +from TOGSim.acs.node_mapping import ELEMENTWISE, KernelDesc, NodeDesc, Op, lookup +from TOGSim.onnx_frontend import emit +from TOGSim.onnx_frontend.config import array_dim_from_name +from TOGSim.onnx_frontend.onnx_ops import OPS, resolve +from TOGSim.onnx_frontend.tiling import Hardware, check_systolic, systolic_tile + +EXAMPLE = os.path.join(os.path.dirname(os.path.dirname( + os.path.dirname(os.path.abspath(__file__)))), "example") +HW = Hardware(sa=8, vlane=8, vlen_bits=256) + + +def test_retarget_rewrites_only_the_named_constant(): + src = "static const int64_t M = 128;\nstatic const int64_t TILES = (M + 1) / 2;\n" + out = emit.retarget(src, {"M": 512}) + assert "M = 512;" in out + assert "(M + 1) / 2" in out # derived constants follow, untouched + + +def test_retarget_refuses_a_constant_the_kernel_does_not_declare(): + # silently skipping leaves the kernel at the example's size, and the run + # still reports a cycle count -- for a shape nobody asked about + with pytest.raises(KeyError): + emit.retarget("static const int64_t M = 128;\n", {"TM": 32}) + + +@pytest.mark.parametrize("op,kernels", sorted(OPS.items())) + + +def test_every_mapped_kernel_exists_and_declares_its_constants(op, kernels): + for kernel in kernels: + path = os.path.join(EXAMPLE, kernel.directory, kernel.stem + ".cpp") + assert os.path.exists(path), f"{op}: no kernel at {path}" + source = open(path).read() + for name in kernel.size_names + kernel.tile_names + kernel.hw_names: + # the kernels align their '=' signs, so match on the declaration + assert re.search(rf"\b{re.escape(name)}\s*=", source), \ + f"{op}/{kernel.stem}: no constant {name}" + + +@pytest.mark.parametrize("op,kernels", sorted(OPS.items())) + + +def test_every_mapped_kernel_has_a_node_table(op, kernels): + for kernel in kernels: + lookup(kernel.stem) # raises with a message if absent + + +def test_a_fused_node_decomposes_to_kernels_that_exist(): + assert [k.stem for k in resolve("BiasGelu")] == ["bias_act", "bias_gelu"] + assert [k.stem for k in resolve("SkipLayerNormalization")] == \ + ["bias_act", "layernorm"] + + +def test_a_fused_node_with_its_own_kernel_maps_straight_to_it(): + assert [k.stem for k in resolve("Attention")] == ["attention"] + + +def test_an_unsupported_op_is_reported_not_guessed(): + assert resolve("Einsum") is None # None -> unmapped, never costed + + +def test_metadata_ops_carry_no_compute(): + assert resolve("Reshape") == [] + + +def test_systolic_tile_is_a_multiple_of_the_array_geometry(): + tm, tn, tk = systolic_tile(512, 1024, 512, HW) + assert tm % HW.sa == 0 + # under SA*VLANE the kernel's STEPS_N truncates to 0 and issues no matmul + assert tn % (HW.sa * HW.vlane) == 0 and tn >= HW.sa * HW.vlane + + +def test_too_narrow_for_the_array_is_refused(): + check_systolic(64, HW) # 64 == SA*VLANE, the minimum + with pytest.raises(ValueError): + check_systolic(63, HW) + + +def test_array_dim_comes_from_the_config_name(): + assert array_dim_from_name("systolic_ws_8x8_c1_simple_noc.yml") == 8 + assert array_dim_from_name("systolic_ws_128x128_c2_booksim_tpuv3.yml") == 128 + assert array_dim_from_name("no_array_here.yml") is None + + +def test_the_table_column_is_an_overlap_not_an_interval(): + rows = acs_cycles.build_table(lookup("gemm"), + {"rows": 32, "elems": 32 * 256}, HW) + preload_cycles, preload_overlap = rows[1] + assert preload_cycles - preload_overlap == 0 + + matmul_cycles, matmul_overlap = rows[2] + assert matmul_cycles - matmul_overlap == min(HW.sa, 32) + + +def test_vector_work_is_charged_in_full(): + rows = acs_cycles.build_table(lookup("bias_act"), {"elems": 4096}, HW) + cycles, overlapping = rows[0] + assert overlapping == 0 + assert cycles == (4096 // HW.throughput) * 2 # + bias, then max(x, 0) + + +def test_gelu_costs_more_than_relu_on_the_same_tile(): + tile = {"elems": 4096} + relu = acs_cycles.build_table(lookup("bias_act"), tile, HW)[0][0] + gelu = acs_cycles.build_table(lookup("bias_gelu"), tile, HW)[0][0] + assert gelu == 3 * relu # six instructions against two + + +def test_the_cost_function_can_be_replaced(): + desc = KernelDesc("x", (NodeDesc("n", "vector", (Op(ELEMENTWISE, 1, "op"),)),)) + try: + acs_cycles.set_cost_function(lambda node, tile, hw: (99, 0)) + assert acs_cycles.build_table(desc, {"elems": 1}, HW) == [(99, 99)] + finally: + acs_cycles.set_cost_function(acs_cycles.default_cost) + + +def _kernel_source(stem): + import glob + hits = glob.glob(os.path.join(EXAMPLE, "*", stem + ".cpp")) + return hits[0] if hits else None + + +@pytest.mark.parametrize("stem", sorted( + k for k in __import__("TOGSim.acs.node_mapping", fromlist=["KERNELS"]).KERNELS)) + + +def test_table_rows_match_the_kernel_and_the_node_list(stem): + import re + from TOGSim.acs.node_mapping import KERNELS + + source = _kernel_source(stem) + if source is None: + pytest.skip(f"{stem} has no kernel in example/") + + text = open(source).read() + tids = set(re.findall(r"static const uint64_t (TID_\w+)\s*=", text)) + + nodes = KERNELS[stem].nodes + assert len(tids) == len(nodes), ( + f"{stem}: {len(tids)} tile_ids in the .cpp but {len(nodes)} in " + f"node_mapping.py") + + # a table is only present when one was generated or measured; when it is, + # its row count has to match too + table = os.path.join(os.path.dirname(source), "cycles.tsv") + if os.path.exists(table): + rows = [l for l in open(table) + if l.strip() and not l.startswith("#")] + assert len(rows) == len(nodes), ( + f"{stem}: cycles.tsv has {len(rows)} rows but {len(nodes)} nodes") + + +@pytest.mark.parametrize("op,kernels", sorted(OPS.items())) + + +def test_one_operator_maps_to_one_kernel(op, kernels): + assert len(kernels) == 1, ( + f"{op} maps to {[k.stem for k in kernels]}; merge them into one kernel") diff --git a/TOGSim/onnx_frontend/tiling.py b/TOGSim/onnx_frontend/tiling.py new file mode 100644 index 000000000..7a6188d3b --- /dev/null +++ b/TOGSim/onnx_frontend/tiling.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) + + +class Hardware: + sa: int # systolic array dimension (see note in config.py) + vlane: int # vpu_num_lanes + vlen_bits: int # vpu_vector_length_bits + elem_bits: int = 32 + + @property + def lane_elems(self) -> int: + return self.vlen_bits // self.elem_bits + + @property + def throughput(self) -> int: + return self.vlane * self.lane_elems + + +def _round_down(value: int, multiple: int) -> int: + return max(multiple, (value // multiple) * multiple) + + +def systolic_tile(m: int, n: int, k: int, hw: Hardware) -> tuple[int, int, int]: + col_unit = hw.sa * hw.vlane + tm = _round_down(min(m, 32), hw.sa) + tn = _round_down(min(n, col_unit * 4), col_unit) + tk = _round_down(min(k, 64), hw.sa) + return tm, tn, tk + + +def check_systolic(n: int, hw: Hardware) -> None: + col_unit = hw.sa * hw.vlane + if n < col_unit: + raise ValueError( + f"{n} columns is under one matmul instruction (SA*VLANE = {col_unit}); " + f"the array cannot be driven at this shape") + + +def reducing_tile(rows: int, cols: int, hw: Hardware) -> tuple[int, int]: + tn = min(cols, max(hw.throughput, hw.lane_elems * hw.vlane)) + tm = max(1, min(rows, max(1, (hw.throughput * 8) // max(1, tn)))) + return tm, tn + + +def pointwise_tile(rows: int, cols: int, hw: Hardware) -> tuple[int, int]: + budget = hw.throughput * 32 + tn = min(cols, budget) + tm = max(1, min(rows, max(1, budget // max(1, tn)))) + return tm, tn