JIT: Propagate flow into finally regions correctly in synthesis (#114016)

Part of #107749. The first profile synthesis run happens before importation, so we don't model flow in and out of finally regions with flow edges yet. As a workaround, synthesis gives each finally region the same weight as its corresponding try region. When call-finally pairs are created, the tail inherits the weight of its head under the (faulty) assumption that all flow into a call-finally will return to the same pair. Once we have flow edges, we can compute the flow out of finally regions the same way as we compute flow elsewhere. It's important that synthesis models flow through finally regions via flow edges once we have them, or else flow through a loop that executes a finally might be lost, messing up the cyclic probability computation and flattening the loop's weight.

Author: amanasifkhalid

src/coreclr/jit/fgbasic.cpp

@@ -372,6 +372,19 @@ void Compiler::fgRemoveEhfSuccessor(BasicBlock* block, const unsigned succIndex)
                   (succCount - succIndex - 1) * sizeof(FlowEdge*));
     }
 
+    // Recompute the likelihoods of the block's other successor edges.
+    const weight_t removedLikelihood = succEdge->getLikelihood();
+    const unsigned newSuccCount      = succCount - 1;
+
+    for (unsigned i = 0; i < newSuccCount; i++)
+    {
+        // If we removed all of the flow out of 'block', distribute flow among the remaining edges evenly.
+        const weight_t currLikelihood = succTab[i]->getLikelihood();
+        const weight_t newLikelihood =
+            (removedLikelihood == 1.0) ? (1.0 / newSuccCount) : (currLikelihood / (1.0 - removedLikelihood));
+        succTab[i]->setLikelihood(min(1.0, newLikelihood));
+    }
+
 #ifdef DEBUG
     // We only expect to see a successor once in the table.
     for (unsigned i = succIndex; i < (succCount - 1); i++)
@@ -431,6 +444,19 @@ void Compiler::fgRemoveEhfSuccessor(FlowEdge* succEdge)
         }
     }
 
+    // Recompute the likelihoods of the block's other successor edges.
+    const weight_t removedLikelihood = succEdge->getLikelihood();
+    const unsigned newSuccCount      = succCount - 1;
+
+    for (unsigned i = 0; i < newSuccCount; i++)
+    {
+        // If we removed all of the flow out of 'block', distribute flow among the remaining edges evenly.
+        const weight_t currLikelihood = succTab[i]->getLikelihood();
+        const weight_t newLikelihood =
+            (removedLikelihood == 1.0) ? (1.0 / newSuccCount) : (currLikelihood / (1.0 - removedLikelihood));
+        succTab[i]->setLikelihood(min(1.0, newLikelihood));
+    }
+
     assert(found);
     ehfDesc->bbeCount--;
 }

src/coreclr/jit/fgehopt.cpp

@@ -1644,27 +1644,28 @@ PhaseStatus Compiler::fgCloneFinally()
 
             for (BasicBlock* const block : Blocks(firstBlock, lastBlock))
             {
-                if (block->hasProfileWeight())
-                {
-                    weight_t const blockWeight = block->bbWeight;
-                    block->setBBProfileWeight(blockWeight * originalScale);
-                    JITDUMP("Set weight of " FMT_BB " to " FMT_WT "\n", block->bbNum, block->bbWeight);
+                weight_t const blockWeight = block->bbWeight;
+                block->setBBProfileWeight(blockWeight * originalScale);
+                JITDUMP("Set weight of " FMT_BB " to " FMT_WT "\n", block->bbNum, block->bbWeight);
 
-                    BasicBlock* const clonedBlock = blockMap[block];
-                    clonedBlock->setBBProfileWeight(blockWeight * clonedScale);
-                    JITDUMP("Set weight of " FMT_BB " to " FMT_WT "\n", clonedBlock->bbNum, clonedBlock->bbWeight);
-                }
+                BasicBlock* const clonedBlock = blockMap[block];
+                clonedBlock->setBBProfileWeight(blockWeight * clonedScale);
+                JITDUMP("Set weight of " FMT_BB " to " FMT_WT "\n", clonedBlock->bbNum, clonedBlock->bbWeight);
+            }
+
+            if (!retargetedAllCalls)
+            {
+                JITDUMP(
+                    "Reduced flow out of EH%u needs to be propagated to continuation block(s). Data %s inconsistent.\n",
+                    XTnum, fgPgoConsistent ? "is now" : "was already");
+                fgPgoConsistent = false;
             }
         }
 
         // Update flow into normalCallFinallyReturn
         if (normalCallFinallyReturn->hasProfileWeight())
         {
-            normalCallFinallyReturn->bbWeight = BB_ZERO_WEIGHT;
-            for (FlowEdge* const predEdge : normalCallFinallyReturn->PredEdges())
-            {
-                normalCallFinallyReturn->increaseBBProfileWeight(predEdge->getLikelyWeight());
-            }
+            normalCallFinallyReturn->setBBProfileWeight(normalCallFinallyReturn->computeIncomingWeight());
         }
 
         // Done!
@@ -2185,33 +2186,13 @@ bool Compiler::fgRetargetBranchesToCanonicalCallFinally(BasicBlock*      block,
     //
     if (block->hasProfileWeight())
     {
-        // Add weight to the canonical call finally pair.
+        // Add weight to the canonical call-finally.
         //
-        weight_t const canonicalWeight =
-            canonicalCallFinally->hasProfileWeight() ? canonicalCallFinally->bbWeight : BB_ZERO_WEIGHT;
-        weight_t const newCanonicalWeight = block->bbWeight + canonicalWeight;
-
-        canonicalCallFinally->setBBProfileWeight(newCanonicalWeight);
+        canonicalCallFinally->increaseBBProfileWeight(block->bbWeight);
 
-        BasicBlock* const canonicalLeaveBlock = canonicalCallFinally->Next();
-
-        weight_t const canonicalLeaveWeight =
-            canonicalLeaveBlock->hasProfileWeight() ? canonicalLeaveBlock->bbWeight : BB_ZERO_WEIGHT;
-        weight_t const newLeaveWeight = block->bbWeight + canonicalLeaveWeight;
-
-        canonicalLeaveBlock->setBBProfileWeight(newLeaveWeight);
-
-        // Remove weight from the old call finally pair.
+        // Remove weight from the old call-finally.
         //
-        if (callFinally->hasProfileWeight())
-        {
-            callFinally->decreaseBBProfileWeight(block->bbWeight);
-        }
-
-        if (leaveBlock->hasProfileWeight())
-        {
-            leaveBlock->decreaseBBProfileWeight(block->bbWeight);
-        }
+        callFinally->decreaseBBProfileWeight(block->bbWeight);
     }
 
     return true;

src/coreclr/jit/fgprofile.cpp

@@ -4787,20 +4787,12 @@ bool Compiler::fgDebugCheckIncomingProfileData(BasicBlock* block, ProfileChecks
     weight_t       incomingLikelyWeight = 0;
     unsigned       missingLikelyWeight  = 0;
     bool           foundPreds           = false;
-    bool           foundEHPreds         = false;
 
     for (FlowEdge* const predEdge : block->PredEdges())
     {
         if (predEdge->hasLikelihood())
         {
-            if (BasicBlock::sameHndRegion(block, predEdge->getSourceBlock()))
-            {
-                incomingLikelyWeight += predEdge->getLikelyWeight();
-            }
-            else
-            {
-                foundEHPreds = true;
-            }
+            incomingLikelyWeight += predEdge->getLikelyWeight();
         }
         else
         {
@@ -4812,29 +4804,11 @@ bool Compiler::fgDebugCheckIncomingProfileData(BasicBlock* block, ProfileChecks
         foundPreds = true;
     }
 
-    // We almost certainly won't get the likelihoods on a BBJ_EHFINALLYRET right,
-    // so special-case BBJ_CALLFINALLYRET incoming flow.
-    //
-    if (block->isBBCallFinallyPairTail())
-    {
-        incomingLikelyWeight = block->Prev()->bbWeight;
-        foundEHPreds         = false;
-    }
-
-    // We almost certainly won't get the likelihoods on a BBJ_EHFINALLYRET right,
-    // so special-case BBJ_CALLFINALLYRET incoming flow.
-    //
-    if (block->isBBCallFinallyPairTail())
-    {
-        incomingLikelyWeight = block->Prev()->bbWeight;
-        foundEHPreds         = false;
-    }
-
     bool likelyWeightsValid = true;
 
     // If we have EH preds we may not have consistent incoming flow.
     //
-    if (foundPreds && !foundEHPreds)
+    if (foundPreds)
     {
         if (verifyLikelyWeights)
         {
@@ -4890,7 +4864,7 @@ bool Compiler::fgDebugCheckOutgoingProfileData(BasicBlock* block, ProfileChecks
     //
     const unsigned numSuccs = block->NumSucc(this);
 
-    if ((numSuccs > 0) && !block->KindIs(BBJ_EHFINALLYRET, BBJ_EHFAULTRET, BBJ_EHFILTERRET))
+    if ((numSuccs > 0) && !block->KindIs(BBJ_EHFAULTRET, BBJ_EHFILTERRET))
     {
         weight_t const blockWeight        = block->bbWeight;
         weight_t       outgoingLikelihood = 0;

src/coreclr/jit/fgprofilesynthesis.cpp

@@ -41,6 +41,11 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
         assert(m_loops != nullptr);
     }
 
+    if (m_loops->NumLoops() > 0)
+    {
+        m_cyclicProbabilities = new (m_comp, CMK_Pgo) weight_t[m_loops->NumLoops()];
+    }
+
     // 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;
@@ -168,7 +173,10 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
     if (m_approximate)
     {
         JITDUMP("Profile is inconsistent. Bypassing post-phase consistency checks.\n");
-        m_comp->Metrics.ProfileInconsistentInitially++;
+        if (!m_comp->fgImportDone)
+        {
+            m_comp->Metrics.ProfileInconsistentInitially++;
+        }
     }
 
     // Derive the method's call count from the entry block's weight
@@ -208,7 +216,7 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
         // Leave a note so we can bypass the post-phase check, and
         // instead assert at the end of fgImport, if we make it that far.
         //
-        if (!isConsistent)
+        if (!isConsistent && !m_comp->fgImportDone)
         {
             m_comp->fgPgoDeferredInconsistency = true;
             JITDUMP("Will defer asserting until after importation\n");
@@ -745,13 +753,6 @@ void ProfileSynthesis::RandomizeLikelihoods()
 //
 void ProfileSynthesis::ComputeCyclicProbabilities()
 {
-    m_cyclicProbabilities = nullptr;
-    if (m_loops->NumLoops() == 0)
-    {
-        return;
-    }
-
-    m_cyclicProbabilities = new (m_comp, CMK_Pgo) weight_t[m_loops->NumLoops()];
     // Walk loops in post order to visit inner loops before outer loops.
     for (FlowGraphNaturalLoop* loop : m_loops->InPostOrder())
     {
@@ -828,10 +829,7 @@ void ProfileSynthesis::ComputeCyclicProbabilities(FlowGraphNaturalLoop* loop)
 
                 for (FlowEdge* const edge : nestedLoop->EntryEdges())
                 {
-                    if (BasicBlock::sameHndRegion(block, edge->getSourceBlock()))
-                    {
-                        newWeight += edge->getLikelyWeight();
-                    }
+                    newWeight += edge->getLikelyWeight();
                 }
 
                 newWeight *= m_cyclicProbabilities[nestedLoop->GetIndex()];
@@ -847,10 +845,10 @@ void ProfileSynthesis::ComputeCyclicProbabilities(FlowGraphNaturalLoop* loop)
                 {
                     BasicBlock* const sourceBlock = edge->getSourceBlock();
 
-                    // Ignore flow across EH, or from preds not in the loop.
-                    // Latter can happen if there are unreachable blocks that flow into the loop.
+                    // Ignore flow from preds not in the loop.
+                    // This can happen if there are unreachable blocks that flow into the loop.
                     //
-                    if (BasicBlock::sameHndRegion(block, sourceBlock) && loop->ContainsBlock(sourceBlock))
+                    if (loop->ContainsBlock(sourceBlock))
                     {
                         newWeight += edge->getLikelyWeight();
                     }
@@ -1214,7 +1212,8 @@ void ProfileSynthesis::GaussSeidelSolver()
     const FlowGraphDfsTree* const dfs                  = m_loops->GetDfsTree();
     unsigned const                blockCount           = dfs->GetPostOrderCount();
     bool                          checkEntryExitWeight = true;
-    bool                          showDetails          = false;
+    bool const                    showDetails          = false;
+    bool const                    callFinalliesCreated = m_comp->fgImportDone;
 
     JITDUMP("Synthesis solver: flow graph has %u improper loop headers\n", m_improperLoopHeaders);
 
@@ -1290,9 +1289,10 @@ void ProfileSynthesis::GaussSeidelSolver()
                     {
                         newWeight = block->bbWeight;
 
-                        // Finallies also add in the weight of their try.
+                        // If we haven't added flow edges into/out of finallies yet,
+                        // add in the weight of their corresponding try regions.
                         //
-                        if (ehDsc->HasFinallyHandler())
+                        if (!callFinalliesCreated && ehDsc->HasFinallyHandler())
                         {
                             newWeight += countVector[ehDsc->ebdTryBeg->bbNum];
                         }
@@ -1318,11 +1318,7 @@ void ProfileSynthesis::GaussSeidelSolver()
                     for (FlowEdge* const edge : loop->EntryEdges())
                     {
                         BasicBlock* const predBlock = edge->getSourceBlock();
-
-                        if (BasicBlock::sameHndRegion(block, predBlock))
-                        {
-                            newWeight += edge->getLikelihood() * countVector[predBlock->bbNum];
-                        }
+                        newWeight += edge->getLikelihood() * countVector[predBlock->bbNum];
                     }
 
                     // Scale by cyclic probability
@@ -1354,10 +1350,7 @@ void ProfileSynthesis::GaussSeidelSolver()
                             continue;
                         }
 
-                        if (BasicBlock::sameHndRegion(block, predBlock))
-                        {
-                            newWeight += edge->getLikelihood() * countVector[predBlock->bbNum];
-                        }
+                        newWeight += edge->getLikelihood() * countVector[predBlock->bbNum];
                     }
 
                     if (selfEdge != nullptr)
@@ -1444,10 +1437,12 @@ void ProfileSynthesis::GaussSeidelSolver()
             }
         }
 
-        // If there were no improper headers, we will have converged in one pass.
+        // If there were no improper headers, we will have converged in one pass
         // (profile may still be inconsistent, if there were capped cyclic probabilities).
+        // After the importer runs, we require that synthesis achieves profile consistency
+        // unless the resultant profile is approximate, so don't skip the below checks.
         //
-        if (m_improperLoopHeaders == 0)
+        if ((m_improperLoopHeaders == 0) && !m_comp->fgImportDone)
         {
             converged = true;
             break;

src/coreclr/jit/importer.cpp

@@ -12576,8 +12576,9 @@ void Compiler::impImport()
 //
 void Compiler::impFixPredLists()
 {
-    unsigned XTnum = 0;
-    bool     added = false;
+    unsigned   XTnum               = 0;
+    bool       added               = false;
+    const bool usingProfileWeights = fgIsUsingProfileWeights();
 
     for (EHblkDsc* HBtab = compHndBBtab; XTnum < compHndBBtabCount; XTnum++, HBtab++)
     {
@@ -12586,6 +12587,7 @@ void Compiler::impFixPredLists()
             BasicBlock* const finallyBegBlock  = HBtab->ebdHndBeg;
             BasicBlock* const finallyLastBlock = HBtab->ebdHndLast;
             unsigned          predCount        = (unsigned)-1;
+            const weight_t    finallyWeight    = finallyBegBlock->bbWeight;
 
             for (BasicBlock* const finallyBlock : BasicBlockRangeList(finallyBegBlock, finallyLastBlock))
             {
@@ -12629,7 +12631,8 @@ void Compiler::impFixPredLists()
                     jumpEhf->bbeCount = predCount;
                     jumpEhf->bbeSuccs = new (this, CMK_FlowEdge) FlowEdge*[predCount];
 
-                    unsigned predNum = 0;
+                    unsigned predNum             = 0;
+                    weight_t remainingLikelihood = 1.0;
                     for (BasicBlock* const predBlock : finallyBegBlock->PredBlocks())
                     {
                         // We only care about preds that are callfinallies.
@@ -12643,6 +12646,22 @@ void Compiler::impFixPredLists()
                         FlowEdge* const   newEdge      = fgAddRefPred(continuation, finallyBlock);
                         newEdge->setLikelihood(1.0 / predCount);
 
+                        if (usingProfileWeights && (finallyWeight != BB_ZERO_WEIGHT))
+                        {
+                            // Derive edge likelihood from the entry block's weight relative to other entries.
+                            //
+                            const weight_t callFinallyWeight = predBlock->bbWeight;
+                            const weight_t likelihood        = min(callFinallyWeight / finallyWeight, 1.0);
+                            newEdge->setLikelihood(min(likelihood, remainingLikelihood));
+                            remainingLikelihood = max(BB_ZERO_WEIGHT, remainingLikelihood - likelihood);
+                        }
+                        else
+                        {
+                            // If we don't have profile data, evenly distribute the likelihoods.
+                            //
+                            newEdge->setLikelihood(1.0 / predCount);
+                        }
+
                         jumpEhf->bbeSuccs[predNum] = newEdge;
                         ++predNum;
 
@@ -12657,6 +12676,26 @@ void Compiler::impFixPredLists()
 
                 finallyBlock->SetEhfTargets(jumpEhf);
             }
+
+            if (usingProfileWeights)
+            {
+                // Compute new flow into the finally region's continuation successors.
+                //
+                bool profileConsistent = true;
+                for (BasicBlock* const callFinally : finallyBegBlock->PredBlocks())
+                {
+                    BasicBlock* const callFinallyRet = callFinally->Next();
+                    callFinallyRet->setBBProfileWeight(callFinallyRet->computeIncomingWeight());
+                    profileConsistent &= fgProfileWeightsEqual(callFinally->bbWeight, callFinallyRet->bbWeight);
+                }
+
+                if (!profileConsistent)
+                {
+                    JITDUMP("Flow into finally handler EH%u does not match outgoing flow. Data %s inconsistent.\n",
+                            XTnum, fgPgoConsistent ? "is now" : "was already");
+                    fgPgoConsistent = false;
+                }
+            }
         }
     }
 }