JIT: Reduce heuristics-derived edge likelihoods into throw blocks found by morph (#116637)

Author: amanasifkhalid

src/coreclr/jit/block.h

@@ -594,6 +594,9 @@ struct FlowEdge
     // Convenience flag for phases that need to track edge visitation
     bool m_visited;
 
+    // Indicates if m_likelihood was determined using profile synthesis's heuristics
+    bool m_heuristicBasedLikelihood;
+
     // True if likelihood has been set
     INDEBUG(bool m_likelihoodSet);
 
@@ -605,6 +608,7 @@ struct FlowEdge
         , m_likelihood(0)
         , m_dupCount(0)
         , m_visited(false)
+        , m_heuristicBasedLikelihood(false)
 #ifdef DEBUG
         , m_likelihoodSet(false)
 #endif // DEBUG
@@ -661,7 +665,8 @@ struct FlowEdge
 
     void clearLikelihood()
     {
-        m_likelihood = 0.0;
+        m_likelihood               = 0.0;
+        m_heuristicBasedLikelihood = false;
         INDEBUG(m_likelihoodSet = false);
     }
 
@@ -706,6 +711,16 @@ struct FlowEdge
         assert(visited());
         m_visited = false;
     }
+
+    bool isHeuristicBased() const
+    {
+        return m_heuristicBasedLikelihood;
+    }
+
+    void setHeuristicBased(bool isHeuristicBased)
+    {
+        m_heuristicBasedLikelihood = isHeuristicBased;
+    }
 };
 
 //------------------------------------------------------------------------

src/coreclr/jit/compiler.cpp

@@ -21,6 +21,7 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 #include "lower.h"
 #include "stacklevelsetter.h"
 #include "patchpointinfo.h"
+#include "fgprofilesynthesis.h"
 #include "jitstd/algorithm.h"
 #include "minipal/time.h"
 
@@ -4592,6 +4593,14 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         //
         DoPhase(this, PHASE_DFS_BLOCKS3, &Compiler::fgDfsBlocksAndRemove);
 
+        auto adjustThrowEdgeLikelihoods = [this]() -> PhaseStatus {
+            return ProfileSynthesis::AdjustThrowEdgeLikelihoods(this);
+        };
+
+        // Adjust heuristic-derived edge likelihoods into paths that are known to throw.
+        //
+        DoPhase(this, PHASE_ADJUST_THROW_LIKELIHOODS, adjustThrowEdgeLikelihoods);
+
         // Discover and classify natural loops (e.g. mark iterative loops as such).
         //
         DoPhase(this, PHASE_FIND_LOOPS, &Compiler::optFindLoopsPhase);

src/coreclr/jit/compphases.h

@@ -72,6 +72,7 @@ CompPhaseNameMacro(PHASE_COMPUTE_DOMINATORS,         "Compute dominators",
 CompPhaseNameMacro(PHASE_CANONICALIZE_ENTRY,         "Canonicalize entry",             false, -1, false)
 CompPhaseNameMacro(PHASE_SET_BLOCK_WEIGHTS,          "Set block weights",              false, -1, false)
 CompPhaseNameMacro(PHASE_ZERO_INITS,                 "Redundant zero Inits",           false, -1, false)
+CompPhaseNameMacro(PHASE_ADJUST_THROW_LIKELIHOODS,   "Adjust throw edge likelihoods",  false, -1, false)
 CompPhaseNameMacro(PHASE_FIND_LOOPS,                 "Find loops",                     false, -1, false)
 CompPhaseNameMacro(PHASE_CLONE_LOOPS,                "Clone loops",                    false, -1, false)
 CompPhaseNameMacro(PHASE_UNROLL_LOOPS,               "Unroll loops",                   false, -1, false)

src/coreclr/jit/fgdiagnostic.cpp

@@ -3106,6 +3106,11 @@ void Compiler::fgDebugCheckBBlist(bool checkBBNum /* = false */, bool checkBBRef
             assert(block->HasInitializedTarget());
         }
 
+        if (block->KindIs(BBJ_COND))
+        {
+            assert(block->GetTrueEdge()->isHeuristicBased() == block->GetFalseEdge()->isHeuristicBased());
+        }
+
         // A branch or fall-through to a BBJ_CALLFINALLY block must come from the `try` region associated
         // with the finally block the BBJ_CALLFINALLY is targeting. There is one special case: if the
         // BBJ_CALLFINALLY is the first block of a `try`, then its predecessor can be outside the `try`:

src/coreclr/jit/fgflow.cpp

@@ -182,6 +182,7 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
             // Copy likelihood from old edge.
             //
             flow->setLikelihood(oldEdge->getLikelihood());
+            flow->setHeuristicBased(oldEdge->isHeuristicBased());
         }
     }
 

src/coreclr/jit/fgopt.cpp

@@ -2247,6 +2247,7 @@ bool Compiler::fgOptimizeUncondBranchToSimpleCond(BasicBlock* block, BasicBlock*
     //
     fgRedirectTargetEdge(block, target->GetTrueTarget());
     block->GetTargetEdge()->setLikelihood(target->GetTrueEdge()->getLikelihood());
+    block->GetTargetEdge()->setHeuristicBased(target->GetTrueEdge()->isHeuristicBased());
 
     FlowEdge* const falseEdge = fgAddRefPred(target->GetFalseTarget(), block, target->GetFalseEdge());
     block->SetCond(block->GetTargetEdge(), falseEdge);
@@ -2741,6 +2742,7 @@ bool Compiler::fgOptimizeBranch(BasicBlock* bJump)
 
     fgRedirectTargetEdge(bJump, falseTarget);
     bJump->GetTargetEdge()->setLikelihood(falseEdge->getLikelihood());
+    bJump->GetTargetEdge()->setHeuristicBased(falseEdge->isHeuristicBased());
 
     FlowEdge* const newTrueEdge = fgAddRefPred(trueTarget, bJump, trueEdge);
 

src/coreclr/jit/fgprofilesynthesis.cpp

@@ -9,6 +9,107 @@
 
 #include "fgprofilesynthesis.h"
 
+//------------------------------------------------------------------------
+// AdjustThrowEdgeLikelihoods: Find throw blocks in the flowgraph, and propagate
+// their throwable state through their predecessors in postorder.
+// Then, adjust heuristic-based likelihoods of edges into paths known to throw.
+//
+// Arguments:
+//   compiler - the Compiler object
+//
+// Returns:
+//   Suitable phase status
+//
+/* static */ PhaseStatus ProfileSynthesis::AdjustThrowEdgeLikelihoods(Compiler* compiler)
+{
+    const FlowGraphDfsTree* dfsTree = compiler->m_dfsTree;
+    assert(dfsTree != nullptr);
+    BitVecTraits traits = dfsTree->PostOrderTraits();
+    BitVec       willThrow(BitVecOps::MakeEmpty(&traits));
+
+    // Adjusts the likelihoods out of a block that conditionally flows into a path that throws
+    auto tweakLikelihoods = [&](BasicBlock* block) {
+        assert(block->KindIs(BBJ_COND));
+        FlowEdge *throwEdge, *normalEdge;
+
+        if (BitVecOps::IsMember(&traits, willThrow, block->GetTrueTarget()->bbPostorderNum))
+        {
+            throwEdge  = block->GetTrueEdge();
+            normalEdge = block->GetFalseEdge();
+        }
+        else
+        {
+            throwEdge  = block->GetFalseEdge();
+            normalEdge = block->GetTrueEdge();
+        }
+
+        throwEdge->setLikelihood(throwLikelihood);
+        normalEdge->setLikelihood(1.0 - throwLikelihood);
+    };
+
+    bool modified = false;
+
+    // Walk the flowgraph in postorder, propagating throw state backwards
+    for (unsigned i = 0; i < dfsTree->GetPostOrderCount(); i++)
+    {
+        BasicBlock* const block = dfsTree->GetPostOrder(i);
+        if (block->KindIs(BBJ_THROW))
+        {
+            JITDUMP(FMT_BB " will throw.\n", block->bbNum);
+            BitVecOps::AddElemD(&traits, willThrow, i);
+        }
+        // Avoid slightly more expensive successor iteration for blocks with one outgoing edge
+        else if ((block->GetUniqueSucc() != nullptr) &&
+                 BitVecOps::IsMember(&traits, willThrow, block->GetUniqueSucc()->bbPostorderNum))
+        {
+            JITDUMP(FMT_BB " flows into a throw block.\n", block->bbNum);
+            BitVecOps::AddElemD(&traits, willThrow, i);
+        }
+        else
+        {
+            bool anyPathThrows = false;
+            bool allPathsThrow = true;
+
+            for (BasicBlock* const succBlock : block->Succs(compiler))
+            {
+                if (BitVecOps::IsMember(&traits, willThrow, succBlock->bbPostorderNum))
+                {
+                    anyPathThrows = true;
+                }
+                else
+                {
+                    allPathsThrow = false;
+                }
+            }
+
+            if (anyPathThrows)
+            {
+                if (allPathsThrow)
+                {
+                    JITDUMP(FMT_BB " flows into a throw block.\n", block->bbNum);
+                    BitVecOps::AddElemD(&traits, willThrow, i);
+                }
+                else if (block->KindIs(BBJ_COND) && block->GetTrueEdge()->isHeuristicBased())
+                {
+                    JITDUMP(FMT_BB " can flow into a throw block.\n", block->bbNum);
+                    assert(block->GetFalseEdge()->isHeuristicBased());
+                    tweakLikelihoods(block);
+                    modified = true;
+                }
+            }
+        }
+    }
+
+    if (modified && compiler->fgIsUsingProfileWeights())
+    {
+        JITDUMP("Modified edge likelihoods. Data %s inconsistent.\n",
+                compiler->fgPgoConsistent ? "is now" : "was already");
+        compiler->fgPgoConsistent = false;
+    }
+
+    return modified ? PhaseStatus::MODIFIED_EVERYTHING : PhaseStatus::MODIFIED_NOTHING;
+}
+
 // TODO
 //
 // * vet against some real data
@@ -252,6 +353,8 @@ void ProfileSynthesis::AssignLikelihoods()
 
             case BBJ_COND:
                 // Two successor cases
+                block->GetTrueEdge()->setHeuristicBased(true);
+                block->GetFalseEdge()->setHeuristicBased(true);
                 AssignLikelihoodCond(block);
                 break;
 

src/coreclr/jit/fgprofilesynthesis.h

@@ -37,6 +37,8 @@ class ProfileSynthesis
         p.Run(option);
     }
 
+    static PhaseStatus AdjustThrowEdgeLikelihoods(Compiler* compiler);
+
     static constexpr weight_t epsilon = 0.001;
 
 private:

src/coreclr/jit/optimizebools.cpp

@@ -824,6 +824,7 @@ bool OptBoolsDsc::optOptimizeRangeTests()
     FlowEdge* const newEdge      = m_comp->fgAddRefPred(inRangeBb, m_b1);
     FlowEdge* const oldFalseEdge = m_b1->GetFalseEdge();
     FlowEdge* const oldTrueEdge  = m_b1->GetTrueEdge();
+    newEdge->setHeuristicBased(oldTrueEdge->isHeuristicBased());
 
     if (!cmp2IsReversed)
     {
@@ -1243,6 +1244,7 @@ void OptBoolsDsc::optOptimizeBoolsUpdateTrees()
             // Modify flow for true side of B1
             //
             m_comp->fgRedirectTrueEdge(m_b1, m_b2->GetTrueTarget());
+            origB1TrueEdge->setHeuristicBased(origB2TrueEdge->isHeuristicBased());
 
             newB1TrueLikelihood =
                 (1.0 - origB1TrueLikelihood) + origB1TrueLikelihood * origB2FalseEdge->getLikelihood();
@@ -1267,6 +1269,7 @@ void OptBoolsDsc::optOptimizeBoolsUpdateTrees()
         // Fix B1 false edge likelihood
         //
         newB1FalseEdge->setLikelihood(1.0 - newB1TrueLikelihood);
+        newB1FalseEdge->setHeuristicBased(origB1TrueEdge->isHeuristicBased());
 
         // Update profile
         if (m_b1->hasProfileWeight())