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

[Jueon Park]

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

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.

>From 67c127a891bb1e2d72c9158b9f4224d8248e50e4 Mon Sep 17 00:00:00 2001
From: rebel-jueonpark <jueonpark at rebellions.ai>
Date: Thu, 12 Feb 2026 23:41:59 +0900
Subject: [PATCH] [mlir] Fix remove-dead-values crash on non-call symbol users
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 23 +++++++++++---------
 mlir/test/Transforms/remove-dead-values.mlir | 16 ++++++++++++++
 2 files changed, 29 insertions(+), 10 deletions(-)

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