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

[rriddle]

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

[rriddle]

Fix typo.

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61162 tests passed, 0 failed and 728 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

Rebase

[merge_guards_bot]

Unit tests: pass. 61735 tests passed, 0 failed and 779 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[mehdi_amini]

mlir/lib/IR/Module.cpp

85

Can you update the comment?

[rriddle]

Address comments.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

mlir/include/mlir/IR/SymbolTable.h

public:
   void insert(Operation *symbol, Block::iterator insertPt = {});
 
   /// Return the name of the attribute used for symbol names.
   static StringRef getSymbolAttrName() { return "sym_name"; }
 
   /// 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);
 
   /// Returns the name of the given symbol operation.
   static StringRef getSymbolName(Operation *symbol);
   /// 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.
   static Operation *lookupSymbolIn(Operation *op, StringRef symbol);
   static Operation *lookupSymbolIn(Operation *op, SymbolRefAttr symbol);
 
   /// Returns the operation registered with the given symbol name within the
   /// closest parent operation of, or including, 'from' with the

 };
 
 /// A trait used to define a symbol that can be used on operations within a
 /// symbol table. Operations using this trait must adhere to the following:
 ///   * Have a StringAttr attribute named 'SymbolTable::getSymbolAttrName()'.
 template <typename ConcreteType>
 class Symbol : public TraitBase<ConcreteType, Symbol> {
 public:
+  using Visibility = mlir::SymbolTable::Visibility;
+
   static LogicalResult verifyTrait(Operation *op) {
     return impl::verifySymbol(op);
   }
 
   /// Returns the name of this symbol.
   StringRef getName() {
     return this->getOperation()
         ->template getAttrOfType<StringAttr>(
             mlir::SymbolTable::getSymbolAttrName())
         .getValue();
   }
 
   /// Set the name of this symbol.
   void setName(StringRef name) {
     this->getOperation()->setAttr(
         mlir::SymbolTable::getSymbolAttrName(),
         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.
   Optional<::mlir::SymbolTable::UseRange> getSymbolUses(Operation *from) {
     return ::mlir::SymbolTable::getSymbolUses(this->getOperation(), from);
   }
 
   /// Return if the current symbol is known to have no uses that are nested

mlir/lib/IR/Module.cpp

if (!has_single_element(bodyRegion))
     return emitOpError("expected body region to have a single block");
 
   // Check that the body has no block arguments.
   auto *body = &bodyRegion.front();
   if (body->getNumArguments() != 0)
     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 << "'";
   }
 
   return success();
 }
 

mlir/lib/IR/SymbolTable.cpp

 //===- SymbolTable.cpp - MLIR Symbol Table Class --------------------------===//
 //
 // Part of the MLIR 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
 /// Return true if the given operation is unknown and may potentially define a
 /// symbol table.
 static bool isPotentiallyUnknownSymbolTable(Operation *op) {
   return !op->getDialect() && op->getNumRegions() == 1;
 }
StringRef SymbolTable::getSymbolName(Operation *symbol) {
   return *name;
 }
 /// Sets the name of the given symbol operation.
 void SymbolTable::setSymbolName(Operation *symbol, StringRef name) {
   symbol->setAttr(getSymbolAttrName(),
                   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
 /// was found.
 Operation *SymbolTable::lookupSymbolIn(Operation *symbolTableOp,
                                        StringRef symbol) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>());
 
for (auto &op : block) {
             .attachNote(it.first->second)
             .append("see existing symbol definition here");
     }
   }
   return success();
 }
 
 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();
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Lists
 //===----------------------------------------------------------------------===//
 
 /// Walk all of the symbol references within the given operation, invoking the

mlir/test/IR/traits.mlir

func @failedSingleBlockImplicitTerminator_missing_terminator() {
   "test.SingleBlockImplicitTerminator"() ({
   ^entry:
     "test.non_existent_op"() : () -> ()
   }) : () -> ()
 }
 
 // -----
 
+// 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() {
   }
   "test.finish" () : () -> ()
 }) : () -> ()
 func @foo() {
 }

mlir/test/lib/TestDialect/TestOps.td

 
 def MultiTensorRankOf : TEST_Op<"multi_tensor_rank_of"> {
   let arguments = (ins
     TensorRankOf<[I8, I32, F32], [0, 1]>:$arg0
   );
 }
 
 //===----------------------------------------------------------------------===//
-// 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";
   let regions = (region SizedRegion<1>:$region);
 }
 
 def SymbolTableRegionOp : TEST_Op<"symbol_table_region", [SymbolTable]> {
   let summary =  "operation which defines a new symbol table without a "