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

[Aman Singh]

https://github.com/amanyagami updated https://github.com/llvm/llvm-project/pull/217515

>From 6e9b43fc3ae65c9563aa84a740ccf88c21c1e6ac Mon Sep 17 00:00:00 2001
From: amanyagami <2amansingh2 at gmail.com>
Date: Wed, 19 Aug 2026 19:45:53 -0700
Subject: [PATCH] [mlir][Bufferization] Don't assert in foldMemRefCasts on a
 func body without func.return

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::ReturnOps, mirroring how the rest of the file already
tolerates an empty getReturnOps() result.

Fixes https://github.com/llvm/llvm-project/issues/217227
---
 .../Transforms/OneShotModuleBufferize.cpp       |  5 +++++
 ...e-shot-module-bufferize-non-func-return.mlir | 17 +++++++++++++++++
 2 files changed, 22 insertions(+)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-non-func-return.mlir

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 different 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-non-func-return.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-non-func-return.mlir
new file mode 100644
index 0000000000000..1996770de50c0
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-non-func-return.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s -one-shot-bufferize='bufferize-function-boundaries=1 copy-before-write=1' | FileCheck %s
+
+// 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 @m(
+// CHECK-SAME:      %{{.*}}: memref<4xi32>)
+// CHECK-NEXT:    spirv.Return
+module {
+  func.func private @m(%arg0: memref<4xi32>) {
+    spirv.Return
+  }
+}