TOGSim C++ trace-generation pipeline (P0-P3): explicit dataflow producer + barriers#267
Merged
Conversation
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
3 times, most recently
from
1151f6a to
7f70bbb
Compare
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
5 times, most recently
from
c166abd to
ed5c747
Compare
Compile a kernel's post-vcix MLIR into a shape-parametric C++ trace producer (build_skeleton -> dep_analysis -> lower_to_emitc -> EmitC -> .so). TOGSim dlopens it, records the emitted instruction stream, and feeds the existing timing Core through togsim_trace_bridge. An async DMA and the memory_barrier that waits on its data pair at runtime by (tag_id, tag_slot): one static dma op runs once per loop iteration, so a compile-time id cannot express the pairing. That also covers multi-tile-K and conv, whose reduction is the kh*kw*C nest.
togsim_dispatch with TILE_BEGIN/TILE_END; outline each work-item into togsim_kernel_tile.
DMA-capacity throttle and frozen-state guard, per-core VMEM in the configs, and the SA weight-buffer throttle.
trace_timeline.py with per-work-item grouping and resource-centric DMA lanes; the trace logs the first DRAM response and the assigned systolic array, and scopes the compute barrier to its dispatch.
The trace .so is now picked up automatically next to the kernel's tile_graph.onnx; TORCHSIM_LEGACY_TOG=1 falls back to the deprecated ONNX producer. Also make the codegen cache replay safe: FxGraphCache can skip PyTorchSim codegen entirely, so the wrapper must carry what the .so needs, and concurrent compiles of one source must not observe a half-written .mlir.
… runtime model Replace the trace bridge's accumulated special cases with one dataflow rule and clean up the runtime that consumes it. Dependency rule: per SRAM buffer keep a writers SET; a reader depends on all current writers (occupancy=ISSUE when both are systolic-array ops, else latency=DONE); a writer REPLACEs the set. The only exception is is_mm_accum (a matmul that reads and writes the same buffer = a commutative accumulator): skip its read edge and UNION its write, waiting only the non-matmul init seed and not ordering co-matmuls. This drops the matmul-accumulator chain that deadlocked the SA weight-slot pipeline while keeping the init->matmul edge, and lets a vector epilogue or the store wait every K matmul (fixes the pure-vector store that an empty COMPUTE_BAR let slip). Remove COMPUTE_BAR entirely: a matmul is its own DONE-handle (finish == SA drain), so the store JOINs the matmul writers directly. The whole emit/loader chain is gone -- build_skeleton, lower_to_emitc, togsim.compute_barrier, the runtime symbol, the Opcode/case/_fence_finish, and TraceRec::COMPUTE_BAR -- so a stale producer fails loudly instead of emitting records the bridge would drop. Only MEMORY_BAR remains (an async load's DONE is its data arrival, not issue). Model compute-output spad footprint in the SRAM version/capacity machinery so buffer reuse (WAR) is capacity-modeled, not a hard edge. The output size comes from the DMA records that touch the same buffer (a buf_bytes pre-pass); an in-place buffer (accumulator, relu) is version-transparent so footprint is not double-counted. The occupy gate and version release sit in the MOVIN/MOVOUT/COMP issue points (release before the COMP skip path so a skipped matmul still frees). Runtime: collapse child_inst / _pipeline_children into one event-indexed _deps[ISSUE|DONE] with add_dep(c, on) and fire(e); collapse the weight-slot release queue and the async-load wakeup into one _due_events timed-effect table drained by process_due_events. Both are behavior-preserving (byte-identical). Require the weight-slot model: sa_weight_buffer_depth must be > 0 (errors at init), and the round-robin disable mode is removed. Degenerate traces (a consumer-less preload, an unpinned matmul) hit explicit error+exit guards rather than asserts that vanish under NDEBUG. Mark the legacy ONNX TOG path deprecated: it is superseded by the trace path, so TileGraphParser logs a deprecation warning and the TORCHSIM_LEGACY_TOG=1 opt-in warns at command build.
… spad/2
The validation-binary spad-overflow check sat inside `if functional_mode:`, so in
timing-only / autotune (non-functional) runs an over-spad tiling was never
rejected and reached TOGSim, which wedged ("spad too small") and crashed the
compile via assert 0. Move the compile + check out of the functional gate (the
Spike execution itself stays gated, run_spike below) and budget per dispatch at
spad/2 -- two work-items run concurrently (double buffer), so each must fit half
the spad or they deadlock competing for it. This matches the GEMM tiling gate
(max_spad_size = spad/2), which pointwise ops lacked. Fixes the resnet /
test_scheduler wedge where a fused BatchNorm+ReLU tile exceeded spad/2.
Every generated trace.cpp now opens with a "GENERATED ... DO NOT EDIT" banner carrying the TOGSim ABI version and the togsim_* call formats, so a dumped trace is self-documenting. Both are read from togsim_runtime.h at codegen time -- the version from the TOGSIM_ABI_VERSION define, the call-format text from a marked block next to the declarations -- so they never drift when the ABI changes.
…rint The bridge sums each work-item's distinct-buffer footprint (each buffer once, so a reduction's reloads of the same section do not inflate it) onto its Tile. can_issue then admits two concurrent dispatches only when each fits half the spad; a dispatch larger than spad/2 runs alone with the whole spad, so two work-items never compete for the shared spad and deadlock -- a runtime safety net beneath the codegen spad/2 gate. The footprint and resulting max_dispatch are logged on the TILE_SCHEDULED trace line for debugging.
Count the fused-epilogue buffers a tiling actually needs, and stop charging the kernel's stack frame against the spad budget -- it lives in DRAM, not the scratchpad, so including it rejected tilings that fit.
Spike v1.0.3 zero-inits the MVIN/MVOUT DMA address buffer so ROUNDUP padding entries are skipped instead of dereferencing uninitialized garbage, fixing the host store segfault on 8-lane (wrapper3) configs where a tile's split axis is not a multiple of vlane_stride*n_vu.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
4 times, most recently
from
ae26ede to
2463f82
Compare
YWHyuk
commented
Jul 10, 2026
| from . import dma_fine_grained | ||
| from . import lower_to_vcix | ||
| from . import peel_transfer | ||
| from . import peel_transfer |
Collaborator
Author
There was a problem hiding this comment.
remove this doc. clear history
Collaborator
Author
There was a problem hiding this comment.
Cleanup this docs series
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
2463f82 to
6908a2c
Compare
pytorchsim_functional_verify_per_kernel compares every realized Spike buffer against a CPU golden and stops at the first divergent kernel, naming the op and the offending indices.
The GEMM path already budgeted its fused epilogue; BMM's fused prologue was unaccounted for, so a tiling that overflowed the scratchpad could be chosen.
Skip the ready-queue scan when nothing can newly issue, skip building the per-instruction trace string when trace logging is off, and fix a use-after-erase in the zero-cycle COMP skip path.
…X TOG The main compile/sim path no longer generates or selects the legacy ONNX Tile-Operation-Graph. extension_codecache emits only trace.so + trace_cycles.tsv (the build-skip now keys on trace.so), and TOGSimulator.run_standalone always drives TOGSim with --trace_so. The TORCHSIM_LEGACY_TOG opt-in is removed from the frontend. The ONNX --models_list branch is kept solely for the STONNE sparse path (extension_op.py); TOGSim's C++ ONNX parser is untouched (separate PR). origins (which FX nodes a kernel came from) is preserved: logged per kernel run and recorded as a trailing "# origins:" line in trace_cycles.tsv -- the legacy ONNX TOG carried this as node metadata, and the C++ cycle-table loader stops at the comment so the current parser is unaffected. Also drop the dead tog_file param from mlir_gem5_compile_command, migrate scripts/chiplet.sh to --trace_so/--cycle_table (the trace path stubs per-tensor addresses and --attributes_list is no longer a Simulator option), and refresh the CLAUDE.md TOG-generation notes.
…iptor Represent a kernel body as ordered steps so a multi-step indirect access can be emitted, and hand the gather/scatter offset to the DMA as an explicit operand instead of inferring it from the operand count.
…emmini togsim.transfer is unregistered, so it can carry every runtime descriptor as an operand -- including the masked-clamp low/high vectors -- which a registered memref.dma_start cannot. The trace path (build_skeleton) reads togsim.transfer/togsim.wait, and the indirect offset operand is detected by its `indirect` attribute rather than by counting operands.
…o CONFIG/2/3/4) Instead of packing the transfer params into four CONFIG asm instructions, build a 144-byte DMA descriptor as an 18xi64 memref.global (.data, non-constant so runtime fields can be written) and emit a single CONFIG_DESC that hands its address to the MVIN/MVOUT (which read the struct, matching the Spike-side torchsim_config_desc). This is the vehicle for the masked-DMA low/high clamp: dim_low/dim_high/fill are descriptor fields (currently defaulted, so behavior is unchanged), not new special registers. - _pack_desc encodes the fields into the struct's i64 slots (dim_size/low/high/mm_stride/ spad_stride/element_size/vlane/config_type/flags/indirect stride+esize). - descriptor globals are deduped by content; the indirect offset spad address is a runtime value, stored into slot 15 before the transfer (flags.bit0 = indirect). Requires the Spike torchsim_config_desc instruction (riscv-isa-sim). Validated functional (Spike) + trace (TOGSim): add, matmul (subtile+async), softmax, gather/scatter (indirect runtime store), constant_pad all pass.
Clamp each tile axis to its real DRAM extent in the DMA descriptor, so a tile that runs past the extent is filled with the consuming reduction's identity instead of MAC-ing garbage, and a store skips the ragged tail instead of writing the systolic pad region. Covers pointwise, pad, reduction, gather/scatter, matmul and conv, plus the fused-epilogue output store. Adds a golden unit test for the lower_transfer_to_gemmini lowering and a non-dividing regression suite to CI.
…ixel_shuffle trace) Both the Gemmini descriptor and the TOGSim DMA model cap at 4 tile dims, but a logical tile can exceed 4D (e.g. pixel_shuffle splits two spatial axes -> a 5D tile [1,1,2,4,2]). lower_transfer_to_gemmini reduces >4D as part of emitting Gemmini asm, but that runs only on the Spike/gem5 lowering path. The trace producer (build_tog) reads togsim.transfer BEFORE that, so it emitted a 5D DMA and TOGSim rejected it: "issued tile is not supported format.. tile.size: 5". Exposed by the pixel_shuffle case newly added to the wrapper1 CI suite. Add a peel_transfer POST_OPT pass that reduces every >4D togsim.transfer up front (before build_tog) while KEEPING it a togsim.transfer, so both consumers see <=4D. It mirrors lower_transfer_to_gemmini's two reductions: collapse unit dims when the effective rank is <=4 (memref.collapse_shape), else peel the outer effective dims into an affine.for nest with the lane-banked physical SRAM offset. The gemmini lowering keeps its own reduction as a now-defensive no-op. Validated on 128x128: test_floormod_axis_split (pixel_shuffle + 9 others), test_masked_nondividing, test_matmul, test_conv2d all pass.
Three MVIN over-reads on shapes whose tiles do not divide the extent: a cat concat-dim split, a vlane index_expr row stride on multi-row splits, and a 1-D reduction tile that exceeded its extent.
torch.sort's frontend lowering exposes the sorted VALUES as the sole scheduler-visible template output; the INDICES buffer is written in place by the same kernel (make_inplace) but is not registered as an output. For argsort the values output is dead, so Inductor DCEs the whole sort kernel and the indices buffer is returned uninitialized (all-zeros), silently corrupting any argsort result (e.g. DeepSeek MoE expert routing). Register the in-place indices write as a MutationOutput owned by the sort operation so the scheduler keeps the kernel alive whenever indices is consumed. OperationBuffer.get_outputs() hardcodes [self], so add a small MLIRTemplateBuffer subclass that can own extra MutationOutputs and have MLIRTemplateCaller.output_node build it; the sort lowering attaches the indices mutation via add_mutation_output. For templates with no mutation this is behaviorally identical to the base (get_outputs == [self]). Then gate the values MVOUT on the values output being live: in argsort it is pruned from the kernel args, so emitting its MVOUT referenced an undeclared %YV. The values scratchpad is still produced internally for the bitonic compare; only the dead DRAM store is skipped. Verified: t.argsort() now returns correct indices (was all-zeros); torch.sort with both outputs live is unchanged; test_sort, test_matmul, test_cat, test_indirect_access unaffected by the output_node change. Note: x[argsort] end to end additionally needs the indirect-gather codegen fix (separate).
…oad clamp A bare indirect index (x[idx]) left the loop-local symbol in the DMA base offset. Separately, reverse-engineering per-dim padding from a single flat offset is ill-posed -- under channels_last the stride-1 axis absorbs the remainder and yields an empty [low, high) box that zeroed the whole load.
codegen_epilogue_body() decided whether to emit the template buffer's MVOUT
by asking "did a fused epilogue already store something?":
if len(self.stores._lines) == 0:
template_store()
That is a proxy for the real question, "does anything outside this kernel still
read the template buffer?", and it is wrong in two of the four cases:
- epilogue + live template buffer -> store skipped, buffer never written
- no epilogue + dead template buffer -> store emitted for a pruned DRAM arg
The first case corrupts DeepSeek-V3's MoE gate. The gate is
sigmoid(mm(hidden, gate_w) + bias), so add+sigmoid fuse onto the GEMM template,
but the raw mm result is also read by a later gather kernel. Its buffer was
declared as a kernel output and never stored, so the gather read an
uninitialized buffer (all zeros), scrambling expert routing. Every config was
affected; only 8x8 crossed the test's loose atol=2e-1 (final Max abs diff
0.2546) while 32x32 landed just under it and passed.
Ask the real question instead. Inductor already answers it:
Kernel.remove_kernel_local_buffers() drops a buffer only when
Scheduler.can_buffer_be_removed_through_fusion() finds every user inside this
fused kernel. A buffer that survived removal therefore has an outside user and
must be stored. Record the template buffer's name in def_kernel() and
def_conv_kernel() (conv has its own implementation) and gate template_store()
on it, in both the main and the reduction_fusion branches. The proxy is gone.
Use the kernel-local self.removed_buffers, not V.graph.removed_buffers: the
former is what remove_buffer() populates, and the latter is still empty when
the store hook runs.
Note this makes a previously missing MVOUT issue, so DRAM traffic and cycle
counts rise for any model with a live template buffer under a fused epilogue.
That is the correct cost, not a regression, but reported cycles will shift.
Verified on the 8x8 config: test_deepseek_v3_base goes from Max abs diff 0.2546
to passing; a dead template buffer (relu(mm)) still prunes its DRAM arg and
emits a single MVOUT; the trace/TOGSim timing path runs with the extra store.
Regressions pass: test_bmm_reduction, test_matmul, test_bmm, test_conv2d,
test_cnn, test_group_conv, test_pool, test_reduce, test_softmax, test_layernorm,
test_mlp.
Lower the remaining pointwise ops, route math.log/atan to VCIX in the Python pass, and add tests/ops/elementwise/test_pointwise.py.
Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
Make int max/min reductions dtype-aware, lower torch.roll to narrow+cat and realize it so the consumer reads a plain affine index, flatten nested single-variable floor/mod in axis-split, define empty_strided_cpu in the generated wrapper header, and place the reduction axis outermost in the per-lane tile layout.
SwinV2 batch>1 used to fail codegen with "Unlinearized floor/mod in DMA index" because the shifted-window path (torch.roll composed with window partition/reverse) produced views axis-split could not linearize. With the roll->slice+cat lowering, the nested/shifted mod handling, and the reduction-axis layout fix, the whole backbone now compiles and matches CPU end to end (max diff ~5e-6 for image 64 / window 8 / batch 2). Add tests/models/test_swinv2.py, wire it as a self-hosted CI job next to the other model tests, and list SwinV2 in the README model coverage.
conv_multi_tile_mapping (used for batch > 1 convs) collects every tile that fits SPAD into tile_candidates and returns them sorted, but it guarded the "no mapping" error on max_used_spad_size instead of on tile_candidates. max_used_spad_size is only bumped when a candidate also satisfies max_k_h_w <= k_h, and max_k_h_w is initialized to K_W, so it only ever fires for the full-kernel (k_h == K_H) tile. When that tile overflows SPAD -- e.g. a CLIP ViT-B/32 patch conv (3->768, 32x32 stride 32) at batch > 1 -- the guard raised "Cannot find a valid mapping" even though smaller-k_h tiles fit and were already in tile_candidates. The mapping variable it tracks is never returned. Guard on `not tile_candidates` instead, so the conv is rejected only when no tile fits at all. The returned candidate list is unchanged, so convs that already worked are unaffected. Adds a batched patch-conv case to test_conv2d.py. Fixes issue #252.
With the conv-mapping fix, CLIP's vision transformer runs with batch > 1 and matches CPU end to end (max diff ~1e-5). Add tests/models/test_clip.py (CLIPVisionModel, batch 2, patch 32), wire it as a self-hosted CI job, and list CLIP in the README model coverage.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
2 times, most recently
from
d9131ba to
1ff2ebb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Replaces the timing-path TOG producer (
MLIR -> Python dict -> ONNX -> C++ TileGraphParser) with a compiled, shape-parametric trace producer:post-vcix MLIR -> skeleton -> EmitC -> C++ -> .so. TOGSim dlopens the.so, runs it to record an instruction trace, and feeds it into the existing Simulator/Core (timing core unchanged). Driven by a new--trace_somode; the legacy ONNX-TOG path is kept and marked DEPRECATED, so nothing existing breaks.Pipeline
Dependency model (no in-order, no runtime tag-hash, no op heuristics)
Dependencies are derived from two sources available pre-collapse:
SA_WEIGHTSbuffer that folds preload->matmul.MEMORY_BAR(renamed fromBAR): the DMA/tag memory fence; an async load -> compute waits the data's resp-complete.COMPUTE_BAR(new): the compute fence; a store waits all systolic-array pipelines to drain.Both barriers are first-class trace ops (
togsim.compute_barrier-> ABItogsim_compute_barrier) visible in the trace dump and the instruction stream.Status
--trace_so.build_togpath on the same kernel + gem5cycle_list: compute work and DRAM traffic match; matmuls pipeline on 2 SAs; the memory fence correctly delays compute until the weight load arrives.docs/design/togsim_cpp_trace.mdsec 10): preload-concurrency cap (needs non-zero preload occupancy), parallel output tiles (dispatch granularity), broader op coverage (conv/SDPA/vector).Testing
tests/test_togsim_skeleton.py,test_togsim_emitc.py,test_togsim_runtime.py(7 tests).--trace_soGEMM through TOGSim.Design of record:
docs/design/togsim_cpp_trace.md(sec 9-10).🤖 Generated with Claude Code