JIT: Support general inlining in runtime async

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

@@ -223,6 +223,8 @@ private ref struct RuntimeAsyncStackState
             public INotifyCompletion? Notifier;
             public ValueTaskSourceContinuation? ValueTaskSourceContinuation;
             public RuntimeAsyncTaskContinuation? TaskContinuation;
+            public delegate*<Continuation, int, Action, void> AwaiterContinuation;
+            public int AwaiterOffset;
 
             // When we suspend in the leaf, the contexts are captured into these fields.
             public ExecutionContext? LeafExecutionContext;
@@ -265,6 +267,8 @@ public void CaptureContexts()
             [NonVersionable]
             public void Push(RuntimeAsyncStackState* stackState)
             {
+                stackState->AwaiterContinuation = null;
+                stackState->AwaiterOffset = 0;
                 stackState->Next = StackState;
                 StackState = stackState;
                 CurrentThread ??= Thread.CurrentThread;
@@ -805,7 +809,15 @@ internal unsafe bool HandleSuspended(ref RuntimeAsyncAwaitState state)
 
                 try
                 {
-                    if (stackState->CriticalNotifier is { } critNotifier)
+                    if (stackState->AwaiterContinuation != null)
+                    {
+                        // The awaiter is stored in the continuation for the caller of
+                        // AwaitAwaiterInContinuation or UnsafeAwaitAwaiterInContinuation.
+                        Debug.Assert((headContinuation.Flags & ContinuationFlags.AllContinuationFlags) == 0);
+                        stackState->AwaiterContinuation(
+                            headContinuation, stackState->AwaiterOffset, GetContinuationAction());
+                    }
+                    else if (stackState->CriticalNotifier is { } critNotifier)
                     {
                         // Result of async call to AwaitAwaiter or UnsafeAwaitAwaiter.
                         // These never have special continuation context handling.
@@ -1477,6 +1489,131 @@ private static void CaptureContinuationContextFlags(ref ContinuationFlags flags,
             flags |= ContinuationFlags.ContinueOnThreadPool;
         }
 
+        // Restore the contexts that an inlined async frame captured when it logically returned to
+        // its caller, after that frame was resumed inside its own body.
+        //
+        // Used by the JIT when inlining runtime async calls. The JIT emits the check of whether
+        // the frame was resumed at all and calls this when it was; everything the async
+        // infrastructure would otherwise have done at that frame boundary happens here.
+        //
+        // Unlike the synchronous restore at the end of a method, the ExecutionContext restore runs
+        // only after a resumption, so it must target the thread we were resumed on rather than the
+        // one whose contexts were captured on entry.
+        //
+        // The continuation context check determines whether resuming the caller's continuation here
+        // would be dispatched inline. It mirrors the "can inline" conditions in
+        // RuntimeAsyncTaskContinuation.QueueIfNecessary; the two must be kept in sync.
+        //
+        // 'flags' must contain only ContinuationFlags.AllContinuationFlags bits.
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.Async)]
+        private static void RestoreInlinedFrameContexts(ExecutionContext? previousExecCtx, object? continuationContext, ContinuationFlags flags)
+        {
+            Debug.Assert((flags & ~ContinuationFlags.AllContinuationFlags) == 0);
+
+            // We are inside a runtime async chain, so the thread has already been cached. Use it
+            // instead of Thread.CurrentThreadAssumedInitialized to keep this to one TLS lookup.
+            ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState;
+            Thread? currentThread = state.CurrentThread;
+            Debug.Assert(currentThread != null);
+
+            RestoreExecutionContext(currentThread, previousExecCtx);
+
+            if ((flags & ContinuationFlags.ContinueOnThreadPool) != 0)
+            {
+                SynchronizationContext? syncCtx = currentThread._synchronizationContext;
+                if (syncCtx is null || syncCtx.GetType() == typeof(SynchronizationContext))
+                {
+                    TaskScheduler? sched = TaskScheduler.InternalCurrent;
+                    if (sched is null || sched == TaskScheduler.Default)
+                    {
+                        return;
+                    }
+                }
+            }
+            else if ((flags & ContinuationFlags.ContinueOnCapturedSynchronizationContext) != 0)
+            {
+                Debug.Assert(continuationContext is SynchronizationContext);
+                if (continuationContext == currentThread._synchronizationContext)
+                {
+                    return;
+                }
+            }
+            else if ((flags & ContinuationFlags.ContinueOnCapturedTaskScheduler) != 0)
+            {
+                Debug.Assert(continuationContext is TaskScheduler);
+                if (continuationContext == TaskScheduler.InternalCurrent)
+                {
+                    return;
+                }
+            }
+            else
+            {
+                // No continuation context was captured, so there is nothing to switch to.
+                return;
+            }
+
+            TailAwait();
+            SwitchToContinuationContext(ref state, continuationContext, flags);
+        }
+
+        // Suspend and resume in the specified continuation context.
+        //
+        // Suspending on an already completed task makes the dispatcher re-dispatch the continuation
+        // immediately. Because that dispatch happens with canInline: false, it always posts or
+        // schedules onto the requested context rather than running inline here, which is what we
+        // want -- we only get here when we are known to be on the wrong context.
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.Async)]
+        private static unsafe void SwitchToContinuationContext(ref RuntimeAsyncAwaitState state, object? continuationContext, ContinuationFlags flags)
+        {
+            Continuation? sentinelContinuation = state.SentinelContinuation ??= new Continuation();
+
+            RuntimeAsyncTaskContinuation? taskCont = state.CachedTaskContinuation;
+            if (taskCont != null)
+            {
+                state.CachedTaskContinuation = null;
+            }
+            else
+            {
+                taskCont = new RuntimeAsyncTaskContinuation();
+            }
+
+            taskCont.Initialize(Task.CompletedTask);
+            taskCont.ContinuationContext = continuationContext;
+            taskCont.Flags |= flags;
+
+            sentinelContinuation.Next = taskCont;
+            state.StackState->TaskContinuation = taskCont;
+
+            state.CaptureContexts();
+            AsyncSuspend(taskCont);
+        }
+
+        // Capture the contexts an inlined async frame hands to its caller when it logically
+        // returns during a suspension, i.e. what the caller's continuation would have captured
+        // had the callee's frame been physically present.
+        //
+        // No-ops when the frame has already resumed: in that case the caller's continuation
+        // already exists and keeps the values it captured when the frame first suspended.
+        //
+        // The suspension walks the inlined frames outward, and a frame having resumed implies
+        // its caller has too, so these can be emitted as a straight line: once one frame has
+        // resumed, this and every subsequent restore for the frames outside it no-op.
+        private static void CaptureInlinedFrameTransition(bool resumed,
+                                                          ref object? continuationContext,
+                                                          ref ContinuationFlags flags,
+                                                          ref ExecutionContext? execContext)
+        {
+            if (resumed)
+            {
+                return;
+            }
+
+            CaptureContinuationContext(ref continuationContext, ref flags);
+            execContext = CaptureExecutionContext();
+        }
+
         // Finish suspension in the common case of a custom await or for a ConfigureAwait(false) task await:
         // - Capture current ExecutionContext into the continuation
         // - Restore ExecutionContext and SynchronizationContext to the current Thread object

src/coreclr/inc/corinfo.h

@@ -1836,6 +1836,12 @@ struct CORINFO_ASYNC_INFO
     CORINFO_METHOD_HANDLE restoreContextsMethHnd;
     // Method handle for AsyncHelpers.RestoreContextsOnSuspension, used before suspending in async methods
     CORINFO_METHOD_HANDLE restoreContextsOnSuspensionMethHnd;
+    // Method handle for AsyncHelpers.RestoreInlinedFrameContexts, used when an inlined
+    // async callee logically returns to its caller after having been resumed
+    CORINFO_METHOD_HANDLE restoreInlinedFrameContextsMethHnd;
+    // Method handle for AsyncHelpers.CaptureInlinedFrameTransition, used on suspension to capture
+    // the contexts each inlined async frame hands to its caller
+    CORINFO_METHOD_HANDLE captureInlinedFrameTransitionMethHnd;
     // Finish suspension without saving continuation context (i.e. custom awaiter or ConfigureAwait(false))
     CORINFO_METHOD_HANDLE finishSuspensionNoContinuationContextMethHnd;
     // Finish suspension with saving continuation context (i.e. normal task await)
@@ -3168,6 +3174,14 @@ class ICorStaticInfo
     // instantiation argument that must be passed to the await call.
     virtual CORINFO_METHOD_HANDLE getAwaitReturnCall(CORINFO_METHOD_HANDLE callerHandle, CORINFO_CONTEXT_HANDLE* contextHandle, CORINFO_LOOKUP* instArg) = 0;
 
+    virtual CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall(
+        CORINFO_METHOD_HANDLE callerHandle,
+        CORINFO_SIG_INFO* callSig,
+        bool isUnsafe,
+        CORINFO_CONTEXT_HANDLE* contextHandle,
+        CORINFO_LOOKUP* instArg
+    ) = 0;
+
     /*********************************************************************************/
     //
     // Diagnostic methods

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -506,6 +506,13 @@ CORINFO_METHOD_HANDLE getAwaitReturnCall(
           CORINFO_CONTEXT_HANDLE* contextHandle,
           CORINFO_LOOKUP* instArg) override;
 
+CORINFO_METHOD_HANDLE getAwaitAwaiterInContinuationCall(
+          CORINFO_METHOD_HANDLE callerHandle,
+          CORINFO_SIG_INFO* callSig,
+          bool isUnsafe,
+          CORINFO_CONTEXT_HANDLE* contextHandle,
+          CORINFO_LOOKUP* instArg) override;
+
 mdMethodDef getMethodDefFromMethod(
           CORINFO_METHOD_HANDLE hMethod) override;
 

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* 134b051a-e1ec-4f52-a1bf-f919908aa33c */
-    0x134b051a,
-    0xe1ec,
-    0x4f52,
-    {0xa1, 0xbf, 0xf9, 0x19, 0x90, 0x8a, 0xa3, 0x3c}
+constexpr GUID JITEEVersionIdentifier = { /* a41245a2-74c0-47ac-87d7-9375cb833125 */
+    0xa41245a2,
+    0x74c0,
+    0x47ac,
+    {0x87, 0xd7, 0x93, 0x75, 0xcb, 0x83, 0x31, 0x25}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -126,6 +126,7 @@ DEF_CLR_API(runWithSPMIErrorTrap)
 DEF_CLR_API(getEEInfo)
 DEF_CLR_API(getAsyncInfo)
 DEF_CLR_API(getAwaitReturnCall)
+DEF_CLR_API(getAwaitAwaiterInContinuationCall)
 DEF_CLR_API(getMethodDefFromMethod)
 DEF_CLR_API(printMethodName)
 DEF_CLR_API(getMethodNameFromMetadata)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -1199,6 +1199,19 @@ CORINFO_METHOD_HANDLE WrapICorJitInfo::getAwaitReturnCall(
     return temp;
 }
 
+CORINFO_METHOD_HANDLE WrapICorJitInfo::getAwaitAwaiterInContinuationCall(
+          CORINFO_METHOD_HANDLE callerHandle,
+          CORINFO_SIG_INFO* callSig,
+          bool isUnsafe,
+          CORINFO_CONTEXT_HANDLE* contextHandle,
+          CORINFO_LOOKUP* instArg)
+{
+    API_ENTER(getAwaitAwaiterInContinuationCall);
+    CORINFO_METHOD_HANDLE temp = wrapHnd->getAwaitAwaiterInContinuationCall(callerHandle, callSig, isUnsafe, contextHandle, instArg);
+    API_LEAVE(getAwaitAwaiterInContinuationCall);
+    return temp;
+}
+
 mdMethodDef WrapICorJitInfo::getMethodDefFromMethod(
           CORINFO_METHOD_HANDLE hMethod)
 {