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).

diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
index d533082..b9c0e7c 100644 (file)
--- a/src/coreclr/jit/compiler.h
+++ b/src/coreclr/jit/compiler.h
@@ -4420,7 +4420,9 @@ public:
     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.

diff --git a/src/coreclr/jit/fgbasic.cpp b/src/coreclr/jit/fgbasic.cpp
index f05a6f9..50ee88c 100644 (file)
--- a/src/coreclr/jit/fgbasic.cpp
+++ b/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;

diff --git a/src/coreclr/jit/fgdiagnostic.cpp b/src/coreclr/jit/fgdiagnostic.cpp
index 15ccb9e..884e3e9 100644 (file)
--- a/src/coreclr/jit/fgdiagnostic.cpp
+++ b/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 */

diff --git a/src/coreclr/jit/importer.cpp b/src/coreclr/jit/importer.cpp
index b05f31b..81f6d91 100644 (file)
--- a/src/coreclr/jit/importer.cpp
+++ b/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)

diff --git a/src/coreclr/jit/importercalls.cpp b/src/coreclr/jit/importercalls.cpp
index cb75aef..09ab191 100644 (file)
--- a/src/coreclr/jit/importercalls.cpp
+++ b/src/coreclr/jit/importercalls.cpp
@@ -1295,34 +1295,6 @@ DONE:
                 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);
-                }
-            }
         }
     }
 

diff --git a/src/coreclr/jit/lclvars.cpp b/src/coreclr/jit/lclvars.cpp
index 59fb7ba..eb8dde8 100644 (file)
--- a/src/coreclr/jit/lclvars.cpp
+++ b/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));
-                }
-            }
         }
     }
 

diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp
index 3f5ecf9..b578928 100644 (file)
--- a/src/coreclr/jit/morph.cpp
+++ b/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;
     }
 

diff --git a/src/tests/JIT/opt/OSR/exposure1.cs b/src/tests/JIT/opt/OSR/exposure1.cs
new file mode 100644 (file)
index 0000000..a21459b
--- /dev/null
+++ b/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;
+    }
+}

diff --git a/src/tests/JIT/opt/OSR/exposure1.csproj b/src/tests/JIT/opt/OSR/exposure1.csproj
new file mode 100644 (file)
index 0000000..fe316c4
--- /dev/null
+++ b/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>

diff --git a/src/tests/JIT/opt/OSR/exposure2.cs b/src/tests/JIT/opt/OSR/exposure2.cs
new file mode 100644 (file)
index 0000000..12abe16
--- /dev/null
+++ b/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;
+    }
+}

diff --git a/src/tests/JIT/opt/OSR/exposure2.csproj b/src/tests/JIT/opt/OSR/exposure2.csproj
new file mode 100644 (file)
index 0000000..fe316c4
--- /dev/null
+++ b/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>