Implement a change in IL API to use RuntimeHelpers.Await<T>(Task<T>) and similar helpers.

buildroslynnugets.cmd

@@ -1,7 +1,7 @@
 setlocal ENABLEEXTENSIONS
 pushd %~dp0
-set ASYNC_ROSLYN_COMMIT=10a5611cb10cd64876a2638664f0740255197e1b
-set ASYNC_SUFFIX=async-11
+set ASYNC_ROSLYN_COMMIT=c1d47ed4bdbd28137fce8fbb350771677aa4cf15
+set ASYNC_SUFFIX=async-12
 set ASYNC_ROSLYN_BRANCH=demos/async2-experiment1
 
 cd ..

eng/Versions.props

@@ -44,9 +44,9 @@
       Any tools that contribute to the design-time experience should use the MicrosoftCodeAnalysisVersion_LatestVS property above to ensure
       they do not break the local dev experience.
     -->
-    <MicrosoftCodeAnalysisCSharpVersion>4.13.0-async-11</MicrosoftCodeAnalysisCSharpVersion>
-    <MicrosoftCodeAnalysisVersion>4.13.0-async-11</MicrosoftCodeAnalysisVersion>
-    <MicrosoftNetCompilersToolsetVersion>4.13.0-async-11</MicrosoftNetCompilersToolsetVersion>
+    <MicrosoftCodeAnalysisCSharpVersion>4.13.0-async-12</MicrosoftCodeAnalysisCSharpVersion>
+    <MicrosoftCodeAnalysisVersion>4.13.0-async-12</MicrosoftCodeAnalysisVersion>
+    <MicrosoftNetCompilersToolsetVersion>4.13.0-async-12</MicrosoftNetCompilersToolsetVersion>
   </PropertyGroup>
   <!--
     For source generator support we need to target multiple versions of Roslyn in order to be able to run on older versions of Roslyn.

src/coreclr/inc/corinfo.h

@@ -1408,6 +1408,9 @@ enum CorInfoTokenKind
 
     // token comes from devirtualizing a method
     CORINFO_TOKENKIND_DevirtualizedMethod = 0x800 | CORINFO_TOKENKIND_Method,
+
+    // token comes from runtime async awaiting pattern
+    CORINFO_TOKENKIND_Await = 0x2000 | CORINFO_TOKENKIND_Method,
 };
 
 struct CORINFO_RESOLVED_TOKEN

src/coreclr/jit/importer.cpp

@@ -8947,7 +8947,52 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 // many other places.  We unfortunately embed that knowledge here.
                 if (opcode != CEE_CALLI)
                 {
-                    _impResolveToken(CORINFO_TOKENKIND_Method);
+                    bool isAwait = false;
+                    if (JitConfig.JitOptimizeAwait())
+                    {
+                        // If we see the following code pattern in runtime async methods:
+                        //
+                        //    call[virt] <Method>
+                        //    call       <Await>
+                        //
+                        //  we emit an eqivalent of
+                        //
+                        //    call[virt] <RtMethod>
+                        //
+                        //  where "RtMethod" is the runtime-async counterpart of a Task-returning method.
+                        //
+                        //  NOTE: we could potentially check if Method is not a thunk and, in cases when we can tell,
+                        //        bypass this optimization. Otherwise in a non-thunk case we would be
+                        //        replacing the pattern with a call to a thunk, which contains roughly the same code.
+
+                        const BYTE* nextOpcode = codeAddr + sizeof(mdToken);
+                        if (compIsAsync2() && (nextOpcode + sizeof(mdToken) < codeEndp) &&
+                            (getU1LittleEndian(nextOpcode) == CEE_CALL))
+                        {
+                            // resolve the next token
+                            CORINFO_RESOLVED_TOKEN nextCallTok;
+                            impResolveToken(nextOpcode + 1, &nextCallTok, CORINFO_TOKENKIND_Method);
+
+                            // check if it is an Await intrinsic
+                            if (eeIsIntrinsic(nextCallTok.hMethod) &&
+                                lookupNamedIntrinsic(nextCallTok.hMethod) == NI_System_Runtime_CompilerServices_RuntimeHelpers_Await)
+                            {
+                                // yes, this is an Await
+                                isAwait = true;
+                            }
+                        }
+                    }
+
+                    if (isAwait)
+                    {
+                        _impResolveToken(CORINFO_TOKENKIND_Await);
+                        // consume the extra call
+                        codeAddr += 1 + sizeof(mdToken);
+                    }
+                    else
+                    {
+                        _impResolveToken(CORINFO_TOKENKIND_Method);
+                    }
 
                     eeGetCallInfo(&resolvedToken,
                                   (prefixFlags & PREFIX_CONSTRAINED) ? &constrainedResolvedToken : nullptr,

src/coreclr/jit/importercalls.cpp

@@ -528,7 +528,8 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
         // calls in JIT generated state machines only.
         if (compIsAsync2() &&
             ((ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_AwaitAwaiterFromRuntimeAsync) ||
-             (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync)))
+             (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync) ||
+             (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_Await)))
         {
             assert((call != nullptr) && call->OperIs(GT_CALL));
             call->AsCall()->gtIsAsyncCall = true;
@@ -913,7 +914,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
     impPopCallArgs(sig, call->AsCall());
 
     // Extra args
-    if ((instParam != nullptr) || sig->isAsyncCall() || (varArgsCookie != nullptr))
+    if ((instParam != nullptr) || call->AsCall()->IsAsync2() || (varArgsCookie != nullptr))
     {
         if (Target::g_tgtArgOrder == Target::ARG_ORDER_R2L)
         {
@@ -923,7 +924,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
                                                            .WellKnown(WellKnownArg::VarArgsCookie));
             }
 
-            if (sig->isAsyncCall())
+            if (call->AsCall()->IsAsync2())
             {
                 call->AsCall()->gtArgs.PushFront(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
                                                            .WellKnown(WellKnownArg::AsyncContinuation));
@@ -943,7 +944,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
                                                 NewCallArg::Primitive(instParam).WellKnown(WellKnownArg::InstParam));
             }
 
-            if (sig->isAsyncCall())
+            if (call->AsCall()->IsAsync2())
             {
                 call->AsCall()->gtArgs.PushBack(this, NewCallArg::Primitive(gtNewNull(), TYP_REF)
                                                           .WellKnown(WellKnownArg::AsyncContinuation));
@@ -3376,7 +3377,8 @@ GenTree* Compiler::impIntrinsic(CORINFO_CLASS_HANDLE    clsHnd,
     }
 
     if ((ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_AwaitAwaiterFromRuntimeAsync) ||
-        (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync))
+        (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync) ||
+        (ni == NI_System_Runtime_CompilerServices_RuntimeHelpers_Await))
     {
         // These are marked intrinsics simply to mark the call node as async,
         // which the caller will do. Make sure we keep pIntrinsicName assigned
@@ -10840,6 +10842,11 @@ NamedIntrinsic Compiler::lookupNamedIntrinsic(CORINFO_METHOD_HANDLE method)
                                 result =
                                     NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync;
                             }
+                            else if (strcmp(methodName, "Await") == 0)
+                            {
+                                result =
+                                    NI_System_Runtime_CompilerServices_RuntimeHelpers_Await;
+                            }
                             else if (strcmp(methodName, "SuspendAsync2") == 0)
                             {
                                 result = NI_System_Runtime_CompilerServices_RuntimeHelpers_SuspendAsync2;

src/coreclr/jit/jitconfigvalues.h

@@ -586,6 +586,8 @@ OPT_CONFIG_INTEGER(JitDoIfConversion, "JitDoIfConversion", 1)
 OPT_CONFIG_INTEGER(JitDoOptimizeMaskConversions, "JitDoOptimizeMaskConversions", 1) // Perform optimization of mask
                                                                                     // conversions
 
+RELEASE_CONFIG_INTEGER(JitOptimizeAwait, "JitOptimizeAwait", 1) // Perform optimization of Await intrinsics
+
 RELEASE_CONFIG_INTEGER(JitEnableOptRepeat, "JitEnableOptRepeat", 1) // If zero, do not allow JitOptRepeat
 RELEASE_CONFIG_METHODSET(JitOptRepeat, "JitOptRepeat")            // Runs optimizer multiple times on specified methods
 RELEASE_CONFIG_INTEGER(JitOptRepeatCount, "JitOptRepeatCount", 2) // Number of times to repeat opts when repeating

src/coreclr/jit/namedintrinsiclist.h

@@ -120,6 +120,7 @@ enum NamedIntrinsic : unsigned short
     NI_System_Runtime_CompilerServices_RuntimeHelpers_GetMethodTable,
     NI_System_Runtime_CompilerServices_RuntimeHelpers_AwaitAwaiterFromRuntimeAsync,
     NI_System_Runtime_CompilerServices_RuntimeHelpers_UnsafeAwaitAwaiterFromRuntimeAsync,
+    NI_System_Runtime_CompilerServices_RuntimeHelpers_Await,
     NI_System_Runtime_CompilerServices_RuntimeHelpers_SuspendAsync2,
     NI_System_Runtime_CompilerServices_RuntimeHelpers_get_RuntimeAsyncViaJitGeneratedStateMachines,
 

src/coreclr/nativeaot/System.Private.CoreLib/src/CompatibilitySuppressions.xml

@@ -857,6 +857,30 @@
     <DiagnosticId>CP0002</DiagnosticId>
     <Target>M:System.Threading.Lock.#ctor(System.Boolean)</Target>
   </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:System.Runtime.CompilerServices.RuntimeHelpers.Await(System.Threading.Tasks.Task)</Target>
+    <Left>ref/net10.0/System.Private.CoreLib.dll</Left>
+    <Right>lib/net10.0/System.Private.CoreLib.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:System.Runtime.CompilerServices.RuntimeHelpers.Await(System.Threading.Tasks.ValueTask)</Target>
+    <Left>ref/net10.0/System.Private.CoreLib.dll</Left>
+    <Right>lib/net10.0/System.Private.CoreLib.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:System.Runtime.CompilerServices.RuntimeHelpers.Await``1(System.Threading.Tasks.Task{``0})</Target>
+    <Left>ref/net10.0/System.Private.CoreLib.dll</Left>
+    <Right>lib/net10.0/System.Private.CoreLib.dll</Right>
+  </Suppression>
+  <Suppression>
+    <DiagnosticId>CP0002</DiagnosticId>
+    <Target>M:System.Runtime.CompilerServices.RuntimeHelpers.Await``1(System.Threading.Tasks.ValueTask{``0})</Target>
+    <Left>ref/net10.0/System.Private.CoreLib.dll</Left>
+    <Right>lib/net10.0/System.Private.CoreLib.dll</Right>
+  </Suppression>
   <Suppression>
     <DiagnosticId>CP0002</DiagnosticId>
     <Target>M:System.Runtime.CompilerServices.RuntimeHelpers.AwaitAwaiterFromRuntimeAsync``1(``0)</Target>

src/coreclr/vm/jitinterface.cpp

@@ -1100,6 +1100,10 @@ void CEEInfo::resolveToken(/* IN, OUT */ CORINFO_RESOLVED_TOKEN * pResolvedToken
             th = ClassLoader::LoadArrayTypeThrowing(th);
             break;
 
+        case CORINFO_TOKENKIND_Await:
+            pMD = pMD->GetAsyncOtherVariant();
+            break;
+
         default:
             // Disallow ELEMENT_TYPE_BYREF and ELEMENT_TYPE_VOID
             if (et == ELEMENT_TYPE_BYREF || et == ELEMENT_TYPE_VOID)

src/coreclr/vm/method.hpp

@@ -70,7 +70,12 @@ enum class AsyncMethodKind
     AsyncImplHelper,
 
     // Synthetic Async2 method that forwards to the NotAsync Task-returning method
-    AsyncThunkHelper
+    AsyncThunkHelper,
+
+    // Actual IL method that is explicitly declared as Async2 and thus compiled into a state machine.
+    // Such methods do not get Async thunks and can only be called from another Async2 method using Async2 call convention.
+    // This is used in a few infrastructure methods like `Await`
+    AsyncImplExplicit,
 };
 
 struct AsyncMethodData
@@ -1829,10 +1834,13 @@ class MethodDesc
     // CONSIDER: We probably need a better name for the concept, but it is hard to beat shortness of "async2"
     inline bool IsAsync2Method() const
     {
-        // Right now the only Async2 methods that exist are synthetic helpers.
-        // It is possible to declare an Async2 method directly in IL/Metadata,
-        // but we do not have a scenario for that.
-        return IsAsyncHelperMethod();
+        LIMITED_METHOD_DAC_CONTRACT;
+        if (!HasAsyncMethodData())
+            return false;
+        auto asyncKind = GetAddrOfAsyncMethodData()->kind;
+        return asyncKind == AsyncMethodKind::AsyncThunkHelper ||
+            asyncKind == AsyncMethodKind::AsyncImplHelper ||
+            asyncKind == AsyncMethodKind::AsyncImplExplicit;
     }
 
     inline bool IsStructMethodOperatingOnCopy()

src/coreclr/vm/methodtablebuilder.cpp

@@ -3605,7 +3605,17 @@ MethodTableBuilder::EnumerateClassMethods()
                 else
                 {
                     _ASSERTE(IsAsyncTaskMethodNormal(asyncMethodType));
-                    pNewMethod->SetAsyncMethodKind(AsyncMethodKind::NotAsync);
+
+                    if (IsMiAsync(dwImplFlags))
+                    {
+                        // TODO: explicitly-async methods have special semantics that is useful in the implementation of runtime async itself.
+                        //       It should not be valid to declare this outside of runtime infrastructure methods. (exact criteria TBD)
+                        pNewMethod->SetAsyncMethodKind(AsyncMethodKind::AsyncImplExplicit);
+                    }
+                    else
+                    {
+                        pNewMethod->SetAsyncMethodKind(AsyncMethodKind::NotAsync);
+                    }
                 }
 
                 pDeclaredMethod = pNewMethod;

src/coreclr/vm/threadsuspend.cpp

@@ -4909,15 +4909,6 @@ void STDCALL OnHijackWorker(HijackArgs * pArgs)
 #endif // HIJACK_NONINTERRUPTIBLE_THREADS
 }
 
-bool IsSpecialCaseAsyncRet(MethodDesc* pMD)
-{
-    // TODO: What's the right way to do this through CoreLibBinder without
-    // causing loading to happen? Also, can we just mark them as async2 in SPC,
-    // or force them to be fully interruptible?
-    LPCUTF8 name = pMD->GetName();
-    return strcmp(name, "UnsafeAwaitAwaiterFromRuntimeAsync") == 0 || strcmp(name, "AwaitAwaiterFromRuntimeAsync") == 0;
-}
-
 static bool GetReturnAddressHijackInfo(EECodeInfo *pCodeInfo, ReturnKind *pReturnKind, bool* hasAsyncRet)
 {
     GCInfoToken gcInfoToken = pCodeInfo->GetGCInfoToken();
@@ -4928,8 +4919,7 @@ static bool GetReturnAddressHijackInfo(EECodeInfo *pCodeInfo, ReturnKind *pRetur
     *hasAsyncRet = false;
 
     MethodDesc* pMD = pCodeInfo->GetMethodDesc();
-    *hasAsyncRet = pMD->IsAsync2Method() ||
-        (pMD->IsIntrinsic() && IsSpecialCaseAsyncRet(pMD));
+    *hasAsyncRet = pMD->IsAsync2Method();
 
     return true;
 }

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeHelpers.cs

@@ -178,8 +178,8 @@ public static unsafe ReadOnlySpan<T> CreateSpan<T>(RuntimeFieldHandle fldHandle)
 
 #if !NATIVEAOT
         [Intrinsic]
-        [MethodImpl(MethodImplOptions.NoInlining)]
         [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
         public static void AwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) where TAwaiter : INotifyCompletion
         {
             ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState;
@@ -193,10 +193,10 @@ public static void AwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) wher
         }
 
         // Marked intrinsic since for JIT state machines this needs to be
-        // recognizes as an async2 call.
+        // recognized as an async2 call.
         [Intrinsic]
         [BypassReadyToRun]
-        [MethodImpl(MethodImplOptions.NoInlining)]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
         public static void UnsafeAwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) where TAwaiter : ICriticalNotifyCompletion
         {
             ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState;
@@ -208,6 +208,73 @@ public static void UnsafeAwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter
             SuspendAsync2(sentinelContinuation);
             return;
         }
+
+        // Marked intrinsic since this needs to be
+        // recognized as an async2 call.
+        [Intrinsic]
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
+        public static T Await<T>(Task<T> task)
+        {
+            TaskAwaiter<T> awaiter = task.GetAwaiter();
+            if (!awaiter.IsCompleted)
+            {
+                UnsafeAwaitAwaiterFromRuntimeAsync(awaiter);
+            }
+
+            return awaiter.GetResult();
+        }
+
+        // Marked intrinsic since this needs to be
+        // recognized as an async2 call.
+        [Intrinsic]
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
+        public static void Await(Task task)
+        {
+            TaskAwaiter awaiter = task.GetAwaiter();
+            if (!awaiter.IsCompleted)
+            {
+                UnsafeAwaitAwaiterFromRuntimeAsync(awaiter);
+            }
+
+            awaiter.GetResult();
+            return;
+        }
+
+        // Marked intrinsic since this needs to be
+        // recognized as an async2 call.
+        [Intrinsic]
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
+        public static T Await<T>(ValueTask<T> task)
+        {
+            ValueTaskAwaiter<T> awaiter = task.GetAwaiter();
+            if (!awaiter.IsCompleted)
+            {
+                UnsafeAwaitAwaiterFromRuntimeAsync(awaiter);
+            }
+
+            return awaiter.GetResult();
+        }
+
+        // Marked intrinsic since this needs to be
+        // recognized as an async2 call.
+        [Intrinsic]
+        [BypassReadyToRun]
+        [MethodImpl(MethodImplOptions.NoInlining | (MethodImplOptions)0x0400)]  // NoInlining | Async
+        public static void Await(ValueTask task)
+        {
+            ValueTaskAwaiter awaiter = task.GetAwaiter();
+            if (!awaiter.IsCompleted)
+            {
+                UnsafeAwaitAwaiterFromRuntimeAsync(awaiter);
+            }
+
+            awaiter.GetResult();
+            return;
+        }
+
 #endif
     }
 }

src/libraries/System.Runtime/ref/System.Runtime.cs

@@ -13708,6 +13708,10 @@ public static void RunModuleConstructor(System.ModuleHandle module) { }
         public delegate void TryCode(object? userData);
         public static void UnsafeAwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) where TAwaiter : ICriticalNotifyCompletion { }
         public static void AwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) where TAwaiter : INotifyCompletion { }
+        public static void Await(System.Threading.Tasks.Task task) { throw null; }
+        public static T Await<T>(System.Threading.Tasks.Task<T> task) { throw null; }
+        public static void Await(System.Threading.Tasks.ValueTask task) { throw null; }
+        public static T Await<T>(System.Threading.Tasks.ValueTask<T> task) { throw null; }
     }
     public sealed partial class RuntimeWrappedException : System.Exception
     {