JIT: Handle async calls implemented via ldvirtftn properly

src/coreclr/jit/compiler.h

@@ -4599,6 +4599,10 @@ class Compiler
                             CORINFO_CALL_INFO* callInfo,
                             IL_OFFSET          rawILOffset);
 
+    void impSetupAndSpillForAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags);
+
+    void impInsertAsyncContinuationForLdvirtftnCall(GenTreeCall* call);
+
     CORINFO_CLASS_HANDLE impGetSpecialIntrinsicExactReturnType(GenTreeCall* call);
 
     GenTree* impFixupCallStructReturn(GenTreeCall* call, CORINFO_CLASS_HANDLE retClsHnd);

src/coreclr/jit/importercalls.cpp

@@ -384,8 +384,18 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
                 // take the call now....
                 call = gtNewIndCallNode(nullptr, callRetTyp, di);
 
+                if (sig->isAsyncCall())
+                {
+                    impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags);
+                }
+
                 impPopCallArgs(sig, call->AsCall());
 
+                if (call->AsCall()->IsAsync())
+                {
+                    impInsertAsyncContinuationForLdvirtftnCall(call->AsCall());
+                }
+
                 GenTree* thisPtr = impPopStack().val;
                 thisPtr          = impTransformThis(thisPtr, pConstrainedResolvedToken, callInfo->thisTransform);
                 assert(thisPtr != nullptr);
@@ -681,68 +691,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
 
     if (sig->isAsyncCall())
     {
-        AsyncCallInfo asyncInfo;
-
-        if ((prefixFlags & PREFIX_IS_TASK_AWAIT) != 0)
-        {
-            JITDUMP("Call is an async task await\n");
-
-            asyncInfo.ExecutionContextHandling                  = ExecutionContextHandling::SaveAndRestore;
-            asyncInfo.SaveAndRestoreSynchronizationContextField = true;
-
-            if ((prefixFlags & PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT) != 0)
-            {
-                asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnCapturedContext;
-                JITDUMP("  Continuation continues on captured context\n");
-            }
-            else
-            {
-                asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnThreadPool;
-                JITDUMP("  Continuation continues on thread pool\n");
-            }
-        }
-        else if (opcode == CEE_CALLI)
-        {
-            // Used for unboxing/instantiating stubs
-            JITDUMP("Call is an async calli\n");
-        }
-        else
-        {
-            JITDUMP("Call is an async non-task await\n");
-            // Only expected non-task await to see in IL is one of the AsyncHelpers.AwaitAwaiter variants.
-            // These are awaits of custom awaitables, and they come with the behavior that the execution context
-            // is captured and restored on suspension/resumption.
-            // We could perhaps skip this for AwaitAwaiter (but not for UnsafeAwaitAwaiter) since it is expected
-            // that the safe INotifyCompletion will take care of flowing ExecutionContext.
-            asyncInfo.ExecutionContextHandling = ExecutionContextHandling::AsyncSaveAndRestore;
-        }
-
-        // For tailcalls the contexts does not need saving/restoring: they will be
-        // overwritten by the caller anyway.
-        //
-        // More specifically, if we can show that
-        // Thread.CurrentThread._executionContext is not accessed between the
-        // call and returning then we can omit save/restore of the execution
-        // context. We do not do that optimization yet.
-        if (tailCallFlags != 0)
-        {
-            asyncInfo.ExecutionContextHandling                  = ExecutionContextHandling::None;
-            asyncInfo.ContinuationContextHandling               = ContinuationContextHandling::None;
-            asyncInfo.SaveAndRestoreSynchronizationContextField = false;
-        }
-
-        call->AsCall()->SetIsAsync(new (this, CMK_Async) AsyncCallInfo(asyncInfo));
-
-        if (asyncInfo.ExecutionContextHandling == ExecutionContextHandling::SaveAndRestore)
-        {
-            compMustSaveAsyncContexts = true;
-
-            // In this case we will need to save the context after the arguments are evaluated.
-            // Spill the arguments to accomplish that.
-            // (We could do this via splitting in SaveAsyncContexts, but since we need to
-            //  handle inline candidates we won't gain much.)
-            impSpillSideEffects(true, CHECK_SPILL_ALL DEBUGARG("Async await with execution context save and restore"));
-        }
+        impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags);
     }
 
     // Now create the argument list.
@@ -6729,7 +6678,7 @@ bool Compiler::impCanPInvokeInlineCallSite(BasicBlock* block)
 //   call passes a combination of legality and profitability checks.
 //
 //   If GTF_CALL_UNMANAGED is set, increments info.compUnmanagedCallCountWithGCTransition
-
+//
 void Compiler::impCheckForPInvokeCall(
     GenTreeCall* call, CORINFO_METHOD_HANDLE methHnd, CORINFO_SIG_INFO* sig, unsigned mflags, BasicBlock* block)
 {
@@ -6850,6 +6799,110 @@ void Compiler::impCheckForPInvokeCall(
     }
 }
 
+//------------------------------------------------------------------------
+// impSetupAndSpillForAsyncCall:
+//   Register a call as being async and set up context handling information depending on the IL.
+//   Also spill IL arguments if necessary.
+//
+// Arguments:
+//    call        - The call
+//    opcode      - The IL opcode for the call
+//    prefixFlags - Flags containing context handling information from IL
+//
+void Compiler::impSetupAndSpillForAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags)
+{
+    AsyncCallInfo asyncInfo;
+
+    if ((prefixFlags & PREFIX_IS_TASK_AWAIT) != 0)
+    {
+        JITDUMP("Call is an async task await\n");
+
+        asyncInfo.ExecutionContextHandling                  = ExecutionContextHandling::SaveAndRestore;
+        asyncInfo.SaveAndRestoreSynchronizationContextField = true;
+
+        if ((prefixFlags & PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT) != 0)
+        {
+            asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnCapturedContext;
+            JITDUMP("  Continuation continues on captured context\n");
+        }
+        else
+        {
+            asyncInfo.ContinuationContextHandling = ContinuationContextHandling::ContinueOnThreadPool;
+            JITDUMP("  Continuation continues on thread pool\n");
+        }
+    }
+    else if (opcode == CEE_CALLI)
+    {
+        // Used for unboxing/instantiating stubs
+        JITDUMP("Call is an async calli\n");
+    }
+    else
+    {
+        JITDUMP("Call is an async non-task await\n");
+        // Only expected non-task await to see in IL is one of the AsyncHelpers.AwaitAwaiter variants.
+        // These are awaits of custom awaitables, and they come with the behavior that the execution context
+        // is captured and restored on suspension/resumption.
+        // We could perhaps skip this for AwaitAwaiter (but not for UnsafeAwaitAwaiter) since it is expected
+        // that the safe INotifyCompletion will take care of flowing ExecutionContext.
+        asyncInfo.ExecutionContextHandling = ExecutionContextHandling::AsyncSaveAndRestore;
+    }
+
+    // For tailcalls the contexts does not need saving/restoring: they will be
+    // overwritten by the caller anyway.
+    //
+    // More specifically, if we can show that
+    // Thread.CurrentThread._executionContext is not accessed between the
+    // call and returning then we can omit save/restore of the execution
+    // context. We do not do that optimization yet.
+    if ((prefixFlags & PREFIX_TAILCALL) != 0)
+    {
+        asyncInfo.ExecutionContextHandling                  = ExecutionContextHandling::None;
+        asyncInfo.ContinuationContextHandling               = ContinuationContextHandling::None;
+        asyncInfo.SaveAndRestoreSynchronizationContextField = false;
+    }
+
+    call->AsCall()->SetIsAsync(new (this, CMK_Async) AsyncCallInfo(asyncInfo));
+
+    if (asyncInfo.ExecutionContextHandling == ExecutionContextHandling::SaveAndRestore)
+    {
+        compMustSaveAsyncContexts = true;
+
+        // In this case we will need to save the context after the arguments are evaluated.
+        // Spill the arguments to accomplish that.
+        // (We could do this via splitting in SaveAsyncContexts, but since we need to
+        //  handle inline candidates we won't gain much.)
+        impSpillSideEffects(true, CHECK_SPILL_ALL DEBUGARG("Async await with execution context save and restore"));
+    }
+}
+
+//------------------------------------------------------------------------
+// impInsertAsyncContinuationForLdvirtftnCall:
+//   Insert the async continuation argument for a call the EE asked to be
+//   performed via ldvirtftn.
+//
+// Arguments:
+//    call - The call
+//
+// Remarks:
+//   Should be called before the 'this' arg is inserted, but after other IL args
+//   have been inserted.
+//
+void Compiler::impInsertAsyncContinuationForLdvirtftnCall(GenTreeCall* call)
+{
+    assert(call->AsCall()->IsAsync());
+
+    if (Target::g_tgtArgOrder == Target::ARG_ORDER_R2L)
+    {
+        call->AsCall()->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
+                                                   .WellKnown(WellKnownArg::AsyncContinuation));
+    }
+    else
+    {
+        call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
+                                                  .WellKnown(WellKnownArg::AsyncContinuation));
+    }
+}
+
 //------------------------------------------------------------------------
 // SpillRetExprHelper: iterate through arguments tree and spill ret_expr to local variables.
 //

src/tests/async/inst-unbox-thunks/inst-unbox-thunks.cs

@@ -137,4 +137,50 @@ public static void ManyArgGenericInstantiating()
     {
         Assert.Equal("System.String", CallStruct1M1b().Result);
     }
+
+    interface I2
+    {
+        Task<string> M0<T>();
+        Task<string> M1<T>(object a0, object a1, object a2, object a3, object a4, object a5, object a6, object a7, object a8);
+    }
+
+    class Class2 : I2
+    {
+        public async Task<string> M0<T>()
+        {
+            await Task.Yield();
+            return typeof(T).ToString();
+        }
+
+        public async Task<string> M1<T>(object a0, object a1, object a2, object a3, object a4, object a5, object a6, object a7, object a8)
+        {
+            await Task.Yield();
+            return typeof(T).ToString();
+        }
+    }
+
+    static I2 o2;
+    static async Task<string> CallClass2M0()
+    {
+        o2 = new Class2();
+        return await o2.M0<string>();
+    }
+
+    static async Task<string> CallClass2M1()
+    {
+        o2 = new Class2();
+        return await o2.M1<string>(default, default, default, default, default, default, default, default, default);
+    }
+
+    [Fact]
+    public static void NoArgGVM()
+    {
+        Assert.Equal("System.String", CallClass2M0().Result);
+    }
+
+    [Fact]
+    public static void ManyArgGVM()
+    {
+        Assert.Equal("System.String", CallClass2M1().Result);
+    }
 }