[mlir] Add support for attaching a visibility to symbols.

Summary:
The visibility defines the structural reachability of the symbol within the IR. Symbols can define one of three visibilities:

* Public
The symbol \may be accessed from outside of the visible IR. We cannot assume that we can observe all of the uses of this symbol.

* Private
The symbol may only be referenced from within the operations in the current symbol table, via SymbolRefAttr.

* Nested
The symbol may be referenced by operations in symbol tables above the current symbol table, as long as each symbol table parent also defines a non-private symbol. This allows or referencing the symbol from outside of the defining symbol table, while retaining the ability for the compiler to see all uses.

These properties help to reason about the properties of a symbol, and will be used in a follow up to implement a dce pass on dead symbols.

A few examples of what this would look like in the IR are shown below:

  module @public_module {
    // This function can be accessed by 'live.user'
    func @nested_function() attributes { sym_visibility = "nested" }

    // This function cannot be accessed outside of 'public_module'
   func @private_function() attributes { sym_visibility = "private" }
  }

  // This function can only be accessed from within this module.
  func @private_function() attributes { sym_visibility = "private" }

  // This function may be referenced externally.
  func @public_function()

  "live.user"() {uses = [@public_module::@nested_function,
                                      @private_function,
                                      @public_function]} : () -> ()

Depends On D72043

Reviewed By: mehdi_amini

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

1 files changed, 38 insertions, 0 deletions

diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index 2df39ea1b73..1b93b66b687 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -46,10 +46,31 @@ public:
   /// Returns the associated operation.
   Operation *getOp() const { return symbolTableOp; }
 
+  /// Return the name of the attribute used for symbol visibility.
+  static StringRef getVisibilityAttrName() { return "sym_visibility"; }
+
   //===--------------------------------------------------------------------===//
   // Symbol Utilities
   //===--------------------------------------------------------------------===//
 
+  /// An enumeration detailing the different visibility types that a symbol may
+  /// have.
+  enum class Visibility {
+    /// The symbol is public and may be referenced anywhere internal or external
+    /// to the visible references in the IR.
+    Public,
+
+    /// The symbol is private and may only be referenced by SymbolRefAttrs local
+    /// to the operations within the current symbol table.
+    Private,
+
+    /// The symbol is visible to the current IR, which may include operations in
+    /// symbol tables above the one that owns the current symbol. `Nested`
+    /// visibility allows for referencing a symbol outside of its current symbol
+    /// table, while retaining the ability to observe all uses.
+    Nested,
+  };
+
   /// Returns true if the given operation defines a symbol.
   static bool isSymbol(Operation *op);
 
@@ -58,6 +79,11 @@ public:
   /// Sets the name of the given symbol operation.
   static void setSymbolName(Operation *symbol, StringRef name);
 
+  /// Returns the visibility of the given symbol operation.
+  static Visibility getSymbolVisibility(Operation *symbol);
+  /// 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.
@@ -200,6 +226,8 @@ public:
 template <typename ConcreteType>
 class Symbol : public TraitBase<ConcreteType, Symbol> {
 public:
+  using Visibility = mlir::SymbolTable::Visibility;
+
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySymbol(op);
   }
@@ -219,6 +247,16 @@ public:
         StringAttr::get(name, this->getOperation()->getContext()));
   }
 
+  /// Returns the visibility of the current symbol.
+  Visibility getVisibility() {
+    return mlir::SymbolTable::getSymbolVisibility(this->getOperation());
+  }
+
+  /// Sets the visibility of the current symbol.
+  void setVisibility(Visibility vis) {
+    mlir::SymbolTable::setSymbolVisibility(this->getOperation(), vis);
+  }
+
   /// Get all of the uses of the current symbol that are nested within the given
   /// operation 'from'.
   /// Note: See mlir::SymbolTable::getSymbolUses for more details.