[Mlir-commits] [mlir] [mlir][ptr] Preserve broadcasts when folding zero pointer offsets (PR #231077)

[Alexander Shaposhnikov]

https://github.com/alexander-shaposhnikov created https://github.com/llvm/llvm-project/pull/231077

ptr.ptr_add accepts a scalar base with a shaped offset and returns a shaped pointer value. Its zero-offset fold currently returns the scalar base even when the offset is a zero splat, producing a value with the wrong type and triggering an assertion during canonicalization.

Only fold to the base when its type matches the result type. Add regression coverage for fixed vectors, scalable vectors, and tensors, and check that type-preserving zero-offset folds still apply to shaped bases.

>From af9a53f78822b1fe55b0417f4d9a7895065f7110 Mon Sep 17 00:00:00 2001
From: Alexander Shaposhnikov <alexander.v.shaposhnikov at gmail.com>
Date: Sun, 11 Oct 2026 23:22:14 -0700
Subject: [PATCH] [mlir][ptr] Preserve broadcasts when folding zero pointer
 offsets

ptr.ptr_add accepts a scalar base with a shaped offset and returns a shaped
pointer value. Its zero-offset fold currently returns the scalar base even
when the offset is a zero splat, producing a value with the wrong type and
triggering an assertion during canonicalization.

Only fold to the base when its type matches the result type. Add regression
coverage for fixed vectors, scalable vectors, and tensors, and check that
type-preserving zero-offset folds still apply to shaped bases.
---
 mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp  |  3 +-
 mlir/test/Dialect/Ptr/canonicalize.mlir | 62 +++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
index 444197d12097f..6d61c11416733 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
@@ -346,7 +346,8 @@ void StoreOp::build(OpBuilder &builder, OperationState &state, Value value,
 /// Fold: ptradd ptr + 0 ->  ptr
 OpFoldResult PtrAddOp::fold(FoldAdaptor adaptor) {
   Attribute attr = adaptor.getOffset();
-  if (!attr)
+  // A shaped offset can broadcast a scalar base to a shaped result.
+  if (!attr || getBase().getType() != getType())
     return nullptr;
   if (llvm::APInt value; m_ConstantInt(&value).match(attr) && value.isZero())
     return getBase();
diff --git a/mlir/test/Dialect/Ptr/canonicalize.mlir b/mlir/test/Dialect/Ptr/canonicalize.mlir
index e50cd1b76caf3..b173c894c5c24 100644
--- a/mlir/test/Dialect/Ptr/canonicalize.mlir
+++ b/mlir/test/Dialect/Ptr/canonicalize.mlir
@@ -14,6 +14,68 @@ func.func @zero_offset(%ptr: !ptr.ptr<#ptr.generic_space>) -> !ptr.ptr<#ptr.gene
   return %res0 : !ptr.ptr<#ptr.generic_space>
 }
 
+// A zero offset must not fold away a broadcast of a scalar base.
+// CHECK-LABEL: func.func @zero_offset_vector_broadcast(
+// CHECK-SAME: %[[PTR:.*]]: !ptr.ptr<#ptr.generic_space>)
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<4xi64>
+// CHECK-NEXT: %[[RESULT:.*]] = ptr.ptr_add %[[PTR]], %[[ZERO]] : !ptr.ptr<#ptr.generic_space>, vector<4xi64>
+// CHECK-NEXT: return %[[RESULT]] : vector<4x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_vector_broadcast(%ptr: !ptr.ptr<#ptr.generic_space>) -> vector<4x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant dense<0> : vector<4xi64>
+  %result = ptr.ptr_add %ptr, %zero : !ptr.ptr<#ptr.generic_space>, vector<4xi64>
+  return %result : vector<4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: func.func @zero_offset_scalable_vector_broadcast(
+// CHECK-SAME: %[[PTR:.*]]: !ptr.ptr<#ptr.generic_space>)
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<[4]xi32>
+// CHECK-NEXT: %[[RESULT:.*]] = ptr.ptr_add %[[PTR]], %[[ZERO]] : !ptr.ptr<#ptr.generic_space>, vector<[4]xi32>
+// CHECK-NEXT: return %[[RESULT]] : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_scalable_vector_broadcast(%ptr: !ptr.ptr<#ptr.generic_space>) -> vector<[4]x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant dense<0> : vector<[4]xi32>
+  %result = ptr.ptr_add %ptr, %zero : !ptr.ptr<#ptr.generic_space>, vector<[4]xi32>
+  return %result : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: func.func @zero_offset_tensor_broadcast(
+// CHECK-SAME: %[[PTR:.*]]: !ptr.ptr<#ptr.generic_space>)
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : tensor<2x4xi64>
+// CHECK-NEXT: %[[RESULT:.*]] = ptr.ptr_add %[[PTR]], %[[ZERO]] : !ptr.ptr<#ptr.generic_space>, tensor<2x4xi64>
+// CHECK-NEXT: return %[[RESULT]] : tensor<2x4x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_tensor_broadcast(%ptr: !ptr.ptr<#ptr.generic_space>) -> tensor<2x4x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant dense<0> : tensor<2x4xi64>
+  %result = ptr.ptr_add %ptr, %zero : !ptr.ptr<#ptr.generic_space>, tensor<2x4xi64>
+  return %result : tensor<2x4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// Type-preserving zero-offset folds still apply to shaped bases.
+// CHECK-LABEL: func.func @zero_offset_vector(
+// CHECK-SAME: %[[PTR:.*]]: vector<4x!ptr.ptr<#ptr.generic_space>>)
+// CHECK-NEXT: return %[[PTR]] : vector<4x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_vector(%ptr: vector<4x!ptr.ptr<#ptr.generic_space>>) -> vector<4x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant dense<0> : vector<4xi64>
+  %result = ptr.ptr_add %ptr, %zero : vector<4x!ptr.ptr<#ptr.generic_space>>, vector<4xi64>
+  return %result : vector<4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: func.func @zero_offset_scalable_vector_scalar_offset(
+// CHECK-SAME: %[[PTR:.*]]: vector<[4]x!ptr.ptr<#ptr.generic_space>>)
+// CHECK-NEXT: return %[[PTR]] : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_scalable_vector_scalar_offset(%ptr: vector<[4]x!ptr.ptr<#ptr.generic_space>>) -> vector<[4]x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant 0 : i32
+  %result = ptr.ptr_add %ptr, %zero : vector<[4]x!ptr.ptr<#ptr.generic_space>>, i32
+  return %result : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: func.func @zero_offset_tensor(
+// CHECK-SAME: %[[PTR:.*]]: tensor<2x4x!ptr.ptr<#ptr.generic_space>>)
+// CHECK-NEXT: return %[[PTR]] : tensor<2x4x!ptr.ptr<#ptr.generic_space>>
+func.func @zero_offset_tensor(%ptr: tensor<2x4x!ptr.ptr<#ptr.generic_space>>) -> tensor<2x4x!ptr.ptr<#ptr.generic_space>> {
+  %zero = arith.constant dense<0> : tensor<2x4xi64>
+  %result = ptr.ptr_add %ptr, %zero : tensor<2x4x!ptr.ptr<#ptr.generic_space>>, tensor<2x4xi64>
+  return %result : tensor<2x4x!ptr.ptr<#ptr.generic_space>>
+}
+
 /// Tests the the `from_ptr` folder.
 // CHECK-LABEL: @test_from_ptr_0
 // CHECK-SAME: (%[[MEM_REF:.*]]: memref<f32, #ptr.generic_space>)