[Mlir-commits] [mlir] cb300c3 - [MLIR][LLVM][SROA] Fix pointer escape through stores bug (#86291)

[llvmlistbot at llvm.org]

Author: Christian Ulmann
Date: 2024-03-22T16:44:06+01:00
New Revision: cb300c33059c1d14f72392ce5dffcf050ad7567d

URL: https://github.com/llvm/llvm-project/commit/cb300c33059c1d14f72392ce5dffcf050ad7567d
DIFF: https://github.com/llvm/llvm-project/commit/cb300c33059c1d14f72392ce5dffcf050ad7567d.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
    mlir/test/Dialect/LLVMIR/sroa.mlir

Removed: 
    


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