[mlir] Add a DCE pass for dead symbols.

[rriddle]

Summary

This pass deletes all symbols that are found to be unreachable. This is done by computing the set of operations that are known to be live, propagating that liveness to other symbols, and then deleting all symbols that are not within this live set.

[merge_guards_bot]

Unit tests: pass. 61736 tests passed, 0 failed and 779 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

Rebase on use list refactoring.

[merge_guards_bot]

Unit tests: pass. 62198 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

ping.

[echristo]

Do we need a full separate pass structure for mlir or is there some sharing we can do here as this (at the IR level) is encompassed by GlobalDCE.

[rriddle]

In D72482#1843629, @echristo wrote:

Do we need a full separate pass structure for mlir or is there some sharing we can do here as this (at the IR level) is encompassed by GlobalDCE.

We do need our own pass in MLIR for a couple of reasons:

1. There are many targets of MLIR that don't directly lower to LLVM(e.g. some non programming language front-ends).
2. Symbols are used at many different abstraction levels, and having the ability to DCE at each is important and powerful.

Note that even if these abstractions do eventually lower to LLVM, sometimes the 'symbol' does not have a direct representation in LLVM. For example, a Dialect in MLIR may have a 'global variable' like operation that is eventually transformed into a buffer when compiling for a specific device. In situations like those, the resultant code wouldn't contain a global in the traditional sense.

3. MLIR supports nested references. A "module" isn't constrained to be a module like in the LLVM sense. Modules(and module like operations) can be arbitrarily nested within each other.

I've looked a bit at GlobalDCE, but at this time I'm not sure how much sharing we can do. This pass is meant to be generic, i.e. agnostic to the specific design choices of a given dialect, and the LLVM definition of globals/linkage/visibility does not generalize across all of the different users.

Did you have any specific ideas in mind for how/what things could be shared?

[mehdi_amini]

mlir/include/mlir/Transforms/Passes.h

130

Typo "have are"

mlir/lib/Transforms/SymbolDCE.cpp

105

typo: linakage

118

I'm not sure what this comment is trying to express here?

135

Could we get the uses after this step? This would reduce slightly the memory pressure during the recursion I think.

[rriddle]

Resolve comments

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/IR/SymbolTable.h

public:
   /// Sets the visibility of the given symbol operation.
   static void setSymbolVisibility(Operation *symbol, Visibility vis);
 
   /// Returns the operation registered with the given symbol name with the
   /// regions of 'symbolTableOp'. 'symbolTableOp' is required to be an operation
   /// with the 'OpTrait::SymbolTable' trait.
   static Operation *lookupSymbolIn(Operation *op, StringRef symbol);
   static Operation *lookupSymbolIn(Operation *op, SymbolRefAttr symbol);
+  /// A variant of 'lookupSymbolIn' that returns all of the symbols referenced
+  /// by a given SymbolRefAttr. Returns failure if any of the nested references
+  /// could not be resolved.
+  static LogicalResult lookupSymbolIn(Operation *op, SymbolRefAttr symbol,
+                                      SmallVectorImpl<Operation *> &symbols);
 
   /// Returns the operation registered with the given symbol name within the
   /// closest parent operation of, or including, 'from' with the
   /// 'OpTrait::SymbolTable' trait. Returns nullptr if no valid symbol was
   /// found.
   static Operation *lookupNearestSymbolFrom(Operation *from, StringRef symbol);
   static Operation *lookupNearestSymbolFrom(Operation *from,
                                             SymbolRefAttr symbol);

mlir/include/mlir/Transforms/Passes.h

 std::unique_ptr<OpPassBase<FuncOp>> createStripDebugInfoPass();
 
 /// Creates a pass which tests loop fusion utilities.
 std::unique_ptr<OpPassBase<FuncOp>> createTestLoopFusionPass();
 
 /// Creates a pass which inlines calls and callable operations as defined by the
 /// CallGraph.
 std::unique_ptr<Pass> createInlinerPass();
+
+/// Creates a pass which delete symbol operations that are unreachable. This
+/// pass may *only* be scheduled on an operation that defines a SymbolTable.
+std::unique_ptr<Pass> createSymbolDCEPass();
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/lib/IR/SymbolTable.cpp

for (auto &block : symbolTableOp->getRegion(0)) {
     for (auto &op : block)
       if (getNameIfSymbol(&op) == symbol)
         return &op;
   }
   return nullptr;
 }
 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        SymbolRefAttr symbol) {
+  SmallVector<Operation *, 4> resolvedSymbols;
+  if (failed(lookupSymbolIn(symbolTableOp, symbol, resolvedSymbols)))
+    return nullptr;
+  return resolvedSymbols.back();
+}
+
+LogicalResult
+SymbolTable::lookupSymbolIn(Operation *symbolTableOp, SymbolRefAttr symbol,
+                            SmallVectorImpl<Operation *> &symbols) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
 
   // Lookup the root reference for this symbol.
   symbolTableOp = lookupSymbolIn(symbolTableOp, symbol.getRootReference());
   if (!symbolTableOp)
-    return nullptr;
+    return failure();
+  symbols.push_back(symbolTableOp);
 
   // If there are no nested references, just return the root symbol directly.
   ArrayRef<FlatSymbolRefAttr> nestedRefs = symbol.getNestedReferences();
   if (nestedRefs.empty())
-    return symbolTableOp;
+    return success();
 
   // Verify that the root is also a symbol table.
   if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>())
-    return nullptr;
+    return failure();
 
   // Otherwise, lookup each of the nested non-leaf references and ensure that
   // each corresponds to a valid symbol table.
   for (FlatSymbolRefAttr ref : nestedRefs.drop_back()) {
     symbolTableOp = lookupSymbolIn(symbolTableOp, ref.getValue());
     if (!symbolTableOp || !symbolTableOp->hasTrait<OpTrait::SymbolTable>())
-      return nullptr;
+      return failure();
+    symbols.push_back(symbolTableOp);
   }
-  return lookupSymbolIn(symbolTableOp, symbol.getLeafReference());
+  symbols.push_back(lookupSymbolIn(symbolTableOp, symbol.getLeafReference()));
+  return success(symbols.back());
 }
 
 /// Returns the operation registered with the given symbol name within the
 /// closes parent operation with the 'OpTrait::SymbolTable' trait. Returns
 /// nullptr if no valid symbol was found.
 Operation *SymbolTable::lookupNearestSymbolFrom(Operation *from,
                                                 StringRef symbol) {
   Operation *symbolTableOp = getNearestSymbolTable(from);

mlir/lib/Transforms/CMakeLists.txt

add_llvm_library(MLIRTransforms
   LoopInvariantCodeMotion.cpp
   LoopTiling.cpp
   LoopUnrollAndJam.cpp
   LoopUnroll.cpp
   MemRefDataFlowOpt.cpp
   PipelineDataTransfer.cpp
   SimplifyAffineStructures.cpp
   StripDebugInfo.cpp
+  SymbolDCE.cpp
   Vectorize.cpp
   ViewOpGraph.cpp
   ViewRegionGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
   )
 

mlir/lib/Transforms/SymbolDCE.cpp

+//===- SymbolDCE.cpp - Pass to delete dead symbols ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements an algorithm for eliminating symbol operations that are
+// known to be dead.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+namespace {
+struct SymbolDCE : public OperationPass<SymbolDCE> {
+  void runOnOperation() override;
+
+  /// Compute the liveness of the symbols within the given symbol table.
+  /// `symbolTableIsHidden` is true if this symbol table is known to be
+  /// unaccessible from operations in its parent regions.
+  LogicalResult computeLiveness(Operation *symbolTableOp,
+                                bool symbolTableIsHidden,
+                                DenseSet<Operation *> &liveSymbols);
+};
+} // end anonymous namespace
+
+void SymbolDCE::runOnOperation() {
+  Operation *symbolTableOp = getOperation();
+
+  // SymbolDCE should only be run on operations that define a symbol table.
+  if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>()) {
+    symbolTableOp->emitOpError()
+        << " was scheduled to run under SymbolDCE, but does not define a "
+           "symbol table";
+    return signalPassFailure();
+  }
+
+  // A flag that signals if the top level symbol table is hidden, i.e. not
+  // accessible from parent scopes.
+  bool symbolTableIsHidden = true;
+  if (symbolTableOp->getParentOp() && SymbolTable::isSymbol(symbolTableOp)) {
+    symbolTableIsHidden = SymbolTable::getSymbolVisibility(symbolTableOp) ==
+                          SymbolTable::Visibility::Private;
+  }
+
+  // Compute the set of live symbols within the symbol table.
+  DenseSet<Operation *> liveSymbols;
+  if (failed(computeLiveness(symbolTableOp, symbolTableIsHidden, liveSymbols)))
+    return signalPassFailure();
+
+  // After computing the liveness, delete all of the symbols that were found to
+  // be dead.
+  symbolTableOp->walk([&](Operation *nestedSymbolTable) {
+    if (!nestedSymbolTable->hasTrait<OpTrait::SymbolTable>())
+      return;
+    for (auto &block : nestedSymbolTable->getRegion(0)) {
+      for (Operation &op :
+           llvm::make_early_inc_range(block.without_terminator())) {
+        if (SymbolTable::isSymbol(&op) && !liveSymbols.count(&op))
+          op.erase();
+      }
+    }
+  });
+}
+
+/// Compute the liveness of the symbols within the given symbol table.
+/// `symbolTableIsHidden` is true if this symbol table is known to be
+/// unaccessible from operations in its parent regions.
+LogicalResult SymbolDCE::computeLiveness(Operation *symbolTableOp,
+                                         bool symbolTableIsHidden,
+                                         DenseSet<Operation *> &liveSymbols) {
+  // A worklist of live operations to propagate uses from.
+  SmallVector<Operation *, 16> worklist;
+
+  // Walk the symbols within the current symbol table, marking the symbols that
+  // are known to be live.
+  for (auto &block : symbolTableOp->getRegion(0)) {
+    for (Operation &op : block.without_terminator()) {
+      // Always add non symbol operations to the worklist.
+      if (!SymbolTable::isSymbol(&op)) {
+        worklist.push_back(&op);
+        continue;
+      }
+
+      // Check the visibility to see if this symbol may be referenced
+      // externally.
+      SymbolTable::Visibility visibility =
+          SymbolTable::getSymbolVisibility(&op);
+
+      // Private symbols are always initially considered dead.
+      if (visibility == mlir::SymbolTable::Visibility::Private)
+        continue;
+      // We only include nested visibility here if the symbol table isn't
+      // hidden.
+      if (symbolTableIsHidden && visibility == SymbolTable::Visibility::Nested)
+        continue;
+
+      // TODO(riverriddle) Add hooks here to allow symbols to provide additional
+      // information, e.g. linkage can be used to drop some symbols that may
+      // otherwise be considered "live".
+      if (liveSymbols.insert(&op).second)
+        worklist.push_back(&op);
+    }
+  }
+
+  // Process the set of symbols that were known to be live, adding new symbols
+  // that are referenced within.
+  while (!worklist.empty()) {
+    Operation *op = worklist.pop_back_val();
+
+    // If this is a symbol table, recursively compute its liveness.
+    if (op->hasTrait<OpTrait::SymbolTable>()) {
+      // The internal symbol table is hidden if the parent is, if its not a
+      // symbol, or if it is a private symbol.
+      bool symbolIsHidden = symbolTableIsHidden || !SymbolTable::isSymbol(op) ||
+                            SymbolTable::getSymbolVisibility(op) ==
+                                SymbolTable::Visibility::Private;
+      if (failed(computeLiveness(op, symbolIsHidden, liveSymbols)))
+        return failure();
+    }
+
+    // Collect the uses held by this operation.
+    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";
+    }
+
+    SmallVector<Operation *, 4> resolvedSymbols;
+    for (const SymbolTable::SymbolUse &use : *uses) {
+      // Lookup the symbols referenced by this use.
+      resolvedSymbols.clear();
+      if (failed(SymbolTable::lookupSymbolIn(
+              op->getParentOp(), use.getSymbolRef(), resolvedSymbols))) {
+        return use.getUser()->emitError()
+               << "unable to resolve reference to symbol "
+               << use.getSymbolRef();
+      }
+
+      // Mark each of the resolved symbols as live.
+      for (Operation *resolvedSymbol : resolvedSymbols)
+        if (liveSymbols.insert(resolvedSymbol).second)
+          worklist.push_back(resolvedSymbol);
+    }
+  }
+
+  return success();
+}
+
+std::unique_ptr<Pass> mlir::createSymbolDCEPass() {
+  return std::make_unique<SymbolDCE>();
+}
+
+static PassRegistration<SymbolDCE> pass("symbol-dce", "Eliminate dead symbols");

mlir/test/IR/test-symbol-dce.mlir

+// RUN: mlir-opt %s -symbol-dce -split-input-file -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline="module(symbol-dce)" -split-input-file | FileCheck %s --check-prefix=NESTED
+
+// Check that trivially dead and trivially live non-nested cases are handled.
+
+// CHECK-LABEL: module attributes {test.simple}
+module attributes {test.simple} {
+  // CHECK-NOT: func @dead_private_function
+  func @dead_private_function() attributes { sym_visibility = "nested" }
+
+  // CHECK-NOT: func @dead_nested_function
+  func @dead_nested_function() attributes { sym_visibility = "nested" }
+
+  // CHECK: func @live_private_function
+  func @live_private_function() attributes { sym_visibility = "nested" }
+
+  // CHECK: func @live_nested_function
+  func @live_nested_function() attributes { sym_visibility = "nested" }
+
+  // CHECK: func @public_function
+  func @public_function() {
+    "foo.return"() {uses = [@live_private_function, @live_nested_function]} : () -> ()
+  }
+
+  // CHECK: func @public_function_explicit
+  func @public_function_explicit() attributes { sym_visibility = "public" }
+}
+
+// -----
+
+// Check that we don't DCE nested symbols if they are used.
+// CHECK-LABEL: module attributes {test.nested}
+module attributes {test.nested} {
+  // CHECK: module @public_module
+  module @public_module {
+    // CHECK-NOT: func @dead_nested_function
+    func @dead_nested_function() attributes { sym_visibility = "nested" }
+
+    // CHECK: func @private_function
+    func @private_function() attributes { sym_visibility = "private" }
+
+    // CHECK: func @nested_function
+    func @nested_function() attributes { sym_visibility = "nested" } {
+      "foo.return"() {uses = [@private_function]} : () -> ()
+    }
+  }
+
+  "live.user"() {uses = [@public_module::@nested_function]} : () -> ()
+}
+
+// -----
+
+// Check that we don't DCE symbols if we can't prove that the top-level symbol
+// table that we are running on is hidden from above.
+// NESTED-LABEL: module attributes {test.no_dce_non_hidden_parent}
+module attributes {test.no_dce_non_hidden_parent} {
+  // NESTED: module @public_module
+  module @public_module {
+    // NESTED: func @nested_function
+    func @nested_function() attributes { sym_visibility = "nested" }
+  }
+  // NESTED: module @nested_module
+  module @nested_module attributes { sym_visibility = "nested" } {
+    // NESTED: func @nested_function
+    func @nested_function() attributes { sym_visibility = "nested" }
+  }
+
+  // Only private modules can be assumed to be hidden.
+  // NESTED: module @private_module
+  module @private_module attributes { sym_visibility = "private" } {
+    // NESTED-NOT: func @nested_function
+    func @nested_function() attributes { sym_visibility = "nested" }
+  }
+
+  "live.user"() {uses = [@nested_module, @private_module]} : () -> ()
+}
+
+// -----
+
+module {
+  func @private_symbol() attributes { sym_visibility = "private" }
+
+  // expected-error@+1 {{contains potentially unknown symbol table}}
+  "foo.possibly_unknown_symbol_table"() ({
+  }) : () -> ()
+}
+
+// -----
+
+module {
+  // expected-error@+1 {{unable to resolve reference to symbol}}
+  "live.user"() {uses = [@unknown_symbol]} : () -> ()
+}