[mlir][SymbolDCE] Check SSA uses for symbols with results

[timnoack]

Fixes a bug in the SymbolDCE pass where operations that are symbols 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)

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Tim Noack (timnoack)

Changes

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)

---

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

3 Files Affected:

• (modified) mlir/lib/Transforms/SymbolDCE.cpp (+1-1)
• (modified) mlir/test/Transforms/test-symbol-dce.mlir (+20)
• (modified) mlir/test/lib/Dialect/Test/TestOps.td (+7)

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"]>,
 ]> {

[timnoack]

I think we should either verify in the SymbolInterface that symbols have no results, or handle potential uses of their results in the SymbolDCE pass (this PR). I vote for the latter approach, because we rely on operations that are both symbols and have result values in our project.

[joker-eph]

I think we should either verify in the SymbolInterface that symbols have no results,

That seems appropriate to me, I believe it is the intent, see the documentation here: https://mlir.llvm.org/docs/SymbolsAndSymbolTables/#defining-or-declaring-a-symbol

A Symbol operation should use the SymbolOpInterface interface to provide the necessary verification and accessors; it also supports operations, such as builtin.module, that conditionally define a symbol. Symbols must have the following properties:
[....]

• No SSA results

[timnoack]

Hmm, I see, I didn't know about that. Then we need to refactor some things. Thanks! I am gonna send a PR for the symbol interface verifier.