[Mlir-commits] [mlir] [mlir][bufferization] Avoid invalid memref.cast in to_buffer folding (PR #208973)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: pzoln

<details>
<summary>Changes</summary>

Fixes #<!-- -->202782

When folding bufferization.to_buffer(bufferization.to_tensor(...)), only create a memref.cast if the source and destination memref types are cast-compatible.

This avoids asserting on unranked memrefs with different memory spaces and leaves the original bufferization ops in place instead.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+5-4) 
- (modified) mlir/test/Dialect/Bufferization/canonicalize.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 45f5b3eaa5aea..80db89aa1b4bc 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -132,10 +132,11 @@ LogicalResult mlir::bufferization::foldToBufferToTensorPair(
   if (unrankedSrcType && rankedDestType)
     return failure();
 
-  // Unranked memref -> unranked memref cast
-  // Ranked memref -> unranked memref cast: No copy needed.
-  assert(memref::CastOp::areCastCompatible(srcType, destType) &&
-         "expected that types are cast compatible");
+  // Unranked/ranked memref -> unranked memref cast: No copy needed if the types
+  // are cast-compatible.
+  if (!memref::CastOp::areCastCompatible(srcType, destType))
+    return failure();
+
   rewriter.replaceOpWithNewOp<memref::CastOp>(toBuffer, destType,
                                               bufferToTensor.getBuffer());
   return success();
diff --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index df07511798b91..ab3d8df892bc1 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -50,6 +50,24 @@ func.func @canonicalize_buffer_cast_of_tensor_load_different_address_space(%arg0
 
 // -----
 
+// If unranked memrefs are not cast-compatible, don't fold them.
+// CHECK-LABEL: func @canonicalize_unranked_buffer_cast_of_tensor_load_different_address_space(
+//  CHECK-SAME:   %[[MEMREF:.*]]: memref<*xi64>)
+//  CHECK-SAME:     -> memref<*xi64, 1> {
+//  CHECK-NOT: memref.cast
+//      CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<*xi64> to tensor<*xi64>
+//      CHECK: %[[BUFFER:.*]] = bufferization.to_buffer %[[TENSOR]] : tensor<*xi64> to memref<*xi64, 1>
+//  CHECK-NOT: memref.cast
+//      CHECK: return %[[BUFFER]] : memref<*xi64, 1>
+func.func @canonicalize_unranked_buffer_cast_of_tensor_load_different_address_space(%arg0: memref<*xi64>)
+    -> memref<*xi64, 1> {
+  %0 = bufferization.to_tensor %arg0 : memref<*xi64> to tensor<*xi64>
+  %1 = bufferization.to_buffer %0 : tensor<*xi64> to memref<*xi64, 1>
+  return %1 : memref<*xi64, 1>
+}
+
+// -----
+
 // If the memrefs are definitely cast-compatible, canonicalize to
 //            cast.
 // CHECK-LABEL: func @canonicalize_buffer_cast_of_tensor_load(

``````````

</details>


https://github.com/llvm/llvm-project/pull/208973