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

[Kohei Yamaguchi]

https://github.com/sott0n created https://github.com/llvm/llvm-project/pull/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

>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] [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 253fcf2525121b..9275a1e4aacee9 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 3edae7827f25f7..b6c0a0e25efe0e 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>) ->