[RuntimeAsync] Fix support for pinvokes

src/coreclr/jit/async.cpp

@@ -60,6 +60,16 @@ class AsyncLiveness
             return true;
         }
 
+        if (lclNum == m_comp->info.compLvFrameListRoot)
+        {
+            return true;
+        }
+
+        if (lclNum == m_comp->lvaInlinedPInvokeFrameVar)
+        {
+            return true;
+        }
+
 #ifdef FEATURE_EH_WINDOWS_X86
         if (lclNum == m_comp->lvaShadowSPslotsVar)
         {
@@ -1296,12 +1306,22 @@ void Async2Transformation::CreateResumptionSwitch()
 {
     BasicBlock* newEntryBB = BasicBlock::New(m_comp, BBJ_ALWAYS);
 
+    bool hadScratchBB = m_comp->fgFirstBBisScratch();
+
     if (m_comp->fgFirstBB->hasProfileWeight())
     {
         newEntryBB->inheritWeight(m_comp->fgFirstBB);
     }
     m_comp->fgFirstBB->bbRefs--;
 
+    if (hadScratchBB)
+    {
+        // If previous first BB was a scratch BB then we have to recreate it
+        // after since later phases may be relying on it.
+        // TODO-Cleanup: The later phases should be creating it if they need it.
+        m_comp->fgFirstBBScratch = nullptr;
+    }
+
     FlowEdge* toPrevEntryEdge = m_comp->fgAddRefPred(m_comp->fgFirstBB, newEntryBB);
     toPrevEntryEdge->setLikelihood(1);
 
@@ -1310,10 +1330,6 @@ void Async2Transformation::CreateResumptionSwitch()
 
     m_comp->fgInsertBBbefore(m_comp->fgFirstBB, newEntryBB);
 
-    // If previous first BB was a scratch BB, then we must add a new scratch BB
-    // to create IR before the switch.
-    m_comp->fgFirstBBScratch = nullptr;
-
     GenTree* continuationArg = m_comp->gtNewLclvNode(m_comp->lvaAsyncContinuationArg, TYP_REF);
     GenTree* null            = m_comp->gtNewNull();
     GenTree* neNull          = m_comp->gtNewOperNode(GT_NE, TYP_INT, continuationArg, null);
@@ -1493,4 +1509,9 @@ void Async2Transformation::CreateResumptionSwitch()
         GenTree* jtrue            = m_comp->gtNewOperNode(GT_JTRUE, TYP_VOID, ltZero);
         LIR::AsRange(checkILOffsetBB).InsertAtEnd(LIR::SeqTree(m_comp, jtrue));
     }
+
+    if (hadScratchBB)
+    {
+        m_comp->fgEnsureFirstBBisScratch();
+    }
 }

src/coreclr/jit/compiler.cpp

@@ -3368,7 +3368,7 @@ void Compiler::compInitOptions(JitFlags* jitFlags)
 
         if (compIsAsync2())
         {
-            printf("OPTIONS: compilation is an async2\n");
+            printf("OPTIONS: compilation is an async2 state machine\n");
         }
     }
 #endif

src/coreclr/jit/lower.cpp

@@ -4888,7 +4888,6 @@ void Lowering::LowerRet(GenTreeOp* ret)
         }
     }
 
-    // Method doing PInvokes has exactly one return block unless it has tail calls.
     if (comp->compMethodRequiresPInvokeFrame())
     {
         InsertPInvokeMethodEpilog(comp->compCurBB DEBUGARG(ret));
@@ -5332,6 +5331,11 @@ void Lowering::LowerReturnSuspend(GenTree* node)
     {
         BlockRange().Remove(BlockRange().LastNode(), true);
     }
+
+    if (comp->compMethodRequiresPInvokeFrame())
+    {
+        InsertPInvokeMethodEpilog(comp->compCurBB DEBUGARG(node));
+    }
 }
 
 //----------------------------------------------------------------------------------------------

src/tests/async/pinvoke.cs

@@ -0,0 +1,44 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using System.Threading.Tasks;
+using Xunit;
+
+public class Async2PInvoke
+{
+    [Fact]
+    public static void TestEntryPoint()
+    {
+        AsyncEntryPoint().Wait();
+    }
+
+    private static async2 Task AsyncEntryPoint()
+    {
+        unsafe
+        {
+            Assert.Equal(5, GetFPtr()());
+        }
+
+        await Task.Yield();
+
+        unsafe
+        {
+            Assert.Equal(5, GetFPtr()());
+        }
+
+        await Task.Yield();
+
+        unsafe
+        {
+            Assert.Equal(5, GetFPtr()());
+        }
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static unsafe delegate* unmanaged<int> GetFPtr() => &GetValue;
+
+    [UnmanagedCallersOnly]
+    private static int GetValue() => 5;
+}

src/tests/async/pinvoke.csproj

@@ -0,0 +1,8 @@
+<Project Sdk="Microsoft.NET.Sdk.IL">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>