JIT: Enable reusing profile-aware DFS trees between phases

src/coreclr/jit/compiler.h

@@ -1935,12 +1935,16 @@ class FlowGraphDfsTree
     // Whether the DFS that produced the tree found any backedges.
     bool m_hasCycle;
 
+    // Whether the DFS that produced the tree used edge likelihoods to influence successor visitation order.
+    bool m_profileAware;
+
 public:
-    FlowGraphDfsTree(Compiler* comp, BasicBlock** postOrder, unsigned postOrderCount, bool hasCycle)
+    FlowGraphDfsTree(Compiler* comp, BasicBlock** postOrder, unsigned postOrderCount, bool hasCycle, bool profileAware)
         : m_comp(comp)
         , m_postOrder(postOrder)
         , m_postOrderCount(postOrderCount)
         , m_hasCycle(hasCycle)
+        , m_profileAware(profileAware)
     {
     }
 
@@ -1975,6 +1979,11 @@ class FlowGraphDfsTree
         return m_hasCycle;
     }
 
+    bool IsProfileAware() const
+    {
+        return m_profileAware;
+    }
+
 #ifdef DEBUG
     void Dump() const;
 #endif // DEBUG

src/coreclr/jit/fgdiagnostic.cpp

@@ -4735,15 +4735,30 @@ void Compiler::fgDebugCheckFlowGraphAnnotations()
         return;
     }
 
-    unsigned count = fgRunDfs(
-        [](BasicBlock* block, unsigned preorderNum) {
+    auto visitPreorder = [](BasicBlock* block, unsigned preorderNum) {
         assert(block->bbPreorderNum == preorderNum);
-    },
-        [=](BasicBlock* block, unsigned postorderNum) {
+    };
+
+    auto visitPostorder = [=](BasicBlock* block, unsigned postorderNum) {
         assert(block->bbPostorderNum == postorderNum);
         assert(m_dfsTree->GetPostOrder(postorderNum) == block);
-    },
-        [](BasicBlock* block, BasicBlock* succ) {});
+    };
+
+    auto visitEdge = [](BasicBlock* block, BasicBlock* succ) {};
+
+    unsigned count;
+    if (m_dfsTree->IsProfileAware())
+    {
+        count = fgRunDfs<decltype(visitPreorder), decltype(visitPostorder), decltype(visitEdge), true>(visitPreorder,
+                                                                                                       visitPostorder,
+                                                                                                       visitEdge);
+    }
+    else
+    {
+        count = fgRunDfs<decltype(visitPreorder), decltype(visitPostorder), decltype(visitEdge), false>(visitPreorder,
+                                                                                                        visitPostorder,
+                                                                                                        visitEdge);
+    }
 
     assert(m_dfsTree->GetPostOrderCount() == count);
 

src/coreclr/jit/fgopt.cpp

@@ -4663,10 +4663,6 @@ void Compiler::fgDoReversePostOrderLayout()
     }
 #endif // DEBUG
 
-    // Compute DFS of all blocks in the method, using profile data to determine the order successors are visited in.
-    //
-    m_dfsTree = fgComputeDfs</* useProfile */ true>();
-
     // If LSRA didn't create any new blocks, we can reuse its loop-aware RPO traversal,
     // which is cached in Compiler::fgBBs.
     // If the cache isn't available, we need to recompute the loop-aware RPO.
@@ -4675,15 +4671,21 @@ void Compiler::fgDoReversePostOrderLayout()
 
     if (rpoSequence == nullptr)
     {
-        rpoSequence                        = new (this, CMK_BasicBlock) BasicBlock*[m_dfsTree->GetPostOrderCount()];
+        assert(m_dfsTree == nullptr);
+        m_dfsTree                          = fgComputeDfs</* useProfile */ true>();
         FlowGraphNaturalLoops* const loops = FlowGraphNaturalLoops::Find(m_dfsTree);
-        unsigned                     index = 0;
-        auto                         addToSequence = [rpoSequence, &index](BasicBlock* block) {
+        rpoSequence                        = new (this, CMK_BasicBlock) BasicBlock*[m_dfsTree->GetPostOrderCount()];
+        unsigned index                     = 0;
+        auto     addToSequence             = [rpoSequence, &index](BasicBlock* block) {
             rpoSequence[index++] = block;
         };
 
         fgVisitBlocksInLoopAwareRPO(m_dfsTree, loops, addToSequence);
     }
+    else
+    {
+        assert(m_dfsTree != nullptr);
+    }
 
     // Fast path: We don't have any EH regions, so just reorder the blocks
     //

src/coreclr/jit/flowgraph.cpp

@@ -4353,7 +4353,7 @@ FlowGraphDfsTree* Compiler::fgComputeDfs()
         fgRunDfs<decltype(visitPreorder), decltype(visitPostorder), decltype(visitEdge), useProfile>(visitPreorder,
                                                                                                      visitPostorder,
                                                                                                      visitEdge);
-    return new (this, CMK_DepthFirstSearch) FlowGraphDfsTree(this, postOrder, numBlocks, hasCycle);
+    return new (this, CMK_DepthFirstSearch) FlowGraphDfsTree(this, postOrder, numBlocks, hasCycle, useProfile);
 }
 
 // Add explicit instantiations.

src/coreclr/jit/lsra.cpp

@@ -947,13 +947,16 @@ void LinearScan::setBlockSequence()
     bbVisitedSet = BitVecOps::MakeEmpty(traits);
 
     assert((blockSequence == nullptr) && (bbSeqCount == 0));
-    FlowGraphDfsTree* const dfsTree = compiler->fgComputeDfs</* useProfile */ true>();
+
+    compiler->m_dfsTree             = compiler->fgComputeDfs</* useProfile */ true>();
+    FlowGraphDfsTree* const dfsTree = compiler->m_dfsTree;
     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);
+        // Ensure loop bodies are compact in the visitation order.
+        compiler->m_loops                  = FlowGraphNaturalLoops::Find(dfsTree);
+        FlowGraphNaturalLoops* const loops = compiler->m_loops;
         unsigned                     index = 0;
 
         auto addToSequence = [this, &index](BasicBlock* block) {
@@ -1319,6 +1322,7 @@ PhaseStatus LinearScan::doLinearScan()
     {
         assert(compiler->fgBBcount > bbSeqCount);
         compiler->fgBBs = nullptr;
+        compiler->fgInvalidateDfsTree();
     }
 
     return PhaseStatus::MODIFIED_EVERYTHING;