[Mlir-commits] [mlir] [mlir][bufferization] Fix crash when returning an allocation twice (PR #220232)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>



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


3 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp (+22-1) 
- (modified) mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir (+28) 
- (modified) mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index 434501b030e4a..ef6ea1fe099be 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -12,6 +12,9 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/WalkResult.h"
 
 namespace mlir {
 namespace bufferization {
@@ -168,7 +171,20 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
     }
     OpBuilder builder(op);
     SmallVector<SmallVector<Value>> dynamicSizes;
+    SetVector<Operation *> toErase;
+    DenseMap<Value, BlockArgument> primaryArgs;
+    DenseMap<Value, SmallVector<Value>> dynSizes;
     for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
+      auto [it, inserted] = primaryArgs.try_emplace(orig, arg);
+      if (!inserted) {
+        if (auto dynIt = dynSizes.find(orig); dynIt != dynSizes.end())
+          dynamicSizes.push_back(dynIt->second);
+        if (it->second != arg &&
+            failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
+          return WalkResult::interrupt();
+        continue;
+      }
+
       bool hoistStaticAllocs =
           options.hoistStaticAllocs &&
           cast<MemRefType>(orig.getType()).hasStaticShape();
@@ -182,13 +198,18 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
         if (hoistDynamicAllocs) {
           SmallVector<Value> dynamicSize = getDynamicSize(orig, func);
           dynamicSizes.push_back(dynamicSize);
+          dynSizes.try_emplace(orig, std::move(dynamicSize));
         }
-        orig.getDefiningOp()->erase();
+        toErase.insert(orig.getDefiningOp());
       } else {
         if (failed(options.memCpyFn(builder, op.getLoc(), orig, arg)))
           return WalkResult::interrupt();
       }
     }
+
+    for (Operation *eraseOp : toErase)
+      eraseOp->erase();
+
     func::ReturnOp::create(builder, op.getLoc(), keepAsReturnOperands);
     op.erase();
     auto dynamicSizePair =
diff --git a/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir b/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
index f33eb8e26fbce..bbe26f0f91daa 100644
--- a/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
+++ b/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
@@ -77,3 +77,31 @@ func.func @complex_alloc_test(%size0 : index, %size1: index) {
 //       CHECK:   call @complex_alloc(%[[size0]], %[[size1]], %[[alloc0]], %[[alloc1]], %[[alloc2]]) : (index, index, memref<?x?xf32>, memref<4xf32>, memref<?xf32>) -> ()
 //       CHECK:   "test.sink"(%[[alloc0]], %[[alloc1]], %[[alloc2]]) : (memref<?x?xf32>, memref<4xf32>, memref<?xf32>) -> ()
 //       CHECK: }
+
+// -----
+
+func.func private @duplicate_dynamic(%size : index)
+    -> (memref<?xf32>, memref<?xf32>) {
+  %alloc = memref.alloc(%size) : memref<?xf32>
+  return %alloc, %alloc : memref<?xf32>, memref<?xf32>
+}
+
+func.func @duplicate_dynamic_test(%size : index) {
+  %alloc0, %alloc1 = call @duplicate_dynamic(%size)
+      : (index) -> (memref<?xf32>, memref<?xf32>)
+  "test.sink"(%alloc0, %alloc1) : (memref<?xf32>, memref<?xf32>) -> ()
+}
+
+// CHECK-LABEL: func.func private @duplicate_dynamic(
+//  CHECK-SAME:   %{{.*}}: index,
+//  CHECK-SAME:   %[[OUT0:.*]]: memref<?xf32>,
+//  CHECK-SAME:   %[[OUT1:.*]]: memref<?xf32>) {
+//       CHECK:   memref.copy %[[OUT0]], %[[OUT1]] : memref<?xf32> to memref<?xf32>
+//       CHECK:   return
+
+// CHECK-LABEL: func.func @duplicate_dynamic_test(
+//  CHECK-SAME:   %[[SIZE:.*]]: index) {
+//       CHECK:   %[[ALLOC0:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+//       CHECK:   %[[ALLOC1:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+//       CHECK:   call @duplicate_dynamic(%[[SIZE]], %[[ALLOC0]], %[[ALLOC1]]) : (index, memref<?xf32>, memref<?xf32>) -> ()
+//       CHECK:   "test.sink"(%[[ALLOC0]], %[[ALLOC1]]) : (memref<?xf32>, memref<?xf32>) -> ()
diff --git a/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir b/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
index 97fb9f45f57f8..5a4a7c7eb5d30 100644
--- a/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
+++ b/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
@@ -49,3 +49,14 @@ func.func private @return_arg(%arg0: memref<128x256xf32>, %arg1: memref<128x256x
   "test.source"(%arg0, %arg1)  : (memref<128x256xf32>, memref<128x256xf32>) -> ()
   return %arg0 : memref<128x256xf32>
 }
+
+// CHECK-LABEL: func.func private @duplicate_return_value(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<4xf32>, %[[ARG1:.*]]: memref<4xf32>) {
+// CHECK-NOT:     memref.alloc
+// CHECK:         memref.copy %[[ARG0]], %[[ARG1]] : memref<4xf32> to memref<4xf32>
+// CHECK:         return
+func.func private @duplicate_return_value()
+    -> (memref<4xf32>, memref<4xf32>) {
+  %a = memref.alloc() : memref<4xf32>
+  return %a, %a : memref<4xf32>, memref<4xf32>
+}

``````````

</details>


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