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

[llvmlistbot at llvm.org]

https://github.com/pzoln created https://github.com/llvm/llvm-project/pull/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.

>From 71ea7d95a7564a43c08e05580f92f8cd39c65768 Mon Sep 17 00:00:00 2001
From: Pavel Zolnikov <pavel at avarivent.com>
Date: Sat, 11 Jul 2026 18:00:45 -0700
Subject: [PATCH] [mlir][bufferization] Avoid invalid memref.cast in to_buffer
 folding

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.

Fixes #202782.
---
 .../Bufferization/IR/BufferizationOps.cpp      |  9 +++++----
 .../Dialect/Bufferization/canonicalize.mlir    | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

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(