[mlir] Expose getNearestSymbolTable as SymbolTable class method

[antiagainst]

Summary

This is a generally useful utility function for interacting with
symbol tables.

[merge_guards_bot]

Unit tests: pass. 62197 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

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

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/IR/SymbolTable.h

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 nearest symbol table from a given operation `from`. Returns
+  /// nullptr if no valid parent symbol table could be found.
+  static Operation *getNearestSymbolTable(Operation *from);
+
   /// 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

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

emitError(loc, "unimplemented builtin variable generation for ")
         << stringifyBuiltIn(builtin);
   }
   return newVarOp;
 }
 
 Value mlir::spirv::getBuiltinVariableValue(Operation *op,
                                            spirv::BuiltIn builtin,
                                            OpBuilder &builder) {
-  Operation *parent = op->getParentOp();
-  while (parent && !parent->hasTrait<OpTrait::SymbolTable>())
-    parent = parent->getParentOp();
+  Operation *parent = SymbolTable::getNearestSymbolTable(op->getParentOp());
   if (!parent) {
     op->emitError("expected operation to be within a module-like op");
     return nullptr;
   }
 
   spirv::GlobalVariableOp varOp = getOrInsertBuiltinVariable(
       *parent->getRegion(0).begin(), op->getLoc(), builtin, builder);
   Value ptr = builder.create<spirv::AddressOfOp>(op->getLoc(), varOp);

mlir/lib/IR/SymbolTable.cpp

 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;
 }
 
-/// Returns the nearest symbol table from a given operation `from`. Returns
-/// nullptr if no valid parent symbol table could be found.
-static Operation *getNearestSymbolTable(Operation *from) {
-  assert(from && "expected valid operation");
-  if (isPotentiallyUnknownSymbolTable(from))
-    return nullptr;
-
-  while (!from->hasTrait<OpTrait::SymbolTable>()) {
-    from = from->getParentOp();
-
-    // Check that this is a valid op and isn't an unknown symbol table.
-    if (!from || isPotentiallyUnknownSymbolTable(from))
-      return nullptr;
-  }
-  return from;
-}
-
 /// Returns the string name of the given symbol, or None if this is not a
 /// symbol.
 static Optional<StringRef> getNameIfSymbol(Operation *symbol) {
   auto nameAttr =
       symbol->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
   return nameAttr ? nameAttr.getValue() : Optional<StringRef>();
 }
 
void SymbolTable::setSymbolVisibility(Operation *symbol, Visibility vis) {
   // 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 nearest symbol table from a given operation `from`. Returns
+/// nullptr if no valid parent symbol table could be found.
+Operation *SymbolTable::getNearestSymbolTable(Operation *from) {
+  assert(from && "expected valid operation");
+  if (isPotentiallyUnknownSymbolTable(from))
+    return nullptr;
+
+  while (!from->hasTrait<OpTrait::SymbolTable>()) {
+    from = from->getParentOp();
+
+    // Check that this is a valid op and isn't an unknown symbol table.
+    if (!from || isPotentiallyUnknownSymbolTable(from))
+      return nullptr;
+  }
+  return from;
+}
+
 /// 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>());
 
llvm::SetVector<Operation *, SmallVector<Operation *, 4>,
                   SmallPtrSet<Operation *, 4>>
       limitAncestors;
   Operation *limitAncestor = limit;
   do {
     // Check to see if 'symbol' is an ancestor of 'limit'.
     if (limitAncestor == symbol) {
       // Check that the nearest symbol table is 'symbol's parent. SymbolRefAttr
       // doesn't support parent references.
-      if (getNearestSymbolTable(limit) != symbol->getParentOp())
+      if (SymbolTable::getNearestSymbolTable(limit) != symbol->getParentOp())
         return WalkResult::advance();
       return callback(SymbolRefAttr::get(symbolName, symbol->getContext()),
                       limit);
     }
 
     limitAncestors.insert(limitAncestor);
   } while ((limitAncestor = limitAncestor->getParentOp()));