[MLIR][LLVM][SROA] Fix pointer escape through stores bug (#86291)

This commit resolves a SROA bug caused by not properly checking if a
llvm store operation writes the pointer to memory or not. Now, we do no
longer consider stores that use a slot pointer as a value to store as
fixable.

2 files changed, 17 insertions, 0 deletions

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 0ef1d10..f171bf7 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -251,6 +251,10 @@ bool LLVM::StoreOp::canRewire(const DestructurableMemorySlot &slot,
   if (getVolatile_())
     return false;
 
+  // Storing the pointer to memory cannot be dealt with.
+  if (getValue() == slot.ptr)
+    return false;
+
   // A store always accesses the first element of the destructured slot.
   auto index = IntegerAttr::get(IntegerType::get(getContext(), 32), 0);
   Type subslotType = getTypeAtIndex(slot, index);
diff --git a/mlir/test/Dialect/LLVMIR/sroa.mlir b/mlir/test/Dialect/LLVMIR/sroa.mlir
index ca49b12..3f4d17c 100644
--- a/mlir/test/Dialect/LLVMIR/sroa.mlir
+++ b/mlir/test/Dialect/LLVMIR/sroa.mlir
@@ -305,3 +305,16 @@ llvm.func @vector_store_type_mismatch(%arg: vector<4xi32>) {
   llvm.store %arg, %1 : vector<4xi32>, !llvm.ptr
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @store_to_memory
+// CHECK-SAME: %[[ARG:.*]]: !llvm.ptr
+llvm.func @store_to_memory(%arg: !llvm.ptr) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (vector<4xf32>)> : (i32) -> !llvm.ptr
+  // CHECK-NEXT: llvm.store %[[ALLOCA]], %[[ARG]]
+  llvm.store %1, %arg : !llvm.ptr, !llvm.ptr
+  llvm.return
+}