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

[Kohei Yamaguchi]

https://github.com/sott0n updated https://github.com/llvm/llvm-project/pull/79098

>From cceb921d9bc7e3a4c1efe7050d6bfe6cb2eb49e7 Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Tue, 23 Jan 2024 14:30:02 +0000
Subject: [PATCH 1/2] [mlir][bufferization] Fix SimplifyClones with dealloc
 before cloneOp

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
---
 .../Dialect/Bufferization/IR/BufferizationOps.cpp    |  5 +++++
 mlir/test/Dialect/Bufferization/canonicalize.mlir    | 12 ++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 253fcf2525121b8..9275a1e4aacee92 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -505,6 +505,11 @@ struct SimplifyClones : public OpRewritePattern<CloneOp> {
     // of the source.
     for (Operation *pos = cloneOp->getNextNode(); pos != redundantDealloc;
          pos = pos->getNextNode()) {
+      // If the next node is nullptr, it indicates there is no redundantDealloc
+      // after cloneOp, means that there is a redundantDealloc preseding
+      // cloneOp.
+      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>) ->

>From 0ecc46a516fc67c84b6d3ba86cbee46592ed2a3f Mon Sep 17 00:00:00 2001
From: Kohei Yamaguchi <fix7211 at gmail.com>
Date: Tue, 23 Jan 2024 17:27:02 +0000
Subject: [PATCH 2/2] addressed a comment

---
 mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 9275a1e4aacee92..eb4a96f35499048 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -505,9 +505,7 @@ struct SimplifyClones : public OpRewritePattern<CloneOp> {
     // of the source.
     for (Operation *pos = cloneOp->getNextNode(); pos != redundantDealloc;
          pos = pos->getNextNode()) {
-      // If the next node is nullptr, it indicates there is no redundantDealloc
-      // after cloneOp, means that there is a redundantDealloc preseding
-      // cloneOp.
+      // Bail if we run out of operations while looking for a deallocation op.
       if (!pos)
         return failure();
       auto effectInterface = dyn_cast<MemoryEffectOpInterface>(pos);