JIT: Move backward jumps to before their successors after RPO-based layout (#102461)

Part of #93020. In #102343, we noticed the RPO-based layout sometimes makes suboptimal decisions in terms of placing a block's hottest predecessor before it -- in particular, this affects loops that aren't entered at the top. To address this, after establishing a baseline RPO layout, fgMoveBackwardJumpsToSuccessors will try to move backward unconditional jumps to right behind their targets to create fallthrough, if the predecessor block is sufficiently hot.

Author: amanasifkhalid

src/coreclr/jit/compiler.h

@@ -6104,6 +6104,9 @@ class Compiler
     void fgDoReversePostOrderLayout();
     void fgMoveColdBlocks();
 
+    template <bool hasEH>
+    void fgMoveBackwardJumpsToSuccessors();
+
     bool fgFuncletsAreCold();
 
     PhaseStatus fgDetermineFirstColdBlock();

src/coreclr/jit/fgbasic.cpp

@@ -6154,7 +6154,7 @@ BasicBlock* Compiler::fgNewBBFromTreeAfter(
  */
 void Compiler::fgInsertBBbefore(BasicBlock* insertBeforeBlk, BasicBlock* newBlk)
 {
-    if (insertBeforeBlk->IsFirst())
+    if (fgFirstBB == insertBeforeBlk)
     {
         newBlk->SetNext(fgFirstBB);
 

src/coreclr/jit/fgopt.cpp

@@ -4535,6 +4535,101 @@ bool Compiler::fgReorderBlocks(bool useProfile)
 #pragma warning(pop)
 #endif
 
+//-----------------------------------------------------------------------------
+// fgMoveBackwardJumpsToSuccessors: Try to move backward unconditional jumps to fall into their successors.
+//
+// Template parameters:
+//    hasEH - If true, method has EH regions, so check that we don't try to move blocks in different regions
+//
+template <bool hasEH>
+void Compiler::fgMoveBackwardJumpsToSuccessors()
+{
+#ifdef DEBUG
+    if (verbose)
+    {
+        printf("*************** In fgMoveBackwardJumpsToSuccessors()\n");
+
+        printf("\nInitial BasicBlocks");
+        fgDispBasicBlocks(verboseTrees);
+        printf("\n");
+    }
+#endif // DEBUG
+
+    EnsureBasicBlockEpoch();
+    BlockSet visitedBlocks(BlockSetOps::MakeEmpty(this));
+    BlockSetOps::AddElemD(this, visitedBlocks, fgFirstBB->bbNum);
+
+    // Don't try to move the first block.
+    // Also, if we have a funclet region, don't bother reordering anything in it.
+    //
+    BasicBlock* next;
+    for (BasicBlock* block = fgFirstBB->Next(); block != fgFirstFuncletBB; block = next)
+    {
+        next = block->Next();
+        BlockSetOps::AddElemD(this, visitedBlocks, block->bbNum);
+
+        // Don't bother trying to move cold blocks
+        //
+        if (!block->KindIs(BBJ_ALWAYS) || block->isRunRarely())
+        {
+            continue;
+        }
+
+        // We will consider moving only backward jumps
+        //
+        BasicBlock* const target = block->GetTarget();
+        if ((block == target) || !BlockSetOps::IsMember(this, visitedBlocks, target->bbNum))
+        {
+            continue;
+        }
+
+        if (hasEH)
+        {
+            // Don't move blocks in different EH regions
+            //
+            if (!BasicBlock::sameEHRegion(block, target))
+            {
+                continue;
+            }
+
+            // block and target are in the same try/handler regions, and target is behind block,
+            // so block cannot possibly be the start of the region.
+            //
+            assert(!bbIsTryBeg(block) && !bbIsHandlerBeg(block));
+
+            // Don't change the entry block of an EH region
+            //
+            if (bbIsTryBeg(target) || bbIsHandlerBeg(target))
+            {
+                continue;
+            }
+        }
+
+        // We don't want to change the first block, so if the jump target is the first block,
+        // don't try moving this block before it.
+        // Also, if the target is cold, don't bother moving this block up to it.
+        //
+        if (target->IsFirst() || target->isRunRarely())
+        {
+            continue;
+        }
+
+        // If moving block will break up existing fallthrough behavior into target, make sure it's worth it
+        //
+        FlowEdge* const fallthroughEdge = fgGetPredForBlock(target, target->Prev());
+        if ((fallthroughEdge != nullptr) &&
+            (fallthroughEdge->getLikelyWeight() >= block->GetTargetEdge()->getLikelyWeight()))
+        {
+            continue;
+        }
+
+        // Move block to before target
+        //
+        fgUnlinkBlock(block);
+        fgInsertBBbefore(target, block);
+    }
+}
+
 //-----------------------------------------------------------------------------
 // fgDoReversePostOrderLayout: Reorder blocks using a greedy RPO traversal.
 //
@@ -4567,6 +4662,13 @@ void Compiler::fgDoReversePostOrderLayout()
             fgInsertBBafter(block, blockToMove);
         }
 
+        // The RPO established a good base layout, but in some cases, it might produce a subpar layout for loops.
+        // In particular, it may place the loop head after the loop exit, creating unnecessary branches.
+        // Fix this by moving unconditional backward jumps up to their targets,
+        // increasing the likelihood that the loop exit block is the last block in the loop.
+        //
+        fgMoveBackwardJumpsToSuccessors</* hasEH */ false>();
+
         return;
     }
 
@@ -4645,6 +4747,8 @@ void Compiler::fgDoReversePostOrderLayout()
         }
     }
 
+    fgMoveBackwardJumpsToSuccessors</* hasEH */ true>();
+
     // Fix up call-finally pairs
     //
     for (int i = 0; i < callFinallyPairs.Height(); i++)