[Mlir-commits] [mlir] [mlir] Walk nested non-symbol table ops in symbol dce (PR #143353)

Jacques Pienaar llvmlistbot at llvm.org
Tue Jun 24 08:35:48 PDT 2025


https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/143353

>From dfa90b57696a6f73baf72ff0de4ccf311a5da9db Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 24 Jun 2025 09:35:26 +0000
Subject: [PATCH 1/6] [mlir] Walk nested tables in symbol dce

The previous positioning was effectively that a symbol is dead if it cannot be addressed from top level. I think that is too strong a requirement: one can have operations that one cannot delete/DCE that refers to symbols which one could delete. This resulted in symbol-dce deleting symbols that are still referenced and the resulting IR being invalid.

This instead all the symbols of top level operations of non-symbol table ops additionally, as those are either dead and DCE would have handled, or alive and we cannot just delete symbols referenced internally. E.g., this treats non-symbol table regioned ops more conservatively.
---
 mlir/lib/Transforms/SymbolDCE.cpp         | 48 ++++++++++++++++++++---
 mlir/test/Transforms/test-symbol-dce.mlir | 19 +++++++++
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Transforms/SymbolDCE.cpp b/mlir/lib/Transforms/SymbolDCE.cpp
index 93d9a6547883a..4e751d7317512 100644
--- a/mlir/lib/Transforms/SymbolDCE.cpp
+++ b/mlir/lib/Transforms/SymbolDCE.cpp
@@ -14,6 +14,7 @@
 #include "mlir/Transforms/Passes.h"
 
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/Support/Debug.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_SYMBOLDCE
@@ -22,6 +23,8 @@ namespace mlir {
 
 using namespace mlir;
 
+#define DEBUG_TYPE "symbol-dce"
+
 namespace {
 struct SymbolDCE : public impl::SymbolDCEBase<SymbolDCE> {
   void runOnOperation() override;
@@ -84,6 +87,8 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
                                          SymbolTableCollection &symbolTable,
                                          bool symbolTableIsHidden,
                                          DenseSet<Operation *> &liveSymbols) {
+  LLVM_DEBUG(llvm::dbgs() << "computeLiveness: " << symbolTableOp->getName()
+                          << "\n");
   // A worklist of live operations to propagate uses from.
   SmallVector<Operation *, 16> worklist;
 
@@ -105,9 +110,13 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
   }
 
   // Process the set of symbols that were known to be live, adding new symbols
-  // that are referenced within.
+  // that are referenced within. For operations that are not symbol tables, it
+  // considers the liveness with respect to the op itself rather than scope of
+  // nested symbol tables by enqueuing all the top level operations for
+  // consideration.
   while (!worklist.empty()) {
     Operation *op = worklist.pop_back_val();
+    LLVM_DEBUG(llvm::dbgs() << "processing: " << op->getName() << "\n");
 
     // If this is a symbol table, recursively compute its liveness.
     if (op->hasTrait<OpTrait::SymbolTable>()) {
@@ -115,26 +124,55 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
       // symbol, or if it is a private symbol.
       SymbolOpInterface symbol = dyn_cast<SymbolOpInterface>(op);
       bool symIsHidden = symbolTableIsHidden || !symbol || symbol.isPrivate();
+      LLVM_DEBUG(llvm::dbgs() << "\tsymbol table: " << op->getName()
+                              << " is hidden: " << symIsHidden << "\n");
       if (failed(computeLiveness(op, symbolTable, symIsHidden, liveSymbols)))
         return failure();
+    } else {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "\tnon-symbol table: " << op->getName() << "\n");
+      // If the op is not a symbol table, then, unless op itself is dead which
+      // would be handled by DCE, we need to check all the regions and blocks
+      // within the op to find the uses (e.g., consider visibility within op as
+      // if top level rather than relying on pure symbol table visibility). This
+      // is more conservative than SymbolTable::walkSymbolTables in the case
+      // where there is again SymbolTable information to take advantage of.
+      for (auto &region : op->getRegions())
+        for (auto &block : region.getBlocks())
+          for (Operation &op : block)
+            worklist.push_back(&op);
     }
 
+    // Get the first parent symbol table op. Note: due to enqueueing of
+    // top-level ops, we may not have a symbol table parent here, but if we do
+    // not, then we also don't have a symbol.
+    Operation *parentOp = op->getParentOp();
+    if (!parentOp->hasTrait<OpTrait::SymbolTable>())
+      continue;
+
     // Collect the uses held by this operation.
     std::optional<SymbolTable::UseRange> uses = SymbolTable::getSymbolUses(op);
     if (!uses) {
       return op->emitError()
-             << "operation contains potentially unknown symbol table, "
-                "meaning that we can't reliable compute symbol uses";
+             << "operation contains potentially unknown symbol table, meaning "
+             << "that we can't reliable compute symbol uses";
     }
 
     SmallVector<Operation *, 4> resolvedSymbols;
+    LLVM_DEBUG(llvm::dbgs() << "uses of " << op->getName() << "\n");
     for (const SymbolTable::SymbolUse &use : *uses) {
+      LLVM_DEBUG(llvm::dbgs() << "\tuse: " << use.getUser() << "\n");
       // Lookup the symbols referenced by this use.
       resolvedSymbols.clear();
-      if (failed(symbolTable.lookupSymbolIn(
-              op->getParentOp(), use.getSymbolRef(), resolvedSymbols)))
+      if (failed(symbolTable.lookupSymbolIn(parentOp, use.getSymbolRef(),
+                                            resolvedSymbols)))
         // Ignore references to unknown symbols.
         continue;
+      LLVM_DEBUG({
+        llvm::dbgs() << "\t\tresolved symbols: ";
+        llvm::interleaveComma(resolvedSymbols, llvm::dbgs());
+        llvm::dbgs() << "\n";
+      });
 
       // Mark each of the resolved symbols as live.
       for (Operation *resolvedSymbol : resolvedSymbols)
diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index 7bd784928e6f3..dc2f0d7dd3b09 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -98,3 +98,22 @@ module {
   // CHECK: "live.user"() {uses = [@unknown_symbol]} : () -> ()
   "live.user"() {uses = [@unknown_symbol]} : () -> ()
 }
+
+// -----
+
+// Check that we don't DCE nested symbols if they are used even if nested inside
+// an unnamed region.
+// CHECK-LABEL: module attributes {test.nested_unnamed_region}
+module attributes {test.nested_unnamed_region} {
+  "test.one_region_op"() ({
+    "test.symbol_scope"() ({
+      // CHECK: func @nfunction
+      func.func @nfunction() {
+        return
+      }
+      func.call @nfunction() : () -> ()
+      "test.finish"() : () -> ()
+    }) : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+}

>From dfe5b3c4323a98110fb6e8c5e0be85af4ca69b0a Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 24 Jun 2025 11:41:26 +0000
Subject: [PATCH 2/6] Simplify test

---
 mlir/test/Transforms/test-symbol-dce.mlir | 26 ++++++++++++-----------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index dc2f0d7dd3b09..d2703dff2f78b 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -101,19 +101,21 @@ module {
 
 // -----
 
-// Check that we don't DCE nested symbols if they are used even if nested inside
-// an unnamed region.
-// CHECK-LABEL: module attributes {test.nested_unnamed_region}
-module attributes {test.nested_unnamed_region} {
+// Check that we don't DCE nested symbols if they are nested inside region
+// without SymbolTable.
+// CHECK-LABEL: module attributes {test.nested_nosymboltable_region}
+module attributes {test.nested_nosymboltable_region} {
   "test.one_region_op"() ({
-    "test.symbol_scope"() ({
-      // CHECK: func @nfunction
-      func.func @nfunction() {
-        return
-      }
-      func.call @nfunction() : () -> ()
-      "test.finish"() : () -> ()
-    }) : () -> ()
+    module {
+        func.func nested @nested() {
+            return
+        }
+
+        func.func @main() {
+            return
+        }
+    }
     "test.finish"() : () -> ()
   }) : () -> ()
 }
+

>From d3134bdd2df31852a01ccd006df27cddac206c63 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jacques+gh at japienaar.info>
Date: Tue, 24 Jun 2025 07:02:02 -0700
Subject: [PATCH 3/6] Update mlir/lib/Transforms/SymbolDCE.cpp

Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
---
 mlir/lib/Transforms/SymbolDCE.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/SymbolDCE.cpp b/mlir/lib/Transforms/SymbolDCE.cpp
index 4e751d7317512..c9cbefc449722 100644
--- a/mlir/lib/Transforms/SymbolDCE.cpp
+++ b/mlir/lib/Transforms/SymbolDCE.cpp
@@ -140,7 +140,8 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
       for (auto &region : op->getRegions())
         for (auto &block : region.getBlocks())
           for (Operation &op : block)
-            worklist.push_back(&op);
+            if (op->getNumRegions())
+              worklist.push_back(&op);
     }
 
     // Get the first parent symbol table op. Note: due to enqueueing of

>From 15d78c98403a392d9f3a971cd9be197e8ecd344a Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 24 Jun 2025 14:07:52 +0000
Subject: [PATCH 4/6] Add test back and fix compile error

---
 mlir/lib/Transforms/SymbolDCE.cpp         |  2 +-
 mlir/test/Transforms/test-symbol-dce.mlir | 23 +++++++++++++++++++++--
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Transforms/SymbolDCE.cpp b/mlir/lib/Transforms/SymbolDCE.cpp
index c9cbefc449722..0a925c47a9e8b 100644
--- a/mlir/lib/Transforms/SymbolDCE.cpp
+++ b/mlir/lib/Transforms/SymbolDCE.cpp
@@ -140,7 +140,7 @@ LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
       for (auto &region : op->getRegions())
         for (auto &block : region.getBlocks())
           for (Operation &op : block)
-            if (op->getNumRegions())
+            if (op.getNumRegions())
               worklist.push_back(&op);
     }
 
diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index d2703dff2f78b..413ea5716faeb 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -103,8 +103,28 @@ module {
 
 // Check that we don't DCE nested symbols if they are nested inside region
 // without SymbolTable.
+
 // CHECK-LABEL: module attributes {test.nested_nosymboltable_region}
-module attributes {test.nested_nosymboltable_region} {
+module attributes { test.nested_nosymboltable_region } {
+  "test.one_region_op"() ({
+    "test.symbol_scope"() ({
+      // CHECK: func @nfunction
+      func.func @nfunction() {
+        return
+      }
+      func.call @nfunction() : () -> ()
+      "test.finish"() : () -> ()
+    }) : () -> ()
+    "test.finish"() : () -> ()
+  }) : () -> ()
+}
+
+// -----
+
+// CHECK-LABEL: module attributes {test.nested_nosymboltable_region_notcalled}
+// CHECK-NOT: @nested
+// CHECK: @main
+module attributes { test.nested_nosymboltable_region_notcalled } {
   "test.one_region_op"() ({
     module {
         func.func nested @nested() {
@@ -118,4 +138,3 @@ module attributes {test.nested_nosymboltable_region} {
     "test.finish"() : () -> ()
   }) : () -> ()
 }
-

>From 96e2e0502fbd66cec83683c5d56b2eaef4c6b70b Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 24 Jun 2025 14:14:27 +0000
Subject: [PATCH 5/6] Change whitespace in test to make more concise

---
 mlir/test/Transforms/test-symbol-dce.mlir | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index 413ea5716faeb..65b9efccd732d 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -127,13 +127,8 @@ module attributes { test.nested_nosymboltable_region } {
 module attributes { test.nested_nosymboltable_region_notcalled } {
   "test.one_region_op"() ({
     module {
-        func.func nested @nested() {
-            return
-        }
-
-        func.func @main() {
-            return
-        }
+      func.func nested @nested() { return }
+      func.func @main() { return }
     }
     "test.finish"() : () -> ()
   }) : () -> ()

>From 5a395bc88ab4b3bdc7a2a35ed305f104995f1e36 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jacques+gh at japienaar.info>
Date: Tue, 24 Jun 2025 08:35:25 -0700
Subject: [PATCH 6/6] Update mlir/test/Transforms/test-symbol-dce.mlir

Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
---
 mlir/test/Transforms/test-symbol-dce.mlir | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Transforms/test-symbol-dce.mlir b/mlir/test/Transforms/test-symbol-dce.mlir
index 65b9efccd732d..90fe93f806d69 100644
--- a/mlir/test/Transforms/test-symbol-dce.mlir
+++ b/mlir/test/Transforms/test-symbol-dce.mlir
@@ -108,8 +108,8 @@ module {
 module attributes { test.nested_nosymboltable_region } {
   "test.one_region_op"() ({
     "test.symbol_scope"() ({
-      // CHECK: func @nfunction
-      func.func @nfunction() {
+      // CHECK: func nested @nfunction
+      func.func nested @nfunction() {
         return
       }
       func.call @nfunction() : () -> ()



More information about the Mlir-commits mailing list