JIT: Set edge likelihood in SetTargetEdge, SetKindAndTargetEdge (#99269)

Part of #93020. When setting the successor edge of a block with one successor, the edge likelihood should always be 1.0. We can enforce this in BasicBlock methods that set bbTargetEdge, and simplify their call sites.

Author: amanasifkhalid

src/coreclr/jit/block.h

@@ -847,6 +847,9 @@ struct BasicBlock : private LIR::Range
         assert(HasInitializedTarget());
         assert(bbTargetEdge->getSourceBlock() == this);
         assert(bbTargetEdge->getDestinationBlock() != nullptr);
+
+        // This is the only successor edge for this block, so likelihood should be 1.0
+        bbTargetEdge->setLikelihood(1.0);
     }
 
     BasicBlock* GetTrueTarget() const
@@ -932,16 +935,23 @@ struct BasicBlock : private LIR::Range
         SetTrueEdge(trueEdge);
     }
 
-    // Set both the block kind and target edge. This can clear `bbTargetEdge` when setting
-    // block kinds that don't use `bbTargetEdge`.
-    void SetKindAndTargetEdge(BBKinds kind, FlowEdge* targetEdge = nullptr)
+    // Set both the block kind and target edge.
+    void SetKindAndTargetEdge(BBKinds kind, FlowEdge* targetEdge)
     {
         bbKind       = kind;
         bbTargetEdge = targetEdge;
+        assert(HasInitializedTarget());
+
+        // This is the only successor edge for this block, so likelihood should be 1.0
+        bbTargetEdge->setLikelihood(1.0);
+    }
 
-        // If bbKind indicates this block has a jump, bbTargetEdge cannot be null.
-        // You shouldn't use this to set a BBJ_COND, BBJ_SWITCH, or BBJ_EHFINALLYRET.
-        assert(HasTarget() ? HasInitializedTarget() : (bbTargetEdge == nullptr));
+    // Set the block kind, and clear bbTargetEdge.
+    void SetKindAndTargetEdge(BBKinds kind)
+    {
+        bbKind       = kind;
+        bbTargetEdge = nullptr;
+        assert(!HasTarget());
     }
 
     bool HasInitializedTarget() const

src/coreclr/jit/fgbasic.cpp

@@ -227,7 +227,6 @@ bool Compiler::fgEnsureFirstBBisScratch()
 
         // The new scratch bb will fall through to the old first bb
         FlowEdge* const edge = fgAddRefPred(fgFirstBB, block);
-        edge->setLikelihood(1.0);
         block->SetKindAndTargetEdge(BBJ_ALWAYS, edge);
         fgInsertBBbefore(fgFirstBB, block);
     }
@@ -678,7 +677,6 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* oldTarget, Bas
 
                 if (block->KindIs(BBJ_ALWAYS))
                 {
-                    newEdge->setLikelihood(1.0);
                     block->SetTargetEdge(newEdge);
                 }
                 else
@@ -2964,7 +2962,6 @@ 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)
@@ -3862,7 +3859,6 @@ void Compiler::fgFindBasicBlocks()
                 {
                     // Mark catch handler as successor.
                     FlowEdge* const newEdge = fgAddRefPred(hndBegBB, block);
-                    newEdge->setLikelihood(1.0);
                     block->SetTargetEdge(newEdge);
                     assert(hndBegBB->bbCatchTyp == BBCT_FILTER_HANDLER);
                     break;
@@ -4198,7 +4194,6 @@ void Compiler::fgFixEntryFlowForOSR()
     assert(fgFirstBB->KindIs(BBJ_ALWAYS) && fgFirstBB->JumpsToNext());
     fgRemoveRefPred(fgFirstBB->GetTargetEdge());
     FlowEdge* const newEdge = fgAddRefPred(fgOSREntryBB, fgFirstBB);
-    newEdge->setLikelihood(1.0);
     fgFirstBB->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
 
     // We don't know the right weight for this block, since
@@ -4790,7 +4785,6 @@ BasicBlock* Compiler::fgSplitBlockAtEnd(BasicBlock* curr)
 
     // Default to fallthrough, and add the arc for that.
     FlowEdge* const newEdge = fgAddRefPred(newBlock, curr);
-    newEdge->setLikelihood(1.0);
 
     // Transfer the kind and target. Do this after the code above, to avoid null-ing out the old targets used by the
     // above code.
@@ -5041,7 +5035,6 @@ BasicBlock* Compiler::fgSplitEdge(BasicBlock* curr, BasicBlock* succ)
 
     // And 'succ' has 'newBlock' as a new predecessor.
     FlowEdge* const newEdge = fgAddRefPred(succ, newBlock);
-    newEdge->setLikelihood(1.0);
     newBlock->SetTargetEdge(newEdge);
 
     // This isn't accurate, but it is complex to compute a reasonable number so just assume that we take the

src/coreclr/jit/fgdiagnostic.cpp

@@ -2786,6 +2786,7 @@ bool BBPredsChecker::CheckJump(BasicBlock* blockPred, BasicBlock* block)
         case BBJ_EHCATCHRET:
         case BBJ_EHFILTERRET:
             assert(blockPred->TargetIs(block));
+            assert(blockPred->GetTargetEdge()->getLikelihood() == 1.0);
             return true;
 
         case BBJ_EHFINALLYRET:

src/coreclr/jit/fginline.cpp

@@ -682,11 +682,8 @@ class SubstitutePlaceholdersAndDevirtualizeWalker : public GenTreeVisitor<Substi
                 else
                 {
                     m_compiler->fgRemoveRefPred(block->GetFalseEdge());
-                    block->SetKind(BBJ_ALWAYS);
+                    block->SetKindAndTargetEdge(BBJ_ALWAYS, block->GetTrueEdge());
                 }
-
-                FlowEdge* const edge = m_compiler->fgGetPredForBlock(block->GetTarget(), block);
-                edge->setLikelihood(1.0);
             }
         }
         else
@@ -1537,7 +1534,6 @@ void Compiler::fgInsertInlineeBlocks(InlineInfo* pInlineInfo)
                         bottomBlock->bbNum);
 
                 FlowEdge* const newEdge = fgAddRefPred(bottomBlock, block);
-                newEdge->setLikelihood(1.0);
                 block->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
 
                 if (block == InlineeCompiler->fgLastBB)

src/coreclr/jit/fgopt.cpp

@@ -648,7 +648,6 @@ PhaseStatus Compiler::fgPostImportationCleanup()
                         else
                         {
                             FlowEdge* const newEdge = fgAddRefPred(newTryEntry->Next(), newTryEntry);
-                            newEdge->setLikelihood(1.0);
                             newTryEntry->SetFlags(BBF_NONE_QUIRK);
                             newTryEntry->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
                         }
@@ -2634,7 +2633,7 @@ void Compiler::fgRemoveConditionalJump(BasicBlock* block)
 
     /* Conditional is gone - always jump to target */
 
-    block->SetKind(BBJ_ALWAYS);
+    block->SetKindAndTargetEdge(BBJ_ALWAYS, block->GetTrueEdge());
     assert(block->TargetIs(target));
 
     // TODO-NoFallThrough: Set BBF_NONE_QUIRK only when false target is the next block

src/coreclr/jit/fgprofile.cpp

@@ -511,7 +511,6 @@ void BlockCountInstrumentor::RelocateProbes()
             intermediary->SetFlags(BBF_IMPORTED | BBF_MARKED | BBF_NONE_QUIRK);
             intermediary->inheritWeight(block);
             FlowEdge* const newEdge = m_comp->fgAddRefPred(block, intermediary);
-            newEdge->setLikelihood(1.0);
             intermediary->SetTargetEdge(newEdge);
             SetModifiedFlow();
 
@@ -1683,7 +1682,6 @@ void EfficientEdgeCountInstrumentor::RelocateProbes()
             intermediary->SetFlags(BBF_IMPORTED | BBF_NONE_QUIRK);
             intermediary->inheritWeight(block);
             FlowEdge* const newEdge = m_comp->fgAddRefPred(block, intermediary);
-            newEdge->setLikelihood(1.0);
             intermediary->SetTargetEdge(newEdge);
             NewRelocatedProbe(intermediary, probe->source, probe->target, &leader);
             SetModifiedFlow();
@@ -3921,7 +3919,7 @@ void EfficientEdgeCountReconstructor::PropagateOSREntryEdges(BasicBlock* block,
 
         assert(flowEdge != nullptr);
 
-        // Naive likelihood should have been set during pred initialization in fgAddRefPred
+        // Naive likelihood should have been set during pred initialization in fgLinkBasicBlocks
         //
         assert(flowEdge->hasLikelihood());
         weight_t likelihood = 0;
@@ -3980,8 +3978,8 @@ void EfficientEdgeCountReconstructor::PropagateEdges(BasicBlock* block, BlockInf
     // If there is a pseudo edge,
     // There should be only one successor for block. The flow
     // from block to successor will not represent real flow.
-    // We set likelihood anyways so we can assert later
-    // that all flow edges have known likelihood.
+    // Likelihood should be set to 1.0 already, as we already know
+    // this block has only one successor.
     //
     // Note the flowEdge target may not be the same as the pseudo edge target.
     //
@@ -3992,7 +3990,7 @@ void EfficientEdgeCountReconstructor::PropagateEdges(BasicBlock* block, BlockInf
         assert(block->HasInitializedTarget());
         FlowEdge* const flowEdge = block->GetTargetEdge();
         assert(flowEdge != nullptr);
-        flowEdge->setLikelihood(1.0);
+        assert(flowEdge->getLikelihood() == 1.0);
         return;
     }
 

src/coreclr/jit/flowgraph.cpp

@@ -1629,7 +1629,6 @@ void Compiler::fgConvertSyncReturnToLeave(BasicBlock* block)
 
     // Convert the BBJ_RETURN to BBJ_ALWAYS, jumping to genReturnBB.
     FlowEdge* const newEdge = fgAddRefPred(genReturnBB, block);
-    newEdge->setLikelihood(1.0);
     block->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
 
 #ifdef DEBUG
@@ -2101,7 +2100,6 @@ class MergedReturns
                     // Change BBJ_RETURN to BBJ_ALWAYS targeting const return block.
                     assert((comp->info.compFlags & CORINFO_FLG_SYNCH) == 0);
                     FlowEdge* const newEdge = comp->fgAddRefPred(constReturnBlock, returnBlock);
-                    newEdge->setLikelihood(1.0);
                     returnBlock->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
 
                     // Remove GT_RETURN since constReturnBlock returns the constant.

src/coreclr/jit/importer.cpp

@@ -2026,7 +2026,6 @@ BasicBlock* Compiler::impPushCatchArgOnStack(BasicBlock* hndBlk, CORINFO_CLASS_H
         newBlk->bbCodeOffs = hndBlk->bbCodeOffs;
 
         FlowEdge* const newEdge = fgAddRefPred(hndBlk, newBlk);
-        newEdge->setLikelihood(1.0);
         newBlk->SetTargetEdge(newEdge);
 
         // Spill into a temp.
@@ -4437,7 +4436,6 @@ void Compiler::impImportLeave(BasicBlock* block)
 
                 // the previous call to a finally returns to this call (to the next finally in the chain)
                 FlowEdge* const newEdge = fgAddRefPred(callBlock, step);
-                newEdge->setLikelihood(1.0);
                 step->SetTargetEdge(newEdge);
 
                 // The new block will inherit this block's weight.
@@ -4487,7 +4485,6 @@ void Compiler::impImportLeave(BasicBlock* block)
             assert(finallyNesting <= compHndBBtabCount);
 
             FlowEdge* const newEdge = fgAddRefPred(HBtab->ebdHndBeg, callBlock);
-            newEdge->setLikelihood(1.0);
             callBlock->SetKindAndTargetEdge(BBJ_CALLFINALLY, newEdge);
 
             GenTree* endLFin = new (this, GT_END_LFIN) GenTreeVal(GT_END_LFIN, TYP_VOID, finallyNesting);
@@ -4539,7 +4536,6 @@ void Compiler::impImportLeave(BasicBlock* block)
 
         {
             FlowEdge* const newEdge = fgAddRefPred(finalStep, step);
-            newEdge->setLikelihood(1.0);
             step->SetTargetEdge(newEdge);
         }
 
@@ -4572,7 +4568,6 @@ void Compiler::impImportLeave(BasicBlock* block)
         // this is the ultimate destination of the LEAVE
         {
             FlowEdge* const newEdge = fgAddRefPred(leaveTarget, finalStep);
-            newEdge->setLikelihood(1.0);
             finalStep->SetTargetEdge(newEdge);
         }
 
@@ -4692,7 +4687,6 @@ void Compiler::impImportLeave(BasicBlock* block)
                 }
 
                 FlowEdge* const newEdge = fgAddRefPred(exitBlock, step);
-                newEdge->setLikelihood(1.0);
                 step->SetTargetEdge(newEdge); // the previous step (maybe a call to a nested finally, or a nested catch
                                               // exit) returns to this block
 
@@ -4737,7 +4731,6 @@ void Compiler::impImportLeave(BasicBlock* block)
                 // next block, and flow optimizations will remove it.
                 fgRemoveRefPred(block->GetTargetEdge());
                 FlowEdge* const newEdge = fgAddRefPred(callBlock, block);
-                newEdge->setLikelihood(1.0);
                 block->SetKindAndTargetEdge(BBJ_ALWAYS, newEdge);
 
                 // The new block will inherit this block's weight.
@@ -4807,7 +4800,6 @@ void Compiler::impImportLeave(BasicBlock* block)
                     }
 
                     FlowEdge* const newEdge = fgAddRefPred(step2, step);
-                    newEdge->setLikelihood(1.0);
                     step->SetTargetEdge(newEdge);
                     step2->inheritWeight(block);
                     step2->CopyFlags(block, BBF_RUN_RARELY);
@@ -4848,7 +4840,6 @@ void Compiler::impImportLeave(BasicBlock* block)
                 }
 
                 FlowEdge* const newEdge = fgAddRefPred(callBlock, step);
-                newEdge->setLikelihood(1.0);
                 step->SetTargetEdge(newEdge); // the previous call to a finally returns to this call (to the next
                                               // finally in the chain)
 
@@ -4885,7 +4876,6 @@ void Compiler::impImportLeave(BasicBlock* block)
 #endif
 
             FlowEdge* const newEdge = fgAddRefPred(HBtab->ebdHndBeg, callBlock);
-            newEdge->setLikelihood(1.0);
             callBlock->SetKindAndTargetEdge(BBJ_CALLFINALLY, newEdge);
         }
         else if (HBtab->HasCatchHandler() && jitIsBetween(blkAddr, tryBeg, tryEnd) &&
@@ -4954,7 +4944,6 @@ void Compiler::impImportLeave(BasicBlock* block)
                 }
 
                 FlowEdge* const newEdge = fgAddRefPred(catchStep, step);
-                newEdge->setLikelihood(1.0);
                 step->SetTargetEdge(newEdge);
 
                 // The new block will inherit this block's weight.
@@ -5009,7 +4998,6 @@ void Compiler::impImportLeave(BasicBlock* block)
             fgRemoveRefPred(step->GetTargetEdge());
         }
         FlowEdge* const newEdge = fgAddRefPred(leaveTarget, step);
-        newEdge->setLikelihood(1.0);
         step->SetTargetEdge(newEdge); // this is the ultimate destination of the LEAVE
 
 #ifdef DEBUG
@@ -5072,7 +5060,6 @@ void Compiler::impResetLeaveBlock(BasicBlock* block, unsigned jmpAddr)
         BasicBlock* dupBlock = BasicBlock::New(this);
         dupBlock->CopyFlags(block);
         FlowEdge* const newEdge = fgAddRefPred(block->GetTarget(), dupBlock);
-        newEdge->setLikelihood(1.0);
         dupBlock->SetKindAndTargetEdge(BBJ_CALLFINALLY, newEdge);
         dupBlock->copyEHRegion(block);
         dupBlock->bbCatchTyp = block->bbCatchTyp;
@@ -5104,7 +5091,6 @@ void Compiler::impResetLeaveBlock(BasicBlock* block, unsigned jmpAddr)
 
     fgRemoveRefPred(block->GetTargetEdge());
     FlowEdge* const newEdge = fgAddRefPred(fgLookupBB(jmpAddr), block);
-    newEdge->setLikelihood(1.0);
     block->SetKindAndTargetEdge(BBJ_LEAVE, newEdge);
 
     // We will leave the BBJ_ALWAYS block we introduced. When it's reimported
@@ -7187,7 +7173,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                         JITDUMP(FMT_BB " always branches to " FMT_BB ", changing to BBJ_ALWAYS\n", block->bbNum,
                                 block->GetFalseTarget()->bbNum);
                         fgRemoveRefPred(block->GetFalseEdge());
-                        block->SetKind(BBJ_ALWAYS);
+                        block->SetKindAndTargetEdge(BBJ_ALWAYS, block->GetTrueEdge());
 
                         // TODO-NoFallThrough: Once false target can diverge from bbNext, it may not make sense to
                         // set BBF_NONE_QUIRK
@@ -7262,7 +7248,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                             JITDUMP("\nThe conditional jump becomes an unconditional jump to " FMT_BB "\n",
                                     block->GetTrueTarget()->bbNum);
                             fgRemoveRefPred(block->GetFalseEdge());
-                            block->SetKind(BBJ_ALWAYS);
+                            block->SetKindAndTargetEdge(BBJ_ALWAYS, block->GetTrueEdge());
                         }
                         else
                         {
@@ -7276,9 +7262,6 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                             // to set BBF_NONE_QUIRK
                             block->SetFlags(BBF_NONE_QUIRK);
                         }
-
-                        FlowEdge* const edge = fgGetPredForBlock(block->GetTarget(), block);
-                        edge->setLikelihood(1.0);
                     }
 
                     break;
@@ -7452,7 +7435,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                         JITDUMP(FMT_BB " always branches to " FMT_BB ", changing to BBJ_ALWAYS\n", block->bbNum,
                                 block->GetFalseTarget()->bbNum);
                         fgRemoveRefPred(block->GetFalseEdge());
-                        block->SetKind(BBJ_ALWAYS);
+                        block->SetKindAndTargetEdge(BBJ_ALWAYS, block->GetTrueEdge());
 
                         // TODO-NoFallThrough: Once false target can diverge from bbNext, it may not make sense to
                         // set BBF_NONE_QUIRK
@@ -7540,7 +7523,6 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                         {
                             // transform the basic block into a BBJ_ALWAYS
                             block->SetKindAndTargetEdge(BBJ_ALWAYS, curEdge);
-                            curEdge->setLikelihood(1.0);
                             foundVal = true;
                         }
                         else