[Mlir-commits] [mlir] [mlir][ptr] Lower vector pointer addition to LLVM GEP (PR #231080)

[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>

The ptr.ptr_add conversion only accepts scalar LLVM pointer bases and uses the base type as the GEP result type. This leaves vector-base additions unconverted and produces a scalar GEP with an unresolved conversion cast when a vector offset broadcasts a scalar base.

Accept vector pointer bases and derive the GEP result type from the converted operation result. Require a scalar or vector LLVM pointer result so aggregate types are not used as GEP results.

Add tests for scalar/vector operand combinations with fixed and scalable vectors, including index offsets and preservation of no-wrap flags.

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


2 Files Affected:

- (modified) mlir/lib/Conversion/PtrToLLVM/PtrToLLVM.cpp (+9-3) 
- (added) mlir/test/Conversion/PtrToLLVM/ptr-add.mlir (+59) 


``````````diff
diff --git a/mlir/lib/Conversion/PtrToLLVM/PtrToLLVM.cpp b/mlir/lib/Conversion/PtrToLLVM/PtrToLLVM.cpp
index cc1b3242c2619..b7da737e0c615 100644
--- a/mlir/lib/Conversion/PtrToLLVM/PtrToLLVM.cpp
+++ b/mlir/lib/Conversion/PtrToLLVM/PtrToLLVM.cpp
@@ -299,9 +299,15 @@ PtrAddOpConversion::matchAndRewrite(ptr::PtrAddOp op, OpAdaptor adaptor,
                                     ConversionPatternRewriter &rewriter) const {
   // Get and check the base.
   Value base = adaptor.getBase();
-  if (!isa<LLVM::LLVMPointerType>(base.getType()))
+  if (!isa<LLVM::LLVMPointerType>(getElementTypeOrSelf(base.getType())))
     return rewriter.notifyMatchFailure(op, "Incompatible pointer type");
 
+  // A vector offset can broadcast a scalar base to a vector result.
+  Type resultType = getTypeConverter()->convertType(op.getType());
+  if (!resultType ||
+      !isa<LLVM::LLVMPointerType>(getElementTypeOrSelf(resultType)))
+    return rewriter.notifyMatchFailure(op, "Incompatible result type");
+
   // Get the offset.
   Value offset = adaptor.getOffset();
 
@@ -326,8 +332,8 @@ PtrAddOpConversion::matchAndRewrite(ptr::PtrAddOp op, OpAdaptor adaptor,
   }
 
   // Create the GEP operation with appropriate arguments
-  rewriter.replaceOpWithNewOp<LLVM::GEPOp>(op, base.getType(), elementType,
-                                           base, ValueRange{offset}, flags);
+  rewriter.replaceOpWithNewOp<LLVM::GEPOp>(op, resultType, elementType, base,
+                                           ValueRange{offset}, flags);
   return success();
 }
 
diff --git a/mlir/test/Conversion/PtrToLLVM/ptr-add.mlir b/mlir/test/Conversion/PtrToLLVM/ptr-add.mlir
new file mode 100644
index 0000000000000..07fc6a196c657
--- /dev/null
+++ b/mlir/test/Conversion/PtrToLLVM/ptr-add.mlir
@@ -0,0 +1,59 @@
+// RUN: mlir-opt %s -convert-to-llvm | FileCheck %s
+
+// Broadcast a scalar base to the shape of the offsets.
+// CHECK-LABEL: llvm.func @scalar_base_vector_offset(
+// CHECK-SAME: %[[BASE:.*]]: !llvm.ptr, %[[OFFSET:.*]]: vector<4xi64>) -> vector<4x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr %[[BASE]][%[[OFFSET]]] : (!llvm.ptr, vector<4xi64>) -> vector<4x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<4x!llvm.ptr>
+func.func @scalar_base_vector_offset(%base: !ptr.ptr<#ptr.generic_space>, %offset: vector<4xi64>) -> vector<4x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add %base, %offset : !ptr.ptr<#ptr.generic_space>, vector<4xi64>
+  return %result : vector<4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// Broadcast a scalar offset to the shape of the base; preserve the flags.
+// CHECK-LABEL: llvm.func @vector_base_scalar_offset(
+// CHECK-SAME: %[[BASE:.*]]: vector<4x!llvm.ptr>, %[[OFFSET:.*]]: i32) -> vector<4x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr nusw %[[BASE]][%[[OFFSET]]] : (vector<4x!llvm.ptr>, i32) -> vector<4x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<4x!llvm.ptr>
+func.func @vector_base_scalar_offset(%base: vector<4x!ptr.ptr<#ptr.generic_space>>, %offset: i32) -> vector<4x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add nusw %base, %offset : vector<4x!ptr.ptr<#ptr.generic_space>>, i32
+  return %result : vector<4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// Convert index elements while preserving element-wise pointer addition.
+// CHECK-LABEL: llvm.func @vector_base_vector_offset(
+// CHECK-SAME: %[[BASE:.*]]: vector<4x!llvm.ptr>, %[[OFFSET:.*]]: vector<4xi64>) -> vector<4x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr nuw %[[BASE]][%[[OFFSET]]] : (vector<4x!llvm.ptr>, vector<4xi64>) -> vector<4x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<4x!llvm.ptr>
+func.func @vector_base_vector_offset(%base: vector<4x!ptr.ptr<#ptr.generic_space>>, %offset: vector<4xindex>) -> vector<4x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add nuw %base, %offset : vector<4x!ptr.ptr<#ptr.generic_space>>, vector<4xindex>
+  return %result : vector<4x!ptr.ptr<#ptr.generic_space>>
+}
+
+// The same combinations work for scalable vectors.
+// CHECK-LABEL: llvm.func @scalar_base_scalable_offset(
+// CHECK-SAME: %[[BASE:.*]]: !llvm.ptr, %[[OFFSET:.*]]: vector<[4]xi64>) -> vector<[4]x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr inbounds %[[BASE]][%[[OFFSET]]] : (!llvm.ptr, vector<[4]xi64>) -> vector<[4]x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<[4]x!llvm.ptr>
+func.func @scalar_base_scalable_offset(%base: !ptr.ptr<#ptr.generic_space>, %offset: vector<[4]xindex>) -> vector<[4]x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add inbounds %base, %offset : !ptr.ptr<#ptr.generic_space>, vector<[4]xindex>
+  return %result : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: llvm.func @scalable_base_scalar_offset(
+// CHECK-SAME: %[[BASE:.*]]: vector<[4]x!llvm.ptr>, %[[OFFSET:.*]]: i64) -> vector<[4]x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr %[[BASE]][%[[OFFSET]]] : (vector<[4]x!llvm.ptr>, i64) -> vector<[4]x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<[4]x!llvm.ptr>
+func.func @scalable_base_scalar_offset(%base: vector<[4]x!ptr.ptr<#ptr.generic_space>>, %offset: i64) -> vector<[4]x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add %base, %offset : vector<[4]x!ptr.ptr<#ptr.generic_space>>, i64
+  return %result : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+}
+
+// CHECK-LABEL: llvm.func @scalable_base_scalable_offset(
+// CHECK-SAME: %[[BASE:.*]]: vector<[4]x!llvm.ptr>, %[[OFFSET:.*]]: vector<[4]xi32>) -> vector<[4]x!llvm.ptr>
+// CHECK-NEXT: %[[RESULT:.*]] = llvm.getelementptr %[[BASE]][%[[OFFSET]]] : (vector<[4]x!llvm.ptr>, vector<[4]xi32>) -> vector<[4]x!llvm.ptr>, i8
+// CHECK-NEXT: llvm.return %[[RESULT]] : vector<[4]x!llvm.ptr>
+func.func @scalable_base_scalable_offset(%base: vector<[4]x!ptr.ptr<#ptr.generic_space>>, %offset: vector<[4]xi32>) -> vector<[4]x!ptr.ptr<#ptr.generic_space>> {
+  %result = ptr.ptr_add %base, %offset : vector<[4]x!ptr.ptr<#ptr.generic_space>>, vector<[4]xi32>
+  return %result : vector<[4]x!ptr.ptr<#ptr.generic_space>>
+}

``````````

</details>


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