[Mlir-commits] [mlir] [MLIR][LLVM][SROA] Avoid splitting dynamically indexed allocas (PR #85758)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Christian Ulmann (Dinistro)

<details>
<summary>Changes</summary>

This commit ensures that SROA does no longer attempt to split allocas that are indexed into dynamically. Dynamic indices into arrays are allowed to be negative or out-of-bounds, when the alloca containing the array has memory backing these produced indices.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp (+8) 
- (modified) mlir/test/Dialect/LLVMIR/sroa.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index bf24194d03ddb2..81766b307d3d78 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -344,6 +344,10 @@ LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
     return failure();
   if (!isFirstIndexZero(*this))
     return failure();
+  // Dynamic indices can be out-of-bounds (even negative), so an access with
+  // dynamic indices can never be considered safe.
+  if (!getDynamicIndices().empty())
+    return failure();
   Type reachedType = getResultPtrElementType();
   if (!reachedType)
     return failure();
@@ -362,6 +366,10 @@ bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
     return false;
   if (!isFirstIndexZero(*this))
     return false;
+  // Dynamic indices can be out-of-bounds (even negative), so an access with
+  // dynamic indices can never be properly rewired.
+  if (!getDynamicIndices().empty())
+    return false;
   Type reachedType = getResultPtrElementType();
   if (!reachedType || getIndices().size() < 2)
     return false;
diff --git a/mlir/test/Dialect/LLVMIR/sroa.mlir b/mlir/test/Dialect/LLVMIR/sroa.mlir
index f56ea53ac029e7..02d25f27f978a6 100644
--- a/mlir/test/Dialect/LLVMIR/sroa.mlir
+++ b/mlir/test/Dialect/LLVMIR/sroa.mlir
@@ -197,3 +197,21 @@ llvm.func @no_dynamic_indexing(%arg: i32) -> i32 {
   // CHECK: llvm.return %[[RES]] : i32
   llvm.return %3 : i32
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @no_nested_dynamic_indexing
+// CHECK-SAME: (%[[ARG:.*]]: i32)
+llvm.func @no_nested_dynamic_indexing(%arg: i32) -> i32 {
+  // CHECK: %[[SIZE:.*]] = llvm.mlir.constant(1 : i32)
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[SIZE]] x !llvm.struct<(array<10 x i32>, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  %1 = llvm.alloca %0 x !llvm.struct<(array<10 x i32>, i32)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+  // CHECK-NOT: = llvm.alloca
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0, %[[ARG]]]
+  %2 = llvm.getelementptr %1[0, 0, %arg] : (!llvm.ptr, i32) -> !llvm.ptr, !llvm.struct<(array<10 x i32>, i32)>
+  // CHECK: %[[RES:.*]] = llvm.load %[[GEP]]
+  %3 = llvm.load %2 : !llvm.ptr -> i32
+  // CHECK: llvm.return %[[RES]] : i32
+  llvm.return %3 : i32
+}

``````````

</details>


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