[Mlir-commits] [mlir] [mlir] Fix remove-dead-values crash on non-call symbol users (PR #181155)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Jueon Park (JueonPark)

<details>
<summary>Changes</summary>

Fixes #<!-- -->180416.

processFuncOp in the remove-dead-values pass previously assumed that all symbol users of a function are CallOpInterface ops. This is not always true — for example, spirv.EntryPoint references a function symbol but is not a call-like op. The assumption caused a crash when iterating over symbol users and unconditionally casting them to CallOpInterface.

Instead of bailing out entirely when a non-call symbol user is present, this patch filters the symbol uses to collect only call-like users into a SmallVector<Operation *>. The downstream loops then iterate over this filtered list. This allows the pass to still optimize functions that have both call-like and non-call-like symbol users, in addition to cleaning all assertions.

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


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+13-10) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+16) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 66f369e8a5f65..9d1dd36ad4ca2 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -286,6 +286,16 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
     return;
   }
 
+  // Collect only call-like symbol users. Non-call users (e.g. spirv.EntryPoint)
+  // are skipped because we cannot safely modify the function signature without
+  // knowing how such users depend on it.
+  SmallVector<Operation *> callUses;
+  auto symbolUses = funcOp.getSymbolUses(module);
+  for (SymbolTable::SymbolUse use : *symbolUses) {
+    if (isa<CallOpInterface>(use.getUser()))
+      callUses.push_back(use.getUser());
+  }
+
   // Get the list of unnecessary (non-live) arguments in `nonLiveArgs`.
   SmallVector<Value> arguments(funcOp.getArguments());
   BitVector nonLiveArgs = markLives(arguments, nonLiveSet, la);
@@ -299,10 +309,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // Do (2). (Skip creating generic operand cleanup entries for call ops.
   // Call arguments will be removed in the call-site specific segment-aware
   // cleanup, avoiding generic eraseOperands bitvector mechanics.)
-  SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
-  for (SymbolTable::SymbolUse use : uses) {
-    Operation *callOp = use.getUser();
-    assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
+  for (Operation *callOp : callUses) {
     // Push an empty operand cleanup entry so that call-site specific logic in
     // cleanUpDeadVals runs (it keys off CallOpInterface). The BitVector is
     // intentionally all false to avoid generic erasure.
@@ -336,9 +343,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // since it forwards only to non-live value(s) (%1#1).
   size_t numReturns = funcOp.getNumResults();
   BitVector nonLiveRets(numReturns, true);
-  for (SymbolTable::SymbolUse use : uses) {
-    Operation *callOp = use.getUser();
-    assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
+  for (Operation *callOp : callUses) {
     BitVector liveCallRets = markLives(callOp->getResults(), nonLiveSet, la);
     nonLiveRets &= liveCallRets.flip();
   }
@@ -360,9 +365,7 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
   // Do (5) and (6).
   if (numReturns == 0)
     return;
-  for (SymbolTable::SymbolUse use : uses) {
-    Operation *callOp = use.getUser();
-    assert(isa<CallOpInterface>(callOp) && "expected a call-like user");
+  for (Operation *callOp : callUses) {
     cl.results.push_back({callOp, nonLiveRets});
     collectNonLiveValues(nonLiveSet, callOp->getResults(), nonLiveRets);
   }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index ae83eac0c376f..e6d34881ebd14 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -796,3 +796,19 @@ func.func @scf_while_dead_iter_args() -> i32 {
   }
   return %result#0 : i32
 }
+
+// -----
+
+// Verify that `remove-dead-values` does not crash when a function has a
+// non-call symbol user (e.g. spirv.EntryPoint). The function should be
+// preserved as-is since its signature cannot be safely modified.
+
+// CHECK-LABEL: module
+// CHECK: spirv.EntryPoint "GLCompute" @main_func
+// CHECK: llvm.func @main_func()
+module {
+  spirv.EntryPoint "GLCompute" @main_func
+  llvm.func @main_func() attributes {sym_visibility = "private"} {
+    llvm.return
+  }
+}

``````````

</details>


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