Repository navigation
fix(lowering): issue a window that covers its operand as the operand itself - #273
Merged
Merged
Conversation
FeathBow
marked this pull request as draft
October 26, 2026 17:06
FeathBow
marked this pull request as ready for review
October 26, 2026 17:34
This was referenced Oct 27, 2026
The schedule corpus already checks, analyzes and finalizes hir/sm89_e4m3_gemm.py once. Its case now compiles the reloaded TIR in guard() instead of a separate test repeating those stages.
zhen8838
approved these changes
Oct 27, 2026
zhen8838
left a comment
Collaborator
There was a problem hiding this comment.
Thanks @FeathBow, approving.
Removing the whole-tensor windows in lowering is right: such a window is the operand itself, and the finalized sm_89 GEMM now compiles.
Reviewing this exposed two TileFoundry issues, which I filed separately. Neither blocks this PR:
- #274: CUDA codegen slices only tensors with a plain
Layout. A partial window of a sharded register tile, or of a swizzled smem tile, still fails.sparse_decodeis the swizzled case, not a sharded one as the Risk section says. - #275: the instruction projection is recomputed at several stages, and lowering computes it again instead of using what check produced.
About the test: test_schedule_corpus.py already collects hir/sm89_e4m3_gemm.py and checks, analyzes and finalizes it once, so test_finalize_compile.py repeated those stages through _finalized. I moved the compile into that fixture's corpus case (454a11a): Sm89E4m3Gemm.guard() compiles the reloaded TIR, which is also the path a finalize output file takes. When the kernel can run, the run can go in the same guard(), with a device check inside it.
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
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.
Why
What
schedule finalizewrites for the sm_89 GEMM doesn't compile. Lowering wraps the whole accumulator in a window that starts at zero and spans it,T.ptr_of(acc[0:0 + 16, 0:0 + 8]), and CUDA codegen refuses a pointer slice of a sharded register tile withCUDA pointer slices require a plain source layout. Past that, the MMA reads its register fragments aslhs_frag_dst[0:0 + 16, 0:0 + 32]andrhs_frag_dst[0:0 + 32, 0:0 + 8], and CUDA codegen has no handler for a bareSlice. Each of these windows is the whole operand.The fix belongs in lowering.
StoragePlan.operandalready treats a window that starts at zero and spans its operand as the operand itself, but materialization and the instruction issue loop still built aSlicefor it. MLIR'sTrivialSubViewOpFolderdoes the same thing formemref.subview. It replaces such a subview with its source, or casts the source when the types differ.What
ConvertHIRToTIRdecides once whether a window covers its operand, with static zero starts and sizes equal to the operand's shape, and uses that answer in all three places that build windows.StoragePlan.operandbehaves as before. Materialization takes the operand's own pointer for such a window. The instruction issue loop passes the operand itself when its type is the one the instruction wants, and otherwise takes its pointer and retypes it as before.For the sm_89 GEMM, finalize now writes
T.ptr_of(acc)and the MMA readslhs_frag_dstandrhs_frag_dstdirectly. Across the 49 schedule HIR fixtures that have a TIR twin, 29 finalize outputs change, all in this way, and noTensorViewlayout or shape changes. When I compiled those 29 for their own targets, the sm_89 GEMM andmesh_viewsnow compile. Two wgmma outputs get past the pointer slice refusal and stop at the next view, whereshard_engine_layoutdoesn't accept the nested mode(8, 8)in its shape. That view's layout is identical on main and is refused there in the same way, so the earlier refusal was only hiding it. For the rest, the first line of each error matches main, mostly TMA copies that need tensor maps encoded on the host. I haven't compared the full errors.The new test finalizes the sm_89 GEMM HIR fixture and compiles the result for
nvidia.rtx_6000_ada. It only needs nvcc. Without this change it fails with the pointer slice refusal. The CPU suite gives the same results before and after apart from that test.The schedule tutorial prints the start of the finalize output for its optimal GEMM. Its block now passes
acc_localto the MMA, and prints one line fewer so that it still ends on that MMA.Contract
docs/spec/passes.mdnow says, underConvertHIRToTIR, that a window which starts at zero and spans its operand's shape is the operand itself, and that lowering issues that operand, or its own pointer where the window takes another type, with noSlice.Risk
The finalize output for those 29 fixtures reads differently, so anything comparing it with saved text will see these lines change. The TIR fixtures under
tests/nvidia/cuda/fixtures/schedule/tiraren't snapshots of finalize output (48 of the 49 already differ from it), so I've left them alone. Only static zero starts count, so a window whose start is only known at run time keeps itsSlice. A partial window of a sharded tile still meets the codegen refusal, as insparse_decode, and that needs separate work.