Enable runtime async for native AOT testing

src/coreclr/inc/corinfo.h

@@ -1434,14 +1434,14 @@ enum CORINFO_THIS_TRANSFORM
 
 enum CORINFO_CALLINFO_FLAGS
 {
-    CORINFO_CALLINFO_NONE           = 0x0000,
-    CORINFO_CALLINFO_ALLOWINSTPARAM = 0x0001,   // Can the compiler generate code to pass an instantiation parameters? Simple compilers should not use this flag
-    CORINFO_CALLINFO_CALLVIRT       = 0x0002,   // Is it a virtual call?
-    // UNUSED                       = 0x0004,
-    CORINFO_CALLINFO_DISALLOW_STUB  = 0x0008,   // Do not use a stub for this call, even if it is a virtual call.
-    CORINFO_CALLINFO_SECURITYCHECKS = 0x0010,   // Perform security checks.
-    CORINFO_CALLINFO_LDFTN          = 0x0020,   // Resolving target of LDFTN
-    // UNUSED                       = 0x0040,
+    CORINFO_CALLINFO_NONE              = 0x0000,
+    CORINFO_CALLINFO_ALLOWINSTPARAM    = 0x0001,   // Can the compiler generate code to pass an instantiation parameters? Simple compilers should not use this flag
+    CORINFO_CALLINFO_CALLVIRT          = 0x0002,   // Is it a virtual call?
+    CORINFO_CALLINFO_ALLOWASYNCVARIANT = 0x0004,   // allow resolution to an async variant
+    CORINFO_CALLINFO_DISALLOW_STUB     = 0x0008,   // Do not use a stub for this call, even if it is a virtual call.
+    CORINFO_CALLINFO_SECURITYCHECKS    = 0x0010,   // Perform security checks.
+    CORINFO_CALLINFO_LDFTN             = 0x0020,   // Resolving target of LDFTN
+    // UNUSED                          = 0x0040,
 };
 
 enum CorInfoIsAccessAllowedResult
@@ -1484,10 +1484,6 @@ 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,
-    CORINFO_TOKENKIND_AwaitVirtual = 0x4000 | CORINFO_TOKENKIND_Method,
 };
 
 struct CORINFO_RESOLVED_TOKEN
@@ -1560,6 +1556,12 @@ struct CORINFO_CALL_INFO
     CORINFO_CONST_LOOKUP    instParamLookup;
 
     bool                    wrapperDelegateInvoke;
+
+    // If CORINFO_CALLINFO_ALLOWASYNCVARIANT was passed, this is the resolved
+    // async variant or NULL if no async variant was resolved.
+    // This is the async variant of the token's method and differs from hMethod
+    // of this class in cases of sharing, constrained resolution etc.
+    CORINFO_METHOD_HANDLE   resolvedAsyncVariant;
 };
 
 enum CORINFO_DEVIRTUALIZATION_DETAIL

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* 8f2ee94f-5111-4b7b-97bd-5c689d476385 */
-    0x8f2ee94f,
-    0x5111,
-    0x4b7b,
-    {0x97, 0xbd, 0x5c, 0x68, 0x9d, 0x47, 0x63, 0x85}
+constexpr GUID JITEEVersionIdentifier = { /* 8edbe247-3d6f-43ca-81de-648d358d36f4 */
+    0x8edbe247,
+    0x3d6f,
+    0x43ca,
+    {0x81, 0xde, 0x64, 0x8d, 0x35, 0x8d, 0x36, 0xf4}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/interpreter/compiler.cpp

@@ -4551,10 +4551,13 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re
     }
     else
     {
+        unsigned flags = CORINFO_CALLINFO_ALLOWINSTPARAM | CORINFO_CALLINFO_SECURITYCHECKS | CORINFO_CALLINFO_DISALLOW_STUB;
+        const uint8_t* callIP = m_ip;
         if (!newObj && m_methodInfo->args.isAsyncCall() && AsyncCallPeeps.FindAndApplyPeep(this))
         {
-            resolvedCallToken = m_resolvedAsyncCallToken;
+            ResolveToken(token, CORINFO_TOKENKIND_Method, &resolvedCallToken);
             continuationContextHandling = m_currentContinuationContextHandling;
+            flags |= CORINFO_CALLINFO_ALLOWASYNCVARIANT;
         }
         else
         {
@@ -4570,11 +4573,23 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re
             m_compHnd->getNewHelper(resolvedCallToken.hClass, &hasSideEffects);
         }
 
-        CORINFO_CALLINFO_FLAGS flags = (CORINFO_CALLINFO_FLAGS)(CORINFO_CALLINFO_ALLOWINSTPARAM | CORINFO_CALLINFO_SECURITYCHECKS | CORINFO_CALLINFO_DISALLOW_STUB);
         if (isVirtual)
-            flags = (CORINFO_CALLINFO_FLAGS)(flags | CORINFO_CALLINFO_CALLVIRT);
+            flags |= CORINFO_CALLINFO_CALLVIRT;
 
-        m_compHnd->getCallInfo(&resolvedCallToken, pConstrainedToken, m_methodInfo->ftn, flags, &callInfo);
+        m_compHnd->getCallInfo(&resolvedCallToken, pConstrainedToken, m_methodInfo->ftn, (CORINFO_CALLINFO_FLAGS)flags, &callInfo);
+
+        if (flags & CORINFO_CALLINFO_ALLOWASYNCVARIANT)
+        {
+            if (callInfo.resolvedAsyncVariant != NULL)
+            {
+                resolvedCallToken.hMethod = callInfo.resolvedAsyncVariant;
+            }
+            else
+            {
+                // Undo the IP change from the peep
+                m_ip = callIP + 5;
+            }
+        }
 
         if (callInfo.sig.isVarArg())
         {
@@ -6739,14 +6754,6 @@ int InterpCompiler::ApplyLdftnDelegateCtorPeep(const uint8_t* ip, OpcodePeepElem
     return -1;
 }
 
-bool InterpCompiler::ResolveAsyncCallToken(const uint8_t* ip)
-{
-    CorInfoTokenKind tokenKind =
-        ip[0] == CEE_CALL ? CORINFO_TOKENKIND_Await : CORINFO_TOKENKIND_AwaitVirtual;
-    ResolveToken(getU4LittleEndian(ip + 1), tokenKind, &m_resolvedAsyncCallToken);
-    return m_resolvedAsyncCallToken.hMethod != NULL;
-}
-
 bool InterpCompiler::IsRuntimeAsyncCall(const uint8_t* ip, OpcodePeepElement* pattern, void** ppComputedInfo)
 {
     CORINFO_RESOLVED_TOKEN awaitResolvedToken;
@@ -6757,10 +6764,6 @@ bool InterpCompiler::IsRuntimeAsyncCall(const uint8_t* ip, OpcodePeepElement* pa
         return false;
     }
 
-    if (!ResolveAsyncCallToken(ip))
-    {
-        return false;
-    }
     m_currentContinuationContextHandling = ContinuationContextHandling::ContinueOnCapturedContext;
     return true;
 }
@@ -6825,7 +6828,7 @@ bool InterpCompiler::IsRuntimeAsyncCallConfigureAwaitTask(const uint8_t* ip, Opc
         return false;
     }
 
-    return ResolveAsyncCallToken(ip);
+    return true;
 }
 
 bool InterpCompiler::IsRuntimeAsyncCallConfigureAwaitValueTaskExactStLoc(const uint8_t* ip, OpcodePeepElement* pattern, void** ppComputedInfo)
@@ -6895,7 +6898,7 @@ bool InterpCompiler::IsRuntimeAsyncCallConfigureAwaitValueTaskExactStLoc(const u
         return false;
     }
 
-    return ResolveAsyncCallToken(ip);
+    return true;
 }
 
 bool InterpCompiler::IsRuntimeAsyncCallConfigureAwaitValueTask(const uint8_t* ip, OpcodePeepElement* pattern, void** ppComputedInfo)
@@ -6928,7 +6931,7 @@ bool InterpCompiler::IsRuntimeAsyncCallConfigureAwaitValueTask(const uint8_t* ip
         return false;
     }
 
-    return ResolveAsyncCallToken(ip);
+    return true;
 }
 
 int InterpCompiler::ApplyConvRUnR4Peep(const uint8_t* ip, OpcodePeepElement* peep, void* computedInfo)

src/coreclr/interpreter/compiler.h

@@ -927,7 +927,6 @@ class InterpCompiler
     bool    IsLdftnDelegateCtorPeep(const uint8_t* ip, OpcodePeepElement* peep, void** outComputedInfo);
     int     ApplyLdftnDelegateCtorPeep(const uint8_t* ip, OpcodePeepElement* peep, void* computedInfo);
 
-    bool ResolveAsyncCallToken(const uint8_t* ip);
     enum class ContinuationContextHandling : uint8_t
     {
         ContinueOnCapturedContext,

src/coreclr/jit/ee_il_dll.hpp

@@ -339,7 +339,12 @@ inline var_types Compiler::TypeHandleToVarType(CorInfoType jitType, CORINFO_CLAS
     return type;
 }
 
-inline CORINFO_CALLINFO_FLAGS combine(CORINFO_CALLINFO_FLAGS flag1, CORINFO_CALLINFO_FLAGS flag2)
+constexpr CORINFO_CALLINFO_FLAGS operator|(CORINFO_CALLINFO_FLAGS a, CORINFO_CALLINFO_FLAGS b)
 {
-    return (CORINFO_CALLINFO_FLAGS)(flag1 | flag2);
+    return (CORINFO_CALLINFO_FLAGS)((uint32_t)a | (uint32_t)b);
+}
+
+inline CORINFO_CALLINFO_FLAGS& operator|=(CORINFO_CALLINFO_FLAGS& a, CORINFO_CALLINFO_FLAGS b)
+{
+    return a = (CORINFO_CALLINFO_FLAGS)((uint32_t)a | (uint32_t)b);
 }

src/coreclr/jit/importer.cpp

@@ -8600,7 +8600,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 JITDUMP(" %08X", resolvedToken.token);
 
                 eeGetCallInfo(&resolvedToken, (prefixFlags & PREFIX_CONSTRAINED) ? &constrainedResolvedToken : nullptr,
-                              combine(CORINFO_CALLINFO_SECURITYCHECKS, CORINFO_CALLINFO_LDFTN), &callInfo);
+                              CORINFO_CALLINFO_SECURITYCHECKS | CORINFO_CALLINFO_LDFTN, &callInfo);
 
                 // This check really only applies to intrinsic Array.Address methods
                 if (callInfo.sig.callConv & CORINFO_CALLCONV_PARAMTYPE)
@@ -8639,8 +8639,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 JITDUMP(" %08X", resolvedToken.token);
 
                 eeGetCallInfo(&resolvedToken, nullptr /* constraint typeRef */,
-                              combine(combine(CORINFO_CALLINFO_SECURITYCHECKS, CORINFO_CALLINFO_LDFTN),
-                                      CORINFO_CALLINFO_CALLVIRT),
+                              CORINFO_CALLINFO_SECURITYCHECKS | CORINFO_CALLINFO_LDFTN | CORINFO_CALLINFO_CALLVIRT,
                               &callInfo);
 
                 // This check really only applies to intrinsic Array.Address methods
@@ -8773,7 +8772,7 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 _impResolveToken(CORINFO_TOKENKIND_NewObj);
 
                 eeGetCallInfo(&resolvedToken, nullptr /* constraint typeRef*/,
-                              combine(CORINFO_CALLINFO_SECURITYCHECKS, CORINFO_CALLINFO_ALLOWINSTPARAM), &callInfo);
+                              CORINFO_CALLINFO_SECURITYCHECKS | CORINFO_CALLINFO_ALLOWINSTPARAM, &callInfo);
 
                 mflags = callInfo.methodFlags;
 
@@ -9000,6 +8999,14 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                 // many other places.  We unfortunately embed that knowledge here.
                 if (opcode != CEE_CALLI)
                 {
+                    CORINFO_CALLINFO_FLAGS callInfoFlags =
+                        CORINFO_CALLINFO_ALLOWINSTPARAM | CORINFO_CALLINFO_SECURITYCHECKS;
+
+                    if (opcode == CEE_CALLVIRT)
+                    {
+                        callInfoFlags |= CORINFO_CALLINFO_CALLVIRT;
+                    }
+
                     bool        isAwait            = false;
                     int         configVal          = -1; // -1 not configured, 0/1 configured to false/true
                     const BYTE* codeAddrAfterMatch = nullptr;
@@ -9021,44 +9028,36 @@ void Compiler::impImportBlockCode(BasicBlock* block)
                             {
                                 prefixFlags |= PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT;
                             }
+
+                            callInfoFlags |= CORINFO_CALLINFO_ALLOWASYNCVARIANT;
                         }
                     }
 
+                    _impResolveToken(CORINFO_TOKENKIND_Method);
+
+                    eeGetCallInfo(&resolvedToken,
+                                  (prefixFlags & PREFIX_CONSTRAINED) ? &constrainedResolvedToken : nullptr,
+                                  callInfoFlags, &callInfo);
+
                     if (isAwait)
                     {
-                        _impResolveToken(opcode == CEE_CALLVIRT ? CORINFO_TOKENKIND_AwaitVirtual
-                                                                : CORINFO_TOKENKIND_Await);
-                        if (resolvedToken.hMethod != nullptr)
+                        if (callInfo.resolvedAsyncVariant != NO_METHOD_HANDLE)
                         {
-                            // There is a runtime async variant that is implicitly awaitable, just call that.
-                            // skip the await pattern to the last token.
                             codeAddr   = codeAddrAfterMatch;
                             opcodeOffs = awaitOffset;
+
+                            JITDUMP("Async variant provided by VM\n");
+                            // Update the token as we may need it again for various runtime lookups
+                            resolvedToken.hMethod = callInfo.resolvedAsyncVariant;
                         }
                         else
                         {
-                            // This can happen in cases when the Task-returning method is not a runtime Async
-                            // function. For example "T M1<T>(T arg) => arg" when called with a Task argument.
-                            // It can also happen generally if the VM does not think using the async entry point
-                            // is worth it. Treat these as a regular call that is Awaited.
-                            _impResolveToken(CORINFO_TOKENKIND_Method);
                             prefixFlags &= ~(PREFIX_IS_TASK_AWAIT | PREFIX_TASK_AWAIT_CONTINUE_ON_CAPTURED_CONTEXT);
                             isAwait = false;
 
                             JITDUMP("No async variant provided by VM, treating as regular call that is awaited\n");
                         }
                     }
-                    else
-                    {
-                        _impResolveToken(CORINFO_TOKENKIND_Method);
-                    }
-
-                    eeGetCallInfo(&resolvedToken,
-                                  (prefixFlags & PREFIX_CONSTRAINED) ? &constrainedResolvedToken : nullptr,
-                                  // this is how impImportCall invokes getCallInfo
-                                  combine(combine(CORINFO_CALLINFO_ALLOWINSTPARAM, CORINFO_CALLINFO_SECURITYCHECKS),
-                                          (opcode == CEE_CALLVIRT) ? CORINFO_CALLINFO_CALLVIRT : CORINFO_CALLINFO_NONE),
-                                  &callInfo);
                 }
                 else
                 {

src/coreclr/jit/importercalls.cpp

@@ -6923,11 +6923,6 @@ void Compiler::impSetupAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned pref
             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");

src/coreclr/jit/morph.cpp

@@ -5358,14 +5358,14 @@ GenTree* Compiler::fgMorphTailCallViaHelpers(GenTreeCall* call, CORINFO_TAILCALL
         {
             assert(!call->tailCallInfo->GetSig()->hasTypeArg());
 
-            CORINFO_CALL_INFO callInfo;
-            unsigned          flags = CORINFO_CALLINFO_LDFTN;
+            CORINFO_CALL_INFO      callInfo;
+            CORINFO_CALLINFO_FLAGS flags = CORINFO_CALLINFO_LDFTN;
             if (call->tailCallInfo->IsCallvirt())
             {
                 flags |= CORINFO_CALLINFO_CALLVIRT;
             }
 
-            eeGetCallInfo(call->tailCallInfo->GetToken(), nullptr, (CORINFO_CALLINFO_FLAGS)flags, &callInfo);
+            eeGetCallInfo(call->tailCallInfo->GetToken(), nullptr, flags, &callInfo);
             target = getVirtMethodPointerTree(thisPtrStubArg, call->tailCallInfo->GetToken(), &callInfo);
         }
 

src/coreclr/nativeaot/Runtime/PalLimitedContext.h

@@ -25,6 +25,8 @@ struct PAL_LIMITED_CONTEXT
     // Includes special registers, callee saved registers and general purpose registers used to return values from functions (not floating point return registers)
 #ifdef TARGET_ARM
     uintptr_t  R0;
+    uintptr_t  R1;
+    uintptr_t  R2;
     uintptr_t  R4;
     uintptr_t  R5;
     uintptr_t  R6;

src/coreclr/nativeaot/Runtime/StackFrameIterator.cpp

@@ -441,6 +441,8 @@ void StackFrameIterator::InternalInit(Thread * pThreadToWalk, PTR_PAL_LIMITED_CO
     // scratch regs
     //
     m_RegDisplay.pR0  = (PTR_uintptr_t)PTR_TO_MEMBER_TADDR(PAL_LIMITED_CONTEXT, pCtx, R0);
+    m_RegDisplay.pR1  = (PTR_uintptr_t)PTR_TO_MEMBER_TADDR(PAL_LIMITED_CONTEXT, pCtx, R1);
+    m_RegDisplay.pR2  = (PTR_uintptr_t)PTR_TO_MEMBER_TADDR(PAL_LIMITED_CONTEXT, pCtx, R2);
 
 #elif defined(TARGET_ARM64)
     //

src/coreclr/nativeaot/Runtime/arm/AsmOffsetsCpu.h

@@ -30,20 +30,22 @@ PLAT_ASM_OFFSET(10, StackFrameIterator, m_RegDisplay)
 PLAT_ASM_OFFSET(100, StackFrameIterator, m_OriginalControlPC)
 PLAT_ASM_OFFSET(104, StackFrameIterator, m_pPreviousTransitionFrame)
 
-PLAT_ASM_SIZEOF(70, PAL_LIMITED_CONTEXT)
-PLAT_ASM_OFFSET(24, PAL_LIMITED_CONTEXT, IP)
+PLAT_ASM_SIZEOF(78, PAL_LIMITED_CONTEXT)
+PLAT_ASM_OFFSET(2c, PAL_LIMITED_CONTEXT, IP)
 
 PLAT_ASM_OFFSET(0, PAL_LIMITED_CONTEXT, R0)
-PLAT_ASM_OFFSET(4, PAL_LIMITED_CONTEXT, R4)
-PLAT_ASM_OFFSET(8, PAL_LIMITED_CONTEXT, R5)
-PLAT_ASM_OFFSET(0c, PAL_LIMITED_CONTEXT, R6)
-PLAT_ASM_OFFSET(10, PAL_LIMITED_CONTEXT, R7)
-PLAT_ASM_OFFSET(14, PAL_LIMITED_CONTEXT, R8)
-PLAT_ASM_OFFSET(18, PAL_LIMITED_CONTEXT, R9)
-PLAT_ASM_OFFSET(1c, PAL_LIMITED_CONTEXT, R10)
-PLAT_ASM_OFFSET(20, PAL_LIMITED_CONTEXT, R11)
-PLAT_ASM_OFFSET(28, PAL_LIMITED_CONTEXT, SP)
-PLAT_ASM_OFFSET(2c, PAL_LIMITED_CONTEXT, LR)
+PLAT_ASM_OFFSET(4, PAL_LIMITED_CONTEXT, R1)
+PLAT_ASM_OFFSET(8, PAL_LIMITED_CONTEXT, R2)
+PLAT_ASM_OFFSET(0c, PAL_LIMITED_CONTEXT, R4)
+PLAT_ASM_OFFSET(10, PAL_LIMITED_CONTEXT, R5)
+PLAT_ASM_OFFSET(14, PAL_LIMITED_CONTEXT, R6)
+PLAT_ASM_OFFSET(18, PAL_LIMITED_CONTEXT, R7)
+PLAT_ASM_OFFSET(1c, PAL_LIMITED_CONTEXT, R8)
+PLAT_ASM_OFFSET(20, PAL_LIMITED_CONTEXT, R9)
+PLAT_ASM_OFFSET(24, PAL_LIMITED_CONTEXT, R10)
+PLAT_ASM_OFFSET(28, PAL_LIMITED_CONTEXT, R11)
+PLAT_ASM_OFFSET(30, PAL_LIMITED_CONTEXT, SP)
+PLAT_ASM_OFFSET(34, PAL_LIMITED_CONTEXT, LR)
 
 PLAT_ASM_SIZEOF(80, REGDISPLAY)
 PLAT_ASM_OFFSET(38, REGDISPLAY, SP)

src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj

@@ -1,4 +1,8 @@
 <Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <Features>$(Features);runtime-async=on</Features>
+  </PropertyGroup>
+
   <PropertyGroup>
     <EmitCompilerGeneratedFiles>true</EmitCompilerGeneratedFiles>
   </PropertyGroup>