JIT: Enable profile consistency checks throughout JIT frontend (#111498)

Part of #107749.

src/coreclr/jit/block.h

@@ -1197,6 +1197,18 @@ struct BasicBlock : private LIR::Range
     // getBBWeight -- get the normalized weight of this block
     weight_t getBBWeight(Compiler* comp) const;
 
+    // computeIncomingWeight -- sum the weights of the flow edges into this block
+    weight_t computeIncomingWeight() const
+    {
+        weight_t incomingWeight = BB_ZERO_WEIGHT;
+        for (FlowEdge* const predEdge : PredEdges())
+        {
+            incomingWeight += predEdge->getLikelyWeight();
+        }
+
+        return incomingWeight;
+    }
+
     // hasProfileWeight -- Returns true if this block's weight came from profile data
     bool hasProfileWeight() const
     {

src/coreclr/jit/compiler.cpp

@@ -5153,6 +5153,11 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
         DoPhase(this, PHASE_SWITCH_RECOGNITION, &Compiler::optSwitchRecognition);
     }
 
+    // Drop back to just checking profile likelihoods.
+    //
+    activePhaseChecks &= ~PhaseChecks::CHECK_PROFILE;
+    activePhaseChecks |= PhaseChecks::CHECK_LIKELIHOODS;
+
 #ifdef DEBUG
     // Stash the current estimate of the function's size if necessary.
     if (verbose && opts.OptimizationEnabled())

src/coreclr/jit/helperexpansion.cpp

@@ -2700,8 +2700,29 @@ bool Compiler::fgLateCastExpansionForCall(BasicBlock** pBlock, Statement* stmt,
     if ((call->gtCallMoreFlags & GTF_CALL_M_CAST_OBJ_NONNULL) != 0)
     {
         fgRemoveStmt(nullcheckBb, nullcheckBb->lastStmt());
-        fgRemoveRefPred(nullcheckBb->GetTrueEdge());
+        FlowEdge* const removedEdge = nullcheckBb->GetTrueEdge();
+        fgRemoveRefPred(removedEdge);
         nullcheckBb->SetKindAndTargetEdge(BBJ_ALWAYS, nullcheckBb->GetFalseEdge());
+
+        // Locally repair profile
+        if (nullcheckBb->hasProfileWeight())
+        {
+            BasicBlock* const removedTarget = removedEdge->getDestinationBlock();
+            if (removedTarget->bbPreds == nullptr)
+            {
+                // Unreachable path will be removed by the next flowgraph simplification pass.
+                // For now, mark the profile as inconsistent to skip post-phase checks.
+                JITDUMP("fgLateCastExpansionForCall: " FMT_BB
+                        " is unreachable, and will be removed later. Data %s inconsistent.\n",
+                        removedTarget->bbNum, fgPgoConsistent ? "is now" : "was already");
+                fgPgoConsistent = false;
+            }
+
+            for (BasicBlock* const block : Blocks(nullcheckBb->Next(), lastBb))
+            {
+                block->setBBProfileWeight(block->computeIncomingWeight());
+            }
+        }
     }
 
     // Bonus step: merge prevBb with nullcheckBb as they are likely to be mergeable

src/coreclr/jit/ifconversion.cpp

@@ -740,8 +740,10 @@ bool OptIfConversionDsc::optIfConvert()
     }
 
     // Update the flow from the original block.
-    m_comp->fgRemoveAllRefPreds(m_startBlock->GetFalseTarget(), m_startBlock);
-    m_startBlock->SetKindAndTargetEdge(BBJ_ALWAYS, m_startBlock->GetTrueEdge());
+    FlowEdge* const removedEdge  = m_comp->fgRemoveAllRefPreds(m_startBlock->GetFalseTarget(), m_startBlock);
+    FlowEdge* const retainedEdge = m_startBlock->GetTrueEdge();
+    m_startBlock->SetKindAndTargetEdge(BBJ_ALWAYS, retainedEdge);
+    m_comp->fgRepairProfileCondToUncond(m_startBlock, retainedEdge, removedEdge);
 
 #ifdef DEBUG
     if (m_comp->verbose)

src/coreclr/jit/morph.cpp

@@ -13513,12 +13513,7 @@ PhaseStatus Compiler::fgMorphBlocks()
         //
         if (fgEntryBB->hasProfileWeight())
         {
-            weight_t incomingWeight = BB_ZERO_WEIGHT;
-            for (FlowEdge* const predEdge : fgEntryBB->PredEdges())
-            {
-                incomingWeight += predEdge->getLikelyWeight();
-            }
-
+            const weight_t incomingWeight = fgEntryBB->computeIncomingWeight();
             if (!fgProfileWeightsConsistent(incomingWeight, fgEntryBB->bbWeight))
             {
                 JITDUMP("OSR: Original method entry " FMT_BB " has inconsistent weight. Data %s inconsistent.\n",

src/coreclr/jit/optimizebools.cpp

@@ -866,6 +866,23 @@ bool OptBoolsDsc::optOptimizeRangeTests()
     m_comp->fgRemoveRefPred(oldFalseEdge);
     m_comp->fgRemoveBlock(m_b2, true);
 
+    // Update profile
+    if (m_b1->hasProfileWeight())
+    {
+        BasicBlock* const trueTarget  = m_b1->GetTrueTarget();
+        BasicBlock* const falseTarget = m_b1->GetFalseTarget();
+        trueTarget->setBBProfileWeight(trueTarget->computeIncomingWeight());
+        falseTarget->setBBProfileWeight(falseTarget->computeIncomingWeight());
+
+        if ((trueTarget->NumSucc() > 0) || (falseTarget->NumSucc() > 0))
+        {
+            JITDUMP("optOptimizeRangeTests: Profile needs to be propagated through " FMT_BB
+                    "'s successors. Data %s inconsistent.\n",
+                    m_b1->bbNum, m_comp->fgPgoConsistent ? "is now" : "was already");
+            m_comp->fgPgoConsistent = false;
+        }
+    }
+
     Statement* const stmt = m_b1->lastStmt();
     m_comp->gtSetStmtInfo(stmt);
     m_comp->fgSetStmtSeq(stmt);
@@ -1055,8 +1072,13 @@ bool OptBoolsDsc::optOptimizeCompareChainCondBlock()
     m_comp->fgSetStmtSeq(s2);
 
     // Update the flow.
-    m_comp->fgRemoveRefPred(m_b1->GetTrueEdge());
-    m_b1->SetKindAndTargetEdge(BBJ_ALWAYS, m_b1->GetFalseEdge());
+    FlowEdge* const removedEdge  = m_b1->GetTrueEdge();
+    FlowEdge* const retainedEdge = m_b1->GetFalseEdge();
+    m_comp->fgRemoveRefPred(removedEdge);
+    m_b1->SetKindAndTargetEdge(BBJ_ALWAYS, retainedEdge);
+
+    // Repair profile.
+    m_comp->fgRepairProfileCondToUncond(m_b1, retainedEdge, removedEdge);
 
     // Fixup flags.
     m_b2->CopyFlags(m_b1, BBF_COPY_PROPAGATE);
@@ -1338,6 +1360,23 @@ void OptBoolsDsc::optOptimizeBoolsUpdateTrees()
         // Fix B1 false edge likelihood
         //
         newB1FalseEdge->setLikelihood(1.0 - newB1TrueLikelihood);
+
+        // Update profile
+        if (m_b1->hasProfileWeight())
+        {
+            BasicBlock* const trueTarget  = origB1TrueEdge->getDestinationBlock();
+            BasicBlock* const falseTarget = newB1FalseEdge->getDestinationBlock();
+            trueTarget->setBBProfileWeight(trueTarget->computeIncomingWeight());
+            falseTarget->setBBProfileWeight(falseTarget->computeIncomingWeight());
+
+            if ((trueTarget->NumSucc() > 0) || (falseTarget->NumSucc() > 0))
+            {
+                JITDUMP("optOptimizeRangeTests: Profile needs to be propagated through " FMT_BB
+                        "'s successors. Data %s inconsistent.\n",
+                        m_b1->bbNum, m_comp->fgPgoConsistent ? "is now" : "was already");
+                m_comp->fgPgoConsistent = false;
+            }
+        }
     }
 
     // Get rid of the second block

src/coreclr/jit/redundantbranchopts.cpp

@@ -1652,8 +1652,11 @@ bool Compiler::optJumpThreadCore(JumpThreadInfo& jti)
     // If this pred is in the set that will reuse block, do nothing.
     // Else revise pred to branch directly to the appropriate successor of block.
     //
-    for (BasicBlock* const predBlock : jti.m_block->PredBlocksEditing())
+    bool modifiedProfile = false;
+    for (FlowEdge* const predEdge : jti.m_block->PredEdgesEditing())
     {
+        BasicBlock* const predBlock = predEdge->getSourceBlock();
+
         // If this was an ambiguous pred, skip.
         //
         if (BitVecOps::IsMember(&jti.traits, jti.m_ambiguousPreds, predBlock->bbPostorderNum))
@@ -1670,36 +1673,47 @@ bool Compiler::optJumpThreadCore(JumpThreadInfo& jti)
 
         // Jump to the appropriate successor.
         //
+        BasicBlock* newTarget = nullptr;
         if (isTruePred)
         {
             JITDUMP("Jump flow from pred " FMT_BB " -> " FMT_BB
                     " implies predicate true; we can safely redirect flow to be " FMT_BB " -> " FMT_BB "\n",
                     predBlock->bbNum, jti.m_block->bbNum, predBlock->bbNum, jti.m_trueTarget->bbNum);
-
-            fgReplaceJumpTarget(predBlock, jti.m_block, jti.m_trueTarget);
-
-            if (setNoCseIn && !jti.m_trueTarget->HasFlag(BBF_NO_CSE_IN))
-            {
-                JITDUMP(FMT_BB " => BBF_NO_CSE_IN\n", jti.m_trueTarget->bbNum);
-                jti.m_trueTarget->SetFlags(BBF_NO_CSE_IN);
-            }
+            newTarget = jti.m_trueTarget;
         }
         else
         {
             JITDUMP("Jump flow from pred " FMT_BB " -> " FMT_BB
                     " implies predicate false; we can safely redirect flow to be " FMT_BB " -> " FMT_BB "\n",
                     predBlock->bbNum, jti.m_block->bbNum, predBlock->bbNum, jti.m_falseTarget->bbNum);
+            newTarget = jti.m_falseTarget;
+        }
 
-            fgReplaceJumpTarget(predBlock, jti.m_block, jti.m_falseTarget);
+        fgReplaceJumpTarget(predBlock, jti.m_block, newTarget);
 
-            if (setNoCseIn && !jti.m_falseTarget->HasFlag(BBF_NO_CSE_IN))
-            {
-                JITDUMP(FMT_BB " => BBF_NO_CSE_IN\n", jti.m_falseTarget->bbNum);
-                jti.m_falseTarget->SetFlags(BBF_NO_CSE_IN);
-            }
+        if (setNoCseIn && !newTarget->HasFlag(BBF_NO_CSE_IN))
+        {
+            JITDUMP(FMT_BB " => BBF_NO_CSE_IN\n", newTarget->bbNum);
+            newTarget->SetFlags(BBF_NO_CSE_IN);
+        }
+
+        if (predBlock->hasProfileWeight())
+        {
+            newTarget->increaseBBProfileWeight(predEdge->getLikelyWeight());
+            modifiedProfile = true;
         }
     }
 
+    // jti.m_block is unreachable, but we won't remove it until the next flowgraph simplification pass.
+    // Mark the profile as inconsistent to pass the post-phase checks.
+    if (modifiedProfile)
+    {
+        JITDUMP("RBO: " FMT_BB
+                " is now unreachable, and flow into its successors needs to be removed. Data %s inconsistent.\n",
+                jti.m_block->bbNum, fgPgoConsistent ? "is now" : "was already");
+        fgPgoConsistent = false;
+    }
+
     // If block didn't get fully optimized, and now has just one pred, see if
     // we can sharpen the predicate's VN.
     //