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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Author: Aman Singh (amanyagami)

<details>
<summary>Changes</summary>

`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)

---
Full diff: https://github.com/llvm/llvm-project/pull/217515.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp (+5) 
- (added) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-non-func-return.mlir (+17) 


``````````diff
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
+  }
+}

``````````

</details>


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