[MLIR][LLVM][SROA] Fix pointer escape through stores bug

[Dinistro]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Christian Ulmann (Dinistro)

Changes

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.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp (+4)
• (modified) mlir/test/Dialect/LLVMIR/sroa.mlir (+13)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 0ef1d105aca6cb..f171bf7cc4bec3 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 ca49b1298b0e91..3f4d17c6a43f97 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
+}

[gysit]

LGTM