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

Author: timnoack

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);
     }

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

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