[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

2 files changed, 55 insertions, 2 deletions

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();
 }