Skip to content

fix(lowering,verify): bind each selected ring stage before an instruction uses it - #276

Closed
FeathBow wants to merge 1 commit into
tile-ai:mainfrom
FeathBow:fix/lowering-bind-ring-stages
Closed

FeathBow wants to merge 1 commit into
tile-ai:mainfrom
FeathBow:fix/lowering-bind-ring-stages

Conversation

@FeathBow

Copy link
Copy Markdown
Contributor

Why

Lowering passed a selected ring stage straight to an instruction, as in T.copy_async_tensor(value, lhs_ring_stages[k % 3]). The CUDA emitters read tensor operands by variable name, so each of these failed in name_for at compile time. Finalizing the schedule fixtures gave 55 such operands, from copy, copy_async and copy_async_tensor.

Lowering already binds a view at a constant smem address, or over gmem scratch, in each scope that uses it. A staged buffer's selected stage was the one view it returned unbound, and nothing in TIR verification looked at the form of an operand.

This relates to the #230 item in #255. Lowering now binds each selected stage, so the CUDA emitters need no change for it.

What

  • TirBuilder.declare returns a staged buffer's selected stage as a View with the stage's own shape and layout. Copies, other instructions and the MMA issue loop then bind T.tensor_view(T.ptr_of(stages[i]), ...) in the scope that uses it and pass that variable.
  • The TIR verifier refuses a TensorView, Slice or TupleGetItem call as an operand of an effect Op. Finalize verifies what lowering writes, so this form now fails there.
  • tir §1.4 states the rule, and passes §7.1 adds the selected stage to the views bound where they are used.
  • The schedule tutorial's finalize block shows the new bindings and prints 121 lines, so that it still ends on the first MMA.

Contract

An effect Op no longer accepts a TensorView, Slice or TupleGetItem call as an operand; a LetStmt binds it and the Op takes the variable. Other expressions, such as scalar arithmetic, stay inline. This covers a scalar TupleGetItem as well. The handwritten TIR in the repository already binds all three.

Risk

  • Finalize output gains a binding wherever a staged buffer is used. The 75 schedule fixtures that finalize on main still do, and the 9 that fail there fail as before.
  • When I compiled those 75 outputs for their own targets, 12 compile on main and 13 with this change, the new one being smem_padded_stages. Four others get past this point and stop at other errors: a register output parameter, a copy not split over threads, and a sliced mesh naming several levels. persistent_tile_union_count fails earlier than before: its stage index reads a CTA coordinate that codegen cannot render yet, which it now meets before the TMA copy that has no emitter (feat(codegen): emit T.copy_async_tensor with host-encoded tensor maps #248).
  • Nothing here runs on a GPU. The change is to the form of operands; no staged program from finalize runs on the card until its copies are synchronised.
  • Tests: the verifier refuses each of the three calls, and the schedule corpus case for smem_padded_stages now compiles its finalize output. All four fail with main's lowering and verifier. Apart from them, the CPU suite fails the same tests with either version.

@zhen8838

Copy link
Copy Markdown
Collaborator

Thanks @FeathBow for tracking down these 55 operands.

This PR exposes a deeper problem. The failure comes from how lowering represents a staged smem buffer, as a Tuple of views that TupleGetItem selects, and not from the form of the operand. Fixing that touches lowering, the TIR spec and CUDA codegen together.

So I will take this PR over and rework it. No further change is needed from you here. I will link the result in this thread, so that you can continue your work on top of it.

@zhen8838

zhen8838 commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

Superseded by #277: #277

The replacement binds integer-address ring TensorViews, their tuple and each selected stage, carries analyzed shared-memory peaks into host launches, and shares one target-resolved default pipeline. Final local verification: source 1514 passed / 12 skipped; installed 167 passed, including all 13 tutorial checks and four page reproductions. All differences across the 60 finalized HIR fixtures were attributed to host entries, ring bindings or the host entry point. Follow-ups are tracked in #278 and #255.

zhen8838 added a commit that referenced this pull request Oct 28, 2026
## Why
- Supersedes #276. Unbound ring views and stage selections reach CUDA
operands without SSA names, and compilation entry points assemble
different pass lists.

## What
- Keep ring stages as a tuple of bound TensorViews and bind each
selected stage once per consuming scope.
- Share `default_pipeline` across finalize, compile and scaffold.
Resolve the core host-entry pass through `Target.resolve`, with CUDA
supplying an `@bind` implementation.
- Make `Module.entry` the host starting point and carry analyzed
shared-memory peaks into host launches.
- Link all CPU hosts on the call paths, with the root host first.

## Contract
- Both `InsertHostEntryPass.create(target, options)` and
`ConvertHIRToTIR.create(target, options)` support target-owned
implementations.
- `link_program` removes `entries`; the loader binds the first linked
host, and `build` requires a CPU host root.
- Full lowering requires memory-analysis metadata. `ConvertHIRToTIR` is
a no-op for skeletons or modules without HIR; a CPU host root being
lowered must launch exactly one HIR callee.

## Risk
- #278 tracks Broadcast layout normalization; the remaining
layout-adaptation views are retained.
- #255 tracks SymbolRef/function references; indexed child entries
remain excluded from synthesized hosts, and Launch type repair remains
necessary.
@zhen8838 zhen8838 closed this Oct 28, 2026
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