From e3625337691ac87e30aa201e960943e8c601b0d3 Mon Sep 17 00:00:00 2001
From: Andy Ayers <andya@microsoft.com>
Date: Thu, 4 May 2023 19:02:33 -0700
Subject: [PATCH] JIT: prepare for instrumentation before incorporating profile
 counts

Otherwise the spanning tree we generate may be biased by the profile data
and not match the spanning tree we generated in Tier0.

Fixes #85799.
---
 src/coreclr/jit/block.h       | 12 +++++-------
 src/coreclr/jit/compiler.cpp  | 16 ++++++++--------
 src/coreclr/jit/fgprofile.cpp | 10 ++++++++++
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/src/coreclr/jit/block.h b/src/coreclr/jit/block.h
index 748ac2229e1ac6..7beab20b4af65c 100644
--- a/src/coreclr/jit/block.h
+++ b/src/coreclr/jit/block.h
@@ -1089,15 +1089,13 @@ struct BasicBlock : private LIR::Range
     BlockSet bbReach; // Set of all blocks that can reach this one
 
     union {
-        BasicBlock* bbIDom;   // Represent the closest dominator to this block (called the Immediate
-                              // Dominator) used to compute the dominance tree.
-        FlowEdge* bbLastPred; // Used early on by fgLinkBasicBlock/fgAddRefPred
+        BasicBlock* bbIDom;          // Represent the closest dominator to this block (called the Immediate
+                                     // Dominator) used to compute the dominance tree.
+        FlowEdge* bbLastPred;        // Used early on by fgLinkBasicBlock/fgAddRefPred
+        void*     bbSparseProbeList; // Used early on by fgInstrument
     };
 
-    union {
-        void* bbSparseCountInfo; // Used early on by fgIncorporateEdgeCounts
-        void* bbSparseProbeList; // Used early on by fgInstrument
-    };
+    void* bbSparseCountInfo; // Used early on by fgIncorporateEdgeCounts
 
     unsigned bbPreorderNum;  // the block's  preorder number in the graph (1...fgMaxBBNum]
     unsigned bbPostorderNum; // the block's postorder number in the graph (1...fgMaxBBNum]
diff --git a/src/coreclr/jit/compiler.cpp b/src/coreclr/jit/compiler.cpp
index 1bb68bfb975a1b..8a4a74f89e327b 100644
--- a/src/coreclr/jit/compiler.cpp
+++ b/src/coreclr/jit/compiler.cpp
@@ -4496,6 +4496,14 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     };
     DoPhase(this, PHASE_PRE_IMPORT, preImportPhase);
 
+    // If we're going to instrument code, we may need to prepare before
+    // we import. Also do this before we read in any profile data.
+    //
+    if (compileFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR))
+    {
+        DoPhase(this, PHASE_IBCPREP, &Compiler::fgPrepareToInstrumentMethod);
+    }
+
     // Incorporate profile data.
     //
     // Note: the importer is sensitive to block weights, so this has
@@ -4505,14 +4513,6 @@ void Compiler::compCompile(void** methodCodePtr, uint32_t* methodCodeSize, JitFl
     DoPhase(this, PHASE_INCPROFILE, &Compiler::fgIncorporateProfileData);
     activePhaseChecks &= ~PhaseChecks::CHECK_PROFILE;
 
-    // If we're going to instrument code, we may need to prepare before
-    // we import.
-    //
-    if (compileFlags->IsSet(JitFlags::JIT_FLAG_BBINSTR))
-    {
-        DoPhase(this, PHASE_IBCPREP, &Compiler::fgPrepareToInstrumentMethod);
-    }
-
     // If we are doing OSR, update flow to initially reach the appropriate IL offset.
     //
     if (opts.IsOSR())
diff --git a/src/coreclr/jit/fgprofile.cpp b/src/coreclr/jit/fgprofile.cpp
index 39c16f10314b7b..b391ae99dd4f09 100644
--- a/src/coreclr/jit/fgprofile.cpp
+++ b/src/coreclr/jit/fgprofile.cpp
@@ -1220,6 +1220,16 @@ void Compiler::WalkSpanningTree(SpanningTreeVisitor* visitor)
             break;
         }
     }
+
+    // Notify visitor of remaining blocks
+    //
+    for (BasicBlock* const block : Blocks())
+    {
+        if (!BlockSetOps::IsMember(this, marked, block->bbNum))
+        {
+            visitor->VisitBlock(block);
+        }
+    }
 }
 
 // Map a block into its schema key we will use for storing basic blocks.