Control-flow support for PtrAnalysis

[nhat-nguyen]

We have internally implemented control-flow support for PtrAnalysis. This allows the triton-to-structured pass to generate block-pointers in presence of control-flow operations such as scf.if. We would appreciate any feedbacks!

General control flow support in PtrAnalysis

Background

PtrAnalysis supports loops (scf.for) but fails on conditionals (scf.if) that produce pointers, requiring us to fallback to scalar loads and stores. Where applicable, we want the analysis to identify structured loads and stores and lower them to more efficient DMA operations.

The Problem

Pointer analysis recursively visits operations and produces PtrState for each op: source pointer, offsets, strides, etc.

Challenge: An scf.if produces two pointer states (one per branch) that must be reconciled into a single result state.

%result_ptr = scf.if %cond -> (!tt.ptr<f32>) {
  %then_ptr = ... // PtrStateThen
  scf.yield %then_ptr
} else {
  %else_ptr = ... // PtrStateElse
  scf.yield %else_ptr
}
// What is the PtrState of %result_ptr?

Solution: Leverage existing transformation that exposes the pointer states through tts.get_structured_state to derive the correct pointer states that represent both branches. This allows the analysis to correctly populate the pointer state for scf.if operations.

Solutions

Background on tts.get_structured_state

This operation decouples analysis from IR rewriting:

1. Pre-pass: Inserts tts.get_structured_state ops in strategic places where an operation can produce PtrState.
2. Analysis: Builds PtrState for each operation.
3. Rewrite: Process all tts.get_structured_state ops and replace the uses of their return values with the correct offsets and strides populated during step 2.

Example for scf.for:

%arg_ptr, %arg_offset, %arg_stride = tts.get_structured_state ...
%ptr, %offset, %stride = scf.for ... iter_args(%arg_ptr, %arg_offset, %arg_stride) {
  scf.yield %next_ptr, %computed_offset, %computed_stride
}

1. Reconciling states using scf.if results

Leverage existing tts.get_structured_state infrastructure: make scf.if yield all PtrState components (offsets, strides, and source). The reconciled state then references the scf.if results directly:

%ptr, %offset, %stride, %source = scf.if %cond -> (!tt.ptr<f32>, index, index, !tt.ptr<f32>) {
  scf.yield %then_ptr, %then_offset, %then_stride, %then_source
} else {
  scf.yield %else_ptr, %else_offset, %else_stride, %else_source
}
// PtrState for %result_ptr: { source: %source, offsets: [%offset], strides: [%stride] }

This creates “phi nodes” for each state component and is the key solution to generalized control flow support since the analysis can now recursively visit across arbitrary control-flow ops.

2. tts.make_tptr after each condition

When ops need a structured pointer from scf.if, insert a single tts.make_tptr after the conditional:

%offset, %stride, %source = scf.if %cond {
  ...
}

%structured_ptr = tts.make_tptr %source to sizes: [...], strides: [%stride], offsets: [%offset], ...
%data = tts.load %structured_ptr

Why not inside each branch?

• tts.make_tptr is a descriptor operation that encodes how memory should be accessed, not a value-producing operation.
• Using the tts.make_tptr result in each branch will also confuse the bufferization pass.

3. Handling Mixed Structured/Unstructured Branches

When branches produce incompatible structuredness (e.g.: one is structured when the other isn’t), we will fallback to using unstructured pointers. Handling this problem is tricky in practice, especially in nested scf.if.

Example:

Assume that:

• the first outer-most if branch returns unstructured pointers and
• the second outer-most if branch (and its nested branches) return structured pointers.

%4:12 = scf.if %3 -> {
  %unstructured_0, %offsets_1:2,... = "tts.get_structured_state"
  scf.yield %unstructured_0, %offsets#0,...
} else {
  %7:12 = scf.if %6 -> {
    %structured, %offsets:2,... = "tts.get_structured_state"
    %structured_0, %offsets_1:2,... = "tts.get_structured_state"
    scf.yield %structured, %offsets#0,...
  } else {
    %structured, %offsets:2,... = "tts.get_structured_state"
    %structured_0, %offsets_1:2,... = "tts.get_structured_state"
    scf.yield %structured, %offsets#0,...
  }
  scf.yield %7#0, %7#1, %7#2,...
}

Since both branches are incompatible, we want to use unstructured pointers for this scf.if. However, since the tts.get_structured_state operations are processed independently, marking an operation to be skipped during rewriting because other tts.get_structured_state have failed is tricky.

For such cases, for simplicity, we will mark all nested operations inside the scf.if to be unstructured since it is not likely that users will write such code.

4. Extension to handling pointers returned by other control-flow ops

The above approach also scales to other kinds of control-flow ops that return pointers such as scf.for, scf.while,… We just need to create the appropriate visit and rewrite methods for such operations.

Conclusion

This work provides a framework for supporting arbitrary control-flow within the triton-to-structured pass by leveraging the existing transformation provided by the tts.get_structured_state transformation.