[TritonControlFlowOpt](feat) decouple pointers across SCF boundaries - #1528
Open
elstainniles wants to merge 1 commit into
Open
[TritonControlFlowOpt](feat) decouple pointers across SCF boundaries#1528elstainniles wants to merge 1 commit into
elstainniles wants to merge 1 commit into
Conversation
elstainniles
requested review from
KanuaK,
WuTYSFG and
hongziqi
as code owners
August 13, 2026 02:13
Carry complete block-pointer descriptors across supported for, while, and if boundaries as pointer-free SSA components, and rebuild pointers only at their use sites. Keep tensor-pointer bases outside loop signatures when they are invariant while carrying complete lane offsets. Teach TritonToLinalg and TritonToUnstructure to lower integer pointer carriers, scalar and opaque pointer joins, descriptor loops, lane offsets, and rebased memref layouts. Preserve exact descriptor producer slots, retain legacy conversion for mixed pointer boundaries, and keep externally typed memref boundaries layout-compatible. Remove the module-wide addptr-base restriction so unrelated and local make_tensor_ptr operations retain their previous behavior. Add one end-to-end pytest covering dynamic if/for/while block-pointer descriptors, changing bases, ordinary loop results, scalar-base tensor pointers, and opaque lane-wise tensor pointers.
elstainniles
force-pushed
the
feat/cfo-pointer-boundary-main-dev
branch
from
August 13, 2026 02:31
864f583 to
b4f9ee5
Compare
Contributor
|
✅ OpenCodeReview: Review partially complete: 0 finding(s); 7 of 17 selected item(s) failed. |
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.
Motivation
Triton block pointers and tensors of pointers could cross
scf.for,scf.while, andscf.ifboundaries as Triton pointer-typed SSA values. Thismade the downstream conversion depend on pointer values embedded in SCF
signatures and prevented control-flow lowering from using ordinary scalar or
tensor state consistently.
This change establishes a pointer-free boundary representation for supported
pointer states.
TritonControlFlowOptanalyzes an entire control-flow subtree,expands the selected pointer results into descriptor components, and rebuilds
the original Triton pointer only at use sites. The downstream passes consume
the descriptor contract without changing unrelated pointer code.
Changes
1. Carry complete block-pointer descriptors across SCF
A block pointer is represented as:
base_addressis ani64; shape, stride, and offset components retain theiroriginal integer types. The
orderattribute remains compile-time metadata.All descriptor components are carried for a rewritten block-pointer boundary,
which gives
for,while, andifone stable positional schema.2. Separate read-only control-flow analysis from IR rewriting
The pass first computes a schema for every supported SCF result and backedge.
No operations are created during this phase. Rewriting starts only after the
whole control-flow subtree has a compatible plan, so nested control flow uses
the same component order at every edge.
3. Keep complete descriptors correct for dynamic loop behavior
Loop bounds, steps, and pointer deltas may be dynamic. Pointer advancement is
lowered to loop-carried offset arithmetic instead of an induction-variable
closed form, so zero-trip loops and deltas defined in the loop body preserve
the original SCF semantics.
4. Support branch-selected and addptr-derived block-pointer bases
Different
scf.ifbranches may select different block-pointer bases. Eachbranch yields the corresponding integer address and the pointer is rebuilt
after the join. A
tt.make_tensor_ptrwhose base is produced bytt.addptrisaccepted when it actually crosses a supported SCF boundary. Local block
pointers are not globally rejected or rewritten.
An unrelated local descriptor remains local:
5. Represent tensor-of-pointers as common base plus complete offsets
For
tensor<...x!tt.ptr<T>>, the policy records:A scalar pointer introduced by
tt.splatis retained outside a loop while thecomplete per-lane offsets tensor is carried only when it changes. Chained
tt.addptroperations are flattened into the complete offsets. The componentmodel keeps an explicit base slot so a future analysis can support richer
tensor-base decomposition without changing the shared control-flow rewrite.
6. Mark the exact descriptor slots consumed downstream
Rewritten loops receive
PointerDescriptorBoundary, a dense array containingonly the loop slots that belong to pointer descriptors. When multiple pointer
policies rewrite the same loop, existing slots are remapped into the new
signature and merged with the newly expanded slots.
TritonToLinalguses those exact slots as producer-preservation roots. It doesnot retain unrelated accumulators, masks, bounds, or ordinary tensor state.
Malformed or duplicate slot metadata is rejected.
7. Preserve mixed boundaries and legacy pointer conversion
A loop can contain both CFO-expanded descriptor slots and a residual pointer
slot that a later policy proved invariant. Such a loop is not treated as fully
pointer-free. The descriptor roots are preserved, while the existing legacy
loop conversion remains responsible for the residual pointer.
After
TritonToLinalg, both paths converge on a legal pointer-free boundaryinstead of skipping conversion merely because the marker exists.
8. Lower scalar pointer joins through integer address carriers
Scalar pointer
arith.select, pointerscf.ifjoins,tt.ptr_to_int, andtt.int_to_ptrround trips are lowered through integer addresses. This avoidsrequiring
arith.selector SCF to choose between memrefs whose layouts maydiffer.
Pointer-valued
scf.ifuses the same carrier rule:9. Rebase memref layouts without breaking typed boundaries
When a pointer address absorbs a descriptor offset, the corresponding
memref.reinterpret_castis rebased to offset zero. Result layouts arepropagated through chained and rank-reduced
memref.subviewoperations.Layout-sensitive consumers such as
scf.yield,scf.condition,func.return,and
func.callretain their declared memref type instead of receiving anincompatible rewritten layout.
At a fixed function or SCF boundary, the declared layout remains unchanged:
10. Lower opaque tensor-pointer operations lane by lane
TritonToUnstructurecan now bubbletensor.extractthrough pointerarith.selectand preserve opaque tensor-pointer expressions until the lane isknown. Offset analysis and argument replacement then lower each selected lane
to scalar pointer arithmetic and memory access.
Regression coverage
The change adds or updates focused MLIR coverage for:
scf.for,scf.while,scf.if,nested control flow, dynamic deltas, and zero-trip loops;
tt.addptr-derived block-pointer bases;fixed
return/call/if/for/whilelayout boundaries;control-flow cases.