[Mlir-commits] [mlir] [mlir][bufferization] Clone simplify fails when input and result type not cast compatiable (PR #71310)

[donald chen]

https://github.com/cxy-1993 created https://github.com/llvm/llvm-project/pull/71310

The simplify  of bufferization.clone generates a memref.cast op, but the checks in simplify do not verify whether the operand types and return types of clone op is compatiable, leading to errors. This patch addresses this issue.

>From 36d29eeaaf7e88a1f2874ba9dc8d48c807eca15d Mon Sep 17 00:00:00 2001
From: cxy <chenxunyu1993 at gmail.com>
Date: Sun, 5 Nov 2023 17:31:44 +0800
Subject: [PATCH] [mlir][bufferization] Clone simplify fails when input and
 result type not cast compatiable

---
 .../Bufferization/IR/BufferizationOps.cpp     |  4 +++
 .../Dialect/Bufferization/canonicalize.mlir   | 27 +++++++++----------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index ec5feab1ed0d856..edaa3945b70eb50 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -457,6 +457,10 @@ struct SimplifyClones : public OpRewritePattern<CloneOp> {
     }
 
     Value source = cloneOp.getInput();
+    if (!memref::CastOp::areCastCompatible({source.getType()},
+                                           {cloneOp.getType()}))
+      return failure();
+
     // Aims to find the dealloc op for the canonical source
     // which otherwise could prevent removal of unnecessary allocs.
     Value canonicalSource = source;
diff --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index 3ba283928a83f0e..3e183c7d9449817 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -156,6 +156,18 @@ func.func @clone_and_cast(%arg0: memref<?xf32>) -> memref<32xf32> {
 
 // -----
 
+// CHECK-LABEL: @clone_incompatible
+func.func @clone_incompatible(%arg0: memref<32xf32, strided<[2]>>) -> memref<32xf32> {
+  %0 = bufferization.clone %arg0 : memref<32xf32, strided<[2]>> to memref<32xf32>
+  memref.dealloc %arg0 : memref<32xf32, strided<[2]>>
+  return %0 : memref<32xf32>
+}
+// CHECK-SAME: %[[ARG:.*]]: memref<32xf32, strided<[2]>>
+// CHECK-NEXT: bufferization.clone %[[ARG]] : memref<32xf32, strided<[2]>> to memref<32xf32>
+// CHECK-NOT: memref.cast
+
+// -----
+
 // CHECK-LABEL: @alias_is_freed
 func.func @alias_is_freed(%arg0 : memref<?xf32>) {
   %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>
@@ -267,21 +279,6 @@ func.func @alloc_tensor_canonicalize() -> (tensor<4x5x?xf32>) {
 
 // -----
 
-func.func @dealloc_canonicalize_clone_removal(%arg0: memref<?xindex>) -> memref<*xf32> {
-  %c1 = arith.constant 1 : index
-  %0 = memref.alloc(%c1) : memref<?xf32>
-  %1 = memref.reshape %0(%arg0) : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
-  %2 = bufferization.clone %1 : memref<*xf32> to memref<*xf32>
-  memref.dealloc %0 : memref<?xf32>
-  return %2 : memref<*xf32>
-}
-// CHECK-LABEL: @dealloc_canonicalize_clone_removal
-//   CHECK-NOT:   bufferization.clone
-//   CHECK-NOT:   memref.dealloc
-//       CHECK:   return {{.*}}
-
-// -----
-
 func.func @dealloc_canonicalize_duplicates(%arg0: memref<2xi32>, %arg1: i1, %arg2: i1, %arg3: memref<2xi32>, %arg4: memref<2xi32>, %arg5: memref<2xi32>) -> (i1, i1, i1) {
   %0:3 = bufferization.dealloc (%arg4, %arg0, %arg0 : memref<2xi32>, memref<2xi32>, memref<2xi32>) if (%arg1, %arg1, %arg1) retain (%arg3, %arg5, %arg3 : memref<2xi32>, memref<2xi32>, memref<2xi32>)
   bufferization.dealloc (%arg0, %arg0 : memref<2xi32>, memref<2xi32>) if (%arg1, %arg2)