[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

5 files changed, 124 insertions, 3 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.
diff --git a/mlir/lib/IR/Module.cpp b/mlir/lib/IR/Module.cpp
index c5af227459c..e0caeda5ed8 100644
--- a/mlir/lib/IR/Module.cpp
+++ b/mlir/lib/IR/Module.cpp
@@ -82,10 +82,13 @@ LogicalResult ModuleOp::verify() {
     return emitOpError("expected body to have no arguments");
 
   // Check that none of the attributes are non-dialect attributes, except for
-  // the symbol name attribute.
+  // the symbol related attributes.
   for (auto attr : getOperation()->getAttrList().getAttrs()) {
     if (!attr.first.strref().contains('.') &&
-        attr.first.strref() != mlir::SymbolTable::getSymbolAttrName())
+        !llvm::is_contained(
+            ArrayRef<StringRef>{mlir::SymbolTable::getSymbolAttrName(),
+                                mlir::SymbolTable::getVisibilityAttrName()},
+            attr.first.strref()))
       return emitOpError(
                  "can only contain dialect-specific attributes, found: '")
              << attr.first << "'";
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 057aeded242..ed9de48c400 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
@@ -179,6 +180,38 @@ void SymbolTable::setSymbolName(Operation *symbol, StringRef name) {
                   StringAttr::get(name, symbol->getContext()));
 }
 
+/// Returns the visibility of the given symbol operation.
+SymbolTable::Visibility SymbolTable::getSymbolVisibility(Operation *symbol) {
+  // If the attribute doesn't exist, assume public.
+  StringAttr vis = symbol->getAttrOfType<StringAttr>(getVisibilityAttrName());
+  if (!vis)
+    return Visibility::Public;
+
+  // Otherwise, switch on the string value.
+  return llvm::StringSwitch<Visibility>(vis.getValue())
+      .Case("private", Visibility::Private)
+      .Case("nested", Visibility::Nested)
+      .Case("public", Visibility::Public);
+}
+/// Sets the visibility of the given symbol operation.
+void SymbolTable::setSymbolVisibility(Operation *symbol, Visibility vis) {
+  MLIRContext *ctx = symbol->getContext();
+
+  // If the visibility is public, just drop the attribute as this is the
+  // default.
+  if (vis == Visibility::Public) {
+    symbol->removeAttr(Identifier::get(getVisibilityAttrName(), ctx));
+    return;
+  }
+
+  // Otherwise, update the attribute.
+  assert((vis == Visibility::Private || vis == Visibility::Nested) &&
+         "unknown symbol visibility kind");
+
+  StringRef visName = vis == Visibility::Private ? "private" : "nested";
+  symbol->setAttr(getVisibilityAttrName(), StringAttr::get(visName, ctx));
+}
+
 /// 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. Returns nullptr if no valid symbol
@@ -272,9 +305,26 @@ LogicalResult OpTrait::impl::verifySymbolTable(Operation *op) {
 }
 
 LogicalResult OpTrait::impl::verifySymbol(Operation *op) {
+  // Verify the name attribute.
   if (!op->getAttrOfType<StringAttr>(mlir::SymbolTable::getSymbolAttrName()))
     return op->emitOpError() << "requires string attribute '"
                              << mlir::SymbolTable::getSymbolAttrName() << "'";
+
+  // Verify the visibility attribute.
+  if (Attribute vis = op->getAttr(mlir::SymbolTable::getVisibilityAttrName())) {
+    StringAttr visStrAttr = vis.dyn_cast<StringAttr>();
+    if (!visStrAttr)
+      return op->emitOpError() << "requires visibility attribute '"
+                               << mlir::SymbolTable::getVisibilityAttrName()
+                               << "' to be a string attribute, but got " << vis;
+
+    if (!llvm::is_contained(ArrayRef<StringRef>{"public", "private", "nested"},
+                            visStrAttr.getValue()))
+      return op->emitOpError()
+             << "visibility expected to be one of [\"public\", \"private\", "
+                "\"nested\"], but got "
+             << visStrAttr;
+  }
   return success();
 }
 
diff --git a/mlir/test/IR/traits.mlir b/mlir/test/IR/traits.mlir
index 794ed4cd4f7..42044bde5dc 100644
--- a/mlir/test/IR/traits.mlir
+++ b/mlir/test/IR/traits.mlir
@@ -207,6 +207,30 @@ func @failedSingleBlockImplicitTerminator_missing_terminator() {
 
 // -----
 
+// Test the invariants of operations with the Symbol Trait.
+
+// expected-error@+1 {{requires string attribute 'sym_name'}}
+"test.symbol"() {} : () -> ()
+
+// -----
+
+// expected-error@+1 {{requires visibility attribute 'sym_visibility' to be a string attribute}}
+"test.symbol"() {sym_name = "foo_2", sym_visibility} : () -> ()
+
+// -----
+
+// expected-error@+1 {{visibility expected to be one of ["public", "private", "nested"]}}
+"test.symbol"() {sym_name = "foo_2", sym_visibility = "foo"} : () -> ()
+
+// -----
+
+"test.symbol"() {sym_name = "foo_3", sym_visibility = "nested"} : () -> ()
+"test.symbol"() {sym_name = "foo_4", sym_visibility = "private"} : () -> ()
+"test.symbol"() {sym_name = "foo_5", sym_visibility = "public"} : () -> ()
+"test.symbol"() {sym_name = "foo_6"} : () -> ()
+
+// -----
+
 // Test that operation with the SymbolTable Trait define a new symbol scope.
 "test.symbol_scope"() ({
   func @foo() {
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index f10991dfe5b..0dd32d74a75 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -74,9 +74,15 @@ def MultiTensorRankOf : TEST_Op<"multi_tensor_rank_of"> {
 }
 
 //===----------------------------------------------------------------------===//
-// Test Operands
+// Test Symbols
 //===----------------------------------------------------------------------===//
 
+def SymbolOp : TEST_Op<"symbol", [Symbol]> {
+  let summary =  "operation which defines a new symbol";
+  let arguments = (ins StrAttr:$sym_name,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+}
+
 def SymbolScopeOp : TEST_Op<"symbol_scope",
     [SymbolTable, SingleBlockImplicitTerminator<"TerminatorOp">]> {
   let summary =  "operation which defines a new symbol table";