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

[llvmlistbot at llvm.org]

Author: pzoln
Date: 2026-07-18T11:58:44Z
New Revision: 1f4d6acc935f93eca55cfef67348c8abcc1415db

URL: https://github.com/llvm/llvm-project/commit/1f4d6acc935f93eca55cfef67348c8abcc1415db
DIFF: https://github.com/llvm/llvm-project/commit/1f4d6acc935f93eca55cfef67348c8abcc1415db.diff

LOG: [mlir][bufferization] Avoid invalid memref.cast in to_buffer folding (#208973)

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.

 Assisted-by: OpenAI Codex

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
    mlir/test/Dialect/Bufferization/canonicalize.mlir

Removed: 
    


################################################################################
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_
diff erent_address_space(%arg0
 
 // -----
 
+// If unranked memrefs are not cast-compatible, don't fold them.
+// CHECK-LABEL: func @canonicalize_unranked_buffer_cast_of_tensor_load_
diff erent_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_
diff erent_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(