JIT: import entire method for OSR, prune unneeded parts later (#83910)

For OSR compiles, always import from the original entry point in addition
to the OSR entry point. This gives the OSR compiler a chance
to see all of the method and so properly compute address exposure,
instead of relying on the Tier0
analysis.

Once address exposure has been determined, revoke special protection
for the original entry and try and prune away blocks that are no longer
needed (we actually wait until morph).

Fixes #83783.

May also fix some of the cases where OSR perf is lagging (though
don't expect this to fix them all).

src/coreclr/jit/compiler.h

@@ -4420,7 +4420,9 @@ class Compiler
  DomTreeNode* fgSsaDomTree;
 
  bool fgBBVarSetsInited;
- bool fgOSROriginalEntryBBProtected;
+
+ // Track how many artificial ref counts we've added to fgEntryBB (for OSR)
+ unsigned fgEntryBBExtraRefs;
 
  // Allocate array like T* a = new T[fgBBNumMax + 1];
  // Using helper so we don't keep forgetting +1.

src/coreclr/jit/fgbasic.cpp

@@ -43,12 +43,12 @@ void Compiler::fgInit()
 
  /* Initialize the basic block list */
 
- fgFirstBB  = nullptr;
- fgLastBB  = nullptr;
- fgFirstColdBlock  = nullptr;
- fgEntryBB  = nullptr;
- fgOSREntryBB  = nullptr;
- fgOSROriginalEntryBBProtected = false;
+ fgFirstBB = nullptr;
+ fgLastBB = nullptr;
+ fgFirstColdBlock = nullptr;
+ fgEntryBB = nullptr;
+ fgOSREntryBB = nullptr;
+ fgEntryBBExtraRefs = 0;
 
 #if defined(FEATURE_EH_FUNCLETS)
  fgFirstFuncletBB = nullptr;

src/coreclr/jit/fgdiagnostic.cpp

@@ -2888,11 +2888,12 @@ void Compiler::fgDebugCheckBBlist(bool checkBBNum /* = false */, bool checkBBRef
  blockRefs += 1;
  }
 
- // Under OSR, if we also are keeping the original method entry around,
- // mark that as implicitly referenced as well.
- if (opts.IsOSR() && (block == fgEntryBB) && fgOSROriginalEntryBBProtected)
+ // Under OSR, if we also are keeping the original method entry around
+ // via artifical ref counts, account for those.
+ //
+ if (opts.IsOSR() && (block == fgEntryBB))
  {
- blockRefs += 1;
+ blockRefs += fgEntryBBExtraRefs;
  }
 
  /* Check the bbRefs */

src/coreclr/jit/importer.cpp

@@ -12624,6 +12624,17 @@ void Compiler::impImport()
  // which we should verify over when we find jump targets.
  impImportBlockPending(entryBlock);
 
+ if (opts.IsOSR())
+ {
+ // We now import all the IR and keep it around so we can
+ // analyze address exposure more robustly.
+ //
+ JITDUMP("OSR: protecting original method entry " FMT_BB "\n", fgEntryBB->bbNum);
+ impImportBlockPending(fgEntryBB);
+ fgEntryBB->bbRefs++;
+ fgEntryBBExtraRefs++;
+ }
+
  /* Import blocks in the worker-list until there are no more */
 
  while (impPendingList)

src/coreclr/jit/importercalls.cpp

@@ -1295,34 +1295,6 @@ var_types Compiler::impImportCall(OPCODE opcode,
  successor->bbFlags |= BBF_TAILCALL_SUCCESSOR;
  optMethodFlags |= OMF_HAS_TAILCALL_SUCCESSOR;
  }
-
- // If this call might eventually turn into a loop back to method entry, make sure we
- // import the method entry.
- //
- assert(call->IsCall());
- GenTreeCall* const actualCall = call->AsCall();
- const bool mustImportEntryBlock = gtIsRecursiveCall(methHnd) || actualCall->IsInlineCandidate() ||
- actualCall->IsGuardedDevirtualizationCandidate();
-
- // Only schedule importation if we're not currently importing the entry BB.
- //
- if (opts.IsOSR() && mustImportEntryBlock && (compCurBB != fgEntryBB))
- {
- JITDUMP("\nOSR: inlineable or recursive tail call [%06u] in the method, so scheduling " FMT_BB
- " for importation\n",
- dspTreeID(call), fgEntryBB->bbNum);
- impImportBlockPending(fgEntryBB);
-
- if (!fgOSROriginalEntryBBProtected && (fgEntryBB != fgFirstBB))
- {
- // Protect fgEntryBB from deletion, since it may not have any
- // explicit flow references until morph.
- //
- fgEntryBB->bbRefs += 1;
- fgOSROriginalEntryBBProtected = true;
- JITDUMP(" also protecting original method entry " FMT_BB "\n", fgEntryBB->bbNum);
- }
- }
  }
  }
 

src/coreclr/jit/lclvars.cpp

@@ -323,7 +323,7 @@ void Compiler::lvaInitTypeRef()
  }
  }
 
- // If this is an OSR method, mark all the OSR locals and model OSR exposure.
+ // If this is an OSR method, mark all the OSR locals.
  //
  // Do this before we add the GS Cookie Dummy or Outgoing args to the locals
  // so we don't have to do special checks to exclude them.
@@ -334,19 +334,6 @@ void Compiler::lvaInitTypeRef()
  {
  LclVarDsc* const varDsc = lvaGetDesc(lclNum);
  varDsc->lvIsOSRLocal = true;
-
- if (info.compPatchpointInfo->IsExposed(lclNum))
- {
- JITDUMP("-- V%02u is OSR exposed\n", lclNum);
- varDsc->lvHasLdAddrOp = 1;
-
- // todo: Why does it apply only to non-structs?
- //
- if (!varTypeIsStruct(varDsc) && !varTypeIsSIMD(varDsc))
- {
- lvaSetVarAddrExposed(lclNum DEBUGARG(AddressExposedReason::OSR_EXPOSED));
- }
- }
  }
  }
 

src/coreclr/jit/morph.cpp

@@ -13865,14 +13865,13 @@ void Compiler::fgMorphBlocks()
  //
  if (opts.IsOSR() && (fgEntryBB != nullptr))
  {
- if (fgOSROriginalEntryBBProtected)
- {
- JITDUMP("OSR: un-protecting original method entry " FMT_BB "\n", fgEntryBB->bbNum);
- assert(fgEntryBB->bbRefs > 0);
- fgEntryBB->bbRefs--;
- }
+ JITDUMP("OSR: un-protecting original method entry " FMT_BB "\n", fgEntryBB->bbNum);
+ assert(fgEntryBBExtraRefs == 1);
+ assert(fgEntryBB->bbRefs >= 1);
+ fgEntryBB->bbRefs--;
+ fgEntryBBExtraRefs = 0;
 
- // And we don't need to remember this block anymore.
+ // We don't need to remember this block anymore.
  fgEntryBB = nullptr;
  }
 

src/tests/JIT/opt/OSR/exposure1.cs

@@ -0,0 +1,47 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Runtime.CompilerServices;
+
+// Runtime 83738: need to ensure that 's' in 'Foo'
+// is marked as address exposed during OSR compiles.
+
+public class Exposure1
+{
+ [MethodImpl(MethodImplOptions.NoInlining)]
+ static void Bar()
+ {
+ }
+
+ [MethodImpl(MethodImplOptions.NoInlining)]
+ static bool Foo(int n)
+ {
+ S s = new S { F = 1234 };
+ ref int foo = ref s.F;
+
+ for (int i = 0; i < n; i++)
+ {
+ Bar();
+ }
+
+ int abc = s.F * 3 + 4;
+ foo = 25;
+ int def = s.F * 3 + 4;
+
+ int eabc = 1234 * 3 + 4;
+ int edef = 25 * 3 + 4;
+ Console.WriteLine("abc = {0} (expected {1}), def = {2} (expected {3})", abc, eabc, def, edef);
+ return (abc == eabc && def == edef);
+ }
+
+ public static int Main()
+ {
+ return Foo(50000) ? 100 : -1;
+ }
+
+ public struct S
+ {
+ public int F;
+ }
+}

src/tests/JIT/opt/OSR/exposure1.csproj

@@ -0,0 +1,13 @@
+<Project Sdk="Microsoft.NET.Sdk">
+ <PropertyGroup>
+ <OutputType>Exe</OutputType>
+ <DebugType />
+ <Optimize>True</Optimize>
+ </PropertyGroup>
+ <ItemGroup>
+ <Compile Include="$(MSBuildProjectName).cs" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TieredCompilation" Value="1" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TC_QuickJitForLoops" Value="1" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TC_OnStackReplacement" Value="1" />
+ </ItemGroup>
+</Project>

src/tests/JIT/opt/OSR/exposure2.cs

@@ -0,0 +1,47 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Runtime.CompilerServices;
+
+// Runtime 83738: need to ensure that 's' in 'Foo'
+// is marked as address exposed during OSR compiles.
+
+public class Exposure2
+{
+ [MethodImpl(MethodImplOptions.NoInlining)]
+ static void Bar()
+ {
+ }
+
+ [MethodImpl(MethodImplOptions.NoInlining)]
+ static bool Foo(int n)
+ {
+ S s = new S { F = 1234 };
+ ref S foo = ref s;
+
+ for (int i = 0; i < n; i++)
+ {
+ Bar();
+ }
+
+ int abc = s.F * 3 + 4;
+ foo.F = 25;
+ int def = s.F * 3 + 4;
+
+ int eabc = 1234 * 3 + 4;
+ int edef = 25 * 3 + 4;
+ Console.WriteLine("abc = {0} (expected {1}), def = {2} (expected {3})", abc, eabc, def, edef);
+ return (abc == eabc && def == edef);
+ }
+
+ public static int Main()
+ {
+ return Foo(50000) ? 100 : -1;
+ }
+
+ public struct S
+ {
+ public int F;
+ }
+}

src/tests/JIT/opt/OSR/exposure2.csproj

@@ -0,0 +1,13 @@
+<Project Sdk="Microsoft.NET.Sdk">
+ <PropertyGroup>
+ <OutputType>Exe</OutputType>
+ <DebugType />
+ <Optimize>True</Optimize>
+ </PropertyGroup>
+ <ItemGroup>
+ <Compile Include="$(MSBuildProjectName).cs" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TieredCompilation" Value="1" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TC_QuickJitForLoops" Value="1" />
+ <CLRTestEnvironmentVariable Include="DOTNET_TC_OnStackReplacement" Value="1" />
+ </ItemGroup>
+</Project>