mlir/lib/Transforms/SymbolPrivatize.cpp

   1 //===- SymbolPrivatize.cpp - Pass to mark symbols private -----------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements an pass that marks all symbols as private unless
  10 // excluded.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "mlir/Transforms/Passes.h"
  15
  16 #include "mlir/IR/SymbolTable.h"
  17
  18 namespace mlir {
  19 #define GEN_PASS_DEF_SYMBOLPRIVATIZE
  20 #include "mlir/Transforms/Passes.h.inc"
  21 } // namespace mlir
  22
  23 using namespace mlir;
  24
  25 namespace {
  26 struct SymbolPrivatize : public impl::SymbolPrivatizeBase<SymbolPrivatize> {
  27   explicit SymbolPrivatize(ArrayRef<std::string> excludeSymbols);
  28   LogicalResult initialize(MLIRContext *context) override;
  29   void runOnOperation() override;
  30
  31   /// Symbols whose visibility won't be changed.
  32   DenseSet<StringAttr> excludedSymbols;
  33 };
  34 } // namespace
  35
  36 SymbolPrivatize::SymbolPrivatize(llvm::ArrayRef<std::string> excludeSymbols) {
  37   exclude = excludeSymbols;
  38 }
  39
  40 LogicalResult SymbolPrivatize::initialize(MLIRContext *context) {
  41   for (const std::string &symbol : exclude)
  42     excludedSymbols.insert(StringAttr::get(context, symbol));
  43   return success();
  44 }
  45
  46 void SymbolPrivatize::runOnOperation() {
  47   for (Region &region : getOperation()->getRegions()) {
  48     for (Block &block : region) {
  49       for (Operation &op : block) {
  50         auto symbol = dyn_cast<SymbolOpInterface>(op);
  51         if (!symbol)
  52           continue;
  53         if (!excludedSymbols.contains(symbol.getNameAttr()))
  54           symbol.setVisibility(SymbolTable::Visibility::Private);
  55       }
  56     }
  57   }
  58 }
  59
  60 std::unique_ptr<Pass>
  61 mlir::createSymbolPrivatizePass(ArrayRef<std::string> exclude) {
  62   return std::make_unique<SymbolPrivatize>(exclude);
  63 }