JIT: import entire method for OSR, prune unneeded parts later

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

@@ -13840,14 +13840,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>