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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Alexander Shaposhnikov (alexander-shaposhnikov)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/231077.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp (+2-1) 
- (modified) mlir/test/Dialect/Ptr/canonicalize.mlir (+62) 


``````````diff
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>)

``````````

</details>


https://github.com/llvm/llvm-project/pull/231077