JIT: Try to reuse profile-aware DFS tree throughout backend

src/coreclr/jit/compiler.cpp

@@ -5150,6 +5150,10 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         // Conditional to Switch conversion
         //
         DoPhase(this, PHASE_SWITCH_RECOGNITION, &Compiler::optSwitchRecognition);
+
+        // Run profile repair
+        //
+        DoPhase(this, PHASE_REPAIR_PROFILE, &Compiler::fgRepairProfile);
     }
 
 #ifdef DEBUG

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
@@ -6617,6 +6626,7 @@ class Compiler
     }
 
     void fgRemoveProfileData(const char* reason);
+    PhaseStatus fgRepairProfile();
     void fgRepairProfileCondToUncond(BasicBlock* block, FlowEdge* retainedEdge, FlowEdge* removedEdge, int* metric = nullptr);
 
 //-------- Insert a statement at the start or end of a basic block --------

src/coreclr/jit/compphases.h

@@ -114,6 +114,7 @@ CompPhaseNameMacro(PHASE_INSERT_GC_POLLS,            "Insert GC Polls",
 CompPhaseNameMacro(PHASE_CREATE_THROW_HELPERS,       "Create throw helper blocks",     false, -1, true)
 CompPhaseNameMacro(PHASE_DETERMINE_FIRST_COLD_BLOCK, "Determine first cold block",     false, -1, true)
 CompPhaseNameMacro(PHASE_RATIONALIZE,                "Rationalize IR",                 false, -1, false)
+CompPhaseNameMacro(PHASE_REPAIR_PROFILE,             "Repair profile",                 false, -1, false)
 
 CompPhaseNameMacro(PHASE_LCLVARLIVENESS,             "Local var liveness",             true,  -1, false)
 CompPhaseNameMacro(PHASE_LCLVARLIVENESS_INIT,        "Local var liveness init",        false, PHASE_LCLVARLIVENESS, false)

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/fgprofile.cpp

@@ -4977,9 +4977,43 @@ bool Compiler::fgDebugCheckOutgoingProfileData(BasicBlock* block, ProfileChecks
 
 #endif // DEBUG
 
+//------------------------------------------------------------------------
+// fgRepairProfile: If we have PGO data and the profile is inconsistent,
+//   run synthesis to re-establish consistency.
+//
+// Returns:
+//   PhaseStatus indicating if profile synthesis ran or not.
+//
+PhaseStatus Compiler::fgRepairProfile()
+{
+    if (fgIsUsingProfileWeights())
+    {
+        if (fgPgoConsistent)
+        {
+            JITDUMP("Profile is already consistent.\n");
+        }
+        else
+        {
+            ProfileSynthesis::Run(this, ProfileSynthesisOption::RetainLikelihoods);
+
+            // Profile synthesis computes a profile-aware DFS tree.
+            // Note that we don't invalidate the tree here, as we opted to retain likelihoods,
+            // so the traversal order shouldn't be different after repairing the profile.
+            // Thus, future phases can reuse this tree.
+            return PhaseStatus::MODIFIED_EVERYTHING;
+        }
+    }
+    else
+    {
+        JITDUMP("No PGO data. Skipping profile repair.\n");
+    }
+
+    return PhaseStatus::MODIFIED_NOTHING;
+}
+
 //------------------------------------------------------------------------
 // fgRepairProfileCondToUncond: attempt to repair profile after modifying
-//   a conditinal branch to an unconditional branch.
+//   a conditional branch to an unconditional branch.
 //
 // Arguments:
 //   block        - block that was just altered

src/coreclr/jit/fgprofilesynthesis.cpp

@@ -30,10 +30,20 @@
 //
 void ProfileSynthesis::Run(ProfileSynthesisOption option)
 {
-    m_dfsTree             = m_comp->fgComputeDfs();
-    m_loops               = FlowGraphNaturalLoops::Find(m_dfsTree);
-    m_improperLoopHeaders = m_loops->ImproperLoopHeaders();
-    m_entryBlock          = m_comp->opts.IsOSR() ? m_comp->fgEntryBB : m_comp->fgFirstBB;
+    if (m_dfsTree == nullptr)
+    {
+        m_dfsTree             = m_comp->fgComputeDfs</* useProfile */ true>();
+        m_loops               = FlowGraphNaturalLoops::Find(m_dfsTree);
+        m_improperLoopHeaders = m_loops->ImproperLoopHeaders();
+    }
+    else
+    {
+        assert(m_loops != nullptr);
+    }
+
+    // Profile synthesis can be run before or after morph, so tolerate (non-)canonical method entries
+    //
+    m_entryBlock = (m_comp->opts.IsOSR() && (m_comp->fgEntryBB != nullptr)) ? m_comp->fgEntryBB : m_comp->fgFirstBB;
 
     // Retain or compute edge likelihood information
     //

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,26 @@ void LinearScan::setBlockSequence()
     bbVisitedSet = BitVecOps::MakeEmpty(traits);
 
     assert((blockSequence == nullptr) && (bbSeqCount == 0));
-    FlowGraphDfsTree* const dfsTree = compiler->fgComputeDfs</* useProfile */ true>();
+
+    if (compiler->m_dfsTree == nullptr)
+    {
+        assert(compiler->m_loops == nullptr);
+        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.
+        // We might need to recompute loops if they weren't retained from previous phases.
+        if (compiler->m_loops == nullptr)
+        {
+            compiler->m_loops = FlowGraphNaturalLoops::Find(dfsTree);
+        }
+
+        FlowGraphNaturalLoops* const loops = compiler->m_loops;
         unsigned                     index = 0;
 
         auto addToSequence = [this, &index](BasicBlock* block) {
@@ -1319,6 +1332,7 @@ PhaseStatus LinearScan::doLinearScan()
     {
         assert(compiler->fgBBcount > bbSeqCount);
         compiler->fgBBs = nullptr;
+        compiler->fgInvalidateDfsTree();
     }
 
     return PhaseStatus::MODIFIED_EVERYTHING;