[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 7b83b2f8fca16a6fbbef2233c76535a8e74debbe 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       | 31 ++++++++++++++++---
 .../Transforms/one-shot-module-bufferize.mlir | 23 ++++++++++++++
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index e76b51aa8015db..db29920055539c 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,31 @@ 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 dbe44cba4d98f7..fdc11d152a9df4 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,26 @@ func.func @ranked_return_via_unranked_call(%arg0: tensor<64x20x40xf32>) -> tenso
   return %b : tensor<64x20x40xf32>
 }
 func.func private @relu_unranked(tensor<*xf32>) -> tensor<*xf32>
+
+
+// -----
+
+// CHECK-LABEL: func.func @caller
+// CHECK:       %[[SUBVIEW:.*]] = memref.subview
+// CHECK-SAME:    strided<[1], offset: ?>
+// CHECK-NOT:   memref.cast %[[SUBVIEW]]
+// CHECK:       %[[ALLOC:.*]] = memref.alloc
+// CHECK:       memref.copy %[[SUBVIEW]], %[[ALLOC]]
+// CHECK:       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
+}
+