[Mlir-commits] [mlir] [MLIR][Bufferization] Fix IdentityLayoutMap allocation at function boundaries (PR #227253)

[Aaditya Agarwal]

https://github.com/aaditya8979 updated https://github.com/llvm/llvm-project/pull/227253

>From 024cf836fd88db9549cc03e785a3267a68c84007 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

Resolves silent data corruption where `memref.cast` strips non-unit stride
layouts from subviews during function boundary bufferization under
`IdentityLayoutMap`. Injects explicit allocation and `memref.copy` for
non-identity to identity layout transitions.

Fixes #83276
---
 .../FuncBufferizableOpInterfaceImpl.cpp       | 33 ++++++++++++++++---
 .../Transforms/one-shot-module-bufferize.mlir | 30 +++++++++++++++++
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index e76b51aa8015d..b57e0d6275144 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -13,6 +13,7 @@
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include <optional>
@@ -322,11 +323,33 @@ 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);
-        if (failed(replacement))
-          return failure();
-        buffer = *replacement;
+        auto memrefSrcType = dyn_cast<MemRefType>(buffer.getType());
+        if (memrefSrcType && !memrefSrcType.getLayout().isIdentity() &&
+            memrefDstType.getLayout().isIdentity()) {
+          SmallVector<Value, 4> dynamicOperands;
+          for (int64_t i = 0, e = memrefDstType.getRank(); i < e; ++i) {
+            if (memrefDstType.getShape()[i] != ShapedType::kDynamic)
+              continue;
+            dynamicOperands.push_back(
+                rewriter.create<memref::DimOp>(callOp.getLoc(), buffer, i));
+          }
+          FailureOr<Value> alloc =
+              options.allocationFn(rewriter, callOp.getLoc(), memrefDstType,
+                                   dynamicOperands, options.bufferAlignment);
+          if (failed(alloc))
+            return failure();
+          if (failed(
+                  options.memCpyFn(rewriter, callOp.getLoc(), buffer, *alloc)))
+            return failure();
+          buffer = *alloc;
+        } else {
+          FailureOr<Value> replacement =
+              bufferization::castOrReallocMemRefValue(rewriter, buffer,
+                                                      memrefDstType, options);
+          if (failed(replacement))
+            return failure();
+          buffer = *replacement;
+        }
       }
       newOperands.push_back(buffer);
     }
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
+}
+