[Frontend] DMA transfer op + aligned-only decompose pass for >4D access - #257
Merged
Merged
Conversation
Phase 1 (emission only) of the DMA transfer-op plan. A DMA access whose logical tile exceeds the 4D Gemmini descriptor limit no longer hard-fails; it emits a high-level togsim.transfer op for a later decompose pass to peel into a loop of <=4D memref.dma_start. The decompose pass is deferred. init_tile_size in mlir_common.py is generalized to any rank by separating the logical tile from the physical (<=4D) descriptor: only the innermost dims carry the vectorized tile, all further-outer dims stay 1, no rank cap. The nr_dim>=3 formula reproduces the old 3D/4D values exactly, removing the "dummy tile size fail!" assertion that conflated logical and physical rank. mlir_codegen_backend.py: get_dma_info >4D branch builds the full N-D tile and sets _dma_needs_transfer; load()/store() emit togsim.transfer when the flag is set, otherwise the existing get_dma_code path is unchanged so aligned <=4D DMAs stay bit-identical. docs/dma-transfer-lowering.md: append the alignment-decomposability theory (aligned vs misaligned, modular valid iff y*z|extent, mixed-radix), the one-loop-axis -> several-implicit-axes generalization for complex fusion, a case-handling summary table, and the Phase 1 implementation status. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Scope decision: the decompose pass is a pure mechanical rank peel of an already-affine access. It no longer linearizes floor/mod and no longer does relayout. Those move upstream: aligned floor/mod is removed by axis splitting at the Inductor scheduling layer (axis-split-scheduling.md), misaligned access is resolved by graph-level copy insertion. The pass asserts (fail loud) on any non-affine residue instead of silently inserting a relayout, which would be a hidden perf cliff and a global layout decision made at the wrong layer. Adds a division-of-labor table and updates the contract, cost, placement, migration, risks, and deferred-work sections. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… path) Add passes/decompose_transfer.py: a Python out-of-line MLIR pass that lowers each togsim.transfer to a customized memref.dma_start, the aligned-only mechanical peel from the design doc. This first increment handles the unit-dim-collapse case (descriptor reaches <=4D once extent-1 tile dims are dropped); genuine >4 effective rank still raises NotImplementedError pending the affine.for peel loop. Mechanics: - Drop extent-1 tile dims. Collapse the SRAM spad memref to the effective rank via memref.collapse_shape (the customized dma_start convention requires SRAM rank == #indices == len(sram_stride)); DRAM stays flat rank-1 with its N-D structure in dram_stride. - Remap vlane_split_axis from the original tile-dim index to the collapsed-dim index and rematerialize the const. Supporting changes: - emit_transfer carries the SSA operands a dma_start needs (dma_type, vlane_stride) and the vlane_split_axis value as an attr (so the pass can remap it); operand prep mirrors get_dma_code for cache compatibility. - lower_to_llvm.py adds expand-strided-metadata to lower collapse_shape. - register decompose_transfer before lower_vlane_idx in passes/__init__.py. Validated end-to-end (Gem5 + Spike + TOGSim, allclose=True) on the 5D permute x.permute(4,3,2,1,0).contiguous() + 1.0; no regression on 2D/3D/elementwise add. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…bview Make the decompose pass total: when more than 4 non-unit tile dims survive, keep the inner 4 as the <=4D descriptor and peel the outer dims by full unrolling -- one customized memref.dma_start per outer-index combo, the SRAM slice a rank-reduced memref.subview at the static slice offset, the DRAM base dram_idx + constant. Unrolling keeps the slice offsets static so no per-iteration SRAM index arithmetic is needed; the vlane axis is remapped into the inner descriptor. Currently unreachable through the full pipeline: init_tile_size caps non-unit tile dims at 3 (effective rank <= 3 in practice), so this path is implemented for completeness / future tilings and validated only in isolation via lower_text on a synthetic 5-effective transfer (2 descriptors, correct subview offsets 0/24, dram offset +1, inner strides). The unit-collapse fast path and 2D/3D/5D pipeline runs are unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
1 participant
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.
Stacked on top of #256 (feature/mlir-python-bindings). Targets that branch, not develop.
What
Represents a DMA whose access exceeds the 4D Gemmini descriptor limit as a high-level togsim.transfer op, and decomposes it into <=4D memref.dma_start in a Python out-of-line MLIR pass. Removes the hard-fail at >4D in get_dma_info.
Pieces
Scope decision
The pass is a pure mechanical rank peel of already-affine access. It does NOT linearize floor/mod and does NOT relayout; those move upstream (aligned floor/mod -> axis split at the scheduling layer; misaligned -> graph copy insertion). The pass asserts on non-affine residue rather than silently relaying out. See docs/dma-transfer-lowering.md.
Validation
End-to-end (Gem5 + Spike + TOGSim, allclose=True) on 5D permute, 5D broadcast, 5D reduction. No regression on 2D/3D/elementwise add. The >4-effective peel validated in isolation via lower_text on a synthetic transfer.
🤖 Generated with Claude Code