[Build] Ship MLIR Python bindings in the LLVM artifact; add design docs - #256
Closed
YWHyuk wants to merge 26 commits into
Closed
[Build] Ship MLIR Python bindings in the LLVM artifact; add design docs#256YWHyuk wants to merge 26 commits into
YWHyuk wants to merge 26 commits into
Conversation
The llvm-project fork CI now builds the MLIR Python bindings into the riscv-llvm release (v1.0.9). This wires PyTorchSim to consume them so MLIR passes can be written in Python (imperative IR rewriting via the bindings) instead of only as C++ passes in the fork. - Dockerfile.base: PYTHONPATH -> /riscv-llvm/python_packages/mlir_core so `import mlir` works in the container. - github-releases.json: bump llvm_project pin v1.0.8 -> v1.0.9 (first artifact with bindings); triggers a base image rebuild. - build_from_source.sh: enable MLIR_ENABLE_BINDINGS_PYTHON, build against the runtime Python (3.11), pin pybind11 <= 2.10.3 (this fork's bindings use pybind11), and copy python_packages into the install tree. - docs/: mlir-python-bindings (this setup + rollout), dma-transfer-lowering (decompose DMA into a loop of affine descriptors to retire the heuristic recompile/tile-forcing dance), linalg-codegen-migration (deferred full structured-ops rewrite). Verified: the v1.0.9 artifact's bindings import under the runtime conda 3.11 and can parse/rewrite custom ops (togsim.transfer with floordiv/mod maps). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the arith.addi-with-vlane_offset-attribute hack (rewritten by the C++ -global-idx pass) with a dedicated torchsim.vlane_idx op, lowered by a Python out-of-line MLIR pass (mlir.ir bindings) to (vcix.v.i per-lane index * offset). - mlir_ops.py: vlane_offset handler emits "torchsim.vlane_idx" (generic form). - mlir/passes/: run_python_passes orchestrator (parse once / run passes on the shared Module / print once, with a marker fast-path) + the lower_vlane_idx pass. Add future Python passes to PASSES. - extension_codecache.py: run the Python passes on the kernel .mlir before mlir-opt; drop -global-idx from both mlir-opt pipelines. Depends on the riscv-llvm v1.0.9 artifact shipping the MLIR Python bindings, VCIX dialect registration, and the DmaStartOp print/parse fix, so vcix and the customized memref.dma_start round-trip through the bindings. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…assManager
Split the mlir-opt invocation after memref-to-gemmini: the custom passes
(test-loop-padding, dma-fine-grained, test-pytorchsim-to-vcix,
test-tile-operation-graph, test-memref-to-gemmini) still run in mlir-opt, but
the standard upstream lowering (convert-*-to-llvm, lower-affine,
reconcile-unrealized-casts, ...) now runs in-process through the MLIR Python
bindings' PassManager. A step toward an all-in-process flow as the custom passes
are migrated to Python.
- passes/lower_to_llvm.py: run_standard_lowering(). Only lower-vector-multi-reduction
is func.func-scoped and is nested explicitly, since the bindings pass-pipeline
parser does not auto-nest like the mlir-opt CLI; pass order is preserved.
- extension_codecache.py: mlir-opt now writes the post-custom IR to
{name}_custom.mlir; run_standard_lowering produces the LLVM-dialect
{name}_llvm.mlir consumed by mlir-translate. (Drops the standard passes from
both mlir-opt pipelines.)
Validated to produce byte-identical LLVM IR to the previous all-mlir-opt pipeline
and end-to-end (test_add and an arange/vlane_idx kernel pass under
gem5/spike/TOGSim). This makes the MLIR Python bindings a required runtime
dependency of every compile (satisfied by the riscv-llvm v1.0.9 artifact).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
YWHyuk
force-pushed
the
feature/mlir-python-bindings
branch
from
June 15, 2026 13:10
9e449c8 to
47fe19d
Compare
- Decide decomposition by affine rank after linearizing floordiv/mod via split (D<=4 -> one descriptor, D>4 -> peel an outer affine.for), not by the presence of floordiv/mod. Genuinely non-affine (data-dependent/indirect) access is out of scope and stays on the indirect-indexing path. Maps onto the existing apply_divisor/get_dma_info >4D site. - Add the memref-to-gemmini boundary: decompose-transfer stops at memref.dma_start; Gemmini ISA encoding stays in the C++ test-memref-to-gemmini pass (separation of concerns). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…memref-to-gemmini
Port the C++ test-memref-to-gemmini conversion to a Python out-of-line MLIR pass
(passes/lower_dma_to_gemmini.py) and drop -test-memref-to-gemmini from both
mlir-opt pipelines. The pass works at the memref level -- addresses via
memref.extract_aligned_pointer_as_index + arith, Gemmini instructions as
llvm.inline_asm (.insn r CUSTOM_1 ...) -- so it avoids the C++ conversion
framework (LLVMTypeConverter / getStridedElementPtr / MemRefDescriptor); the
existing standard lowering finalizes everything to LLVM.
- Timing semantics preserved: the functional/Spike path emits gemmini config +
mvin/mvout asm; the gem5 cycle path (run_standard_lowering timing=True) erases
dma_start (the TOG already carries DMA timing). dma_wait is erased in both.
This matches the old test-memref-to-gemmini timing=1 behavior.
- Indirect access (gather/scatter): CONFIG4 + the indirect bit in CONFIG; the
index-spad base address is taken via extract_aligned_pointer_as_index after
tracing affine.apply{indirect_access} -> index_cast -> affine.load.
- run_standard_lowering runs this pass (after the custom mlir-opt passes) then
the standard MLIR->LLVM PassManager pipeline.
Validated end-to-end (gem5/Spike/TOGSim allclose) on add, matmul, conv2d,
layernorm, softmax, indirect_access (gather + scatter), and an arange/vlane_idx
kernel; the config-instruction constants are byte-identical to the C++ pass.
docs/dma-transfer-lowering.md: gemmini ISA lowering is now this Python pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The Python MLIR passes (run_python_passes / run_standard_lowering) are now a hard dependency of every compile, but a plain local run may not export PYTHONPATH to the bindings. Derive the bindings dir (python_packages/mlir_core) from TORCHSIM_LLVM_PATH and prepend it to sys.path when `import mlir` would otherwise fail. No-op when PYTHONPATH already provides it (the container/CI case). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Checkpoint before axis-split scheduling prototype. Will be regrouped later. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Removes aligned FloorDiv/ModularIndexing from index expressions before MLIR codegen by splitting loop axes at the Inductor scheduling layer, reusing the LoopBody rebuild machinery (same as revert_group). Env-gated by TORCHSIM_AXIS_SPLIT (dump via TORCHSIM_DEBUG_AXIS_SPLIT). axis_split.py: find_split_plan detects FloorDiv/ModularIndexing on a single iter var whose divisor divides the extent; build_split_body rebuilds the body with v = outer*k + inner so the floor/mod collapses. Validated on group_norm: idx1 = 3*p0 + (p1//2) -> 3*s0 + (s1//1) i.e. FloorDiv eliminated, mean access affine. Known issues (5D blow-up from raw size, ModularIndexing under-split) and full coverage classification recorded in docs/axis-split-scheduling.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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>
build_split_body now reindexes the existing (collapsed/reordered) node._body via LoopBody's copy path (pass the body as fn -> _init_with_copy) instead of re-tracing the raw store function over the un-collapsed size. This keeps already-merged dims merged (no rank blow-up: the prior approach un-collapsed spatial and produced a 5D tile that tripped the <=4D init_tile_size) and lets simplify_with_ranges fold the split floor/mod: v -> outer*k + inner makes FloorDiv(v, k) collapse to the outer axis. indexing_from_args requires exactly one replacement expr per original var (index dims then reduce dims); reduction dims pass through unchanged. Validated on group_norm(num_groups=3): the normalize kernel goes (2,6,16) -> (2,3,2,16) (stays 4D), idx1 = 3*p0 + (p1//2) -> 3*s0 + s1 (the channel FloorDiv is eliminated), and the run is allclose=True end-to-end (Gem5 + Spike + TOGSim). Gated behind TORCHSIM_AXIS_SPLIT; default path unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…dims Build the new split symbols with torch._inductor.utils.sympy_index_symbol (integer, non-negative) instead of bare sympy.Symbol, so simplify_with_ranges actually folds the split floor: idx1 = 3*p0 + (p1//2) now becomes 3*z0 + z1 instead of leaving a 3*z0 + (z1//1) residue. sympy_index_symbol forbids names starting with s (reserved for shape symbols), so index dims use the z prefix. Reduction dims use the r prefix and stay after the index dims so the reduction axis remains innermost (var_ranges ordered iter-then-reduce; LoopBody.sizes splits on len(iter_vars)). LoopBody var names are remapped to index<N> in MLIR codegen, so the prefix is internal but must not collide with the original body names (p/q), which z/r do not. group_norm(num_groups=3) stays allclose=True end-to-end. Reduction-dim split path is convention-correct but not yet exercised (no available test splits an index dim of a reduction kernel). Gated behind TORCHSIM_AXIS_SPLIT. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…radix
Generalize the aligned split from a single FloorDiv divisor to a mixed-radix
boundary chain so ModularIndexing and multiple radices on one axis linearize:
- find_split_plan now returns {axis: boundaries}, an ascending divisibility
chain [1, b1, ..., E] of cut points gathered from the axis terms:
FloorDiv(v,k) -> boundary k; ModularIndexing(v,k,m) -> boundaries k and k*m.
If the boundaries do not form a divisibility chain (incompatible radices,
e.g. floor-by-2 and mod-by-3 on extent 6), the axis is left unsplit.
- build_split_body splits each planned axis into one sub-var per segment
(v = sum_i d_i*b_i), most-significant outermost. FloorDiv/ModularIndexing on
the axis then collapse to affine combinations of the sub-vars.
Also fix a decompose-transfer peel bug surfaced once axis-split makes the peel
path reachable: operandSegmentSizes on memref.subview must be DenseI32ArrayAttr
([1,0,0,0]); the i64 version silently zeroed to [0,0,0,0] and failed
verification (only caught now because the isolation test did parse/print, not
mlir-opt verification).
Validated end-to-end (allclose=True): group_norm (FloorDiv, chain [1,2,6]) and
x.repeat(1,2) (single-axis ModularIndexing, chain [1,8,16]) -> floor/mod fully
eliminated. pixel_shuffle (floor+mod on two axes) linearizes correctly too,
though its 5D tile then exercises high-rank TOG serialization (separate issue).
Gated behind TORCHSIM_AXIS_SPLIT.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Reduction dims are carried through the reindex unchanged (r prefix, kept innermost after the index dims). No natural op produces a floor/mod on a reduction kernel's index axis, so add a TORCHSIM_AXIS_SPLIT_FORCE validation gate: force-split the first even index axis of a reduction kernel even without a floor. A floor-free index split is an identity transform, so allclose must hold -- this exercises the reduce pass-through on a real reduction body. Validated: layernorm (512)->(256,2) and reduce (68)->(34,2) keep their reduction groups (r0 innermost) and pass allclose. Off by default; normal axis-split (TORCHSIM_AXIS_SPLIT) unaffected. docs/axis-split-scheduling.md: record mixed-radix + reduction validation, the incompatible-radices and high-rank-blow-up limitations. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
v1.0.9's artifact shipped the MLIR python bindings as dangling symlinks into the build tree, so import mlir.ir failed at runtime. v1.0.10 is built with cp -rL so the bindings are real files. Bumping the pin also changes the thirdparty base-image PIN, forcing CI to rebuild the base image from the fixed artifact. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
axis_split.ledger(nodes, plan) classifies every FloorDiv/ModularIndexing in a kernel against the split plan and reports the ones axis-split cannot cover, by reason: multi_axis_arg (case 7), non_dividing (case 6), incompatible_radix (case 5), dynamic. Wired into codegen_node behind TORCHSIM_AXIS_LEDGER (prints [AXIS_LEDGER] lines); independent of TORCHSIM_AXIS_SPLIT and behavior-neutral. Used to measure how often the graph-copy cases actually reach codegen across models, so we can decide whether graph-level copy insertion is worth building. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…istic The DeepSeek V3 base test only seeded input_ids; the model weights from from_config used the unseeded global RNG, so every run built a different network. The NPU-vs-CPU worst-element error sits near the (loose) allclose threshold, so it randomly crossed it and the test was flaky. Seed the global RNG before model construction to make runs reproducible. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…floor fold Refactor find_split_plan onto collect_boundaries() + _is_chain() (shared with graph-copy). Add a rank guard that skips a split which would push the index rank past 4 (the >4D peel is not yet numerically correct, so pixel_shuffle falls back to baseline). Add _fold_with_ranges to fold residual FloorDiv/ModularIndexing that simplify_with_ranges misses on a multi-level (>=3) mixed-radix split, proving the bound from the split sub-var ranges via bound_sympy. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…is floor/mod New module gated by TORCHSIM_GRAPH_COPY: wrap the registered lowering entries (so every elementwise consumer is one hook); for each consumer trace the operands' loaders via extract_read_writes and detect either (case 5) two operands with incompatible-radix groupings on a shared axis (a[c//2]+b[c%3]) or (case 7) an operand whose floor/mod argument spans multiple axes ((3*p0+p1)//4 from a transpose+reshape feeding a broadcast/softmax). Replace the cheaper operand with ExternKernel.copy_input (a realized identity Pointwise -- materializes views too, unlike StorageBox.realize() which is a no-op on a ReinterpretView). The consumer then reads it affine and the remaining single grouping is axis-split's job. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…DRAM offset) The >4D peel advanced the per-slice DRAM base with arith.addi(dram_idx, const), which the TOG pass's processDramIndices cannot read (it handles affine.apply / block-arg / constant only) -> empty loop_idx_list -> ONNX serialization failure. Fold the constant offset into affine.apply (d0)->(d0+const) over the original dram_idx instead; processDramIndices recurses through it. The crash is fixed; the peel's SRAM offset is still wrong for the lane-banked scratchpad, so the >4D path stays behind the axis-split rank guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…mpile-dance axis-split and graph-copy are ON by default (disable with TORCHSIM_AXIS_SPLIT=0 / TORCHSIM_GRAPH_COPY=0); wire graph_copy.install() at backend import. Add a TORCHSIM_RECOMPILE_LOG counter in codegen_nodes to measure what still depends on the recompile-dance. Validated default-on across 33 tests (elementwise/gemm/ reduce/conv/view/fusion + mlp/resnet/transformer/vit + cnn/pool/group_conv/sort/ indirect/exponent/conv_fusion): all pass, recompile fires only for the >4D rank-guard fallback (pixel_shuffle). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
tests/ops/view/test_floormod_axis_split.py: group_norm / repeat / repeat_interleave / permute+reshape (axis-split), 3-level mixed-radix, pixel_shuffle (rank-guard fallback), incompat (case 5), reshape+broadcast / softmax(reshape) / layernorm(reshape) (case 7). Self-enables the features (TORCHSIM_AXIS_SPLIT + TORCHSIM_GRAPH_COPY). Not in the CI allowlist (local feature/regression test). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…258 notes 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.
Summary
Make the MLIR Python bindings available so we can author MLIR passes in Python (imperative IR rewriting via
mlir.ir) instead of only as C++ passes in thePSAL-POSTECH/llvm-projectfork.The fork CI now builds the bindings into the
riscv-llvmrelease (v1.0.9); this PR wires PyTorchSim to consume them.Changes
PYTHONPATH -> /riscv-llvm/python_packages/mlir_core(preserves existingPYTHONPATH).llvm_projectpinv1.0.8 -> v1.0.9(first artifact shipping the bindings); triggers a base-image rebuild.MLIR_ENABLE_BINDINGS_PYTHON, build against the runtime Python (3.11), pinpybind11<=2.10.3(this fork's bindings use pybind11, newer fails adef_property/keep_alivestatic assert), copypython_packagesinto the install tree.mlir-python-bindings.md— this setup, cross-repo rollout, ABI notes.dma-transfer-lowering.md— proposal to decompose DMA into a loop of affine descriptors (retire the heuristic recompile/tile-forcing dance).linalg-codegen-migration.md— deferred full structured-ops (linalg) rewrite.Verification
import mlirworks under the runtime conda 3.11.13.togsim.transferop, read itsAffineMap(withfloordiv/mod), buildscf.for+affine.apply+ an inner descriptor op, erase the original, re-verify.Notes
.vscode/PYTHONPATH wiring is local-only (gitignored).cpidempotency fix lands with the next LLVM release; v1.0.9 is functionally correct.🤖 Generated with Claude Code