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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Tim Noack (timnoack)

<details>
<summary>Changes</summary>

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

``````````

</details>


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