[Mlir-commits] [mlir] [mlir][SymbolDCE] Check SSA uses for symbols with results (PR #168376)

[Tim Noack]

https://github.com/timnoack created https://github.com/llvm/llvm-project/pull/168376

Fixes a bug in the SymbolDCE pass where operations that are a symbol **and** produce SSA results could be deleted even when their results were being used.

Previously, SymbolDCE only checked for symbol table references when determining if a symbol operation could be erased. However, dialects may define operations that are both symbols and produce SSA values (e.g., operations producing values that are accessible both via a symbol and via its SSA result. For such operations, the pass would incorrectly delete them if they had no symbol table references, even if their SSA results were actively used. This resulted in invalid IR.

The fix adds an additional check for `op.use_empty()` when determining if a symbol is discardable. A symbol can now only be deleted if:
  1. It's private (or in a hidden symbol table), and
  2. It can be discarded when unused (`canDiscardOnUseEmpty()`), and
  3. It has no symbol table references, and
  4. **it has no SSA value uses** (new check)

>From 44b6f171193d8a3094e4f93e545cc08bf21d2ec2 Mon Sep 17 00:00:00 2001
From: tn <noack at esa.tu-darmstadt.de>
Date: Mon, 17 Nov 2025 15:27:27 +0100
Subject: [PATCH] [mlir][SymbolDCE] Check SSA uses for symbols with results

---
 mlir/lib/Transforms/SymbolDCE.cpp         |  2 +-
 mlir/test/Transforms/test-symbol-dce.mlir | 20 ++++++++++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td     |  7 +++++++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/SymbolDCE.cpp b/mlir/lib/Transforms/SymbolDCE.cpp
index 87885be00daa3..3b025385872e4 100644
--- a/mlir/lib/Transforms/SymbolDCE.cpp
+++ b/mlir/lib/Transforms/SymbolDCE.cpp
@@ -106,7 +106,7 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
         continue;
       }
       bool isDiscardable = (symbolTableIsHidden || symbol.isPrivate()) &&
-                           symbol.canDiscardOnUseEmpty();
+                           symbol.canDiscardOnUseEmpty() && op.use_empty();
       if (!isDiscardable && liveSymbols.insert(&op).second)
         worklist.push_back(&op);
     }
diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index 90fe93f806d69..b5ab42882e02c 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -133,3 +133,23 @@ module attributes { test.nested_nosymboltable_region_notcalled } {
     "test.finish"() : () -> ()
   }) : () -> ()
 }
+
+// -----
+
+// Check that symbols with SSA results are not DCE'd if their result is used.
+// CHECK-LABEL: module attributes {test.symbol_with_ssa_result}
+module attributes {test.symbol_with_ssa_result} {
+  // CHECK: "test.symbol_with_result"() <{sym_name = "used_symbol", sym_visibility = "private"}> : () -> i32
+  %0 = "test.symbol_with_result"() <{sym_name = "used_symbol", sym_visibility = "private"}> : () -> i32
+
+  // CHECK-NOT: test.symbol_with_result
+  // CHECK-NOT: unused_symbol
+  %1 = "test.symbol_with_result"() <{sym_name = "unused_symbol", sym_visibility = "private"}> : () -> i32
+
+  // CHECK-NOT: test.symbol
+  // CHECK-NOT: another_unused_symbol
+  "test.symbol"() <{sym_name = "another_unused_symbol", sym_visibility = "private"}> : () -> ()
+
+  // Use %0, keeping @used_symbol alive via SSA
+  "test.use"(%0) : (i32) -> ()
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 275025978a784..9d88b6f2b64cc 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -120,6 +120,13 @@ def SymbolOp : TEST_Op<"symbol", [NoMemoryEffect, Symbol]> {
                        OptionalAttr<StrAttr>:$sym_visibility);
 }
 
+def SymbolWithResultOp : TEST_Op<"symbol_with_result", [NoMemoryEffect, Symbol]> {
+  let summary = "symbol operation that produces an SSA result";
+  let arguments = (ins StrAttr:$sym_name,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let results = (outs AnyType:$result);
+}
+
 def OverriddenSymbolVisibilityOp : TEST_Op<"overridden_symbol_visibility", [
   DeclareOpInterfaceMethods<Symbol, ["getVisibility", "setVisibility"]>,
 ]> {