JIT: Allow hot/cold splitting between a BBJ_COND block and its false target (#96431)

Next step for #93020. When doing hot/cold splitting, if the first cold block succeeds a BBJ_COND block (meaning the false target is the first cold block), we previously needed to insert a BBJ_ALWAYS block at the end of the hot section to unconditionally jump to the cold section. Since we will need to conditionally generate a jump to the false target depending on its location once bbFalseTarget can diverge from bbNext, this seemed like a nice opportunity to add that logic in, and instead generate a jump to the cold section by checking if a jump is needed to the false target, rather than by appending a BBJ_ALWAYS block to the hot section.

Author: amanasifkhalid

Diff for: src/coreclr/jit/block.cpp

@@ -294,12 +294,27 @@ bool BasicBlock::IsFirstColdBlock(Compiler* compiler) const
 // Returns:
 //    true if block is a BBJ_ALWAYS to the next block that we can fall into
 //
-bool BasicBlock::CanRemoveJumpToNext(Compiler* compiler)
+bool BasicBlock::CanRemoveJumpToNext(Compiler* compiler) const
 {
     assert(KindIs(BBJ_ALWAYS));
     return JumpsToNext() && !hasAlign() && !compiler->fgInDifferentRegions(this, bbTarget);
 }
 
+//------------------------------------------------------------------------
+// CanRemoveJumpToFalseTarget: determine if jump to false target can be omitted
+//
+// Arguments:
+//    compiler - current compiler instance
+//
+// Returns:
+//    true if block is a BBJ_COND that can fall into its false target
+//
+bool BasicBlock::CanRemoveJumpToFalseTarget(Compiler* compiler) const
+{
+    assert(KindIs(BBJ_COND));
+    return NextIs(bbFalseTarget) && !hasAlign() && !compiler->fgInDifferentRegions(this, bbFalseTarget);
+}
+
 //------------------------------------------------------------------------
 // checkPredListOrder: see if pred list is properly ordered
 //

Diff for: src/coreclr/jit/block.h

@@ -616,7 +616,9 @@ struct BasicBlock : private LIR::Range
 
     bool IsFirstColdBlock(Compiler* compiler) const;
 
-    bool CanRemoveJumpToNext(Compiler* compiler);
+    bool CanRemoveJumpToNext(Compiler* compiler) const;
+
+    bool CanRemoveJumpToFalseTarget(Compiler* compiler) const;
 
     unsigned GetTargetOffs() const
     {

Diff for: src/coreclr/jit/codegencommon.cpp

@@ -395,6 +395,13 @@ void CodeGen::genMarkLabelsForCodegen()
             case BBJ_COND:
                 JITDUMP("  " FMT_BB " : branch target\n", block->GetTrueTarget()->bbNum);
                 block->GetTrueTarget()->SetFlags(BBF_HAS_LABEL);
+
+                // If we need a jump to the false target, give it a label
+                if (!block->CanRemoveJumpToFalseTarget(compiler))
+                {
+                    JITDUMP("  " FMT_BB " : branch target\n", block->GetFalseTarget()->bbNum);
+                    block->GetFalseTarget()->SetFlags(BBF_HAS_LABEL);
+                }
                 break;
 
             case BBJ_SWITCH:

Diff for: src/coreclr/jit/codegenlinear.cpp

@@ -314,8 +314,8 @@ void CodeGen::genCodeForBBlist()
                 printf("\nThis is the start of the cold region of the method\n");
             }
 #endif
-            // We should never have a block that falls through into the Cold section
-            noway_assert(!block->Prev()->bbFallsThrough());
+            // We should never split call/finally pairs between hot/cold sections
+            noway_assert(!block->isBBCallFinallyPairTail());
 
             needLabel = true;
         }
@@ -2619,6 +2619,12 @@ void CodeGen::genCodeForJcc(GenTreeCC* jcc)
     assert(jcc->OperIs(GT_JCC));
 
     inst_JCC(jcc->gtCondition, compiler->compCurBB->GetTrueTarget());
+
+    // If we cannot fall into the false target, emit a jump to it
+    if (!compiler->compCurBB->CanRemoveJumpToFalseTarget(compiler))
+    {
+        inst_JMP(EJ_jmp, compiler->compCurBB->GetFalseTarget());
+    }
 }
 
 //------------------------------------------------------------------------

Diff for: src/coreclr/jit/compiler.h

@@ -6406,7 +6406,7 @@ class Compiler
     bool fgIsThrow(GenTree* tree);
 
 public:
-    bool fgInDifferentRegions(BasicBlock* blk1, BasicBlock* blk2);
+    bool fgInDifferentRegions(const BasicBlock* blk1, const BasicBlock* blk2) const;
 
 private:
     bool fgIsBlockCold(BasicBlock* block);

Diff for: src/coreclr/jit/flowgraph.cpp

@@ -635,7 +635,7 @@ bool Compiler::fgIsThrow(GenTree* tree)
  * It returns false when the blocks are both in the same regions
  */
 
-bool Compiler::fgInDifferentRegions(BasicBlock* blk1, BasicBlock* blk2)
+bool Compiler::fgInDifferentRegions(const BasicBlock* blk1, const BasicBlock* blk2) const
 {
     noway_assert(blk1 != nullptr);
     noway_assert(blk2 != nullptr);
@@ -3185,56 +3185,11 @@ PhaseStatus Compiler::fgDetermineFirstColdBlock()
             }
         }
 
-        // When the last Hot block fall through into the Cold section
-        // we may need to add a jump
-        //
-        if (prevToFirstColdBlock->bbFallsThrough())
+        // Don't split up call/finally pairs
+        if (prevToFirstColdBlock->isBBCallFinallyPair())
         {
-            switch (prevToFirstColdBlock->GetKind())
-            {
-                default:
-                    noway_assert(!"Unhandled jumpkind in fgDetermineFirstColdBlock()");
-                    break;
-
-                case BBJ_CALLFINALLY:
-                    // A BBJ_CALLFINALLY that falls through is always followed
-                    // by an empty BBJ_CALLFINALLYRET.
-                    //
-                    assert(prevToFirstColdBlock->isBBCallFinallyPair());
-                    firstColdBlock =
-                        firstColdBlock->Next(); // Note that this assignment could make firstColdBlock == nullptr
-                    break;
-
-                case BBJ_COND:
-                    //
-                    // This is a slightly more complicated case, because we will
-                    // probably need to insert a block to jump to the cold section.
-                    //
-
-                    // TODO-NoFallThrough: Below logic will need additional check once bbFalseTarget can diverge from
-                    // bbNext
-                    assert(prevToFirstColdBlock->FalseTargetIs(firstColdBlock));
-                    if (firstColdBlock->isEmpty() && firstColdBlock->KindIs(BBJ_ALWAYS))
-                    {
-                        // We can just use this block as the transitionBlock
-                        firstColdBlock = firstColdBlock->Next();
-                        // Note that this assignment could make firstColdBlock == NULL
-                    }
-                    else
-                    {
-                        BasicBlock* transitionBlock =
-                            fgNewBBafter(BBJ_ALWAYS, prevToFirstColdBlock, true, firstColdBlock);
-                        transitionBlock->inheritWeight(firstColdBlock);
-                        prevToFirstColdBlock->SetFalseTarget(transitionBlock);
-
-                        // Update the predecessor list for firstColdBlock
-                        fgReplacePred(firstColdBlock, prevToFirstColdBlock, transitionBlock);
-
-                        // Add prevToFirstColdBlock as a predecessor for transitionBlock
-                        fgAddRefPred(transitionBlock, prevToFirstColdBlock);
-                    }
-                    break;
-            }
+            // Note that this assignment could make firstColdBlock == nullptr
+            firstColdBlock = firstColdBlock->Next();
         }
     }