[mlir] Add pass to privatize symbols unless excluded.

[jpienaar]

Summary

Simple pass that changes all symbols to private unless symbol is excluded (and
in which case there is no change to symbol's visibility).

[bondhugula]

Looks too specialized and a bit trivial of a task to justify a pass. But otherwise LGTM.

mlir/test/Transforms/test-symbol-privatize.mlir

2

Don't need -allow-unregistered...

[jpienaar]

Removed unneeded allow unregistered dialects.

[mehdi_amini]

mlir/lib/Transforms/SymbolPrivatize.cpp

41

I would limit this to the current level of nesting and not walk recursively.

[rriddle]

This is essentially the MLIR equivalent of LLVM's Internalize pass, so seems fine as an initial implementation of that. We may grow an interface or something to help drive this though, but that can come later.

mlir/include/mlir/Transforms/Passes.td

220

? Slightly more explicit.

mlir/lib/Transforms/SymbolPrivatize.cpp

22

41

Drop the mlir::

41

re: Mehdi's comment

It would also be good to include this in the documentation.

[jpienaar]

Addressed comments (made explicit, avoided recursing, updated pass doc)

mlir/include/mlir/Transforms/Passes.h

 /// Creates a pass which performs sparse conditional constant propagation over
 /// nested operations.
 std::unique_ptr<Pass> createSCCPPass();
 
 /// 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();
 
+/// Creates a pass which marks symbol operations as `private` unless marked
+/// as excluded and then skipped.
+std::unique_ptr<Pass>
+createSymbolPrivatizePass(ArrayRef<std::string> excludeSymbols = {});
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_PASSES_H

mlir/include/mlir/Transforms/Passes.td

let description = [{
     ```
 
     See [Symbols and SymbolTables](SymbolsAndSymbolTables.md) for more
     information on `Symbols`.
   }];
   let constructor = "mlir::createSymbolDCEPass()";
 }
 
+def SymbolPrivatize : Pass<"symbol-privatize"> {
+  let summary = "Mark symbols private";
+  let description = [{
+    This pass simply marks all symbols not listed in pass options as `private`.
+  }];
+  let options = [
+    ListOption<"exclude", "exclude", "std::string",
+       "Comma separated list of symbols that should not be marked private",
+       "llvm::cl::MiscFlags::CommaSeparated">
+  ];
+  let constructor = "mlir::createSymbolPrivatizePass()";
+}
+
 def ViewOpGraph : Pass<"view-op-graph"> {
   let summary = "Print Graphviz visualization of an operation";
   let description = [{
     This pass prints a Graphviz graph of a module.
 
     - Operations are represented as nodes;
     - Uses (data flow) as edges;
     - Control flow as dashed edges;

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
   ControlFlowSink.cpp
   CSE.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   OpStats.cpp
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
+  SymbolPrivatize.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS
   MLIRTransformsPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRCopyOpInterface
   MLIRLoopLikeInterface
   MLIRPass
   MLIRSupport
   MLIRTransformUtils
   )

mlir/lib/Transforms/SymbolPrivatize.cpp

+//===- SymbolPrivatize.cpp - Pass to mark symbols private -----------------===//
+//
+// 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 pass that marks all symbols as private unless
+// excluded.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+namespace {
+struct SymbolPrivatize : public SymbolPrivatizeBase<SymbolPrivatize> {
+  SymbolPrivatize(llvm::ArrayRef<std::string> excludeSymbols);
+  LogicalResult initialize(MLIRContext *context) override;
+  void runOnOperation() override;
+
+  /// Symbols whose visibility won't be changed.
+  DenseSet<StringAttr> excludedSymbols;
+};
+} // namespace
+
+SymbolPrivatize::SymbolPrivatize(llvm::ArrayRef<std::string> excludeSymbols) {
+  exclude = excludeSymbols;
+}
+
+LogicalResult SymbolPrivatize::initialize(MLIRContext *context) {
+  for (const std::string &symbol : exclude)
+    excludedSymbols.insert(StringAttr::get(context, symbol));
+  return success();
+}
+
+void SymbolPrivatize::runOnOperation() {
+  getOperation()->walk([this](mlir::SymbolOpInterface symbol) {
+    if (!excludedSymbols.contains(symbol.getNameAttr()))
+      symbol.setVisibility(SymbolTable::Visibility::Private);
+  });
+}
+
+std::unique_ptr<Pass>
+mlir::createSymbolPrivatizePass(ArrayRef<std::string> exclude) {
+  return std::make_unique<SymbolPrivatize>(exclude);
+}

mlir/test/Transforms/test-symbol-privatize.mlir

+// RUN: mlir-opt %s -symbol-privatize=exclude="aap" | FileCheck %s
+
+// CHECK-LABEL: module attributes {test.simple}
+module attributes {test.simple} {
+  // CHECK: func @aap
+  func @aap() { return }
+
+  // CHECK: func private @kat
+  func @kat() { return }
+}
+