JIT: Track visited blocks in LSRA with post-order number-based BitVec

src/coreclr/jit/lsra.cpp

@@ -911,8 +911,7 @@ LinearScan::LinearScan(Compiler* theCompiler)
 
     // Block sequencing (the order in which we schedule).
     // Note that we don't initialize the bbVisitedSet until we do the first traversal
-    // This is so that any blocks that are added during the first traversal
-    // are accounted for (and we don't have BasicBlockEpoch issues).
+    // This is so that any blocks that are added during the first traversal are accounted for.
     blockSequencingDone = false;
     blockSequence       = nullptr;
     curBBSeqNum         = 0;
@@ -943,47 +942,46 @@ void LinearScan::setBlockSequence()
 {
     assert(!blockSequencingDone); // The method should be called only once.
 
-    compiler->EnsureBasicBlockEpoch();
-#ifdef DEBUG
-    blockEpoch = compiler->GetCurBasicBlockEpoch();
-#endif // DEBUG
-
     // Initialize the "visited" blocks set.
-    bbVisitedSet = BlockSetOps::MakeEmpty(compiler);
+    traits       = new (compiler, CMK_LSRA) BitVecTraits(compiler->fgBBcount, compiler);
+    bbVisitedSet = BitVecOps::MakeEmpty(traits);
 
     assert((blockSequence == nullptr) && (bbSeqCount == 0));
     FlowGraphDfsTree* const dfsTree = compiler->fgComputeDfs</* useProfile */ true>();
+    blockSequence                   = new (compiler, CMK_LSRA) BasicBlock*[compiler->fgBBcount];
 
     if (compiler->opts.OptimizationEnabled() && dfsTree->HasCycle())
     {
         // Ensure loop bodies are compact in the visitation order
         FlowGraphNaturalLoops* const loops = FlowGraphNaturalLoops::Find(dfsTree);
-        blockSequence                      = new (compiler, CMK_LSRA) BasicBlock*[compiler->fgBBcount];
-        unsigned index                     = dfsTree->GetPostOrderCount();
+        unsigned                     index = dfsTree->GetPostOrderCount();
 
         auto addToSequence = [this, &index](BasicBlock* block) {
             assert(index != 0);
             blockSequence[--index] = block;
         };
 
         compiler->fgVisitBlocksInLoopAwareRPO(dfsTree, loops, addToSequence);
+
+        // Flip the DFS traversal to get the reverse post-order traversal
+        // (this is the order in which blocks will be allocated)
+        for (unsigned left = 0, right = dfsTree->GetPostOrderCount() - 1; left < right; left++, right--)
+        {
+            std::swap(blockSequence[left], blockSequence[right]);
+        }
     }
     else
     {
         // TODO: Just use lexical block order in MinOpts
-        blockSequence = dfsTree->GetPostOrder();
+        for (unsigned i = 0; i < dfsTree->GetPostOrderCount(); i++)
+        {
+            blockSequence[i] = dfsTree->GetPostOrder(dfsTree->GetPostOrderCount() - i - 1);
+        }
     }
 
     bbNumMaxBeforeResolution = compiler->fgBBNumMax;
     blockInfo                = new (compiler, CMK_LSRA) LsraBlockInfo[bbNumMaxBeforeResolution + 1];
 
-    // Flip the DFS traversal to get the reverse post-order traversal
-    // (this is the order in which blocks will be allocated)
-    for (unsigned left = 0, right = dfsTree->GetPostOrderCount() - 1; left < right; left++, right--)
-    {
-        std::swap(blockSequence[left], blockSequence[right]);
-    }
-
     hasCriticalEdges = false;
     // We use a bbNum of 0 for entry RefPositions.
     // The other information in blockInfo[0] will never be used.
@@ -1094,8 +1092,10 @@ void LinearScan::setBlockSequence()
         for (BasicBlock* block = compiler->fgLastBB; i < compiler->fgBBcount; block = block->Prev())
         {
             assert(block != nullptr);
-            if (!isBlockVisited(block))
+            if (!dfsTree->Contains(block))
             {
+                // Give this block a unique post-order number that can be used as a key into bbVisitedSet
+                block->bbPostorderNum = i;
                 visitBlock(block);
                 blockSequence[i++] = block;
             }
@@ -1297,7 +1297,6 @@ PhaseStatus LinearScan::doLinearScan()
     compiler->EndPhase(PHASE_LINEAR_SCAN_RESOLVE);
 
     assert(blockSequencingDone); // Should do at least one traversal.
-    assert(blockEpoch == compiler->GetCurBasicBlockEpoch());
 
 #if TRACK_LSRA_STATS
     if ((JitConfig.DisplayLsraStats() == 1)

src/coreclr/jit/lsra.h

@@ -1611,18 +1611,19 @@ class LinearScan : public LinearScanInterface
     Interval** localVarIntervals;
 
     // Set of blocks that have been visited.
-    BlockSet bbVisitedSet;
-    void     markBlockVisited(BasicBlock* block)
+    BitVecTraits* traits;
+    BitVec        bbVisitedSet;
+    void          markBlockVisited(BasicBlock* block)
     {
-        BlockSetOps::AddElemD(compiler, bbVisitedSet, block->bbNum);
+        BitVecOps::AddElemD(traits, bbVisitedSet, block->bbPostorderNum);
     }
     void clearVisitedBlocks()
     {
-        BlockSetOps::ClearD(compiler, bbVisitedSet);
+        BitVecOps::ClearD(traits, bbVisitedSet);
     }
     bool isBlockVisited(BasicBlock* block)
     {
-        return BlockSetOps::IsMember(compiler, bbVisitedSet, block->bbNum);
+        return BitVecOps::IsMember(traits, bbVisitedSet, block->bbPostorderNum);
     }
 
 #if DOUBLE_ALIGN
@@ -1639,10 +1640,6 @@ class LinearScan : public LinearScanInterface
     BasicBlock** blockSequence;
     void         setBlockSequence();
     bool         blockSequencingDone;
-#ifdef DEBUG
-    // LSRA must not change number of blocks and blockEpoch that it initializes at start.
-    unsigned blockEpoch;
-#endif // DEBUG
 
     // Indicates whether the allocation pass has been completed.
     bool allocationPassComplete;