JIT: Try to retain entry weight during profile synthesis

src/coreclr/jit/fgprofile.cpp

@@ -2942,7 +2942,7 @@ PhaseStatus Compiler::fgIncorporateProfileData()
 
         // We now always run repair, to get consistent initial counts
         //
-        JITDUMP("\n%sRepairing profile...\n", opts.IsOSR() ? "blending" : "repairing");
+        JITDUMP("\nRepairing profile...\n");
         ProfileSynthesis::Run(this, ProfileSynthesisOption::RepairLikelihoods);
     }
 

src/coreclr/jit/fgprofilesynthesis.cpp

@@ -33,6 +33,7 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
     m_dfsTree             = m_comp->fgComputeDfs();
     m_loops               = FlowGraphNaturalLoops::Find(m_dfsTree);
     m_improperLoopHeaders = m_loops->ImproperLoopHeaders();
+    m_entryBlock          = m_comp->opts.IsOSR() ? m_comp->fgEntryBB : m_comp->fgFirstBB;
 
     // Retain or compute edge likelihood information
     //
@@ -71,13 +72,18 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
             break;
     }
 
+    // Save entry block's weight.
+    // If the entry block is a loop header, its weight will be overwritten by ComputeCyclicProbabilities.
+    //
+    weight_t entryBlockWeight = m_entryBlock->bbWeight;
+
     // Determine cyclic probabilities
     //
     ComputeCyclicProbabilities();
 
     // Assign weights to entry points in the flow graph
     //
-    AssignInputWeights(option);
+    AssignInputWeights(entryBlockWeight);
 
     // Compute the block weights given the inputs and edge likelihoods
     //
@@ -108,12 +114,13 @@ void ProfileSynthesis::Run(ProfileSynthesisOption option)
         m_approximate               = false;
         m_overflow                  = false;
         m_cappedCyclicProbabilities = 0;
+        entryBlockWeight            = m_entryBlock->bbWeight;
 
         // Regularize the edge likelihoods...
         //
         BlendLikelihoods();
         ComputeCyclicProbabilities();
-        AssignInputWeights(option);
+        AssignInputWeights(entryBlockWeight);
         ComputeBlockWeights();
 
         // Increase blend factor and decrease synthetic loop likelihoods
@@ -975,7 +982,7 @@ void ProfileSynthesis::ComputeCyclicProbabilities(FlowGraphNaturalLoop* loop)
 // fgAssignInputWeights: provide initial profile weights for all blocks
 //
 // Arguments:
-//   option - profile synthesis option
+//   entryBlockWeight - total flow (including method call count) into the entry block
 //
 // Notes:
 //   For finallys we will pick up new entry weights when we process
@@ -986,51 +993,26 @@ void ProfileSynthesis::ComputeCyclicProbabilities(FlowGraphNaturalLoop* loop)
 //
 //   Some parts of the jit are sensitive to the absolute weights.
 //
-void ProfileSynthesis::AssignInputWeights(ProfileSynthesisOption option)
+void ProfileSynthesis::AssignInputWeights(weight_t entryBlockWeight)
 {
-    // Determine input weight for method entry
+    // Determine input weight for method entry.
+    // Ideally, we'd use fgCalledCount, but it may not be available yet.
     //
-    BasicBlock* const entryBlock  = m_comp->opts.IsOSR() ? m_comp->fgEntryBB : m_comp->fgFirstBB;
-    weight_t          entryWeight = BB_UNITY_WEIGHT;
+    weight_t                    entryWeight = entryBlockWeight;
+    FlowGraphNaturalLoop* const loop        = m_loops->GetLoopByHeader(m_entryBlock);
 
-    switch (option)
+    if (loop != nullptr)
     {
-        case ProfileSynthesisOption::BlendLikelihoods:
-        case ProfileSynthesisOption::RepairLikelihoods:
-        {
-            // Try and retain entryBlock's weight.
-            // Easiest to do when the block has no preds.
-            //
-            if (entryBlock->hasProfileWeight())
-            {
-                weight_t currentEntryWeight = entryBlock->bbWeight;
-
-                if (!Compiler::fgProfileWeightsEqual(currentEntryWeight, 0.0, epsilon))
-                {
-                    if (entryBlock->bbPreds == nullptr)
-                    {
-                        entryWeight = currentEntryWeight;
-                    }
-                    else
-                    {
-                        // TODO: something similar to how we compute fgCalledCount;
-                        // try and sum return weights?
-                    }
-                }
-                else
-                {
-                    // Entry weight was zero or nearly zero, just use default
-                }
-            }
-            else
-            {
-                // Entry was unprofiled, just use default
-            }
-            break;
-        }
+        const weight_t cyclicProbability = m_cyclicProbabilities[loop->GetIndex()];
+        assert(cyclicProbability != BB_ZERO_WEIGHT);
+        entryWeight /= cyclicProbability;
+    }
 
-        default:
-            break;
+    // Fall back to BB_UNITY_WEIGHT if we have zero entry weight
+    //
+    if (Compiler::fgProfileWeightsEqual(entryWeight, BB_ZERO_WEIGHT, epsilon))
+    {
+        entryWeight = BB_UNITY_WEIGHT;
     }
 
     // Reset existing weights
@@ -1042,8 +1024,8 @@ void ProfileSynthesis::AssignInputWeights(ProfileSynthesisOption option)
 
     // Set entry weight
     //
-    JITDUMP("Synthesis: entry " FMT_BB " has input weight " FMT_WT "\n", entryBlock->bbNum, entryWeight);
-    entryBlock->setBBProfileWeight(entryWeight);
+    JITDUMP("Synthesis: entry " FMT_BB " has input weight " FMT_WT "\n", m_entryBlock->bbNum, entryWeight);
+    m_entryBlock->setBBProfileWeight(entryWeight);
 
     // Determine input weight for EH regions, if any.
     //
@@ -1210,9 +1192,6 @@ void ProfileSynthesis::GaussSeidelSolver()
     bool                          checkEntryExitWeight = true;
     bool                          showDetails          = false;
 
-    // Remember the entry block
-    //
-    BasicBlock* const entryBlock = m_comp->opts.IsOSR() ? m_comp->fgEntryBB : m_comp->fgFirstBB;
     JITDUMP("Synthesis solver: flow graph has %u improper loop headers\n", m_improperLoopHeaders);
 
     // This is an iterative solver, and it may require a lot of iterations
@@ -1268,7 +1247,7 @@ void ProfileSynthesis::GaussSeidelSolver()
 
             // Some blocks have additional profile weights that don't come from flow edges.
             //
-            if (block == entryBlock)
+            if (block == m_entryBlock)
             {
                 newWeight   = block->bbWeight;
                 entryWeight = newWeight;
@@ -1459,7 +1438,7 @@ void ProfileSynthesis::GaussSeidelSolver()
             if (entryExitRelResidual > relResidual)
             {
                 relResidual      = entryExitRelResidual;
-                relResidualBlock = entryBlock;
+                relResidualBlock = m_entryBlock;
             }
         }
 

src/coreclr/jit/fgprofilesynthesis.h

@@ -75,7 +75,7 @@ class ProfileSynthesis
     void ComputeCyclicProbabilities();
     void ComputeCyclicProbabilities(FlowGraphNaturalLoop* loop);
 
-    void AssignInputWeights(ProfileSynthesisOption option);
+    void AssignInputWeights(weight_t entryBlockWeight);
 
     void ComputeBlockWeights();
     void ComputeBlockWeight(BasicBlock* block);
@@ -86,6 +86,7 @@ class ProfileSynthesis
     Compiler* const        m_comp;
     FlowGraphDfsTree*      m_dfsTree                   = nullptr;
     FlowGraphNaturalLoops* m_loops                     = nullptr;
+    BasicBlock*            m_entryBlock                = nullptr;
     weight_t*              m_cyclicProbabilities       = nullptr;
     weight_t               m_blendFactor               = initialBlendFactor;
     weight_t               m_loopExitLikelihood        = loopExitLikelihood;