[Mlir-commits] [mlir] f43e67c - [mlir] Allow SymbolTable to update existing symbols

[Tres Popp]

Author: Tres Popp
Date: 2020-12-16T00:44:40+01:00
New Revision: f43e67cc6c6f6e0da925039b28b54e44fc751267

URL: https://github.com/llvm/llvm-project/commit/f43e67cc6c6f6e0da925039b28b54e44fc751267
DIFF: https://github.com/llvm/llvm-project/commit/f43e67cc6c6f6e0da925039b28b54e44fc751267.diff

LOG: [mlir] Allow SymbolTable to update existing symbols

Previous behavior would fail if inserting an operation that already
existed. Now SymbolTable::insert can also be used as a way to make a
symbol's name unique even after insertion.

Further TODOs have been left over naming and consistent behavior
considerations.

Differential Revision: https://reviews.llvm.org/D93349

Added: 
    

Modified: 
    mlir/include/mlir/IR/SymbolTable.h
    mlir/lib/IR/SymbolTable.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index 0e3230672568..edeaab21aabe 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -38,7 +38,8 @@ class SymbolTable {
 
   /// 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.
+  /// 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 = {});
 
   /// Return the name of the attribute used for symbol names.

diff  --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 85c08c3a7e8f..b3bf6e169530 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -151,23 +151,35 @@ void SymbolTable::erase(Operation *symbol) {
   }
 }
 
-/// Insert a new symbol into the table and associated operation, and rename it
-/// as necessary to avoid collisions.
+// 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) {
-  auto &body = symbolTableOp->getRegion(0).front();
-  if (insertPt == Block::iterator() || insertPt == body.end())
-    insertPt = Block::iterator(body.getTerminator());
-
-  assert(insertPt->getParentOp() == symbolTableOp &&
-         "expected insertPt to be in the associated module operation");
-
-  body.getOperations().insert(insertPt, symbol);
+  // 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 == body.end())
+      insertPt = Block::iterator(body.getTerminator());
+
+    assert(insertPt->getParentOp() == symbolTableOp &&
+           "expected insertPt to be in the associated module operation");
+
+    body.getOperations().insert(insertPt, symbol);
+  }
+  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.
   StringRef name = getSymbolName(symbol);
   if (symbolTable.insert({name, symbol}).second)
     return;
+  // If the symbol was already in the table, also return.
+  if (symbolTable.lookup(name) == symbol)
+    return;
   // 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);