[Mlir-commits] [mlir] [mlir][transforms] Fix crash in remove-dead-values when function has non-call users (PR #183655)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/183655

>From 62089ba1e17cfe4859e1ba1f35f2b738998f5c22 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Feb 2026 16:56:33 -0800
Subject: [PATCH] [mlir][transforms] Fix crash in remove-dead-values when
 function has non-call users

`processFuncOp` asserts that all symbol uses of a function are
`CallOpInterface` operations. This is violated when a function is
referenced by a non-call operation such as `spirv.EntryPoint`, which
uses the function symbol for metadata purposes without calling it.

Fix this by replacing the assertion with an early return: if any user of
the function symbol is not a `CallOpInterface`, skip the function
entirely. This is safe because the pass cannot determine the semantics of
arbitrary non-call references, so it should leave such functions alone.

Fixes #180416
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     | 19 ++++++++++++++-----
 mlir/test/Transforms/remove-dead-values.mlir | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 12a47ba2fb65a..4c15fd46d581e 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -285,7 +285,15 @@ static void processFuncOp(FunctionOpInterface funcOp, Operation *module,
            << funcOp.getOperation()->getName();
     return;
   }
-
+  SymbolTable::UseRange uses = *funcOp.getSymbolUses(module);
+  if (llvm::any_of(uses, [](SymbolTable::SymbolUse use) {
+        return !isa<CallOpInterface>(use.getUser());
+      })) {
+    // If a non-call operation references the function (e.g. spirv.EntryPoint),
+    // we cannot safely remove arguments or return values since we don't know
+    // what the user expects. Skip this function entirely.
+    return;
+  }
   // Get the list of unnecessary (non-live) arguments in `nonLiveArgs`.
   SmallVector<Value> arguments(funcOp.getArguments());
   BitVector nonLiveArgs = markLives(arguments, nonLiveSet, la);
@@ -299,10 +307,8 @@ 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");
     // 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.
@@ -614,8 +620,8 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
   // AttrSizedOperandSegments) in the next phase.
   DenseMap<Operation *, BitVector> erasedFuncArgs;
   for (auto &f : list.functions) {
-    LDBG() << "Cleaning up function: " << f.funcOp.getOperation()->getName()
-           << " (" << f.funcOp.getOperation() << ")";
+    LDBG() << "Cleaning up function: " << f.funcOp.getName() << " ("
+           << f.funcOp.getOperation() << ")";
     LDBG_OS([&](raw_ostream &os) {
       os << "  Erasing non-live arguments [";
       llvm::interleaveComma(f.nonLiveArgs.set_bits(), os);
@@ -634,6 +640,9 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
       // Record only if we actually erased something.
       if (f.nonLiveArgs.any())
         erasedFuncArgs.try_emplace(f.funcOp.getOperation(), f.nonLiveArgs);
+    } else {
+      LDBG() << "Failed to erase arguments for function: "
+             << f.funcOp.getName();
     }
     (void)f.funcOp.eraseResults(f.nonLiveRets);
   }
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 87e77b2eb700f..1f8a9fc95d5cd 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -826,3 +826,22 @@ func.func @replace_dead_operation_results_with_poison(%0: vector<1xindex>) -> ve
   }
   return %2 : vector<1xindex>
 }
+
+// -----
+
+// Verify that a referenced by a non-call op (spirv.EntryPoint),
+// while still having another usual call site is preserved as-is
+// since the pass cannot analyse non-call users.
+// CHECK-LABEL: module @func_with_non_call_users
+module @func_with_non_call_users {
+// CHECK: func @callee(%arg0 : i32, %arg1 : i32)
+  func.func private @callee(%arg1 : i32, %arg2 : i32) {
+    func.return
+  }
+  func.func @main_func() {
+    %cst = llvm.mlir.constant(1 : i32) : i32
+    func.call @callee(%cst, %cst) : (i32, i32) -> ()
+    func.return
+  }
+  spirv.EntryPoint "GLCompute" @callee
+}