[SunMMIO] Enable MLIR Codegen and JIT Pipeline#239
Open
xiaoyao-NKU wants to merge 89 commits into
Open
Conversation
* feat: initialize sunmmio codegen and MLIR builder infrastructure * build: update sunmmio source file paths in CMakeLists.txt
* feat: initialize sunmmio codegen and MLIR builder infrastructure * build: update sunmmio source file paths in CMakeLists.txt * style: unify formatting and indentation
…cts in CMakeLists.txt
2。codegen_sunmmio里，增加必要的获取scope函数，注释了非法的alloc入口，对未解决的memory相关的接口，增加了log说明 Co-authored-by: Copilot <copilot@github.com>
…focus SunMMIO skeleton tests on traversal coverage completeness (#138) * [SunMMIO] Refocus skeleton coverage tests with frontend kernel * [SunMMIO] Remove text builder backend and enforce SUVM path * [SunMMIO] Apply pre-commit formatting fixes
…s and conditional statements (#144)
… pre-lowered TIR nodes that should not reach backend codegen. (#145) * WIP: local sunmmio codegen updates before upstream rebase * [SunMMIO] Enforce generic codegen invariants on pre-lowered TIR nodes * [SunMMIO] Treat DeclBuffer as benign wrapper in generic codegen
…Gen) (#158) * 1. 增加新pass hoist_block_annotations_to_func_attr,功能是将block中的信息保存到func_attr中。 1. 合并规则：map类型合并，其他类型以首次为准 2. 增加了kDeviceFuncPropagatedAttrKeys，会将指定的信息保存到device_function 3. pass位置放在HoistNonRestrictParams之后 2. 修改了pass split_host_device,核心修改是处理传入到device_function信息的buffer/var的remap问题 1. 依据func_attr中的kDeviceFuncPropagatedAttrKeys设置，将指定信息存入device_function 2. 由于spilt过程中会新生成buffer/var，所以需要将属性中对应的信息进行remap 1. 基本类型 buffer/var/PrimExpr 直接remap 2. list,map中第一层，如果存在基础类型，也remap 3. 特殊类型，指定key,特殊处理。目前针对 layout_map/globel_layout_map进行专门处理 4. 其他类型保持原样 3. codegen新增了接受layout信息与查询的相关逻辑 1. 在addfunction，获取layout信息放入context 2. 在context中实现了对layout信息的基础操作，包括 push,pop,lookup,apply_to_type 4. codegen完善了memtesor的构建流程 1. 函数参数，会在addfunction的时候，查询到layout信息，构建SunMMIOValue，实现bindVar, 后续在begin_function实现bindMlir 2. 普通的allocate,会在allocateNode的时候，查询到layout信息，EmitAlloc实现bindVar， 在memory的alloc实现bindMlir 3. 使用LookupBuffer 和 LookupMLIRValue 即可实现从tir到mlir的整体查询 * Remove Chinese comments.
* 实现了expr中大部分的功能，还缺少Tile相关的Binary、Broadcast和Ramp * codegen call part * refine region op * fix bug * remove old placeholder * add a test file
* [WIP] codegen tiles * feat: checkpoint sunmmio codegen upstream migration * fix: avoid tail tile predicate var reuse * feat(sunmmio): advance tiles loop and reduce codegen - implement structured tiles loop lowering with tail-mask handling, 1D aligned-load support, and expanded tile op coverage - add reduce-aware tile lowering with register-style intermediates and in-tile reduce plumbing - add initial mlir_codegen smoke tests for tiles, reduce, tile ops, and fill --------- Co-authored-by: Xiaoyao <39077302+xiaoyao-NKU@users.noreply.github.com>
2.新增了对于codegen中allocate,copy，标量的expr的测试 3.临时方式：暂时支持dma_copy和mma的新方案(多了一个参数) 4.其他需要和主分支保持一致的改动
* pileline test * fix format
* implement the critical path schedule algorithm * implementation of inject pipeline * rebase * rebase * code style * refine the algorithm to avoid the read-write loop * Memoization for b_level * optimize algorithm with 3 phases, refactor the sunmmio_pipeline_planning.cc * disable log * update algorithm * for pipeline algorithm * move debug control to python * for pull request * for ci * update pipeline algorithm * withdraw changes in fill.cc * for ci * for ci * code style * for stash * add broadcast_ in ast traverser * 完善了pipeline的逻辑 * 重构了pipeline，等待flatten sunmmio layout transform * update pipeline * ready for pipeline * for ci * ci debug * ci debug * revert sunmmio ci * for ci * add flashmladecode
* fix bug for pipeline * use 1d copy
* fix bug for pipeline * use 1d copy * fix bug for tiled_dims in copy * fix bug
* 新增支持alloc_var操作，包括支持其在 常规位置和for,if,while,tiles中出现。 在tir中alloc_var是作为buffer处理，scope='local.var' 在codegen中视为一般常量处理 新增对应的测试 * 更新测试文件，更贴近实际的分布式gemm的kernel方式。
* [WIP]: tiles loop and reduce codegen: blocked by dialect progress and row-major layout issue * [Reduce] Express tail masking with if_then_else * [Codegen] Lower SunMMIO tile slices with real ops * codegen if_then_else reduce without predicate * codegen: support hybrid tile reduce scopes * codegen: support rank-one broadcast regions * fix(sunmmio): stabilize codegen target tests * [sunmmio] support tail reduce regions Allow SunMMIO reduce operators to accept BufferRegion inputs and outputs, so kernels can pass dynamic valid extents for tail tiles while keeping canonical shared-memory buffer shapes. Lower rank-1-to-scalar reduce finalization as a tile-domain update instead of a scalar store. This keeps singleton RSRAM result buffers in the tile path, allowing the aligned 1D bridge to load, update, and store them correctly. Keep singleton 1D T.fill in tile form for aligned RSRAM buffers, teach layout_transform codegen to accept rank-1 singleton regions, and move aligned 1D access sizing from hard-coded 64 bytes to the target RSRAM alignment config. Extend tiles-loop lowering for aligned 1D read-modify-write stores with explicit masks, nested/hybrid tile-domain scopes, and updated binary operand shaping now that NPU-IR supports broader broadcast forms. Update SunMMIO target tests to use MeshTensor plus explicit DRAM layouts, and add coverage for tail-region reduce, small 1D reduce results, aligned row-major padding, loop-copy tiles, and layout-aware fill/tile-op cases. * codegen: support aligned rank1 side-buffer copies Preserve rank-1 aligned row-major covered extents when deriving layouts, use covered extents for full matched-layout rank-1 SunMMIO DMA copies, and allow tiled side-buffer accesses with a dynamic base offset. This keeps 1D side buffers on the aligned bridge path instead of falling back to small unaligned tile loads. * test: relax aligned store offset assertion Accept dynamic SSA offsets for suvm.tile.insert_slice as long as the aligned 1D store still writes an 8-lane slice at a nonzero offset. --------- Co-authored-by: Xiaoyao <39077302+xiaoyao-NKU@users.noreply.github.com>
* 打通 alloc ping-pong 属性从pipeline到 codegen的通路。 新增内部 TIR alloc ping-pong 属性，并通过 LowerOpaqueBlock 传递到 SunMMIO codegen。仅支持asram,wsram指定ping pang. * 先临时关闭mlir的简化，简化后会丢信息。
* Fix SunMMIO function-local token/barrier state * Fail fast on invalid SunMMIO MLIR modules * [SunMMIO] Fail fast on missing MLIR SSA values * [SunMMIO] Fail fast on unsupported call lowering * [SunMMIO] Fail fast on unsupported type mapping * [SunMMIO] Fail fast on generic memory and vector lowering * [SunMMIO] Replace fake operand lookup with strict lookup * [SunMMIO] Fail fast on fake tile placeholder lowering * [SunMMIO] Fail fast on unbound TIR vars * sunmmio: remove redundant dialect loads * Tighten SunMMIO codegen validation * Simplify SunMMIO BeginFor overloads * Reorganize SunMMIO tests under dedicated package
* [Build] Add Sunmmio toolchain setup * [Runtime] Add Sunmmio JIT adapter plumbing * [Transform] Preserve Sunmmio MeshTensor ABI metadata * [Codegen] Fix Sunmmio tiled layout and async waits * [Test] Add Sunmmio JIT examples and coverage * [CI] Scope Sunmmio tests for codegen lane * [Runtime] Update Sunmmio toolchain target * [Runtime] Clean up Sunmmio JIT cache metadata * [CI] Fix Sunmmio clang-tidy warnings * Improve CI workflow * [SunMMIO] Stabilize JIT cache and reduce tests * [SunMMIO] Update NPU-IR dependency * [SunMMIO] Normalize region view indices in codegen * [SunMMIO] Bind dynamic layout shape values
* test(sunmmio): update reduce layout assertions * fix(sunmmio): carry tile cache and DMA tokens through loops * fix(sunmmio): register reduce tile temps without handles * fix(sunmmio): handle pr223 gqa decode tile lowering * fix(sunmmio): harden tile lowering validation
* 增加 MXFP4 和 MXFP8 数据类型支持，覆盖前端 dtype 解析、布局构造、布局推导以及 SUVM/MLIR 类型映射。 主要改动包括： - 新增 mxfp4/mxfp8 dtype 别名与 dtype 规范化逻辑 - 增加 MX 数据格式相关布局构造与布局导出支持 - 在 SunMMIO 类型映射中支持 MXFP4/MXFP8 以及相关 float4/float8 类型 - 更新 copy、gemm 等路径中的 dtype 与 layout 处理 - 补充 MX dtype codegen 和外部 MX layout 的验证测试 * 更改中文注释
* Add MeshTensor placement shape API * Rename MeshTensor shape query to global_shape * fix 3 bugs --------- Co-authored-by: Xiaoyao <39077302+xiaoyao-NKU@users.noreply.github.com>
* Add MeshTensor placement shape API * Rename MeshTensor shape query to global_shape * fix 3 bugs * Add symbolic Sunmmio mesh intrinsics * Default Sunmmio mesh symbols in DSL * Test Sunmmio mesh symbol lowering * Allow mesh-symbolic comm and GEMM shapes * Restrict Sunmmio kernel blocks to mesh symbols --------- Co-authored-by: Xiaoyao <39077302+xiaoyao-NKU@users.noreply.github.com>
* fix for pre-commit * test: consolidate Sunmmio Python tests under testing/python/sunmmio * test(sunmmio): update kernels to threadless mesh launch style * Use symbolic mesh dims in SunMMIO kernels * fix sunmmio test
Advance the NPU-IR submodule from 78208a2e to c8c6c13 on NPU-IR main, picking up the ODMA and link-sync lowering fixes: - Make ODMA channel assignment transform-function aware (#140) - Reuse one ODMA descriptor block per channel (#161) - Correct A4E ODMA DMA alignment/min-size per A4E spec (#164) - Rename DescPoolCapacity to OdmaDescPoolCapacity (#165) - Fix link sync wait callees (#166) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* fix for pre-commit * test: consolidate Sunmmio Python tests under testing/python/sunmmio * test(sunmmio): update kernels to threadless mesh launch style * Use symbolic mesh dims in SunMMIO kernels * fix sunmmio test * Fix SunMMIO loop-carried wait handling * Fix dynamic SunMMIO barrier mask lowering
… bugs (#246) * [Sunmmio] Reconcile JIT ABI args by name and run online_softmax end-to-end Make the online_softmax example compile and execute through the Sunmmio JIT + sunsim path, and fix the supporting infrastructure it exposed. ABI (abi.py, adapter.py, kernel_cache.py): - The device kernel may order parameters differently from the public Python signature. Track public_param_names and reconcile public args to device order by name in materialize_runtime_args, instead of assuming a positional [public..., scalar...] layout. - Derive device param names from the host kernel call when the device func carries none, and treat any device name absent from the public set as a hidden runtime scalar. - Serialize the full ABI (to_json_dict/from_json_dict) into the cache metadata and restore it on load, so cached kernels reorder arguments correctly. Adapters now accept a pre-built abi. - Always route sunsim args through materialize_runtime_args (drop the full-arg-count short-circuit) so reordering also applies when explicit scalars are passed. Codegen (sunmmio_codegen_tiles_loop.cc, codegen_sunmmio.cc): - Lower if/then/else inside T.Tiles (BeginIf/BeginElse/EndIf with tile-state save/restore). - Handle tl.vector_core_in_tile_reduce, tl.barrier_init and tl.barrier_arrive_and_wait as Evaluate statements in the tiles scope; forward-declare the reduce lambda. - Stop double-registering reduce-register temp buffers and Sunmmio local-var buffers. Sync injection (inject_sunmmio_sync.cc): - Skip re-emitting loop-entry null tokens already initialized by an enclosing loop. - Inherit available tokens into nested SeqStmt scopes and drop pending tokens from parent scopes, so output-DMA waits hoist correctly. JIT (jit/__init__.py): - Apply the resolved target context while inferring lazy/eager mode. Tooling (libgen.py, sunmmio_mlir_function.cc): - Add -suvm-resolve-barriers, -suvm-canonicalize and -suvm-tile-split-unroll to the default npuir-opt pipeline. - Dump the invalid module on MLIR verification failure when TL_SUNMMIO_DUMP_INVALID_MLIR=1. Example/tests: - Rewrite online_softmax.py as a jit-compiled kernel run on sunsim with a numpy reference check. - Add tests for name-based ABI reordering, full ABI cache round-trip, and output-DMA wait hoisting. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [Sunmmio] Resolve mesh symbols in tile annotations and tensor_meta ResolveSunmmioMeshSymbols replaced tl.mesh_nrows/ncols/ncores with the concrete mesh dims in expressions and buffer shapes, but the default StmtExprMutator does not descend into For/Block `annotations` or the tensor_meta func attr. So a loop like for i in T.serial(mesh_ncols(), annotations={"tile.domain": [mesh_ncols(), 128], ...}) had its extent resolved but kept the symbolic `tile.domain` annotation. Codegen reads that annotation directly (Downcast<Array<PrimExpr>>(op->annotations.at("tile.domain")) then EvalExpr), so the unresolved mesh_ncols reached codegen and crashed with "Unsupported SunMMIO call lowering for tl.mesh_ncols". Override VisitStmt_(For) and augment the Block visitor to resolve PrimExprs inside annotations, and extend ResolveTensorMeta to resolve non-Layout PrimExpr/Array<PrimExpr> values (local_shape, mesh_shape, strides). Annotation values are ffi::Any, whose strict `as<>` misses type-erased containers, so the array case uses try_cast<Array<PrimExpr>>. Note: the pass still only covers known positions (not layout_map attrs, AttrStmt nodes, or deeply nested annotation containers); a fail-fast post-condition is a follow-up. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [Example] Run online_softmax on Sunmmio with 256x256 tiles Use a 256x256 block so the horizontal all-gather's mcast_tok transfer is 1KB-aligned: the multicast tile_view is block_M-sized, and 256 f32 = 1KB, which the NPU-IR lowering requires (block_N does not affect it). Raise the sunsim watchdog to 1200s for larger shapes and print a PASS line after the numpy reference check. Validated: 1024x1024 passes end-to-end on sunsim. 2048x2048 compiles and runs but deadlocks in the multi-tile (2x2) all-gather/sync path; tracked separately. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [JIT][Sunmmio] Enforce ABI invariants and document reconciliation The Sunmmio kernel ABI reconciles the device parameter order (which SplitHostDevice sorts by name) against the public signature by name, and derives hidden scalar dims from tensor shape/stride metadata. Several invariants the design relied on were unchecked, so violations surfaced per-call or silently rather than at construction. Enforce them at construction: - SunmmioRuntimeScalar rejects a partial source binding (all three source fields, or none) so is_bound is a true binary. - SunmmioKernelABI checks that public names match public_arg_count, public and scalar names are each unique and disjoint, and device_param_names is a permutation of public + scalars. Runs on both from_modules and from_json_dict. - Document the explicit-scalar ordering contract in materialize_runtime_args, and add stage comments to from_modules. Make stale caches self-heal: _load_abi_metadata treats a ValueError from from_json_dict as a cache miss and recompiles, mirroring the existing missing-keys handling. Add negative tests for each validator and the inconsistent-metadata round trip. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [JIT][Sunmmio] Lower via npuir-compile instead of hand-listed passes Production lowering ran npuir-opt with a hand-listed 14-pass pipeline, then npuir-translate, then a Python GEP-flag rewrite for the LLVM-17 zpu clang. That pass list had already drifted from NPU-IR's own pipeline and was brittle to maintain. Bump NPU-IR to ee880fd (#155), which adds npuir-compile: a one-shot driver that runs the registered suvm-to-llvm-pipeline, translates to LLVM IR, and strips the LLVM-19+ GEP flags in-process. Migrate materialize_llvm_ir to a single `npuir-compile --target=<mcpu> --emit=llvm-ir` call and drop DEFAULT_NPUIR_OPT_ARGS, the opt/translate helpers, and the GEP sanitizer. - NpuirTools now exposes only `compile`; the binary locator is lifted to a module-level find_npuir_tool so the codegen-validation tests can still resolve npuir-opt without the production tool object. - CMake builds and installs npuir-compile alongside npuir-opt/translate. Note: npuir-compile's pipeline includes a SUVMDeviceValidation pass the old hand-listed pipeline omitted. It rejects the 1D suvm.tile values our softmax codegen emits (a4e requires 2D tiles), so 1D-tile kernels now fail to lower until the codegen emits 2D tiles. Simple/2D kernels are unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [Test] Fix softmax JIT test path and output-DMA-wait assertion The online_softmax example loader resolved the path relative to `testing/` (parents[3]) instead of the repo root, causing a FileNotFoundError in CI; the elementwise loaders already use parents[4]. Fixing the path unmasked a stale assertion: it forbade a sync_null_token for the output token anywhere between the outer loop and the output DMA, but the output-DMA pipeline emits a legitimate prologue null-init before the by-loop. Scope that check to the pipelined loop body so it still catches a null-reset that would make the hoisted wait vacuous, while allowing the prologue init. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Update NPU-IR --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
* Revert "Revert "[Feat]Validate Sunmmio T.copy region extents (#197)" (#204)" This reverts commit 937f47e. * [feat] Add copy tile-view validation pipeline integration * [feat] Add copy tile-view validation to Sunmmio pipeline * [fix bug] fix format bugs * restore dynamic bufferload shrink * fix wrong shape * restore the operation for dynamic bufferLoad * update tile view validation pass * update copy frontend * remove dynamic extent fail and add test case * fix dynamic mins validation * rename validate_copy_tile_view -> validate_tile_view_regions and generalized it to fit comm op * fix tile view region validation * Restore NPU-IR submodule pointer * reject non-positive tile view extents * fix sunmmio copy validation tests
* Using SizeVar to define mesh symbols instead built-in op * Harden Sunmmio mesh symbol resolution
2. 修改mxfp4,mxfp7的对齐规则，至少2 3. 针对layout转换，目前对mx类型，也是仅支持zz<->row_major 4. 新增了对应的测试
* feat: support 1d sunmmio tile pick codegen * fix: extend hybrid tile loop codegen * fix: handle predicated 1d tile stores
* Fix SUNMMIO sync analysis for loop-covered buffer accesses * sunmmio: gate legacy barrier wait after token wait * Add loop token reuse wait injection * Fix decl_buffer construction in mesh symbol test
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
7 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.
Motivation
This PR brings the SunMMIO backend to an end-to-end usable state for Tilelang programs. The existing branch had compiler pipeline pieces, but still needed a unified MLIR/NPU-IR codegen path, JIT runtime integration, broader operator lowering coverage, and validation tests before merging into
tilelang_mesh_main.Solution
src/target/sunmmioand replace the legacy codegen path with the MLIR-based backend.testing/python/sunmmio.