[Frontend] Lower togsim.transfer directly to Gemmini (drop memref.dma_start) by YWHyuk · Pull Request #286 · PSAL-POSTECH/PyTorchSim

YWHyuk · 2026-07-06T06:25:22Z

Stacked on #285 (base = feature/indirect-offset-descriptor).

Drops the memref.dma_start intermediate op and lowers togsim.transfer directly to Gemmini. memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them; keeping it alive through the whole pipeline is the substrate for the masked-DMA work.

What changed

New lower_transfer_to_gemmini = merge of decompose_transfer (<=4D handling: unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) + lower_dma_to_gemmini (CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm). A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather.
togsim.transfer gains a tag_idx operand so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices.
memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering). No memref.* DMA ops remain.
decompose_transfer dropped from PRE_OPT; loop-padding runs opaquely on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding -- its DRAM-grow half is what the masked-DMA clamp will replace.
dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer.

Validation

Functional (Spike, timing_mode=0): add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv(padding=0) all allclose.

Known-pending (documented in the commit + `docs/design/transfer-direct-lowering.md`)

Trace/timing path: build_tog (TogBuilder._handle_dma_start) + dep_analysis._dma_buffer still parse memref.dma_start; being ported.
conv padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces. Not a rewire bug (padding=0 passes; wrong values, not a crash).

🤖 Generated with Claude Code

…emmini memref.dma_start is a registered op with a fixed operand list, so it cannot carry the runtime descriptor operands the masked-DMA clamp needs (per-dim low/high are dynamic-shape values, not attributes). togsim.transfer is unregistered, so it carries them. Eliminate the dma_start intermediate: togsim.transfer flows through the whole pipeline and lowers directly to Gemmini. - New lower_transfer_to_gemmini merges decompose_transfer's <=4D handling (unit-dim collapse / >4D affine.for peel + lane-banked SRAM offset) with lower_dma_to_gemmini's CONFIG/CONFIG2/CONFIG3/[CONFIG4]/MVIN|MVOUT asm. A/B-validated (func7 + packed CONFIG rs1/rs2) against the old decompose->lower chain on add/pad/matmul/gather. - togsim.transfer gains a tag_idx operand (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) so the async DMA<->barrier runtime slot survives without memref.dma_start's variadic tag indices. - memref.dma_wait -> togsim.wait (togsim-level counterpart, emitted by the vcix matmul lowering, erased by the merged lowering). No memref.* DMA ops remain. - decompose_transfer dropped from PRE_OPT; loop-padding runs OPAQUELY on togsim.transfer (mlir-opt --allow-unregistered-dialect), doing only its loop-bound rounding (its DRAM-grow half is what the future masked-DMA clamp replaces). - dma_fine_grained, lower_to_vcix._DmaView, build_skeleton ported to read togsim.transfer. Functional path (Spike, timing_mode=0) validated: add, matmul (fine-grained subtile + async + togsim.wait), softmax, gather/scatter (indirect CONFIG4), constant_pad, conv (padding=0) all allclose. Known-pending: conv with padding>0 regresses -- it relied on loop-padding's DRAM-grow, which the masked-DMA clamp (next step) replaces (padding=0 passes; wrong values, not a crash). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

…wait Finishes the memref.dma_start removal for the trace/timing path. build_tog's TogBuilder now decodes togsim.transfer (dram,dram_idx,sram,sram_idx,tag,tag_idx, dma_type,vst[,offset]) and togsim.wait (tag,tag_idx,num_elements) instead of memref.dma_start/dma_wait: op dispatch, _dma_start_fields (positional decode + direction from dma_kind / dma_type), the tag-user base-address scan, and the wait operand reads. TOGDMANode.tile_stride keeps reading the dram_stride attr (unchanged semantics). build_skeleton's wait marker/dispatch move to togsim.wait so barriers still become togsim.memory_barrier. dep_analysis needs no change (the decoder keeps the same field keys). Trace/timing (TOGSim, timing_mode=1) validated: add 5/5, matmul 11/11 (async DMA + tag_idx barrier pairing). No memref.* DMA ops remain in either path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

… not operand count len(operands) > 8 is ambiguous once the masked-DMA clamp adds more optional operands (low/high after offset). Gate on the `indirect` attr instead, so offset vs mask operands are disambiguated by explicit flags. Behavior-preserving: gather/add pass on both the functional (Spike) and trace (TOGSim) paths. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

YWHyuk · 2026-07-06T06:45:17Z

+    """Insert a `togsim.dma` before the original `togsim.transfer`.

    `tag_id` is the identity of this DMA's tag memref. An async DMA pairs with
    its `togsim.memory_barrier` (the original dma_wait) by the RUNTIME tag slot


togsim.memory_barrier?

…gsim.wait) The trace producer now reads togsim.wait (not memref.dma_wait); update the pass docstring and comments to match. No code change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

…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. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

The togsim.transfer descriptor lowering requires the spike torchsim_config_desc instruction (riscv-isa-sim PR #5, tag v1.0.4). Track the new tag so build-from-source picks up the descriptor-capable spike. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

YWHyuk · 2026-07-06T14:30:30Z

+                pairs4 = [(expand + d, lo, hi) for d, lo, hi in desc_masked_pairs]  # 4D-expand axis
+                desc_ptr = desc_ptr_masked(slots, pairs4, masked_fill, ind_addr)
+            else:
+                desc_ptr = desc_ptr_and_store_indirect(slots, ind_addr)


Why two seperated functions. Make a general single function which cover those two features

YWHyuk · 2026-07-06T14:32:38Z

+            init = "-inf"
+        else:
+            init = reduction_init(rtype, dtype)
+        val = {"-inf": float("-inf"), "inf": float("inf")}.get(init, init)


Is it best?

YWHyuk · 2026-07-06T14:33:27Z

+        pad borders fall out naturally. Returns [(tile_axis, low_var, high_var), ...] for
+        the non-trivial axes; low/high are index SSA vars.
+        """
+        local_dims = self._last_local_dims


Does this method can support dynamic dim? We need to support dynamic dim

YWHyuk · 2026-07-06T14:34:07Z

+from .lower_dma_to_gemmini import _i64_signed, _row_major_strides, _elem_bytes, _asm, CONSTRAINTS
+from .decompose_transfer import _int_array, _const_int, _squeeze_reassociation
+
+CONFIG, CONFIG2, CONFIG3, CONFIG4, CONFIG_DESC = 0, 4, 5, 6, 7


CONFIG2, CONFIG3, CONFIG4 remained?

The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high + desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

…-dividing tails Drop the index_expr / indirect tile-divisibility recompile so pad_vlane_tile's lane-aligned tiles are kept instead of shrunk to a divisor of the dim. A non-dividing tile's per-tile remainder is handled by a masked-DMA clamp: - _emit_clamp emits per tile axis a dynamic [low, high) as affine.min/max of that axis' loop iv (the tile base), so the last partial tile and the pad borders fall out per iteration; low = max(0, glo - base), high = min(tile, ghi - base). The pointwise path (_masked_bounds) derives [glo, ghi) per dim (MVOUT [0, output); MVIN of a padded input [pad, pad+input_extent)); the MLIR templates (def_dma_op, e.g. gemm) derive it from the DRAM extent so a matmul tile padding past M/N/K is filled with 0 instead of MAC-ing garbage. - emit_transfer carries the (low, high) index operands + masked_axes; lower_transfer_to_gemmini writes the runtime low/high into the descriptor dim_low/dim_high through an i32-view (byte layout unchanged) + masked flag + desc_fill. One desc_ptr helper builds every descriptor: a fully static one is a deduped i64 global, while any runtime-written field (masked dim_low/dim_high and/or the indirect address, stored as two i32 words) uses a unique i32-view global, so a single descriptor carries both a masked clamp and a gather/scatter, and indirect DMAs drop their divisibility recompile too. - The box-excluded fill is the consuming reduction identity (reduction_init: sum 0, prod 1, max -inf, min +inf; 0 otherwise -- matmul MAC and stores use 0). Log-sum- exp exception (softmax / log_softmax): a sum whose node came from exp reduces exp(input), so the PRIMARY input (indexed by a reduction itervar) is filled with -inf; broadcast operands (the subtracted max) keep their finite identity, else input - (-inf) = +inf -> exp = inf -> NaN. - index_add: instead of the compute gather-add-overwrite (which loses duplicate indices landing in the same non-dividing tile), the MVOUT accumulates out[idx] += val via an accumulate flag; the DMA processes positions sequentially so duplicates accumulate correctly regardless of tile size. acc_float selects float vs integer add (element_size alone is ambiguous). With no caller left, is_dim_dividable / adjust_tile_to_divisible and the TileConstraint.must_divide_dim machinery (trim_large_tail bias, adjust loop) are removed. Fixes elementwise, constant_pad (mobilenet ragged vlane holes), reductions (sum/max/min/prod/softmax/log_softmax/var/std/layernorm/mean), gather/scatter (index_add with duplicates) and matmul/addmm on non-dividing shapes. Requires spike with the descriptor dim_low/dim_high clamp, desc_fill and the MVOUT accumulate flag. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

…ne-grained) Conv on a non-dividing reduction (K = in_ch * kh * kw, e.g. 3*7*7=147 > systolic size) MAC-ed garbage from the tail of the K tile. Two fixes: - def_dma_op now takes the per-dim clamp extent from the kernel's loop_extents {loop_iv_name: loop_bound}, which each conv template declares once (the template's ranges/itervars equivalent). The stride-match extent inference is unreliable for conv because the gemm reads a REPACKED weight (O_C contiguous, stride 1) whose strides do not match the declared node_layout, so the inferred extent is wrong (e.g. O_C=64 read as 8); gemm (no loop_extents) keeps the inference. Only a dim whose index is a single iv is clamped -- an im2col spatial axis (o_h + k_h) has no single base and is left unclamped, which is correct because that input is pre-padded. All four conv templates (base, mt, sb, sbs) declare loop_extents; note the mt template fuses channel and kernel-width so its tile_k loops 0..I_C*K_W. - dma_fine_grained now PROPAGATES the masked (and indirect / accumulate) attrs + the low/high operands when it splits a tile into subtiles, and REMAPS each clamp per subtile: a subtile at offset iv*sub sees high = min(sub, high - iv*sub), low = max(0, low - iv*sub) (out-of-window subtiles get a negative high / low > sub and Spike skips them). Previously it dropped the clamp, so the conv clamp never reached Spike; gemm was unaffected because its DMAs do not go through this pass. Fixes conv2d on non-dividing in_channels / kernels (test_conv2d 13/13), including fused pointwise epilogues (conv + relu/scale/bias/sigmoid). No regression: matmul/addmm, bmm, elementwise, reduce, softmax, index_add all pass. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

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. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

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. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

YWHyuk and others added 2 commits July 6, 2026 15:35

YWHyuk force-pushed the feature/transfer-direct-lowering branch from 032a0e8 to 9884d20 Compare July 6, 2026 06:35

YWHyuk commented Jul 6, 2026

View reviewed changes

YWHyuk and others added 3 commits July 6, 2026 15:57

YWHyuk force-pushed the feature/transfer-direct-lowering branch 3 times, most recently from ceb7e86 to f510f40 Compare July 6, 2026 14:28

YWHyuk commented Jul 6, 2026

View reviewed changes

YWHyuk force-pushed the feature/transfer-direct-lowering branch 12 times, most recently from c9a263e to d996aba Compare July 7, 2026 05:01

YWHyuk added 5 commits July 7, 2026 14:22

[Build] Bump spike to v1.0.5 (masked-DMA clamp + MVOUT accumulate)

debd1e8

The masked-DMA clamp needs spike to read the descriptor dim_low/dim_high + desc_fill, and index_add needs the MVOUT accumulate flag. Pin v1.0.5. Claude-Session: https://claude.ai/code/session_01EEfyUDpkMLRYZ2NAMbb3jN

YWHyuk force-pushed the feature/transfer-direct-lowering branch from 1fd8af4 to 4f8221d Compare July 7, 2026 05:22

YWHyuk merged commit e6633a9 into feature/indirect-offset-descriptor Jul 7, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Frontend] Lower togsim.transfer directly to Gemmini (drop memref.dma_start)#286

[Frontend] Lower togsim.transfer directly to Gemmini (drop memref.dma_start)#286
YWHyuk merged 11 commits into
feature/indirect-offset-descriptorfrom
feature/transfer-direct-lowering

YWHyuk commented Jul 6, 2026

Uh oh!

Uh oh!

YWHyuk Jul 6, 2026

Uh oh!

YWHyuk Jul 6, 2026

Uh oh!

YWHyuk Jul 6, 2026

Uh oh!

YWHyuk Jul 6, 2026

Uh oh!

YWHyuk Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

YWHyuk commented Jul 6, 2026

What changed

Validation

Known-pending (documented in the commit + docs/design/transfer-direct-lowering.md)

Uh oh!

Uh oh!

YWHyuk Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

YWHyuk Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

YWHyuk Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

YWHyuk Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

YWHyuk Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Known-pending (documented in the commit + `docs/design/transfer-direct-lowering.md`)