Prototyping T RuntimeHelpers.Await<T>(Task<T>)

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/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,12 @@
     <DiagnosticId>CP0002</DiagnosticId>
     <Target>M:System.Threading.Lock.#ctor(System.Boolean)</Target>
   </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.AwaitAwaiterFromRuntimeAsync``1(``0)</Target>

src/coreclr/vm/corelib.h

@@ -685,6 +685,7 @@ DEFINE_METHOD(RUNTIME_HELPERS,      FINALIZE_VALUETASK_RETURNING_THUNK, Finalize
 DEFINE_METHOD(RUNTIME_HELPERS,      FINALIZE_VALUETASK_RETURNING_THUNK_1, FinalizeValueTaskReturningThunk, GM_Continuation_RetValueTaskOfT)
 DEFINE_METHOD(RUNTIME_HELPERS,      UNSAFE_AWAIT_AWAITER_FROM_RUNTIME_ASYNC_1, UnsafeAwaitAwaiterFromRuntimeAsync, GM_T_RetVoid)
 DEFINE_METHOD(RUNTIME_HELPERS,      AWAIT_AWAITER_FROM_RUNTIME_ASYNC_1, AwaitAwaiterFromRuntimeAsync, GM_T_RetVoid)
+DEFINE_METHOD(RUNTIME_HELPERS,      AWAIT_1, Await, GM_TaskOfT_RetT)
 
 DEFINE_CLASS(SPAN_HELPERS,          System,                 SpanHelpers)
 DEFINE_METHOD(SPAN_HELPERS,         MEMSET,                 Fill, SM_RefByte_Byte_UIntPtr_RetVoid)

src/coreclr/vm/metasig.h

@@ -571,6 +571,7 @@ DEFINE_METASIG_T(IM(RetValueTask, _, g(VALUETASK)))
 
 DEFINE_METASIG_T(GM(Exception_RetTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(EXCEPTION), GI(C(TASK_1), 1, M(0))))
 DEFINE_METASIG_T(GM(T_RetTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, M(0), GI(C(TASK_1), 1, M(0))))
+DEFINE_METASIG_T(GM(TaskOfT_RetT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, GI(C(TASK_1), 1, M(0)), M(0)))
 DEFINE_METASIG_T(GM(Exception_RetValueTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, C(EXCEPTION), GI(g(VALUETASK_1), 1, M(0))))
 DEFINE_METASIG_T(GM(T_RetValueTaskOfT, IMAGE_CEE_CS_CALLCONV_DEFAULT, 1, M(0), GI(g(VALUETASK_1), 1, M(0))))
 

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,18 @@ MethodTableBuilder::EnumerateClassMethods()
                 else
                 {
                     _ASSERTE(IsAsyncTaskMethodNormal(asyncMethodType));
-                    pNewMethod->SetAsyncMethodKind(AsyncMethodKind::NotAsync);
+
+                    if (IsMiAsync(dwImplFlags))
+                    {
+                        // TODO: VS must validate that only a few special methods can do this.
+                        //       the possibility is useful, but should not become a general
+                        //       feature by accident.
+                        pNewMethod->SetAsyncMethodKind(AsyncMethodKind::AsyncImplExplicit);
+                    }
+                    else
+                    {
+                        pNewMethod->SetAsyncMethodKind(AsyncMethodKind::NotAsync);
+                    }
                 }
 
                 pDeclaredMethod = pNewMethod;

src/coreclr/vm/threadsuspend.cpp

@@ -4915,7 +4915,9 @@ bool IsSpecialCaseAsyncRet(MethodDesc* pMD)
     // 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;
+    return strcmp(name, "UnsafeAwaitAwaiterFromRuntimeAsync") == 0 ||
+        strcmp(name, "AwaitAwaiterFromRuntimeAsync") == 0 ||
+        strcmp(name, "Await") == 0;
 }
 
 static bool GetReturnAddressHijackInfo(EECodeInfo *pCodeInfo, ReturnKind *pReturnKind, bool* hasAsyncRet)

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;
@@ -196,7 +196,7 @@ public static void AwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter) wher
         // recognizes 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,25 @@ public static void UnsafeAwaitAwaiterFromRuntimeAsync<TAwaiter>(TAwaiter awaiter
             SuspendAsync2(sentinelContinuation);
             return;
         }
+
+        // TODO: should this be called "AwaitFromRuntimeAsync" ?  (i.e. same as above, but no "Awaiter")
+        //
+        // Marked intrinsic since this needs to be
+        // recognizes 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();
+        }
+
 #endif
     }
 }

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

@@ -13708,6 +13708,7 @@ 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 T Await<T>(System.Threading.Tasks.Task<T> task) { throw null; }
     }
     public sealed partial class RuntimeWrappedException : System.Exception
     {

src/tests/async/varying-yields-await.csproj

@@ -0,0 +1,9 @@
+<Project Sdk="Microsoft.NET.Sdk.IL">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+    <DefineConstants>AWAIT;$(DefineConstants)</DefineConstants>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="varying-yields.cs" />
+  </ItemGroup>
+</Project>

src/tests/async/varying-yields.cs

@@ -4,6 +4,8 @@
 //#define ASYNC1_TASK
 //#define ASYNC1_VALUETASK
 
+#pragma warning disable 4014, 1998
+
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
@@ -91,9 +93,19 @@ async2 Task<long>
             double liveState3 = _yieldProbability;
 
             if (depth == 0)
+#if AWAIT
+                return RuntimeHelpers.Await(Loop());
+#else
                 return await Loop();
+#endif
+
+            long result =
+#if AWAIT
+                RuntimeHelpers.Await(Run(depth - 1));
+#else
+                await Run(depth - 1);
+#endif
 
-            long result = await Run(depth - 1);
             Sink = (int)liveState1 + (int)liveState2 + (int)(1 / liveState3) + depth;
             return result;
         }
@@ -117,7 +129,12 @@ async2 Task<long>
             {
                 for (int i = 0; i < 20; i++)
                 {
-                    numIters += await DoYields();
+                    numIters +=
+#if AWAIT
+                        RuntimeHelpers.Await(DoYields());
+#else
+                        await DoYields();
+#endif
                 }
             }