Skip to content

[TritonControlFlowOpt](feat) decouple pointers across SCF boundaries - #1528

Open
elstainniles wants to merge 1 commit into
triton-lang:main-devfrom
elstainniles:feat/cfo-pointer-boundary-main-dev
Open

[TritonControlFlowOpt](feat) decouple pointers across SCF boundaries#1528
elstainniles wants to merge 1 commit into
triton-lang:main-devfrom
elstainniles:feat/cfo-pointer-boundary-main-dev

Conversation

@elstainniles

Copy link
Copy Markdown
Contributor

Motivation

Triton block pointers and tensors of pointers could cross scf.for,
scf.while, and scf.if boundaries as Triton pointer-typed SSA values. This
made 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. TritonControlFlowOpt analyzes 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_address, shape..., strides..., offsets...]

base_address is an i64; shape, stride, and offset components retain their
original integer types. The order attribute remains compile-time metadata.
All descriptor components are carried for a rewritten block-pointer boundary,
which gives for, while, and if one stable positional schema.

// Before
%final = scf.for ... iter_args(%ptr = %initial)
    -> (!tt.ptr<tensor<16xf32>>) {
  %next = tt.advance %ptr, [%delta]
  scf.yield %next : !tt.ptr<tensor<16xf32>>
}

// After
%desc:4 = scf.for ...
    iter_args(%base = %base_addr, %shape = %s, %stride = %st, %off = %o)
    -> (i64, i64, i64, i32) {
  %next_off = arith.addi %off, %delta : i32
  scf.yield %base, %shape, %stride, %next_off
      : i64, i64, i64, i32
}
%base = tt.int_to_ptr %desc#0 : i64 -> !tt.ptr<f32>
%ptr = tt.make_tensor_ptr %base, [%desc#1], [%desc#2], [%desc#3]
    {order = array<i32: 0>} : !tt.ptr<tensor<16xf32>>

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.

// One schema is planned for the nested result before either op is rewritten.
%selected = scf.if %outer_cond -> (!tt.ptr<tensor<16xf32>>) {
  %inner = scf.if %inner_cond -> (!tt.ptr<tensor<16xf32>>) {
    scf.yield %ptr0 : !tt.ptr<tensor<16xf32>>
  } else {
    scf.yield %ptr1 : !tt.ptr<tensor<16xf32>>
  }
  scf.yield %inner : !tt.ptr<tensor<16xf32>>
} else {
  scf.yield %ptr2 : !tt.ptr<tensor<16xf32>>
}

// Both nested joins use the same descriptor result schema.
%inner_desc:4 = scf.if ... -> (i64, i64, i64, i32)
%outer_desc:4 = scf.if ... -> (i64, i64, i64, i32)

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.

// Dynamic delta computed in the body.
%delta = arith.addi %iter_delta, %one : i32
%next = tt.advance %ptr, [%delta]

// Rewritten backedge.
%next_offset = arith.addi %offset, %delta : i32
scf.yield %base_addr, %shape, %stride, %next_offset, %delta
    : i64, i64, i64, i32, i32

4. Support branch-selected and addptr-derived block-pointer bases

Different scf.if branches may select different block-pointer bases. Each
branch yields the corresponding integer address and the pointer is rebuilt
after the join. A tt.make_tensor_ptr whose base is produced by tt.addptr is
accepted when it actually crosses a supported SCF boundary. Local block
pointers are not globally rejected or rewritten.

%ptr0 = tt.make_tensor_ptr %base0, ...
%ptr1 = tt.make_tensor_ptr %base1, ...
%selected = scf.if %cond -> (!tt.ptr<tensor<16xf32>>) {
  scf.yield %ptr0 : !tt.ptr<tensor<16xf32>>
} else {
  scf.yield %ptr1 : !tt.ptr<tensor<16xf32>>
}

// Rewritten join.
%addr0 = tt.ptr_to_int %base0 : !tt.ptr<f32> -> i64
%addr1 = tt.ptr_to_int %base1 : !tt.ptr<f32> -> i64
%desc:4 = scf.if %cond -> (i64, i64, i64, i32) {
  scf.yield %addr0, %shape0, %stride0, %offset0
} else {
  scf.yield %addr1, %shape1, %stride1, %offset1
}

An unrelated local descriptor remains local:

%shifted = tt.addptr %base, %offset : !tt.ptr<f32>, i32
%local = tt.make_tensor_ptr %shifted, ...
%value = tt.load %local : !tt.ptr<tensor<16xf32>>

5. Represent tensor-of-pointers as common base plus complete offsets

For tensor<...x!tt.ptr<T>>, the policy records:

[common_base, complete_offsets]

A scalar pointer introduced by tt.splat is retained outside a loop while the
complete per-lane offsets tensor is carried only when it changes. Chained
tt.addptr operations are flattened into the complete offsets. The component
model keeps an explicit base slot so a future analysis can support richer
tensor-base decomposition without changing the shared control-flow rewrite.

%base_tensor = tt.splat %base
%initial = tt.addptr %base_tensor, %range
%final = scf.for ... iter_args(%ptr = %initial)
    -> (tensor<4x!tt.ptr<f32>>) {
  %next = tt.addptr %ptr, %delta
  scf.yield %next : tensor<4x!tt.ptr<f32>>
}

// The invariant scalar base stays outside; only offsets cross the loop.
%final_offsets = scf.for ... iter_args(%offsets = %range)
    -> (tensor<4xi32>) {
  %next_offsets = arith.addi %offsets, %delta : tensor<4xi32>
  scf.yield %next_offsets : tensor<4xi32>
}
%base_tensor = tt.splat %base
%final_ptrs = tt.addptr %base_tensor, %final_offsets

6. Mark the exact descriptor slots consumed downstream

Rewritten loops receive PointerDescriptorBoundary, a dense array containing
only 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.

%results = scf.for ...
    iter_args(%acc = %zero, %base = %addr, %shape = %s,
              %stride = %st, %offset = %o, %ordinary = %value)
    -> (i32, i64, i64, i64, i32, tensor<16xf32>) {
  ...
} {PointerDescriptorBoundary = array<i32: 1, 2, 3, 4>}

TritonToLinalg uses those exact slots as producer-preservation roots. It does
not 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.

// BlockPtr expands, while an invariant tensor-of-pointers remains in slot 4.
%loop:5 = scf.for ...
    -> (i64, i64, i64, i32, tensor<4x!tt.ptr<f32>>) {
  ...
} {PointerDescriptorBoundary = array<i32: 0, 1, 2, 3>}

After TritonToLinalg, both paths converge on a legal pointer-free boundary
instead of skipping conversion merely because the marker exists.

8. Lower scalar pointer joins through integer address carriers

Scalar pointer arith.select, pointer scf.if joins, tt.ptr_to_int, and
tt.int_to_ptr round trips are lowered through integer addresses. This avoids
requiring arith.select or SCF to choose between memrefs whose layouts may
differ.

// Triton input.
%selected = arith.select %cond, %lhs, %rhs : !tt.ptr<f32>
%address = tt.ptr_to_int %selected : !tt.ptr<f32> -> i64

// Lowered address selection.
%lhs_index = memref.extract_aligned_pointer_as_index %lhs
%lhs_address = arith.index_cast %lhs_index : index to i64
%rhs_index = memref.extract_aligned_pointer_as_index %rhs
%rhs_address = arith.index_cast %rhs_index : index to i64
%address = arith.select %cond, %lhs_address, %rhs_address : i64

Pointer-valued scf.if uses the same carrier rule:

%address = scf.if %cond -> (i64) {
  scf.yield %lhs_address : i64
} else {
  scf.yield %rhs_address : i64
}

9. Rebase memref layouts without breaking typed boundaries

When a pointer address absorbs a descriptor offset, the corresponding
memref.reinterpret_cast is rebased to offset zero. Result layouts are
propagated through chained and rank-reduced memref.subview operations.
Layout-sensitive consumers such as scf.yield, scf.condition, func.return,
and func.call retain their declared memref type instead of receiving an
incompatible rewritten layout.

// Before rebasing.
%view = memref.reinterpret_cast %ptr to offset: [%offset],
    sizes: [1, 8], strides: [8, 1]
    : memref<?xf32> to memref<1x8xf32, strided<[8, 1], offset: ?>>

// Address absorbs offset * sizeof(f32); the internal view starts at zero.
%byte_offset = arith.muli %offset_i64, %c4_i64 : i64
%real_address = arith.addi %address, %byte_offset : i64
%rebased = hivm.hir.pointer_cast(%real_address) [%capacity] : memref<?xf32>
%view = memref.reinterpret_cast %rebased to offset: [0],
    sizes: [1, 8], strides: [8, 1]
    : memref<?xf32> to memref<1x8xf32, strided<[8, 1]>>

At a fixed function or SCF boundary, the declared layout remains unchanged:

func.return %view : memref<1xf32, strided<[1], offset: 1>>

10. Lower opaque tensor-pointer operations lane by lane

TritonToUnstructure can now bubble tensor.extract through pointer
arith.select and preserve opaque tensor-pointer expressions until the lane is
known. Offset analysis and argument replacement then lower each selected lane
to scalar pointer arithmetic and memory access.

%selected = arith.select %condition, %lhs_ptrs, %rhs_ptrs
    : tensor<4xi1>, tensor<4x!tt.ptr<f32>>
%advanced = tt.addptr %selected, %delta
%loaded = tt.load %advanced : tensor<4x!tt.ptr<f32>>

// Representative unstructured lane.
%lane_cond = tensor.extract %condition[%i] : tensor<4xi1>
%lhs = tensor.extract %lhs_ptrs[%i] : tensor<4x!tt.ptr<f32>>
%rhs = tensor.extract %rhs_ptrs[%i] : tensor<4x!tt.ptr<f32>>
%lane_base = arith.select %lane_cond, %lhs, %rhs : !tt.ptr<f32>
%lane_ptr = tt.addptr %lane_base, %lane_delta : !tt.ptr<f32>, i64
%value = tt.load %lane_ptr {DiscreteMemAccess} : !tt.ptr<f32>

Regression coverage

The change adds or updates focused MLIR coverage for:

  • full block-pointer descriptors across scf.for, scf.while, scf.if,
    nested control flow, dynamic deltas, and zero-trip loops;
  • different and tt.addptr-derived block-pointer bases;
  • exact descriptor-marker slot remapping and mixed pointer boundaries;
  • scalar pointer selects, integer round trips, memref rebasing, subviews, and
    fixed return/call/if/for/while layout boundaries;
  • tensor-pointer descriptor loops and opaque lane-wise pointer selection;
  • an end-to-end pytest covering dynamic block-pointer and tensor-pointer
    control-flow cases.

@github-actions github-actions Bot added compiler Changes to C/C++ compiler backend (lib/, include/) python Changes to Python runtime or bindings ascend-backend Changes to the Ascend NPU backend labels Aug 13, 2026
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
elstainniles force-pushed the feat/cfo-pointer-boundary-main-dev branch from 864f583 to b4f9ee5 Compare August 13, 2026 02:31
@github-actions

Copy link
Copy Markdown
Contributor

OpenCodeReview: Review partially complete: 0 finding(s); 7 of 17 selected item(s) failed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ascend-backend Changes to the Ascend NPU backend compiler Changes to C/C++ compiler backend (lib/, include/) python Changes to Python runtime or bindings

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant