[mlir][bufferization] Clone simplify fails when input and result type not cast compatiable

[cxy-1993]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: donald chen (cxy-1993)

Changes

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.

---

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

2 Files Affected:

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

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)

[cxy-1993]

@joker-eph @pifon2a Could you please help me review this patch? I do not have access to add reviewers, thanks!

[cxy-1993]

@pifon2a @matthias-springer Could you please help me review this patch? Thanks.

[cxy-1993]

@joker-eph are you available? Could help with a code review on this patch.

[joker-eph]

@matthias-springer already approved, and he knows the Bufferization dialect better than I do, do you need me for something in particular here?

[cxy-1993]

@matthias-springer already approved, and he knows the Bufferization dialect better than I do, do you need me for something in particular here?

I'm not very clear about the merge rules. I would like to ask for your help to review and merge this pull request.

[joker-eph]

This is approved by Matthias, so right now you can click the "squash and merge" button.
It is advised to wait for the pre-merge checks to complete (in particular the CI check "buildkite/github-pull-requests Pending") before merging.

[cxy-1993]

This is approved by Matthias, so right now you can click the "squash and merge" button. It is advised to wait for the pre-merge checks to complete (in particular the CI check "buildkite/github-pull-requests Pending") before merging.

Thank you for your patient explanation. I don't have this button, perhaps because I don't have write access for the LLVM repository.

[joker-eph]

Ah sorry, let me merge this for you (as soon as the CI completes)
In general you can just mention it to the reviewer and they'll happily hit the button for you.

[cxy-1993]

Ah sorry, let me merge this for you (as soon as the CI completes) In general you can just mention it to the reviewer and they'll happily hit the button for you.

Thank you for your help and advice.