From e348e31127a50909c2e95860d41ba45530d7e080 Mon Sep 17 00:00:00 2001
From: "Aman Khalid (from Dev Box)" <amankhalid@microsoft.com>
Date: Fri, 6 Sep 2024 14:37:32 -0400
Subject: [PATCH 1/2] Add late layout pass

---
 src/coreclr/jit/compiler.cpp | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp
index 3f28580aba68a0..5f5e490c7f8590 100644
--- a/src/coreclr/jit/compiler.cpp
+++ b/src/coreclr/jit/compiler.cpp
@@ -5238,11 +5238,12 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     m_pLowering->FinalizeOutgoingArgSpace();
 
     // We can not add any new tracked variables after this point.
-    lvaTrackedFixed = true;
+    lvaTrackedFixed                    = true;
+    const unsigned numBlocksBeforeLSRA = fgBBcount;
 
     // Now that lowering is completed we can proceed to perform register allocation
     //
-    auto linearScanPhase = [this]() {
+    auto linearScanPhase = [this] {
         m_pLinearScan->doLinearScan();
     };
     DoPhase(this, PHASE_LINEAR_SCAN, linearScanPhase);
@@ -5252,8 +5253,25 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
 
     if (opts.OptimizationEnabled())
     {
-        // LSRA and stack level setting can modify the flowgraph.
-        // Now that it won't change, run post-layout optimizations.
+        // LSRA may introduce new blocks. If it does, rerun layout.
+        if (fgBBcount != numBlocksBeforeLSRA)
+        {
+            auto lateLayoutPhase = [this] {
+                fgDoReversePostOrderLayout();
+                fgMoveColdBlocks();
+                return PhaseStatus::MODIFIED_EVERYTHING;
+            };
+
+            DoPhase(this, PHASE_OPTIMIZE_LAYOUT, lateLayoutPhase);
+
+            if (fgFirstColdBlock != nullptr)
+            {
+                fgFirstColdBlock = nullptr;
+                DoPhase(this, PHASE_DETERMINE_FIRST_COLD_BLOCK, &Compiler::fgDetermineFirstColdBlock);
+            }
+        }
+
+        // Now that the flowgraph is finalized, run post-layout optimizations.
         DoPhase(this, PHASE_OPTIMIZE_POST_LAYOUT, &Compiler::optOptimizePostLayout);
     }
 

From dbb3187769f8cd4217d4221510fb01bc19433322 Mon Sep 17 00:00:00 2001
From: "Aman Khalid (from Dev Box)" <amankhalid@microsoft.com>
Date: Fri, 6 Sep 2024 16:41:57 -0400
Subject: [PATCH 2/2] Add config check

---
 src/coreclr/jit/compiler.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp
index 5f5e490c7f8590..54b6b5014e0f59 100644
--- a/src/coreclr/jit/compiler.cpp
+++ b/src/coreclr/jit/compiler.cpp
@@ -5254,7 +5254,7 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     if (opts.OptimizationEnabled())
     {
         // LSRA may introduce new blocks. If it does, rerun layout.
-        if (fgBBcount != numBlocksBeforeLSRA)
+        if ((fgBBcount != numBlocksBeforeLSRA) && JitConfig.JitDoReversePostOrderLayout())
         {
             auto lateLayoutPhase = [this] {
                 fgDoReversePostOrderLayout();