[mlir][bufferization] Fix SimplifyClones with dealloc before cloneOp

[sott0n]

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

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Kohei Yamaguchi (sott0n)

Changes

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

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+5)
• (modified) mlir/test/Dialect/Bufferization/canonicalize.mlir (+12)

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>) ->

[sott0n]

@matthias-springer Could you land it? I don't have commit access.