[Mlir-commits] [mlir] [MLIR][Bufferization] Fix IdentityLayoutMap allocation at function boundaries (PR #227253)
Aaditya Agarwal
llvmlistbot at llvm.org
Wed Sep 30 04:59:35 PDT 2026
https://github.com/aaditya8979 updated https://github.com/llvm/llvm-project/pull/227253
>From 5d3a9c78988af543c1ae939fa942b9d807435d88 Mon Sep 17 00:00:00 2001
From: openhands <openhands at all-hands.dev>
Date: Tue, 29 Sep 2026 15:13:13 +0530
Subject: [PATCH] [MLIR][Bufferization] Fix IdentityLayoutMap allocation at
function boundaries
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Resolves silent data corruption where `memref.cast` strips non-unit stride
layouts from subviews during function boundary bufferization under
`LayoutMapOption::IdentityLayoutMap`.
The root cause is in `castOrReallocMemRefValue`: the `isGuaranteedCastCompatible`
predicate correctly guards against dynamic-to-static offset/stride transitions,
but did not guard against non-identity-to-identity layout transitions. When the
source carries a non-zero or dynamic offset (e.g., from `memref.subview`) and
the destination requires a zero-offset identity layout, `areCastCompatible`
returns true and a `memref.cast` is emitted. This silently strips the offset
metadata while leaving the base pointer unchanged, causing the callee to read
from offset 0 of the original buffer instead of the correct slice.
Fix: extend `isGuaranteedCastCompatible` with an additional guard — if the
source has a non-identity layout and the destination requires identity layout,
return false, causing the function to fall through to the existing
`alloc + copy` path. This correctly covers both `func.call` boundary
materialization and the `to_buffer(to_tensor(x))` canonicalization path.
Added regression tests:
- `mlir/test/Dialect/Bufferization/canonicalize.mlir`: validates that
`to_buffer(to_tensor(x))` with strided/offset source and identity
destination emits alloc+copy, not memref.cast.
- `mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir`:
validates the same for `func.call` boundary bufferization under
`--function-boundary-type-conversion=identity-layout-map`.
Fixes #83276
---
.../Bufferization/IR/BufferizationOps.cpp | 9 +++++-
.../FuncBufferizableOpInterfaceImpl.cpp | 5 ++--
.../Transforms/one-shot-module-bufferize.mlir | 30 +++++++++++++++++++
.../Dialect/Bufferization/canonicalize.mlir | 27 +++++++++++++++++
4 files changed, 68 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 937134dea23b1..9b835c5afff5c 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -52,12 +52,19 @@ FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
for (auto it : zip(sourceStrides, targetStrides))
if (dynamicToStatic(std::get<0>(it), std::get<1>(it)))
return false;
+ // A cast cannot safely zero out a non-zero offset. If the source has a
+ // non-identity layout (e.g., a subview with a dynamic offset) and the
+ // destination requires identity layout (zero offset, unit strides), a cast
+ // would silently strip the offset and cause incorrect loads in the callee.
+ // Fall through to the alloc+copy path instead.
+ if (!source.getLayout().isIdentity() && target.getLayout().isIdentity())
+ return false;
return true;
};
// Note: If `areCastCompatible`, a cast is valid, but may fail at runtime. To
// ensure that we only generate casts that always succeed at runtime, we check
- // a fix extra conditions in `isGuaranteedCastCompatible`.
+ // a few extra conditions in `isGuaranteedCastCompatible`.
if (memref::CastOp::areCastCompatible(srcType, destType) &&
isGuaranteedCastCompatible(srcType, destType)) {
Value casted = *options.castFn(b, value.getLoc(), destType, value);
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index e76b51aa8015d..b3ed8a1eb03c5 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -322,8 +322,9 @@ struct CallOpInterface
auto memrefDstType = dyn_cast<MemRefType>(bufferType);
assert(memrefDstType &&
"buffer layout not supported on unranked tensors");
- FailureOr<Value> replacement = bufferization::castOrReallocMemRefValue(
- rewriter, buffer, memrefDstType, options);
+ FailureOr<Value> replacement =
+ bufferization::castOrReallocMemRefValue(rewriter, buffer,
+ memrefDstType, options);
if (failed(replacement))
return failure();
buffer = *replacement;
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index dbe44cba4d98f..0283837699cc3 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -852,3 +852,33 @@ func.func @ranked_return_via_unranked_call(%arg0: tensor<64x20x40xf32>) -> tenso
return %b : tensor<64x20x40xf32>
}
func.func private @relu_unranked(tensor<*xf32>) -> tensor<*xf32>
+
+
+// -----
+
+// Test that a non-identity layout (strided subview) crossing a function
+// boundary under IdentityLayoutMap triggers an explicit alloc+copy instead
+// of an illegal memref.cast that would silently strip the offset.
+// Reproduces issue #83276.
+//
+// CHECK-NO-LAYOUT-MAP-LABEL: func.func @caller(
+// CHECK-NO-LAYOUT-MAP-SAME: %[[ARG0:.*]]: memref<5xf32>
+// CHECK-NO-LAYOUT-MAP: %[[SUBVIEW:.*]] = memref.subview %[[ARG0]][%{{.*}}] [%{{.*}}] [1]
+// CHECK-NO-LAYOUT-MAP-SAME: : memref<5xf32> to memref<?xf32, strided<[1], offset: ?>>
+// CHECK-NO-LAYOUT-MAP-NOT: memref.cast %[[SUBVIEW]]
+// CHECK-NO-LAYOUT-MAP: %[[DIM:.*]] = memref.dim %[[SUBVIEW]]
+// CHECK-NO-LAYOUT-MAP: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
+// CHECK-NO-LAYOUT-MAP: memref.copy %[[SUBVIEW]], %[[ALLOC]]
+// CHECK-NO-LAYOUT-MAP: call @callee(%[[ALLOC]])
+func.func @caller(%arg0: tensor<5xf32>, %offset: index, %size: index) -> (f32) {
+ %extracted_slice = tensor.extract_slice %arg0[%offset] [%size] [1]
+ : tensor<5xf32> to tensor<?xf32>
+ %result = func.call @callee(%extracted_slice) : (tensor<?xf32>) -> (f32)
+ return %result : f32
+}
+func.func private @callee(%arg0: tensor<?xf32>) -> (f32) {
+ %0 = arith.constant 0 : index
+ %extracted = tensor.extract %arg0[%0] : tensor<?xf32>
+ return %extracted : f32
+}
+
diff --git a/mlir/test/Dialect/Bufferization/canonicalize.mlir b/mlir/test/Dialect/Bufferization/canonicalize.mlir
index d235e673c2336..9d88838f74caa 100644
--- a/mlir/test/Dialect/Bufferization/canonicalize.mlir
+++ b/mlir/test/Dialect/Bufferization/canonicalize.mlir
@@ -112,6 +112,33 @@ func.func @canonicalize_buffer_cast_of_tensor_load_to_copy(
// -----
+// Regression test for #83276: to_buffer(to_tensor(x)) with a strided/offset
+// source and an identity (zero-offset, contiguous) destination must emit an
+// alloc+copy rather than a memref.cast that would silently strip the offset
+// metadata, causing wrong-value loads in the consumer.
+//
+// CHECK-LABEL: func @to_buffer_identity_layout_requires_copy(
+// CHECK-SAME: %[[M:.*]]: memref<?xf32, strided<[1], offset: ?>>)
+// CHECK-SAME: -> memref<?xf32>
+// CHECK-NOT: bufferization.to_tensor
+// CHECK-NOT: bufferization.to_buffer
+// CHECK-NOT: memref.cast
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM:.*]] = memref.dim %[[M]], %[[C0]]
+// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) : memref<?xf32>
+// CHECK: memref.copy %[[M]], %[[ALLOC]]
+// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>> to memref<?xf32>
+// CHECK: return %[[ALLOC]]
+func.func @to_buffer_identity_layout_requires_copy(
+ %arg0: memref<?xf32, strided<[1], offset: ?>>)
+ -> memref<?xf32>
+{
+ %0 = bufferization.to_tensor %arg0 : memref<?xf32, strided<[1], offset: ?>> to tensor<?xf32>
+ %1 = bufferization.to_buffer %0 : tensor<?xf32> to memref<?xf32>
+ return %1 : memref<?xf32>
+}
+
+// -----
// Basic folding of tensor.dim(to_tensor(m)) -> memref.dim(m).
// CHECK-LABEL: func @dim_of_tensor_load(
More information about the Mlir-commits
mailing list