JIT: Cache start of cold section for reuse in 3-opt layout (#111365)

During cold code motion, if fgMoveColdBlocks finds a call-finally pair that begins with a cold block, it will try to move the entire pair, regardless of whether the tail is cold or not. If the tail isn't cold, 3-opt layout's detection of the last hot block may terminate early, causing the cold part of an EH region to also be considered for reordering. Because 3-opt expects each EH subregion being reordered to be contiguous, this split causes asserts. Caching the cold region start point in fgMoveColdBlocks means we don't need to find the first cold block in 3-opt, thus removing the potential of getting it wrong.

While I was here, I decided to enhance fgMoveColdBlocks so that it can move cold call-finally pairs at the end of the hot main method body, and clean up 3-opt's initialization logic to reduce its allocation sizes. Fixes #111207.

Author: amanasifkhalid

src/coreclr/jit/fgopt.cpp

@@ -4828,7 +4828,7 @@ void Compiler::fgMoveColdBlocks()
         }
     };
 
-    BasicBlock* const lastMainBB = fgLastBBInMainFunction();
+    BasicBlock* lastMainBB = fgLastBBInMainFunction();
     if (lastMainBB->IsFirst())
     {
         return;
@@ -4855,15 +4855,41 @@ void Compiler::fgMoveColdBlocks()
 
     // We have moved all cold main blocks before lastMainBB to after lastMainBB.
     // If lastMainBB itself is cold, move it to the end of the method to restore its relative ordering.
+    // But first, we can't move just the tail of a call-finally pair,
+    // so point lastMainBB to the pair's head, if necessary.
     //
-    if (lastMainBB->isBBWeightCold(this) && !lastMainBB->isBBCallFinallyPairTail())
+    if (lastMainBB->isBBCallFinallyPairTail())
     {
+        lastMainBB = lastMainBB->Prev();
+    }
+
+    BasicBlock* lastHotBB = nullptr;
+    if (lastMainBB->isBBWeightCold(this))
+    {
+        // lastMainBB is cold, so the block behind it (if there is one) is the last hot block
+        //
+        lastHotBB = lastMainBB->Prev();
+
+        // Move lastMainBB
+        //
         BasicBlock* const newLastMainBB = fgLastBBInMainFunction();
         if (lastMainBB != newLastMainBB)
         {
             moveBlock(lastMainBB, newLastMainBB);
         }
     }
+    else
+    {
+        // lastMainBB isn't cold, so it (or its call-finally pair tail) the last hot block
+        //
+        lastHotBB = lastMainBB->isBBCallFinallyPair() ? lastMainBB->Next() : lastMainBB;
+    }
+
+    // Save the beginning of the cold section for later.
+    // If lastHotBB is null, there isn't a hot section,
+    // so there's no point in differentiating between sections for layout purposes.
+    //
+    fgFirstColdBlock = (lastHotBB == nullptr) ? nullptr : lastHotBB->Next();
 }
 
 //-----------------------------------------------------------------------------
@@ -5136,10 +5162,8 @@ void Compiler::ThreeOptLayout::ConsiderEdge(FlowEdge* edge)
     const unsigned srcPos = ordinals[srcBlk->bbPostorderNum];
     const unsigned dstPos = ordinals[dstBlk->bbPostorderNum];
 
-    // Don't consider edges from outside the hot range.
-    // If 'srcBlk' has an ordinal of zero and it isn't the first block,
-    // it's not tracked by 'ordinals', so it's not in the hot section.
-    if ((srcPos == 0) && !srcBlk->IsFirst())
+    // Don't consider edges to or from outside the hot range (i.e. ordinal doesn't match 'blockOrder' position).
+    if ((srcBlk != blockOrder[srcPos]) || (dstBlk != blockOrder[dstPos]))
     {
         return;
     }
@@ -5212,19 +5236,19 @@ void Compiler::ThreeOptLayout::AddNonFallthroughPreds(unsigned blockPos)
 //
 void Compiler::ThreeOptLayout::Run()
 {
-    // Walk backwards through the main method body, looking for the last hot block.
     // Since we moved all cold blocks to the end of the method already,
-    // we should have a span of hot blocks to consider reordering at the beginning of the method.
-    // While doing this, try to get as tight an upper bound for the number of hot blocks as possible.
-    // For methods without funclet regions, 'numBlocksUpperBound' is exact.
-    // Otherwise, it's off by the number of handler blocks.
-    BasicBlock* finalBlock;
-    unsigned    numBlocksUpperBound = compiler->fgBBcount;
-    for (finalBlock = compiler->fgLastBBInMainFunction();
-         !finalBlock->IsFirst() && finalBlock->isBBWeightCold(compiler); finalBlock = finalBlock->Prev())
-    {
-        numBlocksUpperBound--;
-    }
+    // we should have a span of hot blocks to consider reordering at the beginning of the method
+    // (unless none of the blocks are cold relative to the rest of the method,
+    // in which case we will reorder the whole main method body).
+    BasicBlock* const finalBlock = (compiler->fgFirstColdBlock != nullptr) ? compiler->fgFirstColdBlock->Prev()
+                                                                           : compiler->fgLastBBInMainFunction();
+
+    // Reset cold section pointer, in case we decide to do hot/cold splitting later
+    compiler->fgFirstColdBlock = nullptr;
+
+    // We better have an end block for the hot section, and it better not be the start of a call-finally pair.
+    assert(finalBlock != nullptr);
+    assert(!finalBlock->isBBCallFinallyPair());
 
     // For methods with fewer than three candidate blocks, we cannot partition anything
     if (finalBlock->IsFirst() || finalBlock->Prev()->IsFirst())
@@ -5233,39 +5257,28 @@ void Compiler::ThreeOptLayout::Run()
         return;
     }
 
-    // If only the first block of a call-finally pair is hot, include the whole pair in the hot section anyway.
-    // This ensures the call-finally pair won't be split up when swapping partitions.
-    if (finalBlock->isBBCallFinallyPair())
-    {
-        finalBlock = finalBlock->Next();
-        numBlocksUpperBound++;
-    }
-
+    // Get an upper bound on the number of hot blocks without walking the whole block list.
+    // We will only consider blocks reachable via normal flow.
+    const unsigned numBlocksUpperBound = compiler->m_dfsTree->GetPostOrderCount();
     assert(numBlocksUpperBound != 0);
-    blockOrder = new (compiler, CMK_BasicBlock) BasicBlock*[numBlocksUpperBound];
-    tempOrder  = new (compiler, CMK_BasicBlock) BasicBlock*[numBlocksUpperBound];
+    blockOrder = new (compiler, CMK_BasicBlock) BasicBlock*[numBlocksUpperBound * 2];
+    tempOrder  = (blockOrder + numBlocksUpperBound);
 
     // Initialize the current block order.
     // Note that we default-initialized 'ordinals' with zeros.
     // Block reordering shouldn't change the method's entry point,
     // so if a block has an ordinal of zero and it's not 'fgFirstBB',
-    // the block wasn't visited below, so it's not in the range of candidate blocks.
-    unsigned nextPostorderNum = compiler->m_dfsTree->GetPostOrderCount();
+    // the block wasn't visited below, meaning it's not in the range of candidate blocks.
     for (BasicBlock* const block : compiler->Blocks(compiler->fgFirstBB, finalBlock))
     {
-        assert(numCandidateBlocks < numBlocksUpperBound);
-        blockOrder[numCandidateBlocks] = tempOrder[numCandidateBlocks] = block;
-
-        // Unreachable blocks should have been pushed out of the candidate set of blocks.
-        // However, the entries of unreachable EH regions are left in-place to facilitate reestablishing contiguity,
-        // so it is possible for us to encounter unreachable blocks.
-        // When we do, assign them postorder numbers that can be used as keys into 'ordinals'.
         if (!compiler->m_dfsTree->Contains(block))
         {
-            assert(nextPostorderNum < compiler->fgBBcount);
-            block->bbPostorderNum = nextPostorderNum++;
+            continue;
         }
 
+        assert(numCandidateBlocks < numBlocksUpperBound);
+        blockOrder[numCandidateBlocks] = tempOrder[numCandidateBlocks] = block;
+
         assert(ordinals[block->bbPostorderNum] == 0);
         ordinals[block->bbPostorderNum] = numCandidateBlocks++;