[Mlir-commits] [mlir] 83af810 - [mlir][Bufferization] Don't assert in foldMemRefCasts on a func body without func.return (#217515)

[llvmlistbot at llvm.org]

Author: Aman Singh
Date: 2026-08-21T10:46:28+02:00
New Revision: 83af810fa5190d779413ed16fb3eec6e80db39db

URL: https://github.com/llvm/llvm-project/commit/83af810fa5190d779413ed16fb3eec6e80db39db
DIFF: https://github.com/llvm/llvm-project/commit/83af810fa5190d779413ed16fb3eec6e80db39db.diff

LOG: [mlir][Bufferization] Don't assert in foldMemRefCasts on a func body without func.return (#217515)

`foldMemRefCasts` (used by One-Shot Module Bufferize when inferring
function result layouts) assumes every `func.func` body with a region
ends in a `func::ReturnOp`, and asserts if `getReturnOps()` finds
none. That assumption does not hold: `func.func`'s verifier only
requires the region to end in some terminator, not specifically
`func::ReturnOp`, so a body ending in e.g. a different dialect's
terminator parses and verifies fine but crashes this pass. This code
path is normally reached through
`insertTensorCopies`/`analyzeModuleOp`, which already reports this
case gracefully via `emitError` — but that analysis is skipped
entirely when `copy-before-write=1` is set, going straight to
`bufferizeModuleOp` and hitting the unguarded assert.

Fix by bailing out of `foldMemRefCasts` (nothing to fold) when there
are no `func::ReturnOp`s, mirroring how the rest of the file already
tolerates an empty `getReturnOps()` result.

Verified: reverting this fix reproduces the reported crash; with the
fix, `mlir-opt -one-shot-bufferize='bufferize-function-boundaries=1
copy-before-write=1'` on the reported reproducer succeeds, and

`mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-non-func-return.mlir`
passes.

Fixes #217227

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
    mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 4d044bbb74df1..5552246427188 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -435,6 +435,11 @@ static void foldMemRefCasts(func::FuncOp funcOp) {
 
   // Compute the common result types of all return ops.
   SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
+  // There is nothing to fold if the function body does not end in a
+  // func::ReturnOp (e.g., the terminator is from a 
diff erent dialect). Bail
+  // out gracefully instead of asserting inside `getReturnTypes`.
+  if (returnOps.empty())
+    return;
   SmallVector<Type> resultTypes = getReturnTypes(returnOps);
 
   // Remove direct casts.

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir
index 7addca2c9d6a5..d44a0b043757d 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir
@@ -14,3 +14,20 @@ func.func @caller(%A : tensor<64xf32>) {
   call @callee(%A) : (tensor<64xf32>) -> ()
   return
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/217227:
+// this function does not have a `func.return` terminator (it uses
+// `spirv.Return` instead, which is a valid terminator for `func.func` since
+// the verifier only requires the region to end in some terminator, not
+// specifically a `func.return`). One-Shot Module Bufferize must not crash on
+// such IR; it has no tensors to bufferize, so it should simply leave the
+// function unchanged.
+
+// CHECK-LABEL: func.func private @non_func_return(
+// CHECK-SAME:      %{{.*}}: memref<4xi32>)
+// CHECK-NEXT:    spirv.Return
+func.func private @non_func_return(%arg0: memref<4xi32>) {
+  spirv.Return
+}