Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[MLIR][MLProgram] Fix crash in mlprogram-pipeline-globals on unresolvable callees (#189244)

Browse source 
The `MLProgramPipelineGlobals` pass crashed with a null pointer dereference
when a `CallOpInterface` operation referred to a callee symbol that could not
be resolved in the IR (e.g. an external function defined outside the module).

Instead  conservatively bail out when a callee symbol cannot be resolved, 
causing the pass to (preserving all loads/stores). This is consistent with
how Value-based callees are handled.

Fixes #109649

Assisted-by: Claude Code
This commit is contained in:
Mehdi Amini 2026-04-14 18:20:29 +02:00 committed by GitHub
parent 4a37b03357
commit 7b744a511e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 33 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp
View file

@ -33,7 +33,7 @@ private:
llvm::DenseMap<SymbolRefAttr, llvm::DenseSet<SymbolRefAttr>> storeSymbolsMap;
};
// Traverses upwards searchign for the operation mapped by the symbol.
// Traverses upwards searching for the operation mapped by the symbol.
static Operation *getFromSymbol(Operation *baseOp, SymbolRefAttr symbol) {
for (auto *op = baseOp; op; op = op->getParentOp()) {
auto *lookup = SymbolTable::lookupNearestSymbolFrom(op, symbol);
@ -57,6 +57,9 @@ LogicalResult MLProgramPipelineGlobals::buildGlobalMap(ModuleOp module) {
auto symbol = mlir::dyn_cast<SymbolRefAttr>(callable);
auto *func = getFromSymbol(op, symbol);
// If the callee cannot be resolved, we cannot safely analyze the IR.
if (!func)
return WalkResult::interrupt();
callableMap[symbol] = func;
}
return WalkResult::advance();
@ -95,8 +98,13 @@ LogicalResult MLProgramPipelineGlobals::buildGlobalMap(ModuleOp module) {
llvm::DenseSet<SymbolRefAttr> storeSymbols;
for (size_t i = 0; i < work.size(); ++i) {
callableMap[work[i]]->walk([&](CallOpInterface call) {
auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
// Defensive: symbols in `work` should always be in `callableMap` since
// buildGlobalMap interrupted on any unresolvable callee, but use find to
// avoid inserting null entries via operator[].
auto it = callableMap.find(work[i]);
assert(it != callableMap.end() && "Expected callable in callableMap");
it->second->walk([&](CallOpInterface call) {
auto symbol = cast<SymbolRefAttr>(call.getCallableForCallee());
if (visited.insert(symbol).second)
work.push_back(symbol);
});

22
mlir/test/Dialect/MLProgram/pipeline-globals.mlir
View file

@ -219,6 +219,28 @@ func.func @call_indirect_load() {
// -----
// Calling via a CallOpInterface op whose callee symbol is not defined in this
// module should not crash - the pass should bail out gracefully.
// See https://github.com/llvm/llvm-project/issues/109649
// CHECK-LABEL: @global_variable
ml_program.global private mutable @global_variable(dense<4> : tensor<4xi32>) : tensor<4xi32>
// CHECK-LABEL: @call_with_unresolvable_callee
func.func @call_with_unresolvable_callee(%arg0: memref<f32>) {
// Both loads must be preserved; the pass conservatively bails out when it
// encounters a call whose callee symbol cannot be resolved.
// CHECK: ml_program.global_load @global_variable
%0 = ml_program.global_load @global_variable : tensor<4xi32>
// @callee is not defined anywhere in this module.
test.call_and_store @callee(%arg0), %arg0 {store_before_call = false} : (memref<f32>, memref<f32>) -> ()
// CHECK: ml_program.global_load @global_variable
%1 = ml_program.global_load @global_variable : tensor<4xi32>
func.return
}
// -----
// CHECK-LABEL: @global_variable
ml_program.global private mutable @global_variable(dense<4> : tensor<4xi32>) : tensor<4xi32>