JIT: Set naive likelihoods when initializing preds (#98192)

Part of #93020. Per discussion on #98054, set "naive" edge likelihoods (in other words, assume every successor edge is equally likely to be taken) in fgAddRefPred when initializing preds.

Author: amanasifkhalid

src/coreclr/jit/fgbasic.cpp

@@ -2904,26 +2904,26 @@ void Compiler::fgLinkBasicBlocks()
     //
     fgFirstBB->bbRefs = 1;
 
-    // Special args to fgAddRefPred so it will use the initialization fast path.
+    // Special arg to fgAddRefPred so it will use the initialization fast path.
     //
-    FlowEdge* const oldEdge           = nullptr;
-    bool const      initializingPreds = true;
+    const bool initializingPreds = true;
 
     for (BasicBlock* const curBBdesc : Blocks())
     {
         switch (curBBdesc->GetKind())
         {
             case BBJ_COND:
             {
-                BasicBlock* const trueTarget = fgLookupBB(curBBdesc->GetTargetOffs());
+                BasicBlock* const trueTarget  = fgLookupBB(curBBdesc->GetTargetOffs());
+                BasicBlock* const falseTarget = curBBdesc->Next();
                 curBBdesc->SetTrueTarget(trueTarget);
-                curBBdesc->SetFalseTarget(curBBdesc->Next());
-                fgAddRefPred<initializingPreds>(trueTarget, curBBdesc, oldEdge);
-                fgAddRefPred<initializingPreds>(curBBdesc->GetFalseTarget(), curBBdesc, oldEdge);
+                curBBdesc->SetFalseTarget(falseTarget);
+                fgAddRefPred<initializingPreds>(trueTarget, curBBdesc);
+                fgAddRefPred<initializingPreds>(falseTarget, curBBdesc);
 
-                if (curBBdesc->GetTrueTarget()->bbNum <= curBBdesc->bbNum)
+                if (trueTarget->bbNum <= curBBdesc->bbNum)
                 {
-                    fgMarkBackwardJump(curBBdesc->GetTrueTarget(), curBBdesc);
+                    fgMarkBackwardJump(trueTarget, curBBdesc);
                 }
 
                 if (curBBdesc->IsLast())
@@ -2945,7 +2945,7 @@ void Compiler::fgLinkBasicBlocks()
                 // Redundantly use SetKindAndTarget() instead of SetTarget() just this once,
                 // so we don't break the HasInitializedTarget() invariant of SetTarget().
                 curBBdesc->SetKindAndTarget(curBBdesc->GetKind(), jumpDest);
-                fgAddRefPred<initializingPreds>(jumpDest, curBBdesc, oldEdge);
+                fgAddRefPred<initializingPreds>(jumpDest, curBBdesc);
 
                 if (curBBdesc->GetTarget()->bbNum <= curBBdesc->bbNum)
                 {
@@ -2978,7 +2978,7 @@ void Compiler::fgLinkBasicBlocks()
                 {
                     BasicBlock* jumpDest = fgLookupBB((unsigned)*(size_t*)jumpPtr);
                     *jumpPtr             = jumpDest;
-                    fgAddRefPred<initializingPreds>(jumpDest, curBBdesc, oldEdge);
+                    fgAddRefPred<initializingPreds>(jumpDest, curBBdesc);
                     if ((*jumpPtr)->bbNum <= curBBdesc->bbNum)
                     {
                         fgMarkBackwardJump(*jumpPtr, curBBdesc);
@@ -3839,7 +3839,8 @@ void Compiler::fgFindBasicBlocks()
                 {
                     // Mark catch handler as successor.
                     block->SetTarget(hndBegBB);
-                    fgAddRefPred(hndBegBB, block);
+                    FlowEdge* const newEdge = fgAddRefPred(hndBegBB, block);
+                    newEdge->setLikelihood(1.0);
                     assert(block->GetTarget()->bbCatchTyp == BBCT_FILTER_HANDLER);
                     break;
                 }

src/coreclr/jit/fgflow.cpp

@@ -136,6 +136,32 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
             if (flowLast->getSourceBlock() == blockPred)
             {
                 flow = flowLast;
+
+                // This edge should have been given a likelihood when it was created.
+                // Since we're increasing its duplicate count, update the likelihood.
+                //
+                assert(flow->hasLikelihood());
+                const unsigned numSucc = blockPred->NumSucc();
+                assert(numSucc > 0);
+
+                if (numSucc == 1)
+                {
+                    // BasicBlock::NumSucc() returns 1 for BBJ_CONDs with the same true/false target.
+                    // For blocks that only ever have one successor (BBJ_ALWAYS, BBJ_LEAVE, etc.),
+                    // their successor edge should never have a duplicate count over 1.
+                    //
+                    assert(blockPred->KindIs(BBJ_COND));
+                    assert(blockPred->TrueTargetIs(blockPred->GetFalseTarget()));
+                    flow->setLikelihood(1.0);
+                }
+                else
+                {
+                    // Duplicate count isn't updated until later, so add 1 for now.
+                    //
+                    const unsigned dupCount = flow->getDupCount() + 1;
+                    assert(dupCount > 1);
+                    flow->setLikelihood((1.0 / numSucc) * dupCount);
+                }
             }
         }
     }
@@ -185,12 +211,19 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
         if (initializingPreds)
         {
             block->bbLastPred = flow;
-        }
 
-        // Copy likelihood from old edge, if any.
-        //
-        if ((oldEdge != nullptr) && oldEdge->hasLikelihood())
+            // When initializing preds, ensure edge likelihood is set,
+            // such that this edge is as likely as any other successor edge
+            //
+            const unsigned numSucc = blockPred->NumSucc();
+            assert(numSucc > 0);
+            assert(flow->getDupCount() == 1);
+            flow->setLikelihood(1.0 / numSucc);
+        }
+        else if ((oldEdge != nullptr) && oldEdge->hasLikelihood())
         {
+            // Copy likelihood from old edge, if any.
+            //
             flow->setLikelihood(oldEdge->getLikelihood());
         }
 
@@ -235,6 +268,10 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
     //
     assert(block->checkPredListOrder());
 
+    // When initializing preds, edge likelihood should always be set.
+    //
+    assert(!initializingPreds || flow->hasLikelihood());
+
     return flow;
 }
 

src/coreclr/jit/fgprofile.cpp

@@ -3916,7 +3916,10 @@ void EfficientEdgeCountReconstructor::PropagateOSREntryEdges(BasicBlock* block,
         FlowEdge* const flowEdge = m_comp->fgGetPredForBlock(edge->m_targetBlock, block);
 
         assert(flowEdge != nullptr);
-        assert(!flowEdge->hasLikelihood());
+
+        // Naive likelihood should have been set during pred initialization in fgAddRefPred
+        //
+        assert(flowEdge->hasLikelihood());
         weight_t likelihood = 0;
 
         if (nEdges == 1)

src/coreclr/jit/importer.cpp

@@ -12280,7 +12280,7 @@ void Compiler::impFixPredLists()
 
                         BasicBlock* const continuation = predBlock->Next();
                         FlowEdge* const   newEdge      = fgAddRefPred(continuation, finallyBlock);
-                        newEdge->setLikelihood(1.0);
+                        newEdge->setLikelihood(1.0 / predCount);
 
                         jumpEhf->bbeSuccs[predNum] = continuation;
                         ++predNum;