[mlir] return the updated symbol table after inserting into SymbolTable

[ftynse]

Summary

Inserting a symbol into a SymbolTable may lead to the name of the symbol being
changed in order to ensure uniqueness of symbol names in the table. Return this
new name to spare the caller the need to extract it from the symbol operation.

Depends On D112700

[mehdi_amini]

Should we return the Attr? It can avoid someone needing to get a attribute from the returned StringRef.

[mehdi_amini]

(LG otherwise)

[ftynse]

Address review.

[ftynse]

Something went wrong with the previous upload...

mlir/include/mlir/IR/SymbolTable.h

public:
   }
 
   /// Erase the given symbol from the table.
   void erase(Operation *symbol);
 
   /// Insert a new symbol into the table, and rename it as necessary to avoid
   /// collisions. Also insert at the specified location in the body of the
   /// associated operation if it is not already there. It is asserted that the
-  /// symbol is not inside another operation.
-  void insert(Operation *symbol, Block::iterator insertPt = {});
+  /// symbol is not inside another operation. Return the name of the symbol
+  /// after insertion.
+  StringRef 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.

mlir/lib/IR/SymbolTable.cpp

void SymbolTable::erase(Operation *symbol) {
   if (it != symbolTable.end() && it->second == symbol) {
     symbolTable.erase(it);
     symbol->erase();
   }
 }
 
 // TODO: Consider if this should be renamed to something like insertOrUpdate
 /// Insert a new symbol into the table and associated operation if not already
-/// there and rename it as necessary to avoid collisions.
-void SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
+/// there and rename it as necessary to avoid collisions. Return the name of
+/// the symbol after insertion.
+StringRef SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
   // The symbol cannot be the child of another op and must be the child of the
   // symbolTableOp after this.
   //
   // TODO: consider if SymbolTable's constructor should behave the same.
   if (!symbol->getParentOp()) {
     auto &body = symbolTableOp->getRegion(0).front();
     if (insertPt == Block::iterator()) {
       insertPt = Block::iterator(body.end());
StringRef SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
   }
   assert(symbol->getParentOp() == symbolTableOp &&
          "symbol is already inserted in another op");
 
   // Add this symbol to the symbol table, uniquing the name if a conflict is
   // detected.
   StringAttr name = getSymbolName(symbol);
   if (symbolTable.insert({name, symbol}).second)
-    return;
+    return name.getValue();
   // If the symbol was already in the table, also return.
   if (symbolTable.lookup(name) == symbol)
-    return;
+    return name.getValue();
   // If a conflict was detected, then the symbol will not have been added to
   // the symbol table. Try suffixes until we get to a unique name that works.
   SmallString<128> nameBuffer(name.getValue());
   unsigned originalLength = nameBuffer.size();
 
   MLIRContext *context = symbol->getContext();
 
   // Iteratively try suffixes until we find one that isn't used.
   do {
     nameBuffer.resize(originalLength);
     nameBuffer += '_';
     nameBuffer += std::to_string(uniquingCounter++);
   } while (!symbolTable.insert({StringAttr::get(context, nameBuffer), symbol})
                 .second);
   setSymbolName(symbol, nameBuffer);
+  return getSymbolName(symbol).getValue();
 }
 
 /// Returns the name of the given symbol operation.
 StringAttr SymbolTable::getSymbolName(Operation *symbol) {
   StringAttr name = getNameIfSymbol(symbol);
   assert(name && "expected valid symbol name");
   return name;
 }