Skip to content

fix(lowering): issue a window that covers its operand as the operand itself - #273

Merged
zhen8838 merged 4 commits into
tile-ai:mainfrom
FeathBow:fix/lowering-whole-windows
Oct 27, 2026
Merged

zhen8838 merged 4 commits into
tile-ai:mainfrom
FeathBow:fix/lowering-whole-windows

Conversation

@FeathBow

@FeathBow FeathBow commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Why

What schedule finalize writes 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 with CUDA pointer slices require a plain source layout. Past that, the MMA reads its register fragments as lhs_frag_dst[0:0 + 16, 0:0 + 32] and rhs_frag_dst[0:0 + 32, 0:0 + 8], and CUDA codegen has no handler for a bare Slice. Each of these windows is the whole operand.

The fix belongs in lowering. StoragePlan.operand already treats a window that starts at zero and spans its operand as the operand itself, but materialization and the instruction issue loop still built a Slice for it. MLIR's TrivialSubViewOpFolder does the same thing for memref.subview. It replaces such a subview with its source, or casts the source when the types differ.

What

ConvertHIRToTIR decides 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.operand behaves 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 reads lhs_frag_dst and rhs_frag_dst directly. Across the 49 schedule HIR fixtures that have a TIR twin, 29 finalize outputs change, all in this way, and no TensorView layout or shape changes. When I compiled those 29 for their own targets, the sm_89 GEMM and mesh_views now compile. Two wgmma outputs get past the pointer slice refusal and stop at the next view, where shard_engine_layout doesn'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_local to the MMA, and prints one line fewer so that it still ends on that MMA.

Contract

docs/spec/passes.md now says, under ConvertHIRToTIR, 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 no Slice.

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/tir aren'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 its Slice. A partial window of a sharded tile still meets the codegen refusal, as in sparse_decode, and that needs separate work.

@FeathBow
FeathBow marked this pull request as draft October 26, 2026 17:06
@FeathBow
FeathBow marked this pull request as ready for review October 26, 2026 17:34
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 zhen8838 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_decode is 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.

@zhen8838
zhen8838 merged commit 41c93eb into tile-ai:main Oct 27, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants