TOGSim C++ trace-generation pipeline (P0-P3): explicit dataflow producer + barriers#267
Open
YWHyuk wants to merge 64 commits into
Open
TOGSim C++ trace-generation pipeline (P0-P3): explicit dataflow producer + barriers#267YWHyuk wants to merge 64 commits into
YWHyuk wants to merge 64 commits into
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
… feed Skeleton + EmitC + cost/dep analysis on the frontend; the trace runtime, loader, bridge, and Core feed on the simulator; shared MLIR pass helpers and the pipeline tests.
Per-record tag key in the bridge plus per-iteration tag alloc in dma-fine-grained so multi-tile-K and conv loads do not collide; strip the reduction accum marker from the memory_barrier slot.
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.
Default to the trace path; fix uninitialized Instruction fields, the matmul accumulator wedge, fused-subtile dedup, nested/fused epilogue dataflow, and dma_wait fusion; bound concurrent dispatches to the spad, round-robin work-items within a partition, benchmark autotune and run the multi-tenant scheduler through the trace path, and emit trace.so for pooling/reduction.
Carry simulator headers through the wrapper for cache-safe replay; drop verbose [P3-trace] logs; fix the key.mlir compile race in load().
… 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.
…election select_tile fed the epilogue buffer count to gemm_combination_mapping as max(n_extra_read - 2, 0), which optimistically assumed the X and W DMA buffers were already freed and reusable by the epilogue. The codegen lays every buffer out as a disjoint .spad global and never reuses freed space, so the estimate undercounted the real footprint: for a fused matmul+relu kernel the budget came to 64512 B/lane (treated as a bare GEMM) while the emitted kernel used 89088 B/lane including the relu output buffer. Under the full-spad guard this was harmless (89088 < 131072), but the spad/2 guard rejected it and crashed test_transformer_fusion with SpadOverflowError, since the heuristic path has no tile-shrink fallback. Pass n_extra_node + n_extra_read instead: one output-tile-sized buffer per epilogue node plus one per extra read operand, matching what the codegen emits. For the matmul+relu kernel the budget now equals the actual footprint exactly, and tile selection picks TILE_M=128 (62976 B/lane) which fits spad/2. Liveness-based spad reuse and broadcast over-allocation are tracked separately as optimizations in issue #275.
The spad-overflow check summed the kernel stack frame into the per-lane scratchpad usage (spad_usage = stack_size + spad_size) and rejected the tiling when that exceeded the spad/2 budget. But only the .spad section actually lives in the scratchpad -- it is pinned there by the --section-start=.spad link option. The kernel stack is in main memory (sp is set up by pk in the -m region, not at the scratchpad vaddr), so it does not consume the scratchpad and must not be charged against it. The scalar frame is also shared across lanes, not per-lane, so adding it double-counted. On small configs (8x8) this falsely rejected feasible tilings: the wrapper3 conv2d/resnet/mistral kernels fit the 32 KB spad with room to spare but were tripped over the spad/2 gate purely by the ~200-800 B stack term, crashing compile with SpadOverflowError. Drop the stack term; the .spad-only check still correctly rejects real buffer overflows (e.g. sparsity, which is fixed separately by f05ac8a).
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.
Add pytorchsim_functional_verify_per_kernel, a sub-option of pytorchsim_functional_mode that localizes the first compiled kernel whose Spike output diverges from a CPU reference. When enabled, the generated wrapper compares every realized buffer (the output of each fused kernel) against a CPU "golden" computed once per call by running the original aten graph (V.graph.module) on CPU with the same inputs, via an fx.Interpreter that records each node's output by name. A buffer is mapped to its originating fx node through V.graph.get_buffer().origin_node, so the check reports the kernel, the originating op, the offending indices and the max abs diff, then raises at the first divergence (stop-at-first). Comparison granularity is the fused-cluster output, the finest observable in a fused pipeline. Auto-disabled when functional mode is off (no Spike values to verify); the config accessor AND-gates the key with functional mode. Codegen bakes the verify_init/verify_check calls into the wrapper only when enabled at compile time, so clear the codegen cache when toggling. Tolerances via TORCHSIM_FUNCTIONAL_VERIFY_RTOL / _ATOL (default 1e-4). extension_functional_verify.py holds the graph registry and the runtime golden/compare logic; mlir_codegen_backend injects the calls at the wrapper level; extension_config reads the new YAML key.
The BMM tile selector fed only n_extra_node to gemm_combination_mapping and dropped the prologue's extra-read operands entirely, so a softmax-fused attention matmul (value^T @ softmax(scores)) was tiled as a bare BMM. The codegen lays the softmax max/sum operands out as their own disjoint weight-tile-sized .spad globals (buf3/buf4) and never reuses freed space, so the estimate undercounted the real footprint: on the 32x32 wrapper2 config (16 KB/lane spad/2 budget) tile selection picked TILE_N=512, whose emitted kernel used 16896 B/lane and overflowed the budget by 512 B, crashing test_transformer with SpadOverflowError. Under the 128x128 full-budget config the slack hid it. This mirrors the epilogue n_extra_read gap fixed for the GEMM template (commit f05ac8a), now on the BMM prologue path. Add an n_prologue_extra_read knob to gemm_combination_mapping that charges each extra prologue-read operand as one weight-tile-sized (TILE_K x TILE_N) .spad buffer, matching what the codegen emits, and have the BMM template count those operands the same way codegen does (the one numel-matching main input reuses the matmul-operand buffer; every other read gets its own global). Tile selection now rejects TILE_N=512 (16896 B/lane) and picks a fitting tile (8704 B/lane), so wrapper2 test_transformer passes the full Spike + allclose check. The new parameter defaults to 0, leaving the GEMM and conv callers unchanged.
Add docs/per-kernel-functional-verify.md (usage, options, mechanism, limitations, code map) and point to it from CLAUDE.md: a one-line entry in the debugging section and in the YAML knobs list.
Keep only a one-line note on what TOGSIM_ABI_VERSION guards; the per-bump v1..v12 history was noise nobody reads.
A fused template epilogue (matmul/conv + relu/add/gelu/...) stores its result via store_epilogue, which -- unlike the plain store() -- emitted no masked-DMA clamp. On a non-dividing output the store then wrote the systolic pad region past the real extent, corrupting the result (fused matmul+relu at 10x10x10 was off by ~7). store_epilogue now clamps each output tile dim to [0, extent): the extent of every loop iv is ranges[k], and the tile dims are indexed by dim_aliasing (tile-dim order), so a dim whose iv extent does not divide the tile is clamped. Keyed by iv name, it is naming-agnostic -- it covers both the index-N matmul epilogue and the c0/tile_n/... conv epilogue -- and _emit_clamp skips dividing dims, so a dividing output is a no-op. Fixes fused matmul/conv epilogues on non-dividing shapes (matmul+relu/add/gelu, conv+relu/scale/bias/sigmoid, non-dividing K and non-dividing output). No regression: test_matmul, matmul_activation, addmm_residual, matmul_scalar all pass.
Hand-write a masked MVOUT togsim.transfer, run lower_transfer_to_gemmini in-process (no torch, no spike; skipped without the MLIR bindings), and compare to an embedded golden -- the exact lowered IR (MLIR SSA numbering is deterministic for a fixed input): the i32-view descriptor global (dim_size / config_type / masked flag packed), the runtime dim_low/dim_high stores at i32 idx 6/10 for masked_axes=[2], and the CONFIG_DESC + MVOUT asm. A second case adds subtile_size=[1,1,4,8] and shows the descriptor dim_size become the subtile (config) shape. Timing mode erases the transfer. Under tests/lowering/ so further per-pass goldens have a home. Regenerate the golden and review the diff on an intended lowering change.
…dims, cleanup Addresses the PR review of the masked-DMA non-dividing work: - _masked_fill_bits detected a log-sum-exp reduction with `"exp" in origin.name`, which also matches expand / expm1 / experimental -- an ordinary non-dividing sum whose origin name contains "exp" would then get a -inf tail fill instead of the sum identity 0 and reduce to -inf. Match the op target instead (_origin_is_exp: the target's overload packet name == "exp"); same fix applied to the get_padding_type twin. Unit test in tests/lowering/test_masked_fill.py pins that expand/expm1 do not match. - Remove the self._last_local_dims temporal coupling: get_dma_info now returns local_dims and load()/store() pass it to _masked_bounds explicitly, so a reordered or skipped get_dma_info can no longer feed stale clamp axes. - Drop the underscore-prefixed locals in def_dma_op / store_epilogue (non-idiomatic), trim the oversized _masked_fill_bits / _masked_bounds docstrings to the invariant plus a link to docs/design/masked-dma-descriptor.md, and clarify the bf16 dtype-table comment (present only for table totality; bf16 is not actually supported).
Dropping the tile-divisibility recompile means non-dividing shapes rely entirely on the masked-DMA clamp/fill; a gap is now a silent wrong answer, not a slow-but-safe recompile. Add tests/ops/misc/test_masked_nondividing.py pinning that path: 2-D constant pad (per-dim greedy pad recovery), sum over a broadcast operand (0 fill, guards the exp-substring bug), softmax (genuine -inf fill), non-dividing elementwise, and non-dividing gather. Register it in the CI allowlist (pytorchsim_test.yml). Also add the missing trailing newline to test_indirect_access.py.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
1e43e9a to
cfcb7a9
Compare
…ilogue The ordered-steps refactor removed the reduction masking mechanism (self.masks / get_mask / mask_cse) but left a compute_body.splice(self.masks) in codegen_epilogue_body, so any matmul/bmm + reduction fusion crashed with AttributeError: no attribute 'masks'. It was masked until now by the stale-spike illegal-instruction CI failure. Remove the dead splice; test_matmul_reduction (matmul+max, var_mean, residual+var_mean) passes. Non-dividing reduction-fusion masking (padded lanes vs the reduction identity) stays a separate gap from that refactor.
…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.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
f280164 to
1e2a3a9
Compare
…-safe) _GoldenInterpreter recorded each node output on CPU but returned the original (npu) tensor to downstream nodes, and cast the recorded value to float32. A device-baked op (e.g. arange on npu) or a consumer that requires operand device agreement (aten.index needs its index tensors on the indexed tensor's device) then saw a mixed CPU/npu pair and the whole golden threw, silently disabling the per-kernel verify for that graph -- e.g. MobileNet's depthwise conv lowered to constant_pad + index gather. run_node now moves EVERY tensor output to CPU (preserving dtype) and returns it, so downstream nodes see CPU operands; only the recorded golden is cast to float32 for the allclose compare, so int index tensors stay usable.
The cat template's _calculate_input_tile_sizes handed each input a concat-dim tile of min(extent, remaining_spad_budget). When that budget did not cover an input's full concat extent, the tile width could fail to divide the extent (e.g. extent 64, tile 62 on the 8-lane config). The copy loop steps by that tile over [0, extent), so a non-dividing tile produces a ragged final iteration with no remainder mask: the MVIN over-reads past the input and the MVOUT over-writes past the output, corrupting the result. This surfaced on the 8x8 (wrapper3) CI config, where the per-input spad budget lands at 126 for two 64-wide inputs, truncating the second tile to 62. It scrambled ~45% of the cat output and broke LlamaDecoderLayer (rotary rotate_half cat) and the diffusion UNet skip-connection concat. Larger arrays give a bigger budget so both inputs got the full 64-wide (dividing) tile, which is why 32x32 and 128x128 passed. Snap each input's concat-dim tile down to the largest divisor of its extent that fits the budget, so every loop iteration is full-width and in-bounds. When the budget already covers the full extent the tile is unchanged, so 32x32/128x128 codegen is a no-op. Verified with a minimal cat repro (64+64 -> 128 on 8x8: fails before, diff 0 after) and the real test_llama LlamaDecoderLayer at 8x8 (now passes).
The vlane index_expr lowering reconstructs each lane element's logical index along the split axis as stride_dim + outer_dim * nr_vector_lane, where outer_dim is the sublane-row index. The next sublane row starts vector_lane * vlane_stride elements later (vector_lane lanes, each holding vlane_stride positions), but the multiplier used nr_vector_lane alone, dropping the vlane_stride factor. Any index tensor (e.g. arange/iota) whose split axis spills into a second sublane row then had every row past the first come out shifted, producing wrong values. This fired on the 8x8 (wrapper3) config: arange(32) at 8 lanes with vlane_stride 2 needs two sublane rows, so positions 16-31 came out as p-8. Larger arrays keep the 32-wide axis in a single row (outer_dim always 0), which is why 32x32/128x128 were unaffected. It broke test_llama's LlamaModel (rotary position ids) at 8x8. Multiply the outer/row term by vector_lane * vlane_stride. This equals the old value whenever vlane_stride == 1, so single-row splits and larger configs are a codegen no-op. Verified with a minimal arange(32) repro at 8x8 (positions 16-31 wrong before, diff 0 after) and the full test_llama LlamaModel at 8x8 (now passes).
init_tile_size gave a 1-D tile of 2*vlane_stride*vector_lane regardless of the data extent. When the extent is smaller than that tile the reduction/pointwise MVIN reads past the end of the DRAM buffer, and the compute loop folds the out-of-range elements into the result. For a sum reduction that means adjacent-DRAM garbage is squared/added into the sum. This broke test_single_perceptron's Loss on the 128x128 config: the MSE mean reduces 128 elements but the tile was 2*2*128 = 512, so the MVIN over-read 384 elements past the 128-wide input and roughly doubled the squared-error sum (loss diff ~4593). It fired on wide-lane configs (tile 512 > extent 128); 32x32 gives tile 128 == extent and 8x8 gives a tile that divides 128, so both were unaffected. Cap the 1-D tile to min(2*vlane_stride*vector_lane, extent) so the tile never exceeds the data. It is a no-op whenever extent >= tile, so larger reductions and other configs are unchanged. Verified with a minimal mse_loss(1-D) repro at 128x128 (over-reads before -> diff 0 after, across extents 64/100/127/200/300) and the real test_single_perceptron at 128x128 (Loss now passes).
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).
convert_indirect_indexing moves an indirect index into an offset spad and
zeros the indirect term in the DRAM base-offset expression, but it only
walked index.args. A BARE indirect index (x[idx], where the load index IS
the symbol so index.args is empty) was left unzeroed, so the DMA base
offset's affine.apply referenced the loop-local index SSA value before it
was defined ("use of undeclared SSA value name"). Any 1-D gather with a
runtime index tensor failed to compile at codegen; it was masked for
argsort only because the sort DCE bug left the indices all-zero (a
degenerate constant load).
Subs any remaining indirect symbol to 0 after the arg loop (mirrors the
store path). The per-position gather is carried by the offset spad, so the
base offset is correctly 0.
Verified: x[idx] (N=32/64) and x[x.argsort()] now compile and match CPU;
test_indirect_access.py (incl. scatter/index_add and multi-dim) and
test_sort.py unchanged.
_masked_bounds reverse-engineered per-dim padding from the single flat index offset and clamped each padded load to [pad, pad + input_extent). That recovery is ill-posed: the offset mixes the tap shift with the padding, and under channels_last the stride-1 channel axis absorbs the offset remainder, producing an impossible clamp (e.g. [4, 8) on a size-2 channel tile). Spike then zeroed the whole load, making channels_last depthwise conv 99.9 percent wrong while NCHW was fine. The clamp was also unnecessary: at pad positions the consumer already yields the correct value (a compute-side select for a padded gather, or the pad op's own fill), so reading OOB is harmless. Keep only the ragged tail clamp (glo=0, ghi=loop_extent), which is genuinely load-bearing for non-dividing tiles. Verified on 32x32: channels_last depthwise conv now matches CPU; NCHW conv and standalone ragged F.pad still pass. Broader regressions (padded reduction loads, full MobileNet) to be addressed separately.
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.
YWHyuk
force-pushed
the
feature/togsim-cpp-trace
branch
from
3fd7ae1 to
79b1aec
Compare
Implement abs, sign, isnan, isinf, fmod, bitwise shifts, copysign, erfc, hypot, cosh, sinh, acos, asin, atan, atan2, acosh, asinh, atanh, and log/ atan (VCIX) in the Inductor MLIR backend. Float-only transcendentals cast non-float inputs to f32 while leaving native float widths (f16/f64) untouched.
Port the C++ MathLog/MathAtanToVCIX patterns to the in-process Python VCIX pass: add math.log (opcode 2, imm 2) and math.atan (opcode 2, imm 3) to _MATH_VIV and MARKERS. Encoding matches the spike rs1 / gem5 VS1 sub-opcodes (sin=0, cos=1, log=2, atan=3). No C++ mlir-opt pass needed. gem5 decodes vlog_v/vatan_v into the CustomVlog/CustomVatan op classes, so give them a functional unit as well. Without one the Minor CPU warns "No functional unit for OpClass CustomVlog" and gem5 exits 1 on any kernel containing torch.log or torch.atan. They belong in SpecialFunctionUnit (opLat 10) next to exp/erf/tanh/sin/cos rather than in a 1-cycle unit, since they are transcendental. Co-Authored-By: Claude Opus 4.8 <[email protected]> Claude-Session: https://claude.ai/code/session_01MTjfq7HAexdCKopgp42rKE
Add the pointwise op test under tests/ops/elementwise and register it in the test workflow allowlist.
Bump the third-party pins to the releases shipping the new pointwise instructions (spike torchsim_vlog/vatan, gem5 CustomVlog/Vatan + decode).
reduction_init returned "-inf"/"inf" for max/min regardless of dtype, so an
integer max/min reduction fed an invalid inf identity into every consumer:
the accumulator init, the template reduction init, and the masked-DMA tail
fill (torch.tensor(inf, dtype=int64) overflowed -- surfaced compiling swinv2
with batch>1). reduction_combine_vec also hardcoded the float multi_reduction
kinds <maximumf>/<minimumf>, which reject integer vectors.
Return the dtype's representable extreme from reduction_init for integer
max/min (max identity -inf -> iinfo.min, min identity +inf -> iinfo.max;
bool -> 0/1), and emit the signed-integer multi_reduction kinds <maxsi>/<minsi>
for integer element types. sum/prod need no change: the <add>/<mul> combining
kinds are polymorphic and ops.add/ops.mul already select arith.add{i,f} by
dtype (as ops.maximum/minimum already select maxsi vs maximumf).
Verified: int amax/amin/sum/prod reductions compile and match CPU; float
test_reduce (ReduceMax) and test_softmax unchanged. Also clears the swinv2
batch>1 masked-fill overflow (it then hits a separate unsupported
shifted-window view, tracked separately).
…lar DMA index
torch.roll has no Inductor lowering; it decomposes to index_select(fmod(...)) -- a
cyclic-shift gather whose ModularIndexing DMA index the affine-only descriptor
cannot express ("Unlinearized floor/mod in DMA index"). Rewrite it as the two
contiguous halves the shift produces, stitched by cat:
roll(x, s, d) == cat([slice(x, N-s, N), slice(x, 0, N-s)], d)
each half is a plain strided slice with no modulo. Do it as a lowering (removing
roll from the decomp table so it is not decomposed first) rather than a
decomposition, so we can also REALIZE the input with copy_input: roll's producer
is often a reshaped view (e.g. swinv2's window_reverse), and the slice offset would
otherwise fuse into that reshape's modular index, re-creating a shifted (p+c)%m
form. A plain .contiguous()/clone does NOT create the buffer boundary (Inductor
inlines it); copy_input does.
Verified: torch.roll (1-D / multi-dim / arbitrary shift) and roll+reduction match
CPU; swinv2 batch>1 now compiles through the shifted-window attention (previously
blocked by the roll's cyclic shift and then by the fused slice offset).
A composition of aligned reshapes on one iteration variable (e.g. swinv2's window partition fused with a later reshape) leaves a nested index like ModularIndexing(ModularIndexing(p, 1, 64), 1, 8) that neither sympy nor simplify_with_ranges reduces. collect_boundaries then skips its cut points (the inner base is not a bare variable) and the affine-only DMA check rejects it. Add a general, pattern-free canonicalizer (flatten_nested_floormod / _as_digit): any nesting of FloorDiv/ModularIndexing over a single symbol is a digit extractor (v // A) % M and collapses to one level by composing divisors, from four divisibility-guarded algebraic identities. Applied in collect_boundaries and in _fold_with_ranges. Multi-variable inner arguments (e.g. a roll shift v+c) are left untouched. Verified: numeric equivalence on the swinv2 window index; test_reduce, test_layernorm, test_cat, test_indirect_access unaffected.
The wrapper header imported empty_strided but not empty_strided_cpu, so a graph that allocates a CPU-side buffer (e.g. swinv2's shifted-window attention mask, which Inductor keeps on CPU) failed at runtime with "NameError: name 'empty_strided_cpu' is not defined". Add the same binding the default Inductor PythonWrapperCodegen header provides.
…yout A per-lane reduction is lowered as a 2-D [reduction | batch] collapse: multi_reduction reshapes the accumulator to <reduction_numel x red_size> and reduces axis 0, which assumes the reduction axis is the outermost in-lane run and every non-reduction (batch) axis is inner. get_dma_info reorders the tile axes for reduction loads to guarantee that, but only for the 2-D and 3-D cases. The 4-D case mis-tested is_reduction (reduction_depth < 3, false once there are three batch dims) and then raised NotImplementedError, and rank 5+ fell through to the default row-major order -- both leaving the reduction axis innermost, so the reduce collapsed a batch axis instead of the reduction axis. This is hit whenever a non-reduction dim-merge is blocked and an extra batch axis stays in-lane: SwinV2 cosine window attention adds a gathered relative-position bias (table[idx[q,k], head]) before the softmax amax, keeping head separate from query (4-D tile, head bled into the max), and its post-attention LayerNorm var_mean splits the token axis into several loop vars (6-D tile, mean off by ~0.005). Generalize the reorder to any tile rank: a tile that carries the reduction axis places that axis-group outermost via range(r, L) + range(r-1, -1, -1), which reduces to the existing 3-D order for L=3. Adds test_reduce_gather_bias (fails max abs diff ~3.3 before the fix).
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
from
ae26ede to
2463f82
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