[CodeGen] Use BasicBlock numbers to map to MBBs (#101883)

Now that basic blocks have numbers, we can replace the BB-to-MBB maps
and the visited set during ISel with vectors for faster lookup.

Author: aengelke

llvm/include/llvm/CodeGen/FunctionLoweringInfo.h

@@ -73,8 +73,8 @@ class FunctionLoweringInfo {
   /// allocated to hold a pointer to the hidden sret parameter.
   Register DemoteRegister;
 
-  /// MBBMap - A mapping from LLVM basic blocks to their machine code entry.
-  DenseMap<const BasicBlock*, MachineBasicBlock *> MBBMap;
+  /// A mapping from LLVM basic block number to their machine block.
+  SmallVector<MachineBasicBlock *> MBBMap;
 
   /// ValueMap - Since we emit code for the function a basic block at a time,
   /// we must remember which virtual registers hold the values for
@@ -172,9 +172,9 @@ class FunctionLoweringInfo {
   /// for a value.
   DenseMap<const Value *, ISD::NodeType> PreferredExtendType;
 
-  /// VisitedBBs - The set of basic blocks visited thus far by instruction
-  /// selection.
-  SmallPtrSet<const BasicBlock*, 4> VisitedBBs;
+  /// The set of basic blocks visited thus far by instruction selection. Indexed
+  /// by basic block number.
+  SmallVector<bool> VisitedBBs;
 
   /// PHINodesToUpdate - A list of phi instructions whose operand list will
   /// be updated after processing the current basic block.
@@ -213,7 +213,8 @@ class FunctionLoweringInfo {
   }
 
   MachineBasicBlock *getMBB(const BasicBlock *BB) const {
-    return MBBMap.lookup(BB);
+    assert(BB->getNumber() < MBBMap.size() && "uninitialized MBBMap?");
+    return MBBMap[BB->getNumber()];
   }
 
   Register CreateReg(MVT VT, bool isDivergent = false);

llvm/include/llvm/CodeGen/GlobalISel/IRTranslator.h

@@ -146,11 +146,6 @@ class IRTranslator : public MachineFunctionPass {
   /// virtual registers and offsets.
   ValueToVRegInfo VMap;
 
-  // N.b. it's not completely obvious that this will be sufficient for every
-  // LLVM IR construct (with "invoke" being the obvious candidate to mess up our
-  // lives.
-  DenseMap<const BasicBlock *, MachineBasicBlock *> BBToMBB;
-
   // One BasicBlock can be translated to multiple MachineBasicBlocks.  For such
   // BasicBlocks translated to multiple MachineBasicBlocks, MachinePreds retains
   // a mapping between the edges arriving at the BasicBlock to the corresponding

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

@@ -287,7 +287,7 @@ Align IRTranslator::getMemOpAlign(const Instruction &I) {
 }
 
 MachineBasicBlock &IRTranslator::getMBB(const BasicBlock &BB) {
-  MachineBasicBlock *&MBB = BBToMBB[&BB];
+  MachineBasicBlock *MBB = FuncInfo.getMBB(&BB);
   assert(MBB && "BasicBlock was not encountered before");
   return *MBB;
 }
@@ -3907,8 +3907,9 @@ bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
   bool HasMustTailInVarArgFn = false;
 
   // Create all blocks, in IR order, to preserve the layout.
+  FuncInfo.MBBMap.resize(F.getMaxBlockNumber());
   for (const BasicBlock &BB: F) {
-    auto *&MBB = BBToMBB[&BB];
+    auto *&MBB = FuncInfo.MBBMap[BB.getNumber()];
 
     MBB = MF->CreateMachineBasicBlock(&BB);
     MF->push_back(MBB);

llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp

@@ -236,6 +236,7 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
   // Create an initial MachineBasicBlock for each LLVM BasicBlock in F.  This
   // also creates the initial PHI MachineInstrs, though none of the input
   // operands are populated.
+  MBBMap.resize(Fn->getMaxBlockNumber());
   for (const BasicBlock &BB : *Fn) {
     // Don't create MachineBasicBlocks for imaginary EH pad blocks. These blocks
     // are really data, and no instructions can live here.
@@ -261,7 +262,7 @@ void FunctionLoweringInfo::set(const Function &fn, MachineFunction &mf,
     }
 
     MachineBasicBlock *MBB = mf.CreateMachineBasicBlock(&BB);
-    MBBMap[&BB] = MBB;
+    MBBMap[BB.getNumber()] = MBB;
     MF->push_back(MBB);
 
     // Transfer the address-taken flag. This is necessary because there could

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

@@ -1643,11 +1643,12 @@ void SelectionDAGISel::SelectAllBasicBlocks(const Function &Fn) {
   }
 
   // Iterate over all basic blocks in the function.
+  FuncInfo->VisitedBBs.assign(Fn.getMaxBlockNumber(), false);
   for (const BasicBlock *LLVMBB : RPOT) {
     if (OptLevel != CodeGenOptLevel::None) {
       bool AllPredsVisited = true;
       for (const BasicBlock *Pred : predecessors(LLVMBB)) {
-        if (!FuncInfo->VisitedBBs.count(Pred)) {
+        if (!FuncInfo->VisitedBBs[Pred->getNumber()]) {
           AllPredsVisited = false;
           break;
         }
@@ -1661,7 +1662,7 @@ void SelectionDAGISel::SelectAllBasicBlocks(const Function &Fn) {
           FuncInfo->InvalidatePHILiveOutRegInfo(&PN);
       }
 
-      FuncInfo->VisitedBBs.insert(LLVMBB);
+      FuncInfo->VisitedBBs[LLVMBB->getNumber()] = true;
     }
 
     BasicBlock::const_iterator const Begin =