diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -36,6 +36,9 @@ #include "llvm/Support/Endian.h" #include "llvm/Support/Regex.h" #include "llvm/Support/SaveAndRestore.h" + +#include + using namespace mlir; using namespace mlir::detail; @@ -835,11 +838,61 @@ SSANameState::SSANameState( Operation *op, DialectInterfaceCollection &interfaces) { - llvm::ScopedHashTable::ScopeTy usedNamesScope(usedNames); + llvm::SaveAndRestore valueIDSaver(nextValueID); + llvm::SaveAndRestore argumentIDSaver(nextArgumentID); + llvm::SaveAndRestore conflictIDSaver(nextConflictID); + + // The context includes nextValueID, nextArgumentID, nextConflictID and scoped + // HashTable. + using hashTableScopeTy = llvm::ScopedHashTable::ScopeTy; + // A namingContext carries the information inherits from parent region. + using namingContext = + std::tuple; + // Allocator for hashTableScopeTy + llvm::BumpPtrAllocator allocator; + + SmallVector regions; + SmallVector nameContext; + for (auto ®ion : op->getRegions()) { + regions.push_back(®ion); + nameContext.push_back( + std::make_tuple(nextValueID, nextArgumentID, nextConflictID, nullptr)); + } + numberValuesInOp(*op, interfaces); - for (auto ®ion : op->getRegions()) - numberValuesInRegion(region, interfaces); + while (!regions.empty()) { + Region *region = regions.pop_back_val(); + hashTableScopeTy *parentScope; + std::tie(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(); + new (curNamesScope) hashTableScopeTy(usedNames); + + numberValuesInRegion(*region, interfaces); + + for (auto &block : *region) + for (auto &op : block) + for (auto ®ion : op.getRegions()) { + regions.push_back(®ion); + nameContext.push_back(std::make_tuple(nextValueID, nextArgumentID, + nextConflictID, curNamesScope)); + } + } + + // Manually remove all the scopes. + while (usedNames.getCurScope() != nullptr) + usedNames.getCurScope()->~hashTableScopeTy(); } void SSANameState::printValueID(Value value, bool printResultNo, @@ -918,15 +971,6 @@ void SSANameState::numberValuesInRegion( Region ®ion, DialectInterfaceCollection &interfaces) { - // Save the current value ids to allow for numbering values in sibling regions - // the same. - llvm::SaveAndRestore valueIDSaver(nextValueID); - llvm::SaveAndRestore argumentIDSaver(nextArgumentID); - llvm::SaveAndRestore conflictIDSaver(nextConflictID); - - // Push a new used names scope. - llvm::ScopedHashTable::ScopeTy usedNamesScope(usedNames); - // Number the values within this region in a breadth-first order. unsigned nextBlockID = 0; for (auto &block : region) { @@ -935,14 +979,6 @@ 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(