[mlir][bufferization] Fix crash with copy-before-write + bufferize-function-boundaries#186446
Merged
Conversation
…nction-boundaries When `copy-before-write=1` is combined with `bufferize-function-boundaries=1`, `bufferizeOp` creates a plain `AnalysisState` (not `OneShotAnalysisState`) and passes it to `insertTensorCopies`. Walking `CallOp`s during conflict resolution called `getCalledFunction(callOp, state)`, which unconditionally cast the `AnalysisState` to `OneShotAnalysisState` via `static_cast`, causing UB and a stack overflow crash. Fix by guarding the cast with `isa<OneShotAnalysisState>()` so that when the state is a plain `AnalysisState`, the function falls through to building a fresh `SymbolTableCollection` — the same safe fallback already present. Fixes llvm#163052 Assisted-by: Claude Code
Member
|
@llvm/pr-subscribers-mlir-bufferization @llvm/pr-subscribers-mlir Author: Mehdi Amini (joker-eph) ChangesWhen Fix by guarding the cast with Fixes #163052 Assisted-by: Claude Code Full diff: https://github.com/llvm/llvm-project/pull/186446.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index e43ab54a048b9..3aaa38272935d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -101,12 +101,14 @@ static FuncOp getCalledFunction(CallOpInterface callOp,
/// Return the FuncOp called by `callOp`.
static FuncOp getCalledFunction(CallOpInterface callOp,
const AnalysisState &state) {
- auto &oneShotAnalysisState = static_cast<const OneShotAnalysisState &>(state);
-
- if (auto *funcAnalysisState =
- oneShotAnalysisState.getExtension<FuncAnalysisState>()) {
- // Use the cached symbol tables.
- return getCalledFunction(callOp, funcAnalysisState->symbolTables);
+ if (isa<OneShotAnalysisState>(state)) {
+ auto &oneShotAnalysisState =
+ static_cast<const OneShotAnalysisState &>(state);
+ if (auto *funcAnalysisState =
+ oneShotAnalysisState.getExtension<FuncAnalysisState>()) {
+ // Use the cached symbol tables.
+ return getCalledFunction(callOp, funcAnalysisState->symbolTables);
+ }
}
SymbolTableCollection symbolTables;
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
new file mode 100644
index 0000000000000..7addca2c9d6a5
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-call-copy-before-write.mlir
@@ -0,0 +1,16 @@
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries=1 copy-before-write=1" | FileCheck %s
+
+// Regression test for https://github.com/llvm/llvm-project/issues/163052
+// copy-before-write=1 + bufferize-function-boundaries=1 with a call to a
+// private (declaration-only) function used to crash with a stack overflow due
+// to an invalid cast of AnalysisState to OneShotAnalysisState inside
+// getCalledFunction().
+
+// CHECK-LABEL: func.func private @callee(memref<64xf32
+// CHECK-LABEL: func.func @caller
+// CHECK: call @callee
+func.func private @callee(tensor<64xf32>)
+func.func @caller(%A : tensor<64xf32>) {
+ call @callee(%A) : (tensor<64xf32>) -> ()
+ return
+}
|
Contributor
Author
|
Ping @matthias-springer |
matthias-springer
approved these changes
Mar 17, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When
copy-before-write=1is combined withbufferize-function-boundaries=1,bufferizeOpcreates a plainAnalysisState(notOneShotAnalysisState) and passes it toinsertTensorCopies. WalkingCallOps during conflict resolution calledgetCalledFunction(callOp, state), which unconditionally cast theAnalysisStatetoOneShotAnalysisStateviastatic_cast, causing UB and a stack overflow crash.Fix by guarding the cast with
isa<OneShotAnalysisState>()so that when the state is a plainAnalysisState, the function falls through to building a freshSymbolTableCollection— the same safe fallback already present.Fixes #163052
Assisted-by: Claude Code