[Mlir-commits] [mlir] bcd14b0 - [mlir][bufferization] Fix SimplifyClones with dealloc before cloneOp (#79098)

[llvmlistbot at llvm.org]

Author: Kohei Yamaguchi
Date: 2024-01-25T11:06:46+01:00
New Revision: bcd14b099dd3d0b9c9336e5cdf7e35279c5cc33b

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

LOG: [mlir][bufferization] Fix SimplifyClones with dealloc before cloneOp (#79098)

The SimplifyClones pass relies on the assumption that the deallocOp
follows the cloneOp. However, a crash occurs when there is a
redundantDealloc preceding the cloneOp. This PR addresses the issue by
ensuring the presence of deallocOp after cloneOp. The verification is
performed by checking if the loop of the sub sequent node of cloneOp
reaches the tail of the list.

Fix #74306

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
    mlir/test/Dialect/Bufferization/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 253fcf2525121b8..eb4a96f35499048 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -505,6 +505,9 @@ struct SimplifyClones : public OpRewritePattern<CloneOp> {
     // of the source.
     for (Operation *pos = cloneOp->getNextNode(); pos != redundantDealloc;
          pos = pos->getNextNode()) {
+      // Bail if we run out of operations while looking for a deallocation op.
+      if (!pos)
+        return failure();
       auto effectInterface = dyn_cast<MemoryEffectOpInterface>(pos);
       if (!effectInterface)
         continue;

diff  --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index 3edae7827f25f71..b6c0a0e25efe0ea 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -235,6 +235,18 @@ func.func @clone_and_realloc(%arg0: memref<?xf32>) {
 
 // -----
 
+// Verify SimplifyClones skips clones with preceding deallocation.
+// CHECK-LABEL: @clone_and_preceding_dealloc
+func.func @clone_and_preceding_dealloc(%arg0: memref<?xf32>) -> memref<32xf32> {
+  memref.dealloc %arg0 : memref<?xf32>
+  %0 = bufferization.clone %arg0 : memref<?xf32> to memref<32xf32>
+  return %0 : memref<32xf32>
+}
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+// CHECK-NOT: %cast = memref.cast %[[ARG]]
+
+// -----
+
 // CHECK-LABEL: func @tensor_cast_to_memref
 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<4x6x16x32xi8>
 func.func @tensor_cast_to_memref(%arg0 : tensor<4x6x16x32xi8>) ->