[MLIR][LLVM][SROA] Avoid splitting dynamically indexed allocas (#85758)

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.

Author: Dinistro

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;

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
+}