JIT: Move edge likelihood initialization to fgLinkBasicBlocks (#99100)

After #98993, the logic for initializing edge likelihoods in fgAddRefPred can potentially read uninitialized memory by calling BasicBlock::NumSucc. Avoid this by moving the edge likelihood logic to fgLinkBasicBlocks, where we can safely and cheaply determine the number of successors. Also, rename AssertionInfo::m_isNextEdgeAssertion based on feedback from #98993.

Author: amanasifkhalid

src/coreclr/jit/assertionprop.cpp

@@ -6146,7 +6146,7 @@ ASSERT_TP* Compiler::optComputeAssertionGen()
                 AssertionIndex valueAssertionIndex;
                 AssertionIndex jumpDestAssertionIndex;
 
-                if (info.IsNextEdgeAssertion())
+                if (info.AssertionHoldsOnFalseEdge())
                 {
                     valueAssertionIndex    = info.GetAssertionIndex();
                     jumpDestAssertionIndex = optFindComplementary(info.GetAssertionIndex());

src/coreclr/jit/fgbasic.cpp

@@ -2979,6 +2979,17 @@ void Compiler::fgLinkBasicBlocks()
                 curBBdesc->SetTrueEdge(trueEdge);
                 curBBdesc->SetFalseEdge(falseEdge);
 
+                if (trueEdge == falseEdge)
+                {
+                    assert(trueEdge->getDupCount() == 2);
+                    trueEdge->setLikelihood(1.0);
+                }
+                else
+                {
+                    trueEdge->setLikelihood(0.5);
+                    falseEdge->setLikelihood(0.5);
+                }
+
                 if (trueTarget->bbNum <= curBBdesc->bbNum)
                 {
                     fgMarkBackwardJump(trueTarget, curBBdesc);
@@ -3003,6 +3014,7 @@ void Compiler::fgLinkBasicBlocks()
                 // Redundantly use SetKindAndTargetEdge() instead of SetTargetEdge() just this once,
                 // so we don't break the HasInitializedTarget() invariant of SetTargetEdge().
                 FlowEdge* const newEdge = fgAddRefPred<initializingPreds>(jumpDest, curBBdesc);
+                newEdge->setLikelihood(1.0);
                 curBBdesc->SetKindAndTargetEdge(curBBdesc->GetKind(), newEdge);
 
                 if (curBBdesc->GetTarget()->bbNum <= curBBdesc->bbNum)
@@ -3029,14 +3041,17 @@ void Compiler::fgLinkBasicBlocks()
 
             case BBJ_SWITCH:
             {
-                unsigned   jumpCnt = curBBdesc->GetSwitchTargets()->bbsCount;
-                FlowEdge** jumpPtr = curBBdesc->GetSwitchTargets()->bbsDstTab;
+                const unsigned numSucc = curBBdesc->GetSwitchTargets()->bbsCount;
+                unsigned       jumpCnt = numSucc;
+                FlowEdge**     jumpPtr = curBBdesc->GetSwitchTargets()->bbsDstTab;
 
                 do
                 {
                     BasicBlock*     jumpDest = fgLookupBB((unsigned)*(size_t*)jumpPtr);
                     FlowEdge* const newEdge  = fgAddRefPred<initializingPreds>(jumpDest, curBBdesc);
-                    *jumpPtr                 = newEdge;
+
+                    newEdge->setLikelihood((1.0 / numSucc) * newEdge->getDupCount());
+                    *jumpPtr = newEdge;
                     if (jumpDest->bbNum <= curBBdesc->bbNum)
                     {
                         fgMarkBackwardJump(jumpDest, curBBdesc);

src/coreclr/jit/fgflow.cpp

@@ -136,32 +136,6 @@ 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->TrueEdgeIs(blockPred->GetFalseEdge()));
-                    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);
-                }
             }
         }
     }
@@ -211,19 +185,6 @@ FlowEdge* Compiler::fgAddRefPred(BasicBlock* block, BasicBlock* blockPred, FlowE
         if (initializingPreds)
         {
             block->bbLastPred = flow;
-
-            // When initializing preds, ensure edge likelihood is set,
-            // such that this edge is as likely as any other successor edge
-            // Note: We probably shouldn't call NumSucc on a new BBJ_COND block.
-            // NumSucc compares bbTrueEdge and bbFalseEdge to determine if this BBJ_COND block has only one successor,
-            // but these members are uninitialized. Aside from the fact that this compares uninitialized memory,
-            // we don't know if the true and false targets are the same in NumSucc until both edges exist.
-            // TODO: Move this edge likelihood logic to fgLinkBasicBlocks.
-            //
-            const unsigned numSucc = blockPred->NumSucc();
-            assert(numSucc > 0);
-            assert(flow->getDupCount() == 1);
-            flow->setLikelihood(1.0 / numSucc);
         }
         else if ((oldEdge != nullptr) && oldEdge->hasLikelihood())
         {
@@ -273,10 +234,6 @@ 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/gentree.h

@@ -195,12 +195,12 @@ inline AssertionIndex GetAssertionIndex(unsigned index)
 class AssertionInfo
 {
     // true if the assertion holds on the false edge instead of the true edge (for GT_JTRUE nodes)
-    unsigned short m_isNextEdgeAssertion : 1;
+    unsigned short m_assertionHoldsOnFalseEdge : 1;
     // 1-based index of the assertion
     unsigned short m_assertionIndex : 15;
 
-    AssertionInfo(bool isNextEdgeAssertion, AssertionIndex assertionIndex)
-        : m_isNextEdgeAssertion(isNextEdgeAssertion), m_assertionIndex(assertionIndex)
+    AssertionInfo(bool assertionHoldsOnFalseEdge, AssertionIndex assertionIndex)
+        : m_assertionHoldsOnFalseEdge(assertionHoldsOnFalseEdge), m_assertionIndex(assertionIndex)
     {
         assert(m_assertionIndex == assertionIndex);
     }
@@ -223,8 +223,8 @@ class AssertionInfo
 
     void Clear()
     {
-        m_isNextEdgeAssertion = 0;
-        m_assertionIndex      = NO_ASSERTION_INDEX;
+        m_assertionHoldsOnFalseEdge = 0;
+        m_assertionIndex            = NO_ASSERTION_INDEX;
     }
 
     bool HasAssertion() const
@@ -237,9 +237,9 @@ class AssertionInfo
         return m_assertionIndex;
     }
 
-    bool IsNextEdgeAssertion() const
+    bool AssertionHoldsOnFalseEdge() const
     {
-        return m_isNextEdgeAssertion;
+        return m_assertionHoldsOnFalseEdge;
     }
 };
 

src/coreclr/jit/morph.cpp

@@ -12993,7 +12993,7 @@ void Compiler::fgAssertionGen(GenTree* tree)
         AssertionIndex ifFalseAssertionIndex;
         AssertionIndex ifTrueAssertionIndex;
 
-        if (info.IsNextEdgeAssertion())
+        if (info.AssertionHoldsOnFalseEdge())
         {
             ifFalseAssertionIndex = info.GetAssertionIndex();
             ifTrueAssertionIndex  = optFindComplementary(ifFalseAssertionIndex);