[Mlir-commits] [mlir] 88c5252 - [mlir] Add pass to privatize symbols unless excluded.

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2022-02-03T20:20:54-08:00
New Revision: 88c525235bc979af9e137d1e9b9d03633b0b0598

URL: https://github.com/llvm/llvm-project/commit/88c525235bc979af9e137d1e9b9d03633b0b0598
DIFF: https://github.com/llvm/llvm-project/commit/88c525235bc979af9e137d1e9b9d03633b0b0598.diff

LOG: [mlir] Add pass to privatize symbols unless excluded.

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

Differential Revision: https://reviews.llvm.org/D118752

Added: 
    mlir/lib/Transforms/SymbolPrivatize.cpp
    mlir/test/Transforms/test-symbol-privatize.mlir

Modified: 
    mlir/include/mlir/Transforms/Passes.h
    mlir/include/mlir/Transforms/Passes.td
    mlir/lib/Transforms/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index e2b5f14d34adf..a74f65a81583f 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -87,6 +87,11 @@ std::unique_ptr<Pass> createSCCPPass();
 /// pass may *only* be scheduled on an operation that defines a SymbolTable.
 std::unique_ptr<Pass> createSymbolDCEPass();
 
+/// Creates a pass which marks top-level symbol operations as `private` unless
+/// listed in `excludeSymbols`.
+std::unique_ptr<Pass>
+createSymbolPrivatizePass(ArrayRef<std::string> excludeSymbols = {});
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 3c3b43d75e440..8b9d2c45b7b57 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -214,6 +214,20 @@ def SymbolDCE : Pass<"symbol-dce"> {
   let constructor = "mlir::createSymbolDCEPass()";
 }
 
+def SymbolPrivatize : Pass<"symbol-privatize"> {
+  let summary = "Mark symbols private";
+  let description = [{
+    This pass marks all top-level symbols of the operation run as `private`
+    except if listed in `exclude` pass option.
+  }];
+  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 = [{

diff  --git a/mlir/lib/Transforms/CMakeLists.txt b/mlir/lib/Transforms/CMakeLists.txt
index 6e59ef2d696ac..fef0feb379e57 100644
--- a/mlir/lib/Transforms/CMakeLists.txt
+++ b/mlir/lib/Transforms/CMakeLists.txt
@@ -11,6 +11,7 @@ add_mlir_library(MLIRTransforms
   SCCP.cpp
   StripDebugInfo.cpp
   SymbolDCE.cpp
+  SymbolPrivatize.cpp
   ViewOpGraph.cpp
 
   ADDITIONAL_HEADER_DIRS

diff  --git a/mlir/lib/Transforms/SymbolPrivatize.cpp b/mlir/lib/Transforms/SymbolPrivatize.cpp
new file mode 100644
index 0000000000000..4aa7897da0ade
--- /dev/null
+++ b/mlir/lib/Transforms/SymbolPrivatize.cpp
@@ -0,0 +1,58 @@
+//===- 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/IR/SymbolTable.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+
+namespace {
+struct SymbolPrivatize : public SymbolPrivatizeBase<SymbolPrivatize> {
+  explicit SymbolPrivatize(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() {
+  for (Region &region : getOperation()->getRegions()) {
+    for (Block &block : region) {
+      for (Operation &op : block) {
+        auto symbol = dyn_cast<SymbolOpInterface>(op);
+        if (!symbol)
+          continue;
+        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);
+}

diff  --git a/mlir/test/Transforms/test-symbol-privatize.mlir b/mlir/test/Transforms/test-symbol-privatize.mlir
new file mode 100644
index 0000000000000..9f67b1ad56ced
--- /dev/null
+++ b/mlir/test/Transforms/test-symbol-privatize.mlir
@@ -0,0 +1,11 @@
+// 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 }
+}
+