[Frontend] Simulate the Triton route's kernel in TOGSim

[YWHyuk]

Stacked on #305 (base is feature/triton-codegen, not develop — review that first). Pairs with PSAL-POSTECH/triton-npu#1, which builds the binary the cycle measurement runs.

#305 got a Triton kernel compiled to a RISC-V ELF. This gets it simulated: one torch.compile of x + y now ends with a TOGSim cycle count.

The one structural difference, and why it is small

A Triton kernel describes a single program instance — the tile loop is not in the kernel, it is the launch grid outside it. PyTorchSim's codegen puts the whole loop nest in the kernel, so the two look incompatible.

They are not. The trace producer already splits at exactly that seam (docs/design/togsim_cpp_trace.md §9.3): togsim_kernel_tile per work-item, enumerated by togsim_kernel, with core assignment owned by togsim_dispatch. A Triton kernel arrives in the shape the outliner otherwise has to produce. What was missing was only the enumeration, and lower_to_emitc.WorkItem supplies it.

What changed

build_tog — four fallbacks, each taken only when the existing condition does not hold, so PyTorchSim's own codegen walks the same path as before:

• _find_kernel accepts 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, and rooting at that one made every DMA a sibling the traversal never reached
• the DMA's tensor identity follows memref.reinterpret_cast back to the block argument — triton-shared types pointers as unranked memref<*xf32>, so the operand is a flat view of an argument, never the argument, which left arg_id at −1 and silently broke the address model
• 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, applied to the trace artifact only. 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. Everything after is untouched — _parallel_loop_chain finds the tagged loop, the outliner threads the induction variable through iv[], and the loop left behind becomes the dispatch enumeration.

triton_backend/timing.py — emit_trace (post-vcix IR → trace.so + trace_cycles.tsv), run_togsim, and measure_tile_cycles, which chains three pieces that already existed: build_tog sample mode annotates the IR, python -m tnpu.cycle lowers it in tnpu's process, gem5 runs it and CycleSimulator reads one numCycles per marker pair.

Measured

x + y, 1024 elements, XBLOCK 128, grid 8:

TOG        root -> DMA(arg0,load) DMA(arg1,load) Compute(vector) DMA(arg2,store)
trace.cpp  togsim_kernel_tile: offset = iv[0]*128, three togsim_dma + one togsim_compute
           togsim_kernel:      for (p = 0; p < 8; ++p) togsim_dispatch(...)
gem5       tile = 19 cycles
TOGSim     650 cycles

Channel-0 DRAM traffic is 16 reads × 32 B × 16 channels = 8192 B, exactly the 8 work-items × 2 loads × 512 B the kernel should move — so every dispatch in the enumeration really ran.

The MLIR route reports 251 cycles on the same computation. The gap is 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. That is the next gap, and it is named in the README.

What this does not do

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.

Matmul timing is also still open: build_tog finds compute nodes by the vcix.iv op name and tnpu emits llvm.riscv.sf.vc.* intrinsics, so a GEMM currently yields no compute node.

Regression

tests/ops/elementwise/test_add.py 5/5 and tests/ops/gemm/test_matmul.py 11/11 on the MLIR route; tnpu baselines add err 0, gemm 1.52588e-05. Verified again after rebasing onto the CI commits in #305.

[YWHyuk]

Replace triton-npu to TOGSim

[YWHyuk]

Done in 4005409 — [Spike]. Same point: this warning is about the functional half, so it should say so.

[YWHyuk]

triton-npu -> Spike

[YWHyuk]

The structural read is right, and it is the part that could have gone wrong: the trace producer already splits at the work-item seam, so a Triton kernel arrives in the shape the outliner would otherwise have to produce, and only the enumeration was missing. _underlying_block_arg also fixes a genuinely silent bug — arg_id = -1 breaks the address model with no diagnostic.

Four things I would change, in priority order. Two of them are about the same failure mode: a wrong number that looks like a measurement.

---

1. _build's new root filter is fail-open

roots = [op for op in block.operations
         if op.operation.name == "affine.for" and _has_loop_role(op)]
if roots: ...
# otherwise: the whole body becomes one work-item

Before, any top-level affine.for was a root (an untagged one fell through to _handle_compute). Now an untagged top-level loop is skipped, and if every top-level loop is untagged the kernel takes the no-loop path and yields a wrong graph — where it previously yielded an empty one, which is an obvious failure.

I checked this is latent rather than active: all eight templates (bmm / cat / conv×3 / gemm / maxpool / sdpa / sort) attach a role, and every top-level affine.for in the generated kernels I have locally (5/5) carries one.

Still, the regression it guards against is silent. Suggest: on the no-loop path, raise (or at minimum warn) if the block still contains an untagged top-level affine.for. A few lines, and it turns "cycle counts quietly drifted" into a message.

2. Placeholder cycles are invisible in the result

measure_tile_cycles returning None is logged loudly, but run_togsim's returned dict is indistinguishable from one backed by a real gem5 measurement. The cycle number is the entire output of this path, and a log line scrolls away.

Suggest carrying the fact into the result — a flag in the dict, or a marker in the TSV header that the reader surfaces — so a consumer cannot mistake PLACEHOLDER_CYCLE for a measurement. The comment on that constant already argues exactly this point; this extends it to the value that actually leaves the module.

3. work_item_for encodes the triton-shared ABI as a constant

pid_x = n_tensor + n_scalar + 3          # after gridX, gridY, gridZ

The layout (pointers → user scalars → grid×3 → pid×3) is documented, but if it ever shifts, some other argument silently becomes the program id and the grid loop drives the wrong value. Suggest asserting the shape rather than only computing from it — e.g. check the trailing six arguments are i32 — so a change announces itself.

Same function assumes a 1-D grid. parallel_args scales with len(grid), but kernel_spec._grid only computes the x extent, so a second axis would make the two disagree. assert len(grid) == 1 pins that until y/z actually lands; the comment already notes x is all that is covered today.

5. _is_address_plumbing whitelists by type

The docstring owns the limitation, and the reasoning is sound for this backend — tile data is vector/float, address math is index/integer. The narrower worry is the op-name side: func.return and memref.cast are matched by name and everything else must be arith.*. If the tnpu lowering later emits an llvm.* or index.* op at top level, it is counted as vector compute instead of plumbing.

Given this file raises (_Fatal) elsewhere rather than guessing, an unknown top-level op might be better treated the same way.

---

Approach and framing both look good to me — stating up front what does not work (output tensors, matmul timing, 650 vs 251 from the missing double buffering) is the right way to land a partial capability. The 8192 B DRAM check is a nice touch: it verifies every dispatch in the enumeration actually ran, which the cycle count alone would not show.

Items 1 and 2 are the ones I would want before merge; 3 and 5 read fine as follow-ups.