[Canonicalize] Transform ptr_to_int->add->int_to_ptr to addptr by yiqian1 · Pull Request #9971 · triton-lang/triton

yiqian1 · 2026-04-08T20:38:11Z

Add canonicalization pattern for IntToPtrOp that recognizes the pattern:
  int_to_ptr(addi(ptr_to_int(ptr), constant_offset))

and transforms it to:
  addptr(ptr, element_offset)

where element_offset = constant_offset / element_size_bytes.

This pattern appears when performing pointer arithmetic via integer
operations (e.g., adding byte offsets to pointers). By canonicalizing
to addptr, AxisInfoAnalysis can correctly track contiguity, enabling
proper vectorization for operations like async_copy_local_to_global.

The pattern only applies when:
- The offset is a compile-time constant (IntegerAttr or SplatElementsAttr)
- The byte offset is evenly divisible by the element size

Added to both standard canonicalize and gluon-canonicalize passes.

Tests added for positive cases (f32, f16, commutative) and negative
cases (non-constant offset, indivisible offset).

Fix AxisInfoAnalysis to preserve contiguity information through tt.int_to_ptr and tt.ptr_to_int operations by registering them as CastOpAxisInfoVisitor. Previously, these operations would reset contiguity to 1, preventing proper vectorization. This fixes async_copy_local_to_global failing to generate wide vector stores (e.g., b128) when pointer arithmetic is performed via integer operations (ptr->int->add->int->ptr pattern). Test case added for async_copy_local_to_global with ptr_to_int/ int_to_ptr casts, verifying 128-bit stores are generated for sizePerThread=[8] with f16 elements. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ThomasRaoux · 2026-04-08T20:49:17Z

+                  CastOpAxisInfoVisitor<triton::IntToPtrOp>,
+                  CastOpAxisInfoVisitor<triton::PtrToIntOp>,


I don't think that's correct since contiguous pointers wouldn't generate contiguous integers.

Also ideally you unit test only running analysis to show the bug and new behavior

I don't think that's correct since contiguous pointers wouldn't generate contiguous integers.

Right. Perhaps I should just handle the specific pattern: PtrToInt -> add -> IntToPtr. because this pattern is basically equivalent to pointer arithmetic and should preserve the contiguity from the original pointer.

can we canonicalize this this into ptr arithmetic?

Add canonicalization pattern for IntToPtrOp that recognizes the pattern: int_to_ptr(addi(ptr_to_int(ptr), constant_offset)) and transforms it to: addptr(ptr, element_offset) where element_offset = constant_offset / element_size_bytes. This pattern appears when performing pointer arithmetic via integer operations (e.g., adding byte offsets to pointers). By canonicalizing to addptr, AxisInfoAnalysis can correctly track contiguity, enabling proper vectorization for operations like async_copy_local_to_global. The pattern only applies when: - The offset is a compile-time constant (IntegerAttr or SplatElementsAttr) - The byte offset is evenly divisible by the element size Added to both standard canonicalize and gluon-canonicalize passes. Tests added for positive cases (f32, f16, commutative) and negative cases (non-constant offset, indivisible offset). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ThomasRaoux · 2026-04-09T02:20:44Z

+      return failure();
+    }
+
+    auto ptrToIntOp = ptrToIntValue.getDefiningOp<PtrToIntOp>();


nit: make ptrToIntValue a PtrToIntOp directly to avoid redundant cast

ThomasRaoux · 2026-04-09T02:23:02Z

+    } else {
+      // Scalar case
+      elementOffsetValue = arith::ConstantOp::create(
+          rewriter, loc, rewriter.getI32IntegerAttr(elementOffset));


I don't think you can assume the offset is a i32.

peterbell10 · 2026-04-09T16:17:29Z

+//-- IntToPtrOp --
+// Canonicalize: int_to_ptr(addi(ptr_to_int(ptr), constant_offset)) ->
+// addptr(ptr, element_offset) Only when offset is constant and divisible by
+// element size.


Wouldn't it be more general to split this into 2 patterns:

int_to_ptr(addi(val, offset)) -> addptr(int_to_ptr(val), element_offset) int_to_ptr(ptr_to_int(ptr)) -> ptr

Good point. Have refactored to 2 patterns.

peterbell10 · 2026-04-09T16:18:27Z

+      offsetValue = addOp.getRhs();
+    } else if (auto rhsPtrToInt = addOp.getRhs().getDefiningOp<PtrToIntOp>()) {
+      ptrToIntOp = rhsPtrToInt;
+      offsetValue = addOp.getLhs();


addi is canonicalized so constants are always the rhs argument.

ThomasRaoux

LGTM

…uity Adding a runtime scalar offset to an existing pointer tensor (output_ptrs + elem_delta) breaks AxisInfoAnalysis contiguity in current Triton, causing fallback to 16-bit async stores which gfx1250 doesn't support (assertion failure). Instead of modifying the existing pointer tensor, build the remote pointer tensor from scratch: (output_ptr + tile_base + elem_delta) + flat_idx. This gives AxisInfoAnalysis the splat(scalar) + arange pattern it trusts, enabling b128 vectorized async stores. Upstream Triton fix: triton-lang/triton#9971 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…n-lang#9971) Add canonicalization pattern for IntToPtrOp that recognizes the pattern: int_to_ptr(addi(ptr_to_int(ptr), constant_offset)) and transforms it to: addptr(ptr, element_offset) where element_offset = constant_offset / element_size_bytes. This pattern appears when performing pointer arithmetic via integer operations (e.g., adding byte offsets to pointers). By canonicalizing to addptr, AxisInfoAnalysis can correctly track contiguity, enabling proper vectorization for operations like async_copy_local_to_global. The pattern only applies when: - The offset is a compile-time constant (IntegerAttr or SplatElementsAttr) - The byte offset is evenly divisible by the element size Added to both standard canonicalize and gluon-canonicalize passes. Tests added for positive cases (f32, f16, commutative) and negative cases (non-constant offset, indivisible offset).

ThomasRaoux reviewed Apr 8, 2026

View reviewed changes

yiqian1 and others added 2 commits April 9, 2026 02:06

Fix file format

2308fac

ThomasRaoux reviewed Apr 9, 2026

View reviewed changes

Preserve offset type in int_to_ptr canonicalization

d48e566

yiqian1 changed the title ~~[AxisInfo] Preserve contiguity through IntToPtr/PtrToInt casts~~ [Canonicalize] Transform ptr_to_int->add->int_to_ptr to addptr Apr 9, 2026

yiqian1 marked this pull request as ready for review April 9, 2026 14:31

yiqian1 requested review from peterbell10 and ptillet as code owners April 9, 2026 14:31

peterbell10 reviewed Apr 9, 2026

View reviewed changes

yiqian1 and others added 3 commits April 9, 2026 18:13

Split int_to_ptr canonicalization into two patterns

f05cb53

Merge branch 'main' into qian/fix-axisInfo-with-cast

4b06044

Merge branch 'main' into qian/fix-axisInfo-with-cast

9d1e148

ThomasRaoux approved these changes Apr 9, 2026

View reviewed changes

Merge branch 'main' into qian/fix-axisInfo-with-cast

8c70d7c

ThomasRaoux enabled auto-merge (squash) April 10, 2026 01:50

ThomasRaoux merged commit 89d154d into triton-lang:main Apr 10, 2026
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Canonicalize] Transform ptr_to_int->add->int_to_ptr to addptr#9971

[Canonicalize] Transform ptr_to_int->add->int_to_ptr to addptr#9971
ThomasRaoux merged 8 commits into
triton-lang:mainfrom
yiqian1:qian/fix-axisInfo-with-cast

yiqian1 commented Apr 8, 2026 •

edited

Loading

Uh oh!

ThomasRaoux Apr 8, 2026

Uh oh!

yiqian1 Apr 8, 2026

Uh oh!

ThomasRaoux Apr 8, 2026

Uh oh!

ThomasRaoux Apr 9, 2026

Uh oh!

ThomasRaoux Apr 9, 2026

Uh oh!

peterbell10 Apr 9, 2026

Uh oh!

yiqian1 Apr 9, 2026

Uh oh!

peterbell10 Apr 9, 2026

Uh oh!

ThomasRaoux left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		CastOpAxisInfoVisitor<triton::IntToPtrOp>,
		CastOpAxisInfoVisitor<triton::PtrToIntOp>,

Conversation

yiqian1 commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ThomasRaoux left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yiqian1 commented Apr 8, 2026 •

edited

Loading