[TritonToLinalg](fix) Restore addptr foading

[CHNJZ]

Background

Upstream Triton defines a standard folder for tt.addptr: tt.addptr(ptr, 0) is canonicalized to the original pointer.

OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) {
  if (matchPattern(adaptor.getOffset(), m_Zero()))
    return getPtr();
  return {};
}

The Ascend backend previously used tt.addptr as the common entry point for pointer analysis and memory lowering. To prevent a synthetic zero-offset tt.addptr from being folded before reaching AddPtrConverter, the backend commented out both TT_AddPtrOp::hasFolder and AddPtrOp::fold() in the shared Triton dialect. This changed upstream canonicalization behavior and introduced an invasive community-code modification.

This PR restores the upstream implementation:

let hasFolder = 1;

and re-enables:

OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) {
  if (matchPattern(adaptor.getOffset(), m_Zero())) {
    return getPtr();
  }
  return {};
}

Once the folder is restored, Ascend pointer lowering can no longer assume that every memory pointer is produced by tt.addptr. This PR therefore also removes the backend dependency on synthetic AddPtr(ptr, 0) anchors.

Problem

The previous LoadStoreCanonicalizer inspected the pointer operands of Load, Store, AtomicRMW, and AtomicCAS operations. When a pointer was not produced by tt.addptr, the canonicalizer inserted a zero-offset tt.addptr to trigger the downstream AddPtrConverter.

Current TTIR:

%ptrs = tt.splat %arg0 : !tt.ptr<f32> -> tensor<4x!tt.ptr<f32>>
%value = tt.load %ptrs : tensor<4x!tt.ptr<f32>>

The previous canonicalizer rewrote it as:

%zero = arith.constant dense<0> : tensor<4xi32>
%addr = tt.addptr %ptrs, %zero : tensor<4x!tt.ptr<f32>>, tensor<4xi32>
%value = tt.load %addr : tensor<4x!tt.ptr<f32>>

After restoring AddPtrOp::fold(), %addr is immediately folded back to %ptrs. Under applyPatternsGreedily(), this creates the following rewrite cycle:

load(ptr) -> load(addptr(ptr, 0)) -> load(ptr) -> load(addptr(ptr, 0)) -> ...

The rewrite never reaches a fixed point, so TritonToLinalg compilation hangs.

The root cause is not the upstream addptr(ptr, 0) -> ptr fold. The actual issue is that the Ascend backend treated a semantically redundant +0 AddPtr as a mandatory pointer-lowering entry point.

Solution

This PR separates pointer lowering into two paths:

1. A real tt.addptr(base, offset) exists: preserve the actual address calculation and continue using AddPtrConverter.
2. No real AddPtr exists: parse the pointer expression directly to obtain its source, offsets, sizes, and strides, then materialize a memref without creating AddPtr(ptr, 0).

Real pointer arithmetic -> AddPtrConverter
Pure pointer expression -> BlockDataParser::materializePointer()

Load, Store, Atomic, and indirect-memory consumers resolve their pointers at the consumption boundary. This also removes the dependency on the execution order of producer and consumer conversion patterns.

Covered scenarios

1. Recover a real AddPtr hidden by Splat

Optimizations such as ReorderBroadcast may transform tensor pointer arithmetic into splat(addptr(base, offset)). The offset in this expression represents a real address calculation and must be preserved.

Current TTIR:

%x = tt.addptr %x_ptr, %pid : !tt.ptr<f32>, i32
%x_3 = tt.splat %x : !tt.ptr<f32> -> tensor<1x!tt.ptr<f32>>
%x_4 = tt.load %x_3 : tensor<1x!tt.ptr<f32>>

It can be rewritten as:

%x_ptrs = tt.splat %x_ptr : !tt.ptr<f32> -> tensor<1x!tt.ptr<f32>>
%pids = tt.splat %pid : i32 -> tensor<1xi32>
%x_3 = tt.addptr %x_ptrs, %pids : tensor<1x!tt.ptr<f32>>, tensor<1xi32>
%x_4 = tt.load %x_3 : tensor<1x!tt.ptr<f32>>

The updated LoadStoreCanonicalizer only recovers this kind of real AddPtr. It no longer inserts a zero offset for ordinary pointers.

2. Materialize a pointer Splat without a real AddPtr

A plain splat(base) does not contain pointer arithmetic, so a synthetic +0 must not be introduced.

Current TTIR:

%ptrs = tt.splat %arg0 : !tt.ptr<f32> -> tensor<4x!tt.ptr<f32>>
%value = tt.load %ptrs : tensor<4x!tt.ptr<f32>>

It is materialized directly as:

%view = memref.reinterpret_cast %arg0 to offset: [0], sizes: [4], strides: [0] : memref<?xf32> to memref<4xf32, strided<[0]>>

The zero stride means that all four lanes access the same base address. BlockDataParser preserves this layout, and the Load, Store, and Atomic converters consume %view directly.

3. Materialize a pointer Bitcast directly

Current TTIR:

%ptrs = tt.splat %arg0 : !tt.ptr<i32> -> tensor<4x!tt.ptr<i32>>
%cast = tt.bitcast %ptrs : tensor<4x!tt.ptr<i32>> -> tensor<4x!tt.ptr<f32>>
%value = tt.load %cast : tensor<4x!tt.ptr<f32>>

It is materialized directly as:

%source = builtin.unrealized_conversion_cast %arg0 : memref<?xi32> to memref<?xf32>
%view = memref.reinterpret_cast %source to offset: [0], sizes: [4], strides: [0] : memref<?xf32> to memref<4xf32, strided<[0]>>

Only the memref element type changes; the original offsets, sizes, and strides are preserved. Pointer Broadcast and ExpandDims expressions are recursively handled through the same materialization entry point.

4. Resolve Load, Store, and Atomic pointers through one entry point

This PR adds a shared pointer-resolution helper:

static FailureOr<Value>
resolveMemoryPointer(Value originalPtr, Value convertedPtr,
                     ConversionPatternRewriter &rewriter);

The helper uses the following order:

1. If the pointer is a Structured CustomOp result with SrcPtrIndex, rematerialize its inherited layout first.
2. If the adaptor pointer is already a memref, reuse it directly.
3. Otherwise, call BlockDataParser::materializePointer() on the original TTIR pointer expression.

The helper is used by:

• LoadConverter
• StoreConverter
• AtomicRMWConverter
• AtomicCASConverter

5. Preserve Structured CustomOp pointer-result layouts

A Structured CustomOp may use SrcPtrIndex to indicate that a pointer result inherits the address layout of a specific input.

Current TTIR:

%custom = hivm.hir.custom {SrcPtrIndex = array<i32: 0>, hivm.is_distributed, symbol = "foo3"} "foo3" ins(%src, %flag : tensor<32x!tt.ptr<f16>>, i32) outs(%empty : tensor<32x!tt.ptr<f16>>) -> tensor<32x!tt.ptr<f16>>
%value = tt.load %custom : tensor<32x!tt.ptr<f16>>

Converted IR:

%custom_memref = hivm.hir.custom {SrcPtrIndex = array<i32: 0>, hivm.is_distributed, symbol = "foo3"} "foo3" ins(%src_view, %flag : memref<32xf16, strided<[1]>>, i32) -> memref<32xf16>
%custom_view = memref.reinterpret_cast %custom_memref to offset: [0], sizes: [32], strides: [1] : memref<32xf16> to memref<32xf16, strided<[1]>>

resolveMemoryPointer() must not return immediately just because the converted CustomOp result already has a memref type. Doing so would lose the layout inherited through SrcPtrIndex. The new implementation gives this case priority and invokes materializePointer().

LoadConverter also no longer assumes that every valid memref is produced by memref.reinterpret_cast. When a pointer comes directly from a CustomOp or a function argument, the converter obtains the static innermost stride from the memref type instead of dereferencing a missing ReinterpretCastOp producer.

6. Support int_to_ptr indirect loads

The previous TritonToUnstructure path wrapped tt.int_to_ptr in a synthetic scalar zero AddPtr:

%base = tt.int_to_ptr %address : i64 -> !tt.ptr<i32>
%zero = arith.constant 0 : i64
%anchor = tt.addptr %base, %zero : !tt.ptr<i32>, i64
%value = ttg.indirect_load %anchor, %offsets : !tt.ptr<i32>, tensor<128xi64>

Because %anchor necessarily disappears after the folder is restored, the updated implementation removes this wrapper and preserves the actual source directly:

%base = tt.int_to_ptr %address : i64 -> !tt.ptr<i32>
%value = ttg.indirect_load %base, %offsets : !tt.ptr<i32>, tensor<128xi64>

In TritonToLinalg, IndirectLoadConverter materializes the pointer at the consumer boundary:

%base_memref = hivm.hir.pointer_cast %address : i64 to memref<?xi32>
%view = memref.reinterpret_cast %base_memref to offset: [0], sizes: [1], strides: [1] : memref<?xi32> to memref<1xi32, strided<[1]>>
%result = call @triton_indirect_load(%view, %offsets, %mask, %other) {isVolatile = true} : (memref<1xi32, strided<[1]>>, tensor<128xi64>, tensor<128xi1>, tensor<128xi32>) -> tensor<128xi32>

requiresVolatileIndirectLoad() still examines the original load pointer, so volatile semantics for unknown integer addresses are preserved.

7. Fold a zero AddPtr used as a CustomOp input

Input TTIR:

%zero = arith.constant 0 : i32
%ptr = tt.addptr %arg4, %zero : !tt.ptr<i64>, i32
%result = hivm.hir.custom {symbol = "foo1"} "foo1" ins(%ptr, %num_programs, %c0_i64 : !tt.ptr<i64>, i32, i64) -> i32

After restoring the folder, the zero AddPtr is removed according to upstream semantics, and the CustomOp receives the converted base memref directly:

%result = hivm.hir.custom {symbol = "foo1"} "foo1" ins(%arg6, %arg7, %c0_i64 : memref<?xi64>, i32, i64) -> i32

The previous memref<1xi64, strided<[1]>> input was a view created only because the synthetic zero AddPtr reached AddPtrConverter. The old and new IR refer to the same underlying address. For this structural conversion test, the FileCheck expectation should be updated instead of recreating a +0 AddPtr merely to preserve the old IR shape.

If a real CustomOp ABI explicitly depends on size or stride descriptor values, the required view should be created explicitly at that ABI-lowering boundary.

8. Use an Atomic result as a dynamic address offset

In the following case, the old value returned by atomic_add is used as a subsequent load or store offset:

%old = tt.atomic_rmw add, acq_rel, gpu, %counter_ptr, %one, %true : (!tt.ptr<i64>, i64, i1) -> i64
%load_ptr = tt.addptr %input_ptr, %old : !tt.ptr<f32>, i64
%value = tt.load %load_ptr : !tt.ptr<f32>

AtomicRMWConverter now invokes the resolver first and obtains MemRefType from the resolved value rather than from the pre-resolution pointer. The dynamic offset is then preserved in memref.reinterpret_cast:

%offset = arith.index_cast %old : i64 to index
%view = memref.reinterpret_cast %input_memref to offset: [%offset], sizes: [1], strides: [1] : memref<?xf32> to memref<1xf32, strided<[1], offset: ?>>
%value = memref.load %view[%c0] : memref<1xf32, strided<[1], offset: ?>>

[github-actions]

✅ OpenCodeReview: Review partially complete: 0 finding(s); 3 of 8 selected item(s) failed.