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

[llvmlistbot at llvm.org]

llvmbot wrote:


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

@llvm/pr-subscribers-mlir-mlprogram

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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 with no definition).

`getFromSymbol` returned nullptr in that case, which was stored in `callableMap` and then unconditionally dereferenced via `->walk()`. Two fixes are applied:

1. In `buildGlobalMap`, interrupt the walk (and bail out) when a callee symbol cannot be resolved, consistent with how Value-based callees are handled.
2. In the transitive-closure loop, use `find` instead of `operator[]` when looking up entries in `callableMap` to avoid creating null entries for symbols not present in the map.

Fixes #<!-- -->109649

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp (+14-3) 
- (modified) mlir/test/Dialect/MLProgram/pipeline-globals.mlir (+25) 


``````````diff
diff --git a/mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp b/mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp
index 54fa15722febe..3ab0931299205 100644
--- a/mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp
+++ b/mlir/lib/Dialect/MLProgram/Transforms/PipelineGlobalOps.cpp
@@ -33,7 +33,7 @@ class MLProgramPipelineGlobals
   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,9 +98,17 @@ LogicalResult MLProgramPipelineGlobals::buildGlobalMap(ModuleOp module) {
     llvm::DenseSet<SymbolRefAttr> storeSymbols;
 
     for (size_t i = 0; i < work.size(); ++i) {
-      callableMap[work[i]]->walk([&](CallOpInterface call) {
+      // 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]);
+      if (it == callableMap.end())
+        continue;
+      it->second->walk([&](CallOpInterface call) {
+        // Skip indirect (Value-based) callees; the outer walk already
+        // interrupted on these, but guard defensively here.
         auto symbol = dyn_cast<SymbolRefAttr>(call.getCallableForCallee());
-        if (visited.insert(symbol).second)
+        if (symbol && visited.insert(symbol).second)
           work.push_back(symbol);
       });
 
diff --git a/mlir/test/Dialect/MLProgram/pipeline-globals.mlir b/mlir/test/Dialect/MLProgram/pipeline-globals.mlir
index a5c9b3e890558..9586cf8247b24 100644
--- a/mlir/test/Dialect/MLProgram/pipeline-globals.mlir
+++ b/mlir/test/Dialect/MLProgram/pipeline-globals.mlir
@@ -219,6 +219,31 @@ func.func @call_indirect_load() {
 
 // -----
 
+// Calling a function that is declared but not defined in this module (no body)
+// 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>
+
+// External function declaration (no body).
+func.func private @external_func() -> ()
+
+// CHECK-LABEL: @call_external
+func.func @call_external() {
+  // Both loads must be preserved; the pass conservatively bails out when it
+  // encounters a call to an unresolvable callee.
+  // CHECK: ml_program.global_load @global_variable
+  %0 = ml_program.global_load @global_variable : tensor<4xi32>
+  // Call an external function with no body - pass cannot resolve callee body.
+  call @external_func() : () -> ()
+  // 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>
 

``````````

</details>


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