Async continuations for native AOT

The same program as in dotnet#121295 still works, but we can newly report async helpers as Intrinsic/Async.

This implements things around suspension/resumption. Most of this change is around handling continuation types.

Continuation types are synthetic types (created in the compiler) that derive from `Continuation` in the CoreLib. The JIT requests these based on the shape of locals it needs to preserve. What we get from the JIT in CorInfoImpl is size of the type and a pointer map (1011001 - non-zero means "GC pointer is here"). What we need to generate is a `MethodTable` for a type that derives from `Continuation` and has the specified GC layout after fields inherited from `Continuation`.

We already have a similar thing in the compiler ("`MethodTable`" we use for GC statics), however because we need to derive from `Continuation` and presumably need to have a working vtable (with Equals/GetHashCode/ToString), we can't use this. So we emit a normal `MethodTable` for a synthetic type. Maybe we could optimize this to the one without vtable in the future.

The synthetic type is a `MetadataType` that reports `Continuation` as the base type. It has instance fields of type `object` or `nint`, based on what we need. Because the standard type layout algorithm (the thing that assigns offsets to fields) would attempt to group GC fields together (it's a layout that works better for the GC), we also have a custom layouting algorithm that lays out these fields sequentially.

The resumption stub is just a stubbed out TODO for someone else to look into for now.

Author: MichalStrehovsky

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

@@ -12,6 +12,10 @@
 using System.Threading;
 using System.Threading.Tasks;
 
+#if NATIVEAOT
+using Internal.Runtime;
+#endif
+
 namespace System.Runtime.CompilerServices
 {
     internal struct ExecutionAndSyncBlockStore
@@ -144,11 +148,16 @@ private struct RuntimeAsyncAwaitState
 
         private static unsafe Continuation AllocContinuation(Continuation prevContinuation, MethodTable* contMT)
         {
-            Continuation newContinuation = (Continuation)RuntimeTypeHandle.InternalAllocNoChecks(contMT);
+#if NATIVEAOT
+            Continuation newContinuation = (Continuation)RuntimeImports.RhNewObject(contMT);
+#else
+            Continuation newContinuation = (Continuation)RuntimeTypeHandle.InternalAllocNoChecks(contMT);       
+#endif
             prevContinuation.Next = newContinuation;
             return newContinuation;
         }
 
+#if !NATIVEAOT
         private static unsafe Continuation AllocContinuationMethod(Continuation prevContinuation, MethodTable* contMT, int keepAliveOffset, MethodDesc* method)
         {
             LoaderAllocator loaderAllocator = RuntimeMethodHandle.GetLoaderAllocator(new RuntimeMethodHandleInternal((IntPtr)method));
@@ -170,6 +179,7 @@ private static unsafe Continuation AllocContinuationClass(Continuation prevConti
             }
             return newContinuation;
         }
+#endif
 
         [BypassReadyToRun]
         [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)]

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

@@ -55,6 +55,10 @@
       <Link>ExceptionStringID.cs</Link>
     </Compile>
   </ItemGroup>
+  <ItemGroup>
+    <!-- TODO: (async) once we know which helpers can actually be shared, move those to libraries partition -->
+    <Compile Include="..\..\..\System.Private.CoreLib\src\System\Runtime\CompilerServices\AsyncHelpers.CoreCLR.cs" />
+  </ItemGroup>
   <ItemGroup>
     <Compile Include="Internal\Runtime\CompilerHelpers\DelegateHelpers.cs" />
     <Compile Include="Internal\Runtime\CompilerHelpers\LibraryInitializer.cs" />

src/coreclr/tools/Common/Compiler/CompilerTypeSystemContext.Async.cs

@@ -1,6 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
+using System.Collections.Generic;
+using System.Text;
+using System.Threading;
+
 using Internal.IL;
 using Internal.TypeSystem;
 using Internal.TypeSystem.Ecma;
@@ -40,5 +45,226 @@ protected override bool CompareValueToValue(AsyncMethodVariant value1, AsyncMeth
             protected override AsyncMethodVariant CreateValueFromKey(EcmaMethod key) => new AsyncMethodVariant(key);
         }
         private AsyncVariantImplHashtable _asyncVariantImplHashtable = new AsyncVariantImplHashtable();
+
+        public MetadataType GetContinuationType(GCPointerMap pointerMap)
+        {
+            return _continuationTypeHashtable.GetOrCreateValue(pointerMap);
+        }
+
+        private sealed class ContinuationTypeHashtable : LockFreeReaderHashtable<GCPointerMap, ContinuationType>
+        {
+            private readonly CompilerTypeSystemContext _parent;
+            private MetadataType _continuationType;
+
+            public ContinuationTypeHashtable(CompilerTypeSystemContext parent)
+                => _parent = parent;
+
+            protected override int GetKeyHashCode(GCPointerMap key) => key.GetHashCode();
+            protected override int GetValueHashCode(ContinuationType value) => value.PointerMap.GetHashCode();
+            protected override bool CompareKeyToValue(GCPointerMap key, ContinuationType value) => key.Equals(value.PointerMap);
+            protected override bool CompareValueToValue(ContinuationType value1, ContinuationType value2)
+                => value1.PointerMap.Equals(value2.PointerMap);
+            protected override ContinuationType CreateValueFromKey(GCPointerMap key)
+            {
+                if (_continuationType == null)
+                    _continuationType = _parent.SystemModule.GetKnownType("System.Runtime.CompilerServices"u8, "Continuation"u8);
+                return new ContinuationType(_continuationType, key);
+            }
+        }
+        private ContinuationTypeHashtable _continuationTypeHashtable;
+
+        /// <summary>
+        /// An async continuation type. The code generator will request this to store local state
+        /// through an async suspension/resumption. We only identify these using a <see cref="GCPointerMap"/>
+        /// since that's all the code generator cares about - size of the type, and where the GC pointers are.
+        /// </summary>
+        private sealed class ContinuationType : MetadataType
+        {
+            private readonly MetadataType _continuationType;
+            private FieldDesc[] _fields;
+            public GCPointerMap PointerMap { get; }
+
+            public override DefType[] ExplicitlyImplementedInterfaces => [];
+            public override ReadOnlySpan<byte> Name => Encoding.UTF8.GetBytes(DiagnosticName);
+            public override ReadOnlySpan<byte> Namespace => [];
+
+            // The layout of the type is "sequential-in-spirit", but since there are GC pointers,
+            // the standard layout algorithm wouldn't respect that. We have a custom layout algorithm.
+            // The following layout-related properties are meaningless.
+            public override bool IsExplicitLayout => false;
+            public override bool IsSequentialLayout => false;
+            public override bool IsExtendedLayout => false;
+            public override bool IsAutoLayout => false;
+            public override ClassLayoutMetadata GetClassLayout() => default;
+
+            public override bool IsBeforeFieldInit => false;
+            public override ModuleDesc Module => _continuationType.Module;
+            public override MetadataType BaseType => _continuationType;
+            public override bool IsSealed => true;
+            public override bool IsAbstract => false;
+            public override MetadataType ContainingType => null;
+            public override PInvokeStringFormat PInvokeStringFormat => default;
+            public override string DiagnosticName => $"ContinuationType_{PointerMap}";
+            public override string DiagnosticNamespace => "";
+            protected override int ClassCode => 0x528741a;
+            public override TypeSystemContext Context => _continuationType.Context;
+
+            public ContinuationType(MetadataType continuationType, GCPointerMap pointerMap)
+                => (_continuationType, PointerMap) = (continuationType, pointerMap);
+
+            public override bool HasCustomAttribute(string attributeNamespace, string attributeName) => false;
+            public override IEnumerable<MetadataType> GetNestedTypes() => [];
+            public override MetadataType GetNestedType(string name) => null;
+            protected override MethodImplRecord[] ComputeVirtualMethodImplsForType() => [];
+            public override MethodImplRecord[] FindMethodsImplWithMatchingDeclName(ReadOnlySpan<byte> name) => [];
+
+            protected override int CompareToImpl(TypeDesc other, TypeSystemComparer comparer)
+            {
+                GCPointerMap otherPointerMap = ((ContinuationType)other).PointerMap;
+                return PointerMap.CompareTo(otherPointerMap);
+            }
+
+            public override int GetHashCode() => PointerMap.GetHashCode();
+
+            protected override TypeFlags ComputeTypeFlags(TypeFlags mask)
+            {
+                TypeFlags flags = 0;
+
+                if ((mask & TypeFlags.HasGenericVarianceComputed) != 0)
+                {
+                    flags |= TypeFlags.HasGenericVarianceComputed;
+                }
+
+                if ((mask & TypeFlags.CategoryMask) != 0)
+                {
+                    flags |= TypeFlags.Class;
+                }
+
+                flags |= TypeFlags.HasFinalizerComputed;
+                flags |= TypeFlags.AttributeCacheComputed;
+
+                return flags;
+            }
+
+            private void InitializeFields()
+            {
+                FieldDesc[] fields = new FieldDesc[PointerMap.Size];
+
+                for (int i = 0; i < PointerMap.Size; i++)
+                    fields[i] = new ContinuationField(this, i);
+
+                Interlocked.CompareExchange(ref _fields, fields, null);
+            }
+            public override IEnumerable<FieldDesc> GetFields()
+            {
+                if (_fields == null)
+                {
+                    InitializeFields();
+                }
+                return _fields;
+            }
+
+            /// <summary>
+            /// A field on a continuation type. The type of the field is determined by consulting the
+            /// associated GC pointer map and it's either `object` or `nint`.
+            /// </summary>
+            private sealed class ContinuationField : FieldDesc
+            {
+                private readonly ContinuationType _owningType;
+                private readonly int _index;
+
+                public ContinuationField(ContinuationType owningType, int index)
+                    => (_owningType, _index) = (owningType, index);
+
+                public override MetadataType OwningType => _owningType;
+                public override TypeDesc FieldType => Context.GetWellKnownType(_owningType.PointerMap[_index] ? WellKnownType.Object : WellKnownType.IntPtr);
+                public override bool HasEmbeddedSignatureData => false;
+                public override bool IsStatic => false;
+                public override bool IsInitOnly => false;
+                public override bool IsThreadStatic => false;
+                public override bool HasRva => false;
+                public override bool IsLiteral => false;
+                public override TypeSystemContext Context => _owningType.Context;
+                public override EmbeddedSignatureData[] GetEmbeddedSignatureData() => null;
+                public override bool HasCustomAttribute(string attributeNamespace, string attributeName) => false;
+
+                protected override int ClassCode => 0xc761a66;
+                protected override int CompareToImpl(FieldDesc other, TypeSystemComparer comparer)
+                {
+                    var otherField = (ContinuationField)other;
+                    int result = _index.CompareTo(otherField._index);
+                    if (result != 0)
+                        return result;
+
+                    return comparer.Compare(_owningType, otherField._owningType);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Layout algorithm that lays out continuation types. It ensures the type has the layout
+        /// that the code generator requested (the GC pointers are where we need them).
+        /// </summary>
+        private sealed class ContinuationTypeFieldLayoutAlgorithm : FieldLayoutAlgorithm
+        {
+            public override bool ComputeContainsGCPointers(DefType type)
+            {
+                // ContainsGCPointers because the base already has some.
+                Debug.Assert(type.BaseType.ContainsGCPointers);
+                return true;
+            }
+
+            public override ComputedInstanceFieldLayout ComputeInstanceLayout(DefType type, InstanceLayoutKind layoutKind)
+            {
+                var continuationType = (ContinuationType)type;
+                var continuationBaseType = (EcmaType)continuationType.BaseType;
+                Debug.Assert(continuationBaseType.Name.SequenceEqual("Continuation"u8));
+
+                LayoutInt pointerSize = continuationType.Context.Target.LayoutPointerSize;
+
+                LayoutInt dataOffset = continuationBaseType.InstanceByteCountUnaligned;
+
+#if DEBUG
+                // Validate we match the expected DataOffset value.
+                // DataOffset is a const int32 field in the Continuation class.
+                EcmaField dataOffsetField = continuationBaseType.GetField("DataOffset"u8);
+                Debug.Assert(dataOffsetField.IsLiteral);
+
+                var reader = dataOffsetField.MetadataReader;
+                var constant = reader.GetConstant(reader.GetFieldDefinition(dataOffsetField.Handle).GetDefaultValue());
+                Debug.Assert(constant.TypeCode == System.Reflection.Metadata.ConstantTypeCode.Int32);
+                int expectedDataOffset = reader.GetBlobReader(constant.Value).ReadInt32()
+                    + pointerSize.AsInt /* + MethodTable field */;
+                Debug.Assert(dataOffset.AsInt == expectedDataOffset);
+#endif
+                FieldAndOffset[] offsets = new FieldAndOffset[continuationType.PointerMap.Size];
+                int i = 0;
+
+                foreach (FieldDesc field in type.GetFields())
+                {
+                    Debug.Assert(field.FieldType.GetElementSize() == pointerSize);
+                    offsets[i++] = new FieldAndOffset(field, dataOffset);
+                    dataOffset += pointerSize;
+                }
+
+                return new ComputedInstanceFieldLayout
+                {
+                    FieldSize = pointerSize,
+                    FieldAlignment = pointerSize,
+                    ByteCountAlignment = pointerSize,
+                    ByteCountUnaligned = dataOffset,
+                    Offsets = offsets,
+                    IsAutoLayoutOrHasAutoLayoutFields = false,
+                    IsInt128OrHasInt128Fields = false,
+                    IsVectorTOrHasVectorTFields = false,
+                    LayoutAbiStable = false,
+                };
+            }
+
+            public override bool ComputeContainsByRefs(DefType type) => throw new NotImplementedException();
+            public override bool ComputeIsUnsafeValueType(DefType type) => throw new NotImplementedException();
+            public override ComputedStaticFieldLayout ComputeStaticFieldLayout(DefType type, StaticLayoutKind layoutKind) => default;
+            public override ValueTypeShapeCharacteristics ComputeValueTypeShapeCharacteristics(DefType type) => throw new NotImplementedException();
+        }
     }
 }

src/coreclr/tools/Common/JitInterface/CorInfoImpl.cs

@@ -685,6 +685,7 @@ private void CompileMethodCleanup()
 
 #if !READYTORUN
             _debugInfo = null;
+            _asyncResumptionStub = null;
 #endif
 
             _debugLocInfos = null;
@@ -3393,7 +3394,19 @@ private void getAsyncInfo(ref CORINFO_ASYNC_INFO pAsyncInfoOut)
         private CORINFO_CLASS_STRUCT_* getContinuationType(nuint dataSize, ref bool objRefs, nuint objRefsSize)
         {
             Debug.Assert(objRefsSize == (dataSize + (nuint)(PointerSize - 1)) / (nuint)PointerSize);
+#if READYTORUN
             throw new NotImplementedException("getContinuationType");
+#else
+            GCPointerMapBuilder gcMapBuilder = new GCPointerMapBuilder((int)dataSize, PointerSize);
+            ReadOnlySpan<bool> bools = MemoryMarshal.CreateReadOnlySpan(ref objRefs, (int)objRefsSize);
+            for (int i = 0; i < bools.Length; i++)
+            {
+                if (bools[i])
+                    gcMapBuilder.MarkGCPointer(i * PointerSize);
+            }
+
+            return ObjectToHandle(_compilation.TypeSystemContext.GetContinuationType(gcMapBuilder.ToGCMap()));
+#endif
         }
 
         private mdToken getMethodDefFromMethod(CORINFO_METHOD_STRUCT_* hMethod)
@@ -3630,9 +3643,6 @@ public static ReadyToRunHelperId GetReadyToRunHelperFromStaticBaseHelper(CorInfo
             return res;
         }
 
-        private void getFunctionFixedEntryPoint(CORINFO_METHOD_STRUCT_* ftn, bool isUnsafeFunctionPointer, ref CORINFO_CONST_LOOKUP pResult)
-        { throw new NotImplementedException("getFunctionFixedEntryPoint"); }
-
 #pragma warning disable CA1822 // Mark members as static
         private CorInfoHelpFunc getLazyStringLiteralHelper(CORINFO_MODULE_STRUCT_* handle)
 #pragma warning restore CA1822 // Mark members as static
@@ -3748,7 +3758,11 @@ private bool getTailCallHelpers(ref CORINFO_RESOLVED_TOKEN callToken, CORINFO_SI
         private CORINFO_METHOD_STRUCT_* getAsyncResumptionStub()
 #pragma warning restore CA1822 // Mark members as static
         {
+#if READYTORUN
             throw new NotImplementedException("Crossgen2 does not support runtime-async yet");
+#else
+            return ObjectToHandle(_asyncResumptionStub ??= new AsyncResumptionStub(MethodBeingCompiled));
+#endif
         }
 
         private byte[] _code;

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncResumptionStub.Mangling.cs

@@ -0,0 +1,15 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Internal.IL.Stubs;
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    public partial class AsyncResumptionStub : ILStubMethod, IPrefixMangledMethod
+    {
+        MethodDesc IPrefixMangledMethod.BaseMethod => _owningMethod;
+
+        string IPrefixMangledMethod.Prefix => "Resume";
+    }
+}

src/coreclr/tools/Common/TypeSystem/IL/Stubs/AsyncResumptionStub.Sorting.cs

@@ -0,0 +1,18 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Internal.IL.Stubs;
+using Internal.TypeSystem;
+
+namespace ILCompiler
+{
+    public partial class AsyncResumptionStub : ILStubMethod
+    {
+        protected override int ClassCode => 0x773ab1;
+
+        protected override int CompareToImpl(MethodDesc other, TypeSystemComparer comparer)
+        {
+            return comparer.Compare(_owningMethod, ((AsyncResumptionStub)other)._owningMethod);
+        }
+    }
+}