Add AsyncHelpers.SetAsyncCallContinuationArg intrinsic

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -156,7 +156,13 @@ public static partial class AsyncHelpers
         //   * a continuation object if the call requires suspension.
         //     In this case the formal result of the call is undefined.
         [Intrinsic]
-        private static Continuation? AsyncCallContinuation() => throw new UnreachableException();
+        internal static Continuation? AsyncCallContinuation() => throw new UnreachableException(); // Unconditionally expanded intrinsic
+
+        // Set call continuation argument for an upcoming call with runtime async calling convention.
+        // Can only be used shortly before the call (with no interfering control flow instructions before the call).
+        // If used anywhere in the method it must be present before _all_ calls with runtime async calling convention.
+        [Intrinsic]
+        internal static void SetAsyncCallContinuationArg(Continuation continuation) => throw new UnreachableException();
 
         // Used during suspensions to hold the continuation chain and on what we are waiting.
         // Methods like FinalizeTaskReturningThunk will unlink the state and wrap into a Task.

src/coreclr/jit/async.cpp

@@ -484,6 +484,11 @@ bool AsyncLiveness::IsLocalCaptureUnnecessary(unsigned lclNum)
         return true;
     }
 
+    if (lclNum == m_comp->lvaNextAsyncCallContArgVar)
+    {
+        return true;
+    }
+
     return false;
 }
 

src/coreclr/jit/compiler.h

@@ -3976,6 +3976,8 @@ class Compiler
 
     bool lvaInlineeReturnSpillTempFreshlyCreated = false; // True if the temp was freshly created for the inlinee return
 
+    unsigned lvaNextAsyncCallContArgVar = BAD_VAR_NUM;
+
 #if FEATURE_FIXED_OUT_ARGS
     unsigned            lvaOutgoingArgSpaceVar = BAD_VAR_NUM;  // var that represents outgoing argument space
     PhasedVar<unsigned> lvaOutgoingArgSpaceSize; // size of fixed outgoing argument space

src/coreclr/jit/importercalls.cpp

@@ -875,7 +875,10 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             // stubs and the VM inserts the arg itself.
             if (call->AsCall()->IsAsync() && (opcode != CEE_CALLI))
             {
-                call->AsCall()->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
+                GenTree* arg = lvaNextAsyncCallContArgVar != BAD_VAR_NUM
+                                   ? (GenTree*)gtNewLclVarNode(lvaNextAsyncCallContArgVar)
+                                   : gtNewNull();
+                call->AsCall()->gtArgs.PushFront(this, NewCallArg::Primitive(arg, TYP_REF)
                                                            .WellKnown(WellKnownArg::AsyncContinuation));
             }
 
@@ -891,7 +894,10 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
             // stubs and the VM inserts the arg itself.
             if (call->AsCall()->IsAsync() && (opcode != CEE_CALLI))
             {
-                call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
+                GenTree* arg = lvaNextAsyncCallContArgVar != BAD_VAR_NUM
+                                   ? (GenTree*)gtNewLclVarNode(lvaNextAsyncCallContArgVar)
+                                   : gtNewNull();
+                call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(arg, TYP_REF)
                                                           .WellKnown(WellKnownArg::AsyncContinuation));
             }
 
@@ -3332,6 +3338,18 @@ GenTree* Compiler::impIntrinsic(CORINFO_CLASS_HANDLE    clsHnd,
         return node;
     }
 
+    if (ni == NI_System_Runtime_CompilerServices_AsyncHelpers_SetAsyncCallContinuationArg)
+    {
+        if (lvaNextAsyncCallContArgVar == BAD_VAR_NUM)
+        {
+            lvaNextAsyncCallContArgVar                     = lvaGrabTemp(false DEBUGARG("Async call continuation arg"));
+            lvaGetDesc(lvaNextAsyncCallContArgVar)->lvType = TYP_REF;
+        }
+
+        GenTree* node = gtNewStoreLclVarNode(lvaNextAsyncCallContArgVar, impPopStack().val);
+        return node;
+    }
+
     if (ni == NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncSuspend)
     {
         GenTree* node = gtNewOperNode(GT_RETURN_SUSPEND, TYP_VOID, impPopStack().val);
@@ -10876,6 +10894,10 @@ NamedIntrinsic Compiler::lookupNamedIntrinsic(CORINFO_METHOD_HANDLE method)
                             {
                                 result = NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncCallContinuation;
                             }
+                            else if (strcmp(methodName, "SetAsyncCallContinuationArg") == 0)
+                            {
+                                result = NI_System_Runtime_CompilerServices_AsyncHelpers_SetAsyncCallContinuationArg;
+                            }
                         }
                         else if (strcmp(className, "StaticsHelpers") == 0)
                         {

src/coreclr/jit/namedintrinsiclist.h

@@ -126,6 +126,7 @@ enum NamedIntrinsic : unsigned short
     NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncSuspend,
     NI_System_Runtime_CompilerServices_AsyncHelpers_Await,
     NI_System_Runtime_CompilerServices_AsyncHelpers_AsyncCallContinuation,
+    NI_System_Runtime_CompilerServices_AsyncHelpers_SetAsyncCallContinuationArg,
 
     NI_System_Runtime_CompilerServices_StaticsHelpers_VolatileReadAsByref,
 

src/coreclr/vm/corelib.h

@@ -750,6 +750,7 @@ DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_CONTINUATION_CONTEXT, CaptureContinuat
 DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_CONTEXTS,          CaptureContexts, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      RESTORE_CONTEXTS,          RestoreContexts, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      ASYNC_CALL_CONTINUATION,   AsyncCallContinuation, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      SET_ASYNC_CALL_CONTINUATION_ARG,   SetAsyncCallContinuationArg, SM_Obj_RetVoid)
 
 #ifdef TARGET_BROWSER
 DEFINE_METHOD(ASYNC_HELPERS,      HANDLE_ASYNC_ENTRYPOINT, HandleAsyncEntryPoint, SM_TaskOfInt_RetInt)

src/coreclr/vm/jitinterface.cpp

@@ -9151,6 +9151,37 @@ void CEEInfo::getFunctionEntryPoint(CORINFO_METHOD_HANDLE  ftnHnd,
             accessType = IAT_PVALUE;
         }
     }
+#ifdef FEATURE_TIERED_COMPILATION
+    else if (ftn->IsAsyncMethod() && m_pMethodBeingCompiled->IsILStub()
+        && m_pMethodBeingCompiled->AsDynamicMethodDesc()->IsAsyncResumptionStub())
+    {
+        // We are looking at a call to an async method from within an async method's
+        // resumption stub. We need to make sure we call the right version of the
+        // async method here.
+
+        // Resumption stubs are uniquely coupled to the code version (since the
+        // continuation is), so we need to make sure we always keep calling the
+        // same version here.
+        PrepareCodeConfig* config = GetThread()->GetCurrentPrepareCodeConfig();
+        NativeCodeVersion ncv = config->GetCodeVersion();
+        if (ncv.GetOptimizationTier() == NativeCodeVersion::OptimizationTier1OSR)
+        {
+#ifdef FEATURE_ON_STACK_REPLACEMENT
+            // The OSR version needs to resume in the tier0 version. The tier0
+            // version will handle setting up the frame that the OSR version
+            // expects and then delegating back into the OSR version (knowing to do
+            // so through information stored in the continuation).
+            unsigned osrIlOffset = 0;
+            PatchpointInfo* patchpointInfo = getOSRInfo(&osrIlOffset);
+            _ASSERTE(patchpointInfo != NULL);
+            ret = (void*)(patchpointInfo->GetTier0EntryPoint());
+            accessType = IAT_VALUE;
+#else // !FEATURE_ON_STACK_REPLACEMENT
+            _ASSERTE(!"Unexpected optimization tier with OSR disabled");
+#endif // FEATURE_ON_STACK_REPLACEMENT
+        }
+    }
+#endif // FEATURE_TIERED_COMPILATION
     else if (ftn->IsVersionableWithPrecode() && (ftn->GetPrecodeType() == PRECODE_FIXUP) && !ftn->IsPointingToStableNativeCode())
     {
         ret = ((FixupPrecode*)ftn->GetOrCreatePrecode())->GetTargetSlot();
@@ -14659,13 +14690,15 @@ CORINFO_METHOD_HANDLE CEEJitInfo::getAsyncResumptionStub(void** entryPoint)
     Signature stubSig = BuildResumptionStubSignature(md->GetLoaderAllocator(), &amTracker);
 
     MetaSig msig(md);
-    Signature calliSig = BuildResumptionStubCalliSignature(msig, md->GetMethodTable(), md->GetLoaderAllocator(), &amTracker);
 
     SigTypeContext emptyCtx;
     ILStubLinker sl(md->GetModule(), stubSig, &emptyCtx, NULL, ILSTUB_LINKER_FLAG_NONE);
 
     ILCodeStream* pCode = sl.NewCodeStream(ILStubLinker::kDispatch);
 
+    pCode->EmitLDARG(0);
+    pCode->EmitCALL(METHOD__ASYNC_HELPERS__SET_ASYNC_CALL_CONTINUATION_ARG, 1, 0);
+
     int numArgs = 0;
 
     if (msig.HasThis())
@@ -14683,18 +14716,6 @@ CORINFO_METHOD_HANDLE CEEJitInfo::getAsyncResumptionStub(void** entryPoint)
         numArgs++;
     }
 
-#ifndef TARGET_X86
-    if (msig.HasGenericContextArg())
-    {
-        pCode->EmitLDC(0);
-        numArgs++;
-    }
-
-    // Continuation
-    pCode->EmitLDARG(0);
-    numArgs++;
-#endif
-
     msig.Reset();
     CorElementType ty;
     while ((ty = msig.NextArg()) != ELEMENT_TYPE_END)
@@ -14707,50 +14728,7 @@ CORINFO_METHOD_HANDLE CEEJitInfo::getAsyncResumptionStub(void** entryPoint)
         numArgs++;
     }
 
-#ifdef TARGET_X86
-    // Continuation
-    pCode->EmitLDARG(0);
-
-    if (msig.HasGenericContextArg())
-    {
-        pCode->EmitLDC(0);
-        numArgs++;
-    }
-
-    numArgs++;
-#endif
-
-#ifdef FEATURE_TIERED_COMPILATION
-    // Resumption stubs are uniquely coupled to the code version (since the
-    // continuation is), so we need to make sure we always keep calling the
-    // same version here.
-    PrepareCodeConfig* config = GetThread()->GetCurrentPrepareCodeConfig();
-    NativeCodeVersion ncv = config->GetCodeVersion();
-    if (ncv.GetOptimizationTier() == NativeCodeVersion::OptimizationTier1OSR)
-    {
-#ifdef FEATURE_ON_STACK_REPLACEMENT
-        // The OSR version needs to resume in the tier0 version. The tier0
-        // version will handle setting up the frame that the OSR version
-        // expects and then delegating back into the OSR version (knowing to do
-        // so through information stored in the continuation).
-        _ASSERTE(m_pPatchpointInfoFromRuntime != NULL);
-        pCode->EmitLDC((DWORD_PTR)m_pPatchpointInfoFromRuntime->GetTier0EntryPoint());
-#else // !FEATURE_ON_STACK_REPLACEMENT
-        _ASSERTE(!"Unexpected optimization tier with OSR disabled");
-#endif // FEATURE_ON_STACK_REPLACEMENT
-    }
-    else
-#endif // FEATURE_TIERED_COMPILATION
-    {
-        {
-            m_finalCodeAddressSlot = (PCODE*)amTracker.Track(m_pMethodBeingCompiled->GetLoaderAllocator()->GetHighFrequencyHeap()->AllocMem(S_SIZE_T(sizeof(PCODE))));
-        }
-
-        pCode->EmitLDC((DWORD_PTR)m_finalCodeAddressSlot);
-        pCode->EmitLDIND_I();
-    }
-
-    pCode->EmitCALLI(pCode->GetSigToken(calliSig.GetRawSig(), calliSig.GetRawSigLen()), numArgs, msig.IsReturnTypeVoid() ? 0 : 1);
+    pCode->EmitCALL(pCode->GetToken(md), numArgs, msig.IsReturnTypeVoid() ? 0 : 1);
 
     DWORD resultLoc = UINT_MAX;
     TypeHandle resultTypeHnd;
@@ -14799,6 +14777,8 @@ CORINFO_METHOD_HANDLE CEEJitInfo::getAsyncResumptionStub(void** entryPoint)
 
     const char* optimizationTierName = "UnknownTier";
 #ifdef FEATURE_TIERED_COMPILATION
+    PrepareCodeConfig* config = GetThread()->GetCurrentPrepareCodeConfig();
+    NativeCodeVersion ncv = config->GetCodeVersion();
     switch (ncv.GetOptimizationTier())
     {
     case NativeCodeVersion::OptimizationTier0: optimizationTierName = "Tier0"; break;

src/coreclr/vm/method.hpp

@@ -2957,6 +2957,12 @@ class DynamicMethodDesc : public StoredSigMethodDesc
         _ASSERTE(IsILStub());
         return GetILStubType() == DynamicMethodDesc::StubDelegateShuffleThunk;
     }
+    bool IsAsyncResumptionStub() const
+    {
+        LIMITED_METHOD_DAC_CONTRACT;
+        _ASSERTE(IsILStub());
+        return GetILStubType() == DynamicMethodDesc::StubAsyncResume;
+    }
 
     // Whether the stub takes a context argument that is an interop MethodDesc.
     // See RequiresMDContextArg() for the non-stub version.