[mlir][AsmPrinter] Remove recursion while SSA naming

[Chia-hungDuan]

Summary

Address the TODO of removing recursion while SSA naming.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

850

What about adding the Region* here? So that you can have a single SmallVector to drive the iterative loop?

883

Please keep auto to places where it improves readability.

[Chia-hungDuan]

Address review comments

1. Use explict type for auto
2. Use one SmallVector

[Chia-hungDuan]

mlir/lib/IR/AsmPrinter.cpp

850

Sure, either way is fine to me. Done!

[rriddle]

mlir/lib/IR/AsmPrinter.cpp

847

Types should be CamelCase.

859

You need a naming context for the top-level operation, otherwise this will crash.

876–877

881–882

mlir/lib/IR/AsmPrinter.cpp

 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SaveAndRestore.h"
+
+#include <tuple>
+
 using namespace mlir;
 using namespace mlir::detail;
 
 void Identifier::print(raw_ostream &os) const { os << str(); }
 
 void Identifier::dump() const { print(llvm::errs()); }
 
 void OperationName::print(raw_ostream &os) const { os << getStringRef(); }
private:
   /// This is the next ID to assign when a name conflict is detected.
   unsigned nextConflictID = 0;
 };
 } // end anonymous namespace
 
 SSANameState::SSANameState(
     Operation *op,
     DialectInterfaceCollection<OpAsmDialectInterface> &interfaces) {
-  llvm::ScopedHashTable<StringRef, char>::ScopeTy usedNamesScope(usedNames);
+  llvm::SaveAndRestore<unsigned> valueIDSaver(nextValueID);
+  llvm::SaveAndRestore<unsigned> argumentIDSaver(nextArgumentID);
+  llvm::SaveAndRestore<unsigned> conflictIDSaver(nextConflictID);
+
+  // The context includes nextValueID, nextArgumentID, nextConflictID and scoped
+  // HashTable.
+  using hashTableScopeTy = llvm::ScopedHashTable<StringRef, char>::ScopeTy;
+  // A namingContext carries the information inherits from parent region.
+  using namingContext =
+      std::tuple<Region *, unsigned, unsigned, unsigned, hashTableScopeTy *>;
+  // Allocator for hashTableScopeTy
+  llvm::BumpPtrAllocator allocator;
+
+  SmallVector<namingContext, 8> nameContext;
+  for (Region &region : op->getRegions())
+    nameContext.push_back(std::make_tuple(&region, nextValueID, nextArgumentID,
+                                          nextConflictID, nullptr));
+
   numberValuesInOp(*op, interfaces);
 
-  for (auto &region : op->getRegions())
-    numberValuesInRegion(region, interfaces);
+  while (!nameContext.empty()) {
+    Region *region;
+    hashTableScopeTy *parentScope;
+    std::tie(region, nextValueID, nextArgumentID, nextConflictID, parentScope) =
+        nameContext.pop_back_val();
+
+    // When we switch from one subtree to another, pop the scopes(needless)
+    // until the parent scope.
+    while (usedNames.getCurScope() != parentScope) {
+      usedNames.getCurScope()->~hashTableScopeTy();
+      assert((usedNames.getCurScope() != nullptr || parentScope == nullptr) &&
+             "top level parentScope must be a nullptr");
+    }
+
+    // Add a scope for the current region.
+    auto *curNamesScope = allocator.Allocate<hashTableScopeTy>();
+    new (curNamesScope) hashTableScopeTy(usedNames);
+
+    numberValuesInRegion(*region, interfaces);
+
+    for (Block &block : *region) {
+      for (Operation &op : block)
+        for (Region &region : op.getRegions())
+          nameContext.push_back(std::make_tuple(&region, nextValueID,
+                                                nextArgumentID, nextConflictID,
+                                                curNamesScope));
+    }
+  }
+
+  // Manually remove all the scopes.
+  while (usedNames.getCurScope() != nullptr)
+    usedNames.getCurScope()->~hashTableScopeTy();
 }
 
 void SSANameState::printValueID(Value value, bool printResultNo,
                                 raw_ostream &stream) const {
   if (!value) {
     stream << "<<NULL>>";
     return;
   }
for (unsigned i = 0, e = namesToUse.size(); i != e; ++i) {
     // Overwrite the name.
     valueNames[nameToReplace] = name.copy(usedNameAllocator);
   }
 }
 
 void SSANameState::numberValuesInRegion(
     Region &region,
     DialectInterfaceCollection<OpAsmDialectInterface> &interfaces) {
-  // Save the current value ids to allow for numbering values in sibling regions
-  // the same.
-  llvm::SaveAndRestore<unsigned> valueIDSaver(nextValueID);
-  llvm::SaveAndRestore<unsigned> argumentIDSaver(nextArgumentID);
-  llvm::SaveAndRestore<unsigned> conflictIDSaver(nextConflictID);
-
-  // Push a new used names scope.
-  llvm::ScopedHashTable<StringRef, char>::ScopeTy usedNamesScope(usedNames);
-
   // Number the values within this region in a breadth-first order.
   unsigned nextBlockID = 0;
   for (auto &block : region) {
     // Each block gets a unique ID, and all of the operations within it get
     // numbered as well.
     blockIDs[&block] = nextBlockID++;
     numberValuesInBlock(block, interfaces);
   }
-
-  // After that we traverse the nested regions.
-  // TODO: Rework this loop to not use recursion.
-  for (auto &block : region) {
-    for (auto &op : block)
-      for (auto &nestedRegion : op.getRegions())
-        numberValuesInRegion(nestedRegion, interfaces);
-  }
 }
 
 void SSANameState::numberValuesInBlock(
     Block &block,
     DialectInterfaceCollection<OpAsmDialectInterface> &interfaces) {
   auto setArgNameFn = [&](Value arg, StringRef name) {
     assert(!valueIDs.count(arg) && "arg numbered multiple times");
     assert(arg.cast<BlockArgument>().getOwner() == &block &&