From a354e901940b196abb97876ba908184b342b2de7 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:09:27 +0900 Subject: [PATCH 1/6] [Frontend] Rename print_operation to visit_operation It prints nothing. The method walks the IR and attaches TOG nodes; `bfs` and `display` do the printing. The name came from the C++ pass this file ports, where one method does both -- the docstring now records that so the correspondence is still findable. Co-Authored-By: Claude Fable 5 --- PyTorchSimFrontend/mlir/passes/build_tog.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index 5a40feec..fba9f5ea 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -568,7 +568,13 @@ def _process_dram_indices(self, value, loop_index_list, indirect_box): loop_index_list.append(("c" + str(c), c)) # ---- main recursion ---- - def print_operation(self, op, node): + def visit_operation(self, op, node): + """Walk `op` and attach the nodes it produces under `node`. + + Builds the graph; it does not print. (The C++ pass this is ported from + does both in one method, `printOperation` -- here `bfs`/`display` own + the printing.) + """ name = _op_name(op) if name in SKIP_OPS: return @@ -605,7 +611,7 @@ def bool_true(k): for region in oper.regions: for block in region.blocks: for inner in block.operations: - self.print_operation(inner, loop_node) + self.visit_operation(inner, loop_node) return if name == "togsim.transfer": @@ -1087,7 +1093,7 @@ def _build(module, builder): continue root = TOGNode("root") builder._reset_matmul_fsm() - builder.print_operation(op, root) + builder.visit_operation(op, root) root.bfs(out) return "".join(out) From 797afde1459bd9a0aab8937a2768f4fea64ced21 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:09:27 +0900 Subject: [PATCH 2/6] [Frontend] Give the trace pipeline a kernel whose body is one work-item A Triton kernel describes a single program instance: the tile loop is not in the kernel, it is the launch grid outside it. The trace producer already wants that same split -- togsim_kernel_tile per work-item, enumerated by togsim_kernel (docs/design/togsim_cpp_trace.md sec 9.3) -- so the two models agree and only the enumeration was missing. This teaches the pipeline to accept a kernel in that shape, rather than requiring the loop nest PyTorchSim's codegen emits. build_tog _find_kernel falls back to the module's only func.func; the name comes from the Triton kernel, not the fixed "kernel". _build roots at top-level loops carrying a ROLE attribute, and treats the whole body as one work-item when there are none. Keying on the attribute matters on its own: bank_vectorize also leaves a bare top-level affine.for for the tile's vector work, and rooting at that one made every DMA a sibling the traversal never reached. A DMA's tensor identity comes from the producer's dram_arg attribute when the operand is a view of the argument rather than the argument. Inferring it would mean chasing memref view ops back, and there are nine of them with no ViewLikeOpInterface in the python bindings to ask generically. DMA nodes are recorded in a list so _collect_dma_nodes can seed from them; it only descended from loop nodes, and a DMA outside any loop was dropped before reaching the skeleton. lower_to_emitc WorkItem + _materialize_grid_loop supply the grid, on the trace artifact only: the body is wrapped in a loop per axis, tagged outer_loop, with each program-id argument replaced by its induction variable. It must run before _rewrite_signature, which erases the arguments and first asserts none are still used. Everything after is unchanged -- _parallel_loop_chain finds the tagged loop, the outliner threads the induction variable through iv[], and the loop left behind becomes the dispatch enumeration. Two things about building that nest are easy to get wrong and only show at rank >= 2: a nested loop is created before the enclosing yield (an InsertionPoint on a block appends, and an scf.for body is already terminated), and every bound is created before the first loop, so that a bound made after an outer loop does not end up below it while an inner loop uses it. A parallel loop may be scf.for as well as affine.for; the role is carried by the attribute, not the dialect. _strip_aux keeps the kernel the caller resolved instead of matching on a name. Every change is a fallback: the existing conditions are tested first, so PyTorchSim's own codegen takes exactly the path it did before. Co-Authored-By: Claude Fable 5 --- .../mlir/passes/build_skeleton.py | 2 + PyTorchSimFrontend/mlir/passes/build_tog.py | 85 ++++++++-- .../mlir/passes/lower_to_emitc.py | 158 ++++++++++++++++-- 3 files changed, 218 insertions(+), 27 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/build_skeleton.py b/PyTorchSimFrontend/mlir/passes/build_skeleton.py index f4ed7d0d..ed52a56d 100644 --- a/PyTorchSimFrontend/mlir/passes/build_skeleton.py +++ b/PyTorchSimFrontend/mlir/passes/build_skeleton.py @@ -366,6 +366,8 @@ def visit(n): for ln in builder.loop_nodes: visit(ln) + for dn in getattr(builder, "dma_nodes", ()): # DMAs outside any tile loop + visit(dn) return by_op diff --git a/PyTorchSimFrontend/mlir/passes/build_tog.py b/PyTorchSimFrontend/mlir/passes/build_tog.py index fba9f5ea..98590d88 100644 --- a/PyTorchSimFrontend/mlir/passes/build_tog.py +++ b/PyTorchSimFrontend/mlir/passes/build_tog.py @@ -414,6 +414,9 @@ def __init__(self): self.loop_var_name = {} # value-identity-key -> loop name self.compute_nodes = [] self.loop_nodes = [] + # `_collect_dma_nodes` descends from the loop nodes, so a DMA hanging off + # the root (no tile loop in the kernel) would be missed. + self.dma_nodes = [] self._reset_matmul_fsm() # ---- matmul FSM ---- @@ -825,9 +828,14 @@ def _handle_dma_start(self, op, node): loop_idx_list.append(key) loop_stride_list.append(reordered[key]) - # base address + # base address: which tensor this DMA touches. The operand is the block + # argument itself in PyTorchSim's codegen; when it is a view of one + # instead, only the producer knows which -- so it says so (`dram_arg`) + # rather than the consumer guessing its way back through view ops. address = "arg" - if _is_block_arg(dram_memref): + if "dram_arg" in oper.attributes: + address += str(ir.IntegerAttr(oper.attributes["dram_arg"]).value) + elif _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) # element size @@ -881,6 +889,7 @@ def _handle_dma_start(self, op, node): tag_stride_list, loop_idx_list, loop_stride_list, indirect_box[0]) dma_node.op = op + self.dma_nodes.append(dma_node) node.add_child(dma_node) dma_node.add_parent(node) @@ -924,7 +933,9 @@ def _handle_dma_wait(self, op, node): dram_memref = f["dst"] elif dst_space == 1 and src_space == 0: dram_memref = f["src"] - if dram_memref is not None and _is_block_arg(dram_memref): + if "dram_arg" in user.attributes: + address += str(ir.IntegerAttr(user.attributes["dram_arg"]).value) + elif dram_memref is not None and _is_block_arg(dram_memref): address += str(ir.BlockArgument(dram_memref).arg_number) if len(tag_stride_list) == 0: @@ -934,6 +945,7 @@ def _handle_dma_wait(self, op, node): wait_node = TOGDMAWaitNode("DMAWaitNode", tag_index_list, tag_stride_list, tag_divider_list, address) wait_node.op = op + self.dma_nodes.append(wait_node) node.add_child(wait_node) wait_node.add_parent(node) @@ -1070,12 +1082,47 @@ def _insert_compute_markers(builder): # Driver. # --------------------------------------------------------------------------- def _find_kernel(module): - for op in module.body.operations: - if op.operation.name != "func.func": - continue + """The kernel function: named `kernel` in PyTorchSim's codegen, else the + module's only func.func (triton-npu carries the Triton kernel's own name). + Declines when there is more than one -- the intent would be a guess.""" + funcs = [op for op in module.body.operations + if op.operation.name == "func.func"] + for op in funcs: if ir.StringAttr(op.operation.attributes["sym_name"]).value == "kernel": return op - return None + return funcs[0] if len(funcs) == 1 else None + + +#: The loop roles (sec 9.1). Without one, a loop is a micro-loop the compute FSM +#: folds into a single node, not a tile loop. +_LOOP_ROLE_ATTRS = ("outer_loop", "accumulation_loop", "inner_loop") + + +def _has_loop_role(op): + attrs = op.operation.attributes + return any(k in attrs and ir.BoolAttr(attrs[k]).value for k in _LOOP_ROLE_ATTRS) + + +def _is_address_plumbing(op): + """Scalar index/integer math (DMA offsets, mask extents) and the terminator. + + Only consulted on the no-top-level-loop path. PyTorchSim's codegen puts this + math in `affine.apply`, which SKIP_OPS drops; triton-npu emits an + arith/index_cast chain that would otherwise count as vector compute. + + Keyed on result type: tile data here is always vector- or float-typed. A + top-level SCALAR arithmetic kernel would be misread, but no path emits one. + """ + name = _op_name(op) + if name in ("func.return", "memref.cast"): + return True + if not name.startswith("arith."): + return False + results = list(op.operation.results) + if not results: + return False + return all(ir.IndexType.isinstance(r.type) or ir.IntegerType.isinstance(r.type) + for r in results) def _build(module, builder): @@ -1088,13 +1135,29 @@ def _build(module, builder): block = func_op.regions[0].blocks[0] out = [] + # A root is a top-level TILE loop, identified by its role attribute (sec + # 9.1) -- not by being an affine.for: bank_vectorize leaves a bare one for + # the tile's vector work, and rooting there orphans every DMA. + roots = [op for op in block.operations + if op.operation.name == "affine.for" and _has_loop_role(op)] + if roots: + for op in roots: + root = TOGNode("root") + builder._reset_matmul_fsm() + builder.visit_operation(op, root) + root.bfs(out) + return "".join(out) + + # No top-level loop: the body is ONE work-item -- the shape a Triton kernel + # arrives in, its grid becoming the trace producer's dispatch loop (sec 9.3). + # PyTorchSim's codegen keeps the tile loops in the kernel and never lands here. + root = TOGNode("root") + builder._reset_matmul_fsm() for op in block.operations: - if op.operation.name != "affine.for": + if _is_address_plumbing(op): continue - root = TOGNode("root") - builder._reset_matmul_fsm() builder.visit_operation(op, root) - root.bfs(out) + root.bfs(out) return "".join(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index 5633769a..bfb94e34 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -119,20 +119,130 @@ def _attr_bool(op, key): # --------------------------------------------------------------------------- # step 1: rewrite signature + togsim.* ops (the unregistered-op glue) # --------------------------------------------------------------------------- -def _strip_aux(module): - """Erase memref.global decls and every func except @kernel (the wrapper).""" +def _strip_aux(module, keep=None): + """Erase memref.global decls and every func except the kernel. + + `keep` is the kernel op: its name is `kernel` only in PyTorchSim's codegen, + so the caller passes what `_find_kernel` resolved. + """ + keep_op = keep.operation if keep is not None else None victims = [] for op in module.body.operations: name = op.operation.name if name == "memref.global": victims.append(op) elif name == "func.func": - if ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": + if keep_op is not None: + if op.operation != keep_op: + victims.append(op) + elif ir.StringAttr(op.operation.attributes["sym_name"]).value != "kernel": victims.append(op) for op in victims: op.operation.erase() +class WorkItem: + """A kernel whose body is ONE work-item, plus the grid over it. + + A Triton kernel describes a single program instance; the grid lives outside + it. The trace producer already splits the same way (design sec 9.3), so only + the enumeration is missing. + + `parallel_args` are the argument positions holding the program ids + (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` + their extents. Both outermost-first. + """ + + def __init__(self, parallel_args, grid): + if len(parallel_args) != len(grid): + raise ValueError( + f"parallel_args {parallel_args} and grid {grid} must have the " + f"same length -- one program-id argument per grid axis") + self.parallel_args = list(parallel_args) + self.grid = [int(g) for g in grid] + + +def _materialize_grid_loop(kernel, work_item, ctx): + """Wrap the body in the grid loop the Triton kernel does not carry: + + func @k(..., %pid: i32) { + scf.for %p = 0 to G { index_cast %p> } {outer_loop} + } + + Downstream is then unchanged: `_parallel_loop_chain` finds the tagged loop, + the outliner threads its induction variable through `iv[]`, and the loop left + behind becomes the dispatch enumeration. `outer_loop` means "independent + work-item" (sec 9.1) -- exactly a Triton program id. + + MUST run before `_rewrite_signature`, which erases the arguments and first + asserts none are still used. + """ + from mlir.dialects import arith, scf + + block = kernel.regions[0].blocks[0] + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + + pid_args = [block.arguments[i] for i in work_item.parallel_args] + body_ops = [o for o in block.operations + if o.operation.name not in _LOOP_TERMINATORS] + terminator = [o for o in block.operations + if o.operation.name in _LOOP_TERMINATORS][0] + + # Every bound first, and all of them before the first loop: each is created + # just before the terminator, so one made after an outer loop would sit + # BELOW it in the block while an inner loop uses it -- which does not + # dominate, and the verifier rejects it (only reachable at rank >= 2). + with ir.InsertionPoint(terminator), loc: + c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result + c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e)).result + for e in work_item.grid] + + loops, inner = [], None + for ub in ubs: + # Nest inside the previous loop, BEFORE its yield: InsertionPoint on a + # block appends, and an scf.for body is already terminated. + ip = ir.InsertionPoint(terminator) if inner is None \ + else ir.InsertionPoint.at_block_terminator(inner.body) + with ip, loc: + loop = scf.ForOp(c0, ub, c1) + # ForOp leaves the body empty here; scf.for needs a terminator, and + # _outline_work_item inserts before it. + if len(loop.body.operations) == 0: + with ir.InsertionPoint(loop.body), loc: + scf.YieldOp([]) + loop.operation.attributes["outer_loop"] = ir.BoolAttr.get(True) + loops.append(loop) + inner = loop + + # Move the tile body inside the innermost loop, ahead of its yield. + inner_block = inner.body + inner_term = inner_block.operations[len(inner_block.operations) - 1] + for op in body_ops: + op.operation.move_before(inner_term) + + # Program ids are i32, induction variables index: cast once, at the top. + with ir.InsertionPoint(inner_block.operations[0]), loc: + casts = [] + for loop, pid in zip(loops, pid_args): + iv = loop.body.arguments[0] + casts.append(arith.IndexCastOp(pid.type, iv).result + if pid.type != idxty else iv) + + for pid, new in zip(pid_args, casts): + _replace_all_uses(pid, new) + + +def _replace_all_uses(old, new): + """The bindings expose no replaceAllUsesWith on a Value.""" + for use in list(old.uses): + owner = use.owner + for i in range(len(owner.operands)): + if owner.operands[i] == old: + owner.operands[i] = new + + def _rewrite_signature(kernel, ctx): """Replace @kernel's memref tensor args with the ABI args (EmitCtx*, int64_t* shape_args, int32_t n) and rename it to togsim_kernel. @@ -196,15 +306,22 @@ def _is_outer(forop): return "outer_loop" in a and ir.BoolAttr(a["outer_loop"]).value +#: The role is carried by the `outer_loop` attribute, not the dialect: +#: PyTorchSim's codegen emits affine.for, _materialize_grid_loop scf.for. Both +#: keep the induction variable in block argument 0. +_LOOP_OPS = ("affine.for", "scf.for") +_LOOP_TERMINATORS = ("affine.yield", "scf.yield", "func.return") + + def _parallel_loop_chain(block): - """The nested chain of `affine.for {outer_loop}` from `block` inward (one + """The nested chain of `{outer_loop}` loops from `block` inward (one work-item's parallel indices). Empty if the kernel has no parallel loop.""" chain = [] cur = block while True: nxt = None for op in cur.operations: - if op.operation.name == "affine.for" and _is_outer(op): + if op.operation.name in _LOOP_OPS and _is_outer(op): nxt = op break if nxt is None: @@ -281,7 +398,7 @@ def _outline_work_item(ctx, kernel, ctx_val): # move the work-item body into the tile fn (terminators stay behind). for op in [o for o in Lbody.operations - if o.operation.name not in ("affine.yield", "func.return")]: + if o.operation.name not in _LOOP_TERMINATORS]: op.operation.move_before(tret) # remap captures (Value `==` is identity): ctx -> ctx2, each parallel IV -> @@ -337,7 +454,7 @@ def _remap(block): # --- the dispatcher: marshal the IVs and hand the tile fn to togsim_dispatch --- term = [o for o in Lbody.operations - if o.operation.name in ("affine.yield", "func.return")][0] + if o.operation.name in _LOOP_TERMINATORS][0] fn_ref = _opaque(ctx, ts.TILE_SYMBOL) # function name -> verbatim pointer in C with ir.InsertionPoint(term): if ivs: @@ -499,15 +616,22 @@ def _add_extern_c(module, ctx): # --------------------------------------------------------------------------- # driver # --------------------------------------------------------------------------- -def lower_to_emitc(skeleton_module): +def lower_to_emitc(skeleton_module, work_item=None): """Lower a skeleton+API module (in place) to an EmitC module with the - `togsim_kernel` entry function. Returns the same module.""" + `togsim_kernel` entry function. Returns the same module. + + `work_item` is for kernels whose body is one work-item with the grid outside + (Triton's shape); None keeps PyTorchSim's, where the tile loops are already + in the kernel. + """ ctx = skeleton_module.context kernel = _find_kernel(skeleton_module) if kernel is None: - raise ValueError("no @kernel found in skeleton module") + raise ValueError("no kernel function found in skeleton module") - _strip_aux(skeleton_module) + _strip_aux(skeleton_module, keep=kernel) + if work_item is not None: + _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch @@ -563,18 +687,20 @@ def _default_include_dir(): return os.path.join(root, "TOGSim", "include") -def skeleton_to_so(skeleton_module, so_path, include_dir=None): +def skeleton_to_so(skeleton_module, so_path, include_dir=None, work_item=None): """skeleton module -> EmitC -> C++ -> compiled trace `.so`. Returns the EmitC module text (for inspection / caching).""" - emitc = lower_to_emitc(skeleton_module) + emitc = lower_to_emitc(skeleton_module, work_item=work_item) inc = include_dir or _default_include_dir() cpp = emitc_to_cpp(emitc, include_dir=inc) compile_so(cpp, so_path, inc) return str(emitc) -def build_trace_so(postvcix_path, so_path, include_dir=None): - """Full P2 path from a post-vcix kernel .mlir to a trace `.so`.""" +def build_trace_so(postvcix_path, so_path, include_dir=None, work_item=None): + """Full P2 path from a post-vcix kernel .mlir to a trace `.so`. + + `work_item` -- see lower_to_emitc.""" from . import build_skeleton as bs ctx = ir.Context() @@ -582,7 +708,7 @@ def build_trace_so(postvcix_path, so_path, include_dir=None): with ctx: module = ir.Module.parse(open(postvcix_path).read(), ctx) bs.build_skeleton(module) - return skeleton_to_so(module, so_path, include_dir) + return skeleton_to_so(module, so_path, include_dir, work_item=work_item) def main(argv): From 3d68dd04e11120f1f229f9d2159c1bf7ba1a2db2 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 3/6] [Frontend] Run the Triton route's kernel through TOGSim Closes the timing loop: one torch.compile now produces a cycle count. The launcher emits the trace producer from tnpu's post-vcix IR and hands it to TOGSim, reusing PyTorchSim's existing trace pipeline unchanged. timing.py emit_trace 04-custom.mlir -> build_skeleton -> trace.so + trace_cycles.tsv run_togsim hand the kernel directory to TOGSimulator.run_standalone work_item_for derives the program-id argument positions from the signature layout (pointers, user scalars, then triton-shared's own gridX,Y,Z / pidX,Y,Z) and the grid from the pinned block sizes codecache persists meta.json beside the artifacts so the timing step can run standalone, and TritonNPULauncher.__call__ simulates instead of raising. kernel_spec._grid becomes grid_of: the timing path needs the same extents to enumerate work-items, so it is computed in one place. The test drives a 2-D grid: that it verifies as MLIR, nests one loop per axis, and dispatches both indices. Checking the module and not only the C++ it becomes is the point -- the emitc lowering hoists constants to a flat scope, so it hides a bound that does not dominate its use. Inductor cannot reach this path here (it uses y/z only when x would overflow), so it gets a test rather than waiting for a kernel to exercise it. Measured on Inductor's `x + y` (1024 elements, XBLOCK 128, grid 8): TOGSim totals 573 cycles, and channel-0 DRAM traffic of 16 reads x 32 B x 16 channels is 8192 B -- exactly the 8 work-items x 2 loads x 512 B the kernel should move, so every dispatch in the enumeration really ran. The cycle table is a placeholder until gem5 sampling lands, and says so on every emit. Output tensors are NOT written; marshalling them through Spike is the remaining functional half. The launcher logs that on every call rather than letting an undefined value pass for a computed one, and the test asserts the timing artifacts exist instead of comparing values. Co-Authored-By: Claude Fable 5 --- .../triton_backend/codecache.py | 28 +++-- .../triton_backend/kernel_spec.py | 10 +- PyTorchSimFrontend/triton_backend/timing.py | 113 ++++++++++++++++++ tests/system/test_triton_codegen.py | 89 +++++++++++++- 4 files changed, 222 insertions(+), 18 deletions(-) create mode 100644 PyTorchSimFrontend/triton_backend/timing.py diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index b4817cc6..69805a87 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -18,7 +18,7 @@ from torch._inductor.codecache import get_hash from PyTorchSimFrontend import extension_config -from . import kernel_spec, tnpu_bridge +from . import kernel_spec, timing, tnpu_bridge logger = extension_config.setup_logger() @@ -43,17 +43,20 @@ def __init__(self, kernel_name, workdir, meta): self.elf = os.path.join(workdir, f"05-{kernel_name}.elf") def __call__(self, *args): - raise NotImplementedError( - f"{self.kernel_name}: compiled to {self.elf}, but the launch is not " - f"wired yet. Two pieces are missing and both are tracked in " - f"triton_backend/README.md:\n" - f" 1. functional -- marshal the caller's tensors into " - f"{self.workdir}/runtime/*.raw, run Spike on the ELF, read the " - f"outputs back into the caller's tensors;\n" - f" 2. timing -- emit trace.so + trace_cycles.tsv from the tnpu IR " - f"and hand them to TOGSim (needs the build_tog adapters).\n" - f"Compilation itself succeeded, so the codegen half of this route " - f"is exercised by getting this far.") + """One launch of the whole grid: simulate, return TOGSim's result. + + Does NOT write the caller's output tensors -- the functional launch is + not wired (README). Logged, so an undefined value cannot pass for a + computed one. + """ + if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): + timing.emit_trace(self.workdir, self.meta) + result = timing.run_togsim(self.workdir) + logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) + logger.warning( + "[Spike] %s: output tensors are NOT written; the functional launch " + "(tensors -> Spike -> tensors) is not wired yet", self.kernel_name) + return result def triton_npu_compile(src_code, meta, kernel_name): @@ -75,6 +78,7 @@ def triton_npu_compile(src_code, meta, kernel_name): tnpu_bridge.tnpu_dir()) with open(os.path.join(write_path, "kernel.py"), "w") as f: f.write(src_code) # the unmodified Inductor source, for diffing + timing.store_meta(write_path, meta) # lets the timing step run standalone tnpu_bridge.run_pipeline(spec_path, write_path, to_stage="binary") logger.info("[triton-npu] %s -> %s", kernel_name, write_path) return TritonNPULauncher(kernel_name, write_path, meta) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index ef779913..d74ea962 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -218,8 +218,12 @@ def strip_for_tnpu(src): return prefix + body -def _grid(meta): - """Sequential launch grid, from the numels and the pinned block sizes.""" +def grid_of(meta): + """Launch grid, from the numels and the pinned block sizes. + + Also read by the timing path, which needs the same extents to enumerate the + work-items -- so it lives here rather than being recomputed per consumer. + """ x = meta["numels"].get("xnumel") xblock = (meta.get("fixed_config") or {}).get("XBLOCK") if x is None or not xblock: @@ -330,7 +334,7 @@ def write_spec_file(src_code, meta, path, tnpu_dir): constexprs=constexprs, args_body=args_body, make_inputs_body=make_inputs_body, - grid=_grid(meta), + grid=grid_of(meta), ) with open(path, "w") as f: f.write(text) diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py new file mode 100644 index 00000000..9a045304 --- /dev/null +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -0,0 +1,113 @@ +"""The timing half of the Triton route: tnpu IR -> trace.so -> TOGSim. + +TOGSim simulates from a compiled trace producer (docs/design/togsim_cpp_trace.md). +PyTorchSim's codegen already emits one; this emits the same from a Triton-shaped +kernel, where the grid must be supplied -- see `lower_to_emitc.WorkItem`. + + emit_trace(workdir, meta) 04-custom.mlir -> trace.so + trace_cycles.tsv + run_togsim(workdir, ...) hand them to TOGSim, return its parsed result +""" + +import json +import os + +from PyTorchSimFrontend import extension_config + +logger = extension_config.setup_logger() + +#: Name TOGSim derives from the kernel directory (Simulator/simulator.py). +TRACE_SO = "trace.so" +CYCLE_TSV = "trace_cycles.tsv" +META_JSON = "meta.json" + +#: Stand-in per-tile cost until gem5 sampling lands. Deliberately not a +#: plausible-looking number: only an obvious non-measurement gets fixed. +PLACEHOLDER_CYCLE = 1 + + +def _runtime_arg_layout(meta): + """(n_tensor_args, n_scalar_args) of the lowered signature. + + triton-shared lays it out as pointers, user scalars, then its own six + (gridX,Y,Z / pidX,Y,Z). constexpr params never become arguments. + """ + sig = meta["signature"] + tensors = [k for k, v in sig.items() if v.startswith("*")] + scalars = [k for k, v in sig.items() + if not v.startswith("*") and v != "constexpr"] + return len(tensors), len(scalars) + + +def work_item_for(meta): + """The WorkItem describing this kernel's program-id args and grid extents.""" + from PyTorchSimFrontend.mlir.passes.lower_to_emitc import WorkItem + from . import kernel_spec + + n_tensor, n_scalar = _runtime_arg_layout(meta) + pid_x = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + grid = list(kernel_spec.grid_of(meta)) + return WorkItem(parallel_args=list(range(pid_x, pid_x + len(grid))), grid=grid) + + +def emit_trace(workdir, meta): + """Build `trace.so` + `trace_cycles.tsv` from tnpu's post-vcix IR. + + Returns the number of compute tiles the cycle table covers. + """ + from PyTorchSimFrontend.mlir.passes import build_skeleton as bs + from PyTorchSimFrontend.mlir.passes import cycle_table as ct + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + postvcix = os.path.join(workdir, "04-custom.mlir") + if not os.path.isfile(postvcix): + raise FileNotFoundError( + f"{postvcix} not found -- tnpu must have run at least to stage 4 " + f"(the post-vcix IR is what the trace is built from)") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(open(postvcix).read(), ctx) + bs.build_skeleton(module) + n_tiles = len(ct._compute_types(module)) + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), + work_item=work_item_for(meta)) + + # Until gem5 sampling lands, every tile costs PLACEHOLDER_CYCLE: TOGSim + # models the DMA and the dependency structure but NOT compute latency. + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) + logger.warning("[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d tiles); " + "compute latency is not modelled yet", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + return n_tiles + + +def run_togsim(workdir, attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict.""" + from Simulator.simulator import TOGSimulator + + so = os.path.join(workdir, TRACE_SO) + if not os.path.isfile(so): + raise FileNotFoundError(f"{so} not found -- call emit_trace first") + + # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its + # DIRECTORY, and reads the file itself only on the STONNE path. + handle = os.path.join(workdir, "tile_graph.onnx") + result_path = TOGSimulator.run_standalone( + handle, attribute_path or os.path.join(workdir, "attribute"), + timeout_sec=timeout_sec) + return TOGSimulator.get_result_from_file(result_path) + + +def store_meta(workdir, meta): + """Persist codegen metadata beside the artifacts, so the timing step can run + standalone.""" + with open(os.path.join(workdir, META_JSON), "w") as f: + json.dump(meta, f, indent=2) + + +def load_meta(workdir): + with open(os.path.join(workdir, META_JSON)) as f: + return json.load(f) diff --git a/tests/system/test_triton_codegen.py b/tests/system/test_triton_codegen.py index a7b0af79..b36715d6 100644 --- a/tests/system/test_triton_codegen.py +++ b/tests/system/test_triton_codegen.py @@ -31,10 +31,73 @@ def fn(x, y): return fn, x, y +def check_multi_axis_grid(): + """A 2-D grid must nest one loop per axis and hand both indices to iv[]. + + Guards the multi-axis path, which the add kernel does not reach: Inductor + only uses y/z when x would overflow, so a 1-D grid exercises just the first + iteration of the nest. + """ + from PyTorchSimFrontend.mlir.passes import lower_to_emitc as l2e + from PyTorchSimFrontend.mlir.passes.build_tog import ir + + src = """ + module { + func.func @k(%arg0: memref<*xf32>, %arg1: i32, %arg2: i32) { + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : i32 + %a = arith.muli %arg1, %c8 : i32 + %b = arith.addi %a, %arg2 : i32 + %o = arith.index_cast %b : i32 to index + "togsim.dma"(%o, %c0) {arg_id = 0 : i32, base = "arg0", dims = [128], + dir = 0 : i32, elem_bits = 32 : i32, is_async = false, read_bufs = [], + strides = [1], tag_id = 0 : i32, write_bufs = [0]} : (index, index) -> () + return + } + } + """ + problems = [] + # Verify the IR the pass itself produces: a bound created after an outer loop + # would not dominate an inner loop's use of it, which only shows at rank >= 2 + # and which the emitc lowering happens to paper over. + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + l2e._materialize_grid_loop( + l2e._find_kernel(module), + l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3]), ctx) + try: + module.operation.verify() + except Exception as e: # noqa: BLE001 + problems.append(f"materialized IR does not verify: {e}") + + ctx = ir.Context() + ctx.allow_unregistered_dialects = True + with ctx: + module = ir.Module.parse(src, ctx) + emitc = l2e.lower_to_emitc( + module, work_item=l2e.WorkItem(parallel_args=[1, 2], grid=[4, 3])) + cpp = l2e.emitc_to_cpp(emitc, include_dir=l2e._default_include_dir()) + + entry = cpp.split("togsim_kernel(EmitCtx*")[-1] + if entry.count("for (") != 2: + problems.append(f"expected 2 nested loops, found {entry.count('for (')}") + if "togsim_dispatch" not in entry: + problems.append("no togsim_dispatch call") + if ", 2);" not in entry: + problems.append("dispatch does not pass 2 indices") + for p in problems: + print(f" multi-axis grid: {p}") + return not problems + + def main(): from PyTorchSimFrontend import extension_config from PyTorchSimFrontend.triton_backend import tnpu_bridge + print(f"multi-axis grid = " + f"{'ok' if check_multi_axis_grid() else 'FAILED'}") print(f"TORCHSIM_TRITON_CODEGEN = {extension_config.CONFIG_TRITON_CODEGEN}") print(f"TNPU_DIR = {extension_config.CONFIG_TNPU_DIR}") ok, _out = tnpu_bridge.doctor() @@ -56,9 +119,29 @@ def main(): "above and README.md's gap list.") return 1 - err = (got.cpu() - expected).abs().max().item() - print(f"max_abs_err = {err}") - return 0 if err < 1e-4 else 1 + # Values are NOT checked: the launch simulates the kernel but does not + # marshal tensors through Spike, so `got` is undefined. What is asserted is + # that the timing path ran end to end -- the two artifacts TOGSim consumes. + del got + import glob + + from PyTorchSimFrontend.triton_backend import timing + + dirs = glob.glob(os.path.join(extension_config.get_dump_path(), "triton_*")) + if not dirs: + print("no kernel directory was produced") + return 1 + workdir = max(dirs, key=os.path.getmtime) + for name in (timing.TRACE_SO, timing.CYCLE_TSV): + path = os.path.join(workdir, name) + if not os.path.isfile(path): + print(f"missing {name} in {workdir}") + return 1 + print(f" {name:18s} {os.path.getsize(path)} bytes") + print(f"\ntiming path OK ({workdir})") + print(f"values NOT verified -- torch would give {expected[:3].tolist()}...; " + f"the functional launch is not wired (triton_backend/README.md)") + return 0 if __name__ == "__main__": From 6f6d6d8fe116aba2a8a9eed04b77a8427fe689d6 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 4/6] [Frontend] Measure the Triton route's tile cost with gem5 The cycle table held a placeholder, so TOGSim modelled the DMA but charged nothing for compute. It is now a measurement. measure_tile_cycles chains three pieces that already existed: build_tog's sample mode annotates the post-vcix IR (inline-asm markers around each compute node, every loop rewritten to one trip, so what runs is one tile), `python -m tnpu.cycle` lowers that to a RISC-V binary in tnpu's own process -- the Gemmini/VCIX lowering and its LLVM live there -- and CycleSimulator runs it under gem5, reading one numCycles per marker pair. build_cycle_table then turns the list into the tsv, keyed by tile_id. Sampling runs before build_skeleton because both consume the same post-vcix IR and build_skeleton rewrites it in place. Failure is not fatal: any step that does not produce a measurement falls back to the placeholder table and says so, since a kernel that simulates with the wrong compute cost is more useful than one that will not simulate -- as long as it announces which it is. Measured on `x + y` (1024 elements, grid 8): the tile samples at 19 cycles and TOGSim's total moves 573 -> 650. The MLIR route reports 251 on the same computation; the remaining gap is double buffering, which tnpu does not emit yet, so nothing overlaps. Co-Authored-By: Claude Fable 5 --- PyTorchSimFrontend/triton_backend/README.md | 47 ++++++++++-- PyTorchSimFrontend/triton_backend/timing.py | 83 ++++++++++++++++++--- 2 files changed, 113 insertions(+), 17 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/README.md b/PyTorchSimFrontend/triton_backend/README.md index 12b01b9d..4b409158 100644 --- a/PyTorchSimFrontend/triton_backend/README.md +++ b/PyTorchSimFrontend/triton_backend/README.md @@ -35,11 +35,19 @@ torch.compile │ a tnpu kernel file (KernelSpec) kernel_spec.py ▼ run.py --to binary (subprocess) tnpu_bridge.py - │ 01-ttir → 02-ttshared → 03-adapted → 04-lowered → 05-*.elf + │ 01-ttir → 02-ttshared → 03-adapted → 04-custom → 05-*.elf ▼ - TritonNPULauncher.__call__ ← NOT WIRED YET + TritonNPULauncher.__call__ codecache.py + ├ timing 04-custom.mlir → trace.so + trace_cycles.tsv → TOGSim + │ cycles measured by gem5 on a one-tile binary + └ functional ← NOT WIRED YET ``` +The timing half reuses PyTorchSim's trace pipeline unchanged. The one structural +difference is that a Triton kernel body is a single program instance, so the grid +that enumerates instances is supplied by `lower_to_emitc.WorkItem` instead of +being read out of the kernel -- see "The grid is not in the kernel" below. + Artifacts land in one directory per source hash under the dump path (`outputs/triton_/`), alongside the unmodified Inductor source (`kernel.py`) so the rewrite is diffable. @@ -54,7 +62,18 @@ Artifacts land in one directory per source hash under the dump path - the lowering is correct in shape: `tl.load/store` became three `togsim.transfer` ops, and Inductor's `xmask` came through as a **masked DMA** (`masked_axes = [0]`, `masked_fill`), which tnpu already supports -- the run stops in `TritonNPULauncher.__call__`, by design +- the trace producer comes out in the shape the design calls for: a + `togsim_kernel_tile` computing `offset = iv[0]*128` around three `togsim_dma` + and one `togsim_compute`, and a `togsim_kernel` looping `p < 8` over + `togsim_dispatch` +- **TOGSim runs it: 650 cycles**, with channel-0 DRAM traffic of 16 reads x 32 B + x 16 channels = 8192 B, exactly the 8 work-items x 2 loads x 512 B the kernel + should move. The MLIR route on the same `x + y` reports 251 cycles -- the same + order, and higher here because tnpu emits synchronous DMA, so nothing overlaps + (gap 2) +- the tile's compute cost is a real measurement: gem5 samples **19 cycles** for + the vector-add tile, via `timing.measure_tile_cycles` +- values are NOT produced: the functional launch is still open (gap 1) ## Gap list, in order @@ -62,17 +81,18 @@ Artifacts land in one directory per source hash under the dump path `runtime/*.raw`, run Spike on the ELF, read outputs back. tnpu's stage 6 does this for its own kernels but generates inputs from the spec; here the tensors come from the caller. -2. **Launch (timing).** Emit `trace.so` + `trace_cycles.tsv` and hand them to - TOGSim. Blocked on the `build_tog` adapters — the tnpu IR is structurally - invisible to it today (no top-level `affine.for`, `scf.for` instead of - `affine.for`, vcix as LLVM intrinsics rather than dialect ops, DMA addresses - as `arith` chains rather than `affine.apply`, grid outside the IR). +2. **Double buffering.** tnpu emits synchronous DMA (`is_async=false`, no + `togsim.wait`), so load → compute → store serialize inside every work-item and + TOGSim has no overlap to model. This is the main remaining gap between the two + routes' cycle counts. 3. **`triton_helpers`.** Any kernel using `triton_helpers.*` (reductions, clamps, `maximum`/`minimum`) cannot compile: the module lives in torch and the tnpu venv has none. `strip_for_tnpu` raises and names the helper. Needs a minimal vendored copy. 4. **Reductions.** Independently blocked in tnpu itself — no lane-aware reduction path; see `triton-npu/kernels/reduce.py`. + Matmul is also still open on the timing side: `build_tog` finds compute nodes + by the `vcix.iv` op name, and tnpu emits `llvm.riscv.sf.vc.*` intrinsics. 5. **Block-size policy.** `fixed_config_for` pins `XBLOCK` to the lane count and deliberately leaves reduction blocks unset. Real tile selection (the MLIR route's autotuner / `codegen_mapping_strategy`) has no equivalent here yet. @@ -103,6 +123,17 @@ GPU, `triton_hash_with_backend()` raises "0 active drivers" because it asks the triton runtime for the current target. We never launch through that runtime, so the value is short-circuited to a deterministic cache key. +**The grid is not in the kernel.** PyTorchSim's codegen puts the tile loops +inside the kernel; a Triton kernel describes one program instance and leaves the +grid to the launch. The trace producer wants that same split already -- +`togsim_kernel_tile` per work-item, enumerated by `togsim_kernel` (design sec +9.3) -- so the models agree and only the enumeration was missing. +`_materialize_grid_loop` supplies it, on the trace artifact only: it wraps the +body in a loop tagged `outer_loop` with each program-id argument replaced by the +induction variable, and everything downstream is unchanged. It runs before +`_rewrite_signature`, which erases the kernel arguments and first asserts none +are still used -- that ordering is what decides where this can live. + ## CI `.github/workflows/triton_npu.yml`, separate from the main CI: this route is WIP, diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index 9a045304..f19423db 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -20,10 +20,56 @@ CYCLE_TSV = "trace_cycles.tsv" META_JSON = "meta.json" -#: Stand-in per-tile cost until gem5 sampling lands. Deliberately not a -#: plausible-looking number: only an obvious non-measurement gets fixed. +#: Used only when gem5 sampling fails. Deliberately not a plausible-looking +#: number: only an obvious non-measurement gets fixed. PLACEHOLDER_CYCLE = 1 +SAMPLE_MLIR = "04-sample.mlir" +CYCLE_BIN = "cycle_bin" + + +def measure_tile_cycles(workdir, meta): + """Per-compute-node cycle counts for ONE tile, measured under gem5. + + build_tog's sample mode marks each compute node and makes every loop a + single trip; tnpu lowers that to a binary (in ITS process -- the Gemmini/VCIX + lowering and its LLVM live there); gem5 runs it. Returns None on any failure, + and the caller falls back to the placeholder table. + """ + from PyTorchSimFrontend.mlir.passes.build_tog import run_tog + + kernel_name = meta["kernel_name"] + spec = os.path.join(workdir, f"{kernel_name}_spec.py") + if not os.path.isfile(spec): + logger.warning("[Gem5] %s not found; cannot sample cycles", spec) + return None + + run_tog(os.path.join(workdir, "04-custom.mlir"), + os.path.join(workdir, "tog_sample.py"), + os.path.join(workdir, SAMPLE_MLIR), sample_mode=True) + + import subprocess + + from . import tnpu_bridge + env = dict(os.environ) + env.pop("PYTHONPATH", None) # keep tnpu on its own MLIR bindings + proc = subprocess.run( + [extension_config.CONFIG_TNPU_PYTHON, "-m", "tnpu.cycle", spec, workdir], + capture_output=True, text=True, cwd=tnpu_bridge.tnpu_dir(), env=env) + if proc.returncode != 0: + logger.warning("[Gem5] cycle binary build failed:\n%s", + (proc.stdout + proc.stderr)[-2000:]) + return None + + from Simulator.simulator import CycleSimulator + try: + return CycleSimulator().compile_and_simulate( + os.path.join(workdir, CYCLE_BIN), int(extension_config.vpu_num_lanes), + silent_mode=True) + except Exception as e: # noqa: BLE001 - fall back to the placeholder table + logger.warning("[Gem5] sampling failed: %s", e) + return None + def _runtime_arg_layout(meta): """(n_tensor_args, n_scalar_args) of the lowered signature. @@ -65,22 +111,41 @@ def emit_trace(workdir, meta): f"{postvcix} not found -- tnpu must have run at least to stage 4 " f"(the post-vcix IR is what the trace is built from)") + # Before build_skeleton: both read the post-vcix IR, which it rewrites in place. + cycles = measure_tile_cycles(workdir, meta) + ctx = ir.Context() ctx.allow_unregistered_dialects = True with ctx: module = ir.Module.parse(open(postvcix).read(), ctx) bs.build_skeleton(module) - n_tiles = len(ct._compute_types(module)) + compute_types = ct._compute_types(module) + n_tiles = len(compute_types) + + if cycles: + # One numCycles per compute node; pad/truncate as the MLIR route does. + cl = list(cycles) + if len(cl) != n_tiles: + logger.warning("[Gem5] returned %d cycle(s) for %d " + "tile(s); padding with the last", len(cl), n_tiles) + cl = (cl + [cl[-1]] * n_tiles)[:n_tiles] + # Systolic-array fill; only matmul tiles use it. + lanes = int(extension_config.vpu_num_lanes) + table = ct.build_cycle_table(module, cl, x_offset=lanes, w_offset=0) + else: + table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles + logger.warning( + "[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d " + "tiles): gem5 sampling did not produce a measurement, so " + "compute latency is NOT modelled", + CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + l2e.skeleton_to_so(module, os.path.join(workdir, TRACE_SO), work_item=work_item_for(meta)) - # Until gem5 sampling lands, every tile costs PLACEHOLDER_CYCLE: TOGSim - # models the DMA and the dependency structure but NOT compute latency. - table = [(PLACEHOLDER_CYCLE, 0)] * n_tiles ct.dump_cycle_table_tsv(table, os.path.join(workdir, CYCLE_TSV)) - logger.warning("[Gem5] %s holds PLACEHOLDER cycles (%d per tile x %d tiles); " - "compute latency is not modelled yet", - CYCLE_TSV, PLACEHOLDER_CYCLE, n_tiles) + if cycles: + logger.info("[Gem5] tile cycles: %s", table) return n_tiles From c2e004f90e05687e87c1eece881b712ca2a0f20f Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:10:01 +0900 Subject: [PATCH 5/6] [Frontend] Let the route ask for a multi-axis grid The grid machinery handles N axes and is tested to, but grid_of only ever computed one from xnumel, so work_item_for could only ever build a 1-D WorkItem -- general plumbing behind a caller that never used it. grid_of now walks every parallel axis, and fixed_config_for pins a block per axis. Parallel vs reduction is Inductor's own test (a prefix starting with "r" is looped inside the kernel, not gridded), and the block name is f"{prefix.upper()}BLOCK", so neither is guessed. Two orderings meet here and they are not the same: grid_of returns axes OUTERMOST first (z, y, x -- x is Inductor's contiguous axis), while triton-shared always appends the program ids as pidX, pidY, pidZ. work_item_for therefore builds the argument list per axis instead of as a range; zipping the two blindly would pair the outermost loop with the wrong id. Block sizes: the outermost axis gets the lane count, because that is the tile dimension bank_vectorize spreads over the lanes. The rest get 1, which leaves the tile exactly the verified [lanes] shape and lets the grid cover everything else. That is correct but pathological -- an inner block of 1 makes each work-item move a strided column -- so a multi-axis kernel logs a warning saying it is not a tiling worth measuring. Choosing real tile sizes is the block-size policy gap in the README. Verified: axes/grid/parallel_args come out ['x']/(8,)/[pidX], ['y','x']/(2,1024)/[pidY,pidX] and ['z','y','x']/[pidZ,pidY,pidX]. End-to-end on a real multi-axis kernel is still untested -- Inductor reaches for y/z only when x would overflow, which the shapes this route handles do not. Co-Authored-By: Claude Fable 5 --- .../triton_backend/kernel_spec.py | 68 +++++++++++++++---- PyTorchSimFrontend/triton_backend/timing.py | 18 ++++- 2 files changed, 68 insertions(+), 18 deletions(-) diff --git a/PyTorchSimFrontend/triton_backend/kernel_spec.py b/PyTorchSimFrontend/triton_backend/kernel_spec.py index d74ea962..fe147286 100644 --- a/PyTorchSimFrontend/triton_backend/kernel_spec.py +++ b/PyTorchSimFrontend/triton_backend/kernel_spec.py @@ -136,21 +136,50 @@ def collect_meta(kernel, kernel_name): } +#: Parallel iteration prefixes, OUTERMOST first. Inductor's `x` is the +#: contiguous axis, so it is innermost; `r*` prefixes are reductions, looped +#: inside the kernel rather than spread over the grid (prefix_is_reduction). +_PARALLEL_PREFIXES = ("z", "y", "x") + + +def _block_name(prefix): + return f"{prefix.upper()}BLOCK" + + +def parallel_axes(numels): + """Grid axes this kernel uses, outermost first.""" + return [p for p in _PARALLEL_PREFIXES if f"{p}numel" in numels] + + def fixed_config_for(kernel): """Block sizes pinned at codegen time. tnpu compiles ONE binary per kernel and the C wrapper walks the grid as a - sequential loop, so there is no autotuner to choose XBLOCK later and no - runtime `grid=` callable. Fixing it here is what makes the launch shape + sequential loop, so there is no autotuner to choose the blocks later and no + runtime `grid=` callable. Fixing them here is what makes the launch shape static. - The lane count is the natural default: `bank_vectorize` distributes tile - dim 0 across the lanes, and a block equal to the lane count gives a per-lane - depth of 1 -- the case every tnpu baseline runs today. + Tile dim 0 is the one `bank_vectorize` spreads over the lanes, so the + OUTERMOST axis gets the lane count -- a per-lane depth of 1, the shape every + tnpu baseline runs. The remaining axes get 1, which leaves the tile exactly + that verified shape and lets the grid cover the rest. It is conservative + rather than fast; choosing real tile sizes is the block-size policy gap in + README, not something to guess at here. """ from PyTorchSimFrontend import extension_config lanes = int(extension_config.vpu_num_lanes) - cfg = {"XBLOCK": lanes} + + axes = parallel_axes(getattr(kernel, "numels", None) or {}) + cfg = {_block_name(p): (lanes if i == 0 else 1) for i, p in enumerate(axes)} + if len(axes) > 1: + # Loud, because the shape is correct but pathological: an inner block of + # 1 makes every work-item move a strided column. Fine for getting a + # multi-axis kernel through the route, misleading to benchmark. + extension_config.setup_logger().warning( + "[triton-npu] %s tiles over %s; inner blocks pinned to 1, which is " + "correct but not a tiling worth measuring", + getattr(kernel, "kernel_name", "kernel"), axes) + cfg.setdefault("XBLOCK", lanes) # a kernel with no tiling info still has x if getattr(kernel, "inside_reduction", False): # A reduction block is NOT free to be the lane count: the reduced axis # has to stay inside a lane (see triton-npu kernels/reduce.py). Left @@ -219,20 +248,29 @@ def strip_for_tnpu(src): def grid_of(meta): - """Launch grid, from the numels and the pinned block sizes. + """Launch grid, from the numels and the pinned block sizes, outermost first. Also read by the timing path, which needs the same extents to enumerate the work-items -- so it lives here rather than being recomputed per consumer. """ - x = meta["numels"].get("xnumel") - xblock = (meta.get("fixed_config") or {}).get("XBLOCK") - if x is None or not xblock: + numels = meta["numels"] + cfg = meta.get("fixed_config") or {} + axes = parallel_axes(numels) + if not axes: raise SpecIncomplete( - f"cannot compute the grid for {meta['kernel_name']}: " - f"xnumel={x!r}, XBLOCK={xblock!r}. Inductor defers the grid to " - f"triton_heuristics at runtime; this route needs it statically " - f"(see fixed_config_for).") - return (int(math.ceil(x / xblock)),) + f"{meta['kernel_name']} has no parallel iteration axis to grid over") + + grid = [] + for prefix in axes: + n, block = numels.get(f"{prefix}numel"), cfg.get(_block_name(prefix)) + if n is None or not block: + raise SpecIncomplete( + f"cannot compute the grid for {meta['kernel_name']} axis " + f"'{prefix}': {prefix}numel={n!r}, {_block_name(prefix)}={block!r}. " + f"Inductor defers the grid to triton_heuristics at runtime; this " + f"route needs it statically (see fixed_config_for).") + grid.append(int(math.ceil(n / block))) + return tuple(grid) SPEC_TEMPLATE = '''\ diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index f19423db..1785e519 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -84,15 +84,27 @@ def _runtime_arg_layout(meta): return len(tensors), len(scalars) +#: triton-shared appends pidX, pidY, pidZ in that order, whatever the tiling is. +_PID_SLOT = {"x": 0, "y": 1, "z": 2} + + def work_item_for(meta): - """The WorkItem describing this kernel's program-id args and grid extents.""" + """The WorkItem describing this kernel's program-id args and grid extents. + + `grid_of` orders axes OUTERMOST first (z, y, x -- x is Inductor's contiguous + one), while the program-id arguments are always laid out x, y, z. The two + are zipped downstream, so the argument list is built per axis rather than as + a range. + """ from PyTorchSimFrontend.mlir.passes.lower_to_emitc import WorkItem from . import kernel_spec n_tensor, n_scalar = _runtime_arg_layout(meta) - pid_x = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + pid_base = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ + axes = kernel_spec.parallel_axes(meta["numels"]) grid = list(kernel_spec.grid_of(meta)) - return WorkItem(parallel_args=list(range(pid_x, pid_x + len(grid))), grid=grid) + return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], + grid=grid) def emit_trace(workdir, meta): From 5fe065fbfd6606af81b71469c27a527398feabea Mon Sep 17 00:00:00 2001 From: root Date: Mon, 27 Jul 2026 22:27:26 +0900 Subject: [PATCH 6/6] [Frontend] Take the grid at run time, so one trace serves every shape Only the NUMBER of grid axes has to be compiled in -- how many loops to nest and how many iv[] slots to fill. The trip counts are just values, and the producer ABI already carries them: togsim_kernel(ctx, shape_args, n). Baking them in was what forced a recompile per shape. A WorkItem extent of None now means "read it from shape_args". The bound cannot be wired when the loop is built -- shape_args does not exist until _rewrite_signature adds it -- so the loop takes a placeholder and _bind_runtime_bounds replaces it once the signature is there. The loops stay in the entry function (the outliner moves only their bodies), so the read is in scope where the bound is used. timing.write_shape computes the grid per launch and writes trace_shape.txt. The launch already knows the real extents: Inductor appends the numels after the tensor arguments, so the trailing values are them. Only the PARALLEL numels ride along -- a reduction axis is looped inside the kernel and never passed, so counting it would misalign the mapping. TOGSim reads that sidecar the same way it reads trace_cycles.tsv, from the kernel directory. main.cc passed nullptr for shape_args; absent file still means nullptr, so a producer with its bounds baked in is unaffected. Measured. One trace.so (md5 identical across all three), torch.compile( dynamic=True), a single kernel directory reused: n=1024 grid 8 650 cycles n=2048 grid 16 1316 cycles n=4096 grid 32 2586 cycles The functional path still compiles per shape -- tnpu's spec bakes the tensor extents into the flat memref view -- so this opens the timing half only. Co-Authored-By: Claude Fable 5 --- .../mlir/passes/lower_to_emitc.py | 52 +++++++++++++++++-- .../triton_backend/codecache.py | 2 +- PyTorchSimFrontend/triton_backend/timing.py | 52 +++++++++++++++++-- TOGSim/src/main.cc | 13 ++++- 4 files changed, 110 insertions(+), 9 deletions(-) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py index bfb94e34..537c8ad0 100644 --- a/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py +++ b/PyTorchSimFrontend/mlir/passes/lower_to_emitc.py @@ -151,6 +151,13 @@ class WorkItem: `parallel_args` are the argument positions holding the program ids (triton-shared appends gridX,Y,Z / pidX,Y,Z after the user scalars); `grid` their extents. Both outermost-first. + + An extent may be None, meaning "read it from shape_args at run time". Only + the NUMBER of axes has to be known when the kernel is compiled -- how many + loops to nest and how many iv[] slots to fill; the trip counts are just + values, and the producer ABI already takes them + (togsim_kernel(ctx, shape_args, n)). That is what lets one compiled trace + serve every shape. """ def __init__(self, parallel_args, grid): @@ -159,7 +166,12 @@ def __init__(self, parallel_args, grid): f"parallel_args {parallel_args} and grid {grid} must have the " f"same length -- one program-id argument per grid axis") self.parallel_args = list(parallel_args) - self.grid = [int(g) for g in grid] + self.grid = [None if g is None else int(g) for g in grid] + + @property + def dynamic_axes(self): + """Indices into `grid` whose extent arrives at run time.""" + return [i for i, g in enumerate(self.grid) if g is None] def _materialize_grid_loop(kernel, work_item, ctx): @@ -196,7 +208,10 @@ def _materialize_grid_loop(kernel, work_item, ctx): with ir.InsertionPoint(terminator), loc: c0 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 0)).result c1 = arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, 1)).result - ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e)).result + # A runtime extent still needs SOMETHING here: shape_args does not exist + # until _rewrite_signature adds it. The placeholder is replaced by + # _bind_runtime_bounds once it does. + ubs = [arith.ConstantOp(idxty, ir.IntegerAttr.get(idxty, e or 1)).result for e in work_item.grid] loops, inner = [], None @@ -233,6 +248,35 @@ def _materialize_grid_loop(kernel, work_item, ctx): for pid, new in zip(pid_args, casts): _replace_all_uses(pid, new) + return [(loops[i], ubs[i]) for i in work_item.dynamic_axes] + + +def _bind_runtime_bounds(pending, shape_arg, ctx): + """Point each runtime loop bound at `shape_args[k]`. + + Runs AFTER _rewrite_signature, which is what creates the shape_args + argument. The loops stay in the entry function (the outliner moves only + their bodies), so the read is in scope where the bound is used. + """ + if not pending: + return + from mlir.dialects import arith + + i64 = ir.IntegerType.get_signless(64) + idxty = ir.IndexType.get() + loc = ir.Location.unknown(ctx) + for k, (loop, placeholder) in enumerate(pending): + with ir.InsertionPoint(placeholder.owner), loc: + kc = ir.Operation.create( + "emitc.constant", results=[i64], + attributes={"value": ir.IntegerAttr.get(i64, k)}).results[0] + elem = ir.Operation.create( + "emitc.subscript", results=[i64], + operands=[shape_arg, kc]).results[0] + bound = arith.IndexCastOp(idxty, elem).result + _replace_all_uses(placeholder, bound) + placeholder.owner.erase() + def _replace_all_uses(old, new): """The bindings expose no replaceAllUsesWith on a Value.""" @@ -630,9 +674,11 @@ def lower_to_emitc(skeleton_module, work_item=None): raise ValueError("no kernel function found in skeleton module") _strip_aux(skeleton_module, keep=kernel) + pending = [] if work_item is not None: - _materialize_grid_loop(kernel, work_item, ctx) + pending = _materialize_grid_loop(kernel, work_item, ctx) ctx_val = _rewrite_signature(kernel, ctx) + _bind_runtime_bounds(pending, kernel.regions[0].blocks[0].arguments[1], ctx) _rewrite_togsim_ops(ctx, kernel, ctx_val) # togsim.* -> emitc.call_opaque _outline_work_item(ctx, kernel, ctx_val) # work-item body -> togsim_kernel_tile + dispatch diff --git a/PyTorchSimFrontend/triton_backend/codecache.py b/PyTorchSimFrontend/triton_backend/codecache.py index 69805a87..76dc264c 100644 --- a/PyTorchSimFrontend/triton_backend/codecache.py +++ b/PyTorchSimFrontend/triton_backend/codecache.py @@ -51,7 +51,7 @@ def __call__(self, *args): """ if not os.path.isfile(os.path.join(self.workdir, timing.TRACE_SO)): timing.emit_trace(self.workdir, self.meta) - result = timing.run_togsim(self.workdir) + result = timing.run_togsim(self.workdir, meta=self.meta, args=args) logger.info("[TOGSim] %s simulated -> %s", self.kernel_name, result) logger.warning( "[Spike] %s: output tensors are NOT written; the functional launch " diff --git a/PyTorchSimFrontend/triton_backend/timing.py b/PyTorchSimFrontend/triton_backend/timing.py index 1785e519..be28060b 100644 --- a/PyTorchSimFrontend/triton_backend/timing.py +++ b/PyTorchSimFrontend/triton_backend/timing.py @@ -18,6 +18,7 @@ #: Name TOGSim derives from the kernel directory (Simulator/simulator.py). TRACE_SO = "trace.so" CYCLE_TSV = "trace_cycles.tsv" +SHAPE_TXT = "trace_shape.txt" META_JSON = "meta.json" #: Used only when gem5 sampling fails. Deliberately not a plausible-looking @@ -102,9 +103,46 @@ def work_item_for(meta): n_tensor, n_scalar = _runtime_arg_layout(meta) pid_base = n_tensor + n_scalar + 3 # after gridX, gridY, gridZ axes = kernel_spec.parallel_axes(meta["numels"]) - grid = list(kernel_spec.grid_of(meta)) + # Extents are left to run time: only the axis COUNT has to be compiled in, + # and the launch knows the real numels. One trace then serves every shape. return WorkItem(parallel_args=[pid_base + _PID_SLOT[p] for p in axes], - grid=grid) + grid=[None] * len(axes)) + + +def write_shape(workdir, meta, args=()): + """Write the grid extents the trace producer reads as shape_args. + + `args` is the launch's positional arguments; Inductor appends the numels + after the tensors, so the trailing values are them, in `meta["numels"]` + order. Falls back to the compile-time hint when they are absent. + """ + from . import kernel_spec + + numels = dict(meta["numels"]) + # Only the PARALLEL numels ride along on the call -- a reduction axis is + # looped inside the kernel, so it is not passed and must not consume one of + # the trailing values. They arrive in kernel order, which is the dict's. + passed = [k for k in numels if not k.startswith("r")] + trailing = [a for a in args if isinstance(a, int) and not isinstance(a, bool)] + if passed and len(trailing) >= len(passed): + for key, val in zip(passed, trailing[-len(passed):]): + numels[key] = val + + axes = kernel_spec.parallel_axes(numels) + + cfg = meta.get("fixed_config") or {} + grid = [] + for p in axes: + n, block = numels.get(f"{p}numel"), cfg.get(f"{p.upper()}BLOCK") + if n is None or not block: + raise ValueError(f"no extent for grid axis '{p}': {n!r} / {block!r}") + grid.append(-(-int(n) // int(block))) # ceil-div + + path = os.path.join(workdir, SHAPE_TXT) + with open(path, "w") as f: + f.write("\n".join(str(g) for g in grid) + "\n") + logger.info("[TOGSim] grid %s -> %s", grid, SHAPE_TXT) + return grid def emit_trace(workdir, meta): @@ -161,13 +199,19 @@ def emit_trace(workdir, meta): return n_tiles -def run_togsim(workdir, attribute_path=None, timeout_sec=None): - """Simulate the emitted trace. Returns TOGSimulator's parsed result dict.""" +def run_togsim(workdir, meta=None, args=(), attribute_path=None, timeout_sec=None): + """Simulate the emitted trace. Returns TOGSimulator's parsed result dict. + + `meta`/`args` supply the grid: the trace producer takes its loop bounds from + shape_args, so they are written out per launch rather than compiled in. + """ from Simulator.simulator import TOGSimulator so = os.path.join(workdir, TRACE_SO) if not os.path.isfile(so): raise FileNotFoundError(f"{so} not found -- call emit_trace first") + if meta is not None: + write_shape(workdir, meta, args) # A handle only: TOGSim derives trace.so / trace_cycles.tsv from its # DIRECTORY, and reads the file itself only on the STONNE path. diff --git a/TOGSim/src/main.cc b/TOGSim/src/main.cc index 0ef98eff..b98d3305 100644 --- a/TOGSim/src/main.cc +++ b/TOGSim/src/main.cc @@ -40,7 +40,18 @@ std::unique_ptr build_trace_tilegraph(Simulator* simulator, while (ct >> c >> o) { cyc.push_back(c); ovl.push_back(o); } } if (cyc.empty()) { cyc.assign(256, 128); ovl.assign(256, 0); } - return trace_to_tilegraph(trace_so_path.c_str(), nullptr, 0, + // Shape args: the producer's grid bounds, one per axis, when the trace was + // compiled without them baked in. Same sidecar convention as the cycle table + // -- absent means the producer carries its own constants. + std::vector shape; + { + std::ifstream sh(fs::path(trace_so_path).parent_path() / "trace_shape.txt"); + int64_t v; + while (sh >> v) shape.push_back(v); + } + return trace_to_tilegraph(trace_so_path.c_str(), + shape.empty() ? nullptr : shape.data(), + (int32_t)shape.size(), bases.data(), (int)bases.size(), cyc.data(), ovl.data(), (int)cyc.size(), partition_cores.data(), (int32_t)partition_cores.size(),