Support dynamically creating tailored continuation layouts

.github/prompts/add-new-jit-ee-api.prompt.md

@@ -1,23 +1,23 @@
 ---
 mode: 'agent'
-tools: ['githubRepo', 'codebase', 'terminalLastCommand']
+tools: ['fetch', 'codebase', 'runCommands', 'usages', 'search', 'think']
 description: 'Add a new API to the JIT-VM (aka JIT-EE) interface in the codebase.'
 ---
 
 #### 1 — Goal
 
-Implement **one** new JIT-VM (also known as JIT-EE) API and all supporting glue. 
+Implement **one** new JIT-VM (also known as JIT-EE) API and all supporting glue.
 The JIT-VM interface defines the APIs through which the JIT compiler communicates with the runtime (VM).
 
 #### 2 — Prerequisites for the model
 
 * You have full repo access
-* You may run scripts (e.g., `.sh` or `.bat`) 
+* You may run scripts (e.g., `.sh` or `.bat`)
 * Ask **clarifying questions** before the first code change if anything (signature, types, platform constraints) is unclear.
 
 #### 3 — Required user inputs
 
-Ask the user for a C-like signature of the new API if it's not provided. 
+Ask the user for a C-like signature of the new API if it's not provided.
 Suggest `<repo_root>/src/coreclr/tools/Common/JitInterface/ThunkGenerator/ThunkInput.txt` file as a reference. Example:
 
 ```
@@ -83,8 +83,8 @@ Use the correct directory for the script to run.
 +}
 ```
 
-6. Now implement the most complex part - SuperPMI. SuperPMI acts as a (de)serializer for JIT-VM queries in order 
-to then replay them without the actual VM to speed up jit-diffs and other scenarios. All parameters and return 
+6. Now implement the most complex part - SuperPMI. SuperPMI acts as a (de)serializer for JIT-VM queries in order
+to then replay them without the actual VM to speed up jit-diffs and other scenarios. All parameters and return
 values recorded/restored using special primitve types and helpers. We need to update the following files:
 
 * `<repo_root>/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h`:
@@ -96,7 +96,7 @@ Go through each of them one by one.
 
 * `<repo_root>/src/coreclr/tools/superpmi/superpmi-shared/agnostic.h`:
 Define two `Agnostic_*` types for input arguments and another one for output parameters (return value, output arguments).
- Do not create them if one of the generics ones can be re-used such as `DLD`, `DD`, `DLDL`, etc. Use `DWORD*` 
+ Do not create them if one of the generics ones can be re-used such as `DLD`, `DD`, `DLDL`, etc. Use `DWORD*`
  like types for integers. Inspect the whole file to see how other APIs are defined.
 
 * `<repo_root>/src/coreclr/tools/superpmi/superpmi-shared/lwmlist.h`:
@@ -126,7 +126,7 @@ Now add a new element to `enum mcPackets` enum in the same file. Example:
 ```
 
 * `<repo_root>/src/coreclr/tools/superpmi/superpmi-shared/methodcontext.cpp`:
-Add the implementation of the 3 methods to `methodcontext.cpp` at the end of it. 
+Add the implementation of the 3 methods to `methodcontext.cpp` at the end of it.
 Consider other similar methods in the file for reference. Do not change implementations of other methods in the file. Example:
 
 ```diff

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

@@ -739,6 +739,9 @@ internal unsafe struct MethodTable
         [FieldOffset(ElementTypeOffset)]
         public MethodTable*** PerInstInfo;
 
+        [FieldOffset(ElementTypeOffset)]
+        public CORINFO_CONTINUATION_DATA_OFFSETS* ContinuationOffsets;
+
         /// <summary>
         /// This interface map used to list out the set of interfaces. Only meaningful if InterfaceCount is non-zero.
         /// </summary>
@@ -897,6 +900,8 @@ public int MultiDimensionalArrayRank
 
         public bool IsArray => (Flags & enum_flag_Category_Array_Mask) == enum_flag_Category_Array;
 
+        public bool IsContinuation => ParentMethodTable == (MethodTable*)typeof(Continuation).TypeHandle.Value;
+
         public bool HasInstantiation => (Flags & enum_flag_HasComponentSize) == 0 && (Flags & enum_flag_GenericsMask) != enum_flag_GenericsMask_NonGeneric;
 
         public bool IsGenericTypeDefinition => (Flags & (enum_flag_HasComponentSize | enum_flag_GenericsMask)) == enum_flag_GenericsMask_TypicalInst;

src/coreclr/inc/corinfo.h

@@ -1700,27 +1700,20 @@ struct CORINFO_EE_INFO
 
 enum CorInfoContinuationFlags
 {
-    // Whether or not the continuation expects the result to be boxed and
-    // placed in the GCData array at index 0. Not set if the callee is void.
-    CORINFO_CONTINUATION_RESULT_IN_GCDATA = 1,
     // If this bit is set the continuation resumes inside a try block and thus
     // if an exception is being propagated, needs to be resumed. The exception
     // should be placed at index 0 or 1 depending on whether the continuation
     // also expects a result.
-    CORINFO_CONTINUATION_NEEDS_EXCEPTION = 2,
-    // If this bit is set the continuation has the IL offset that inspired the
-    // OSR method saved in the beginning of 'Data', or -1 if the continuation
-    // belongs to a tier 0 method.
-    CORINFO_CONTINUATION_OSR_IL_OFFSET_IN_DATA = 4,
+    CORINFO_CONTINUATION_NEEDS_EXCEPTION = 1,
     // If this bit is set the continuation should continue on the thread
     // pool.
-    CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL = 8,
+    CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL = 2,
     // If this bit is set the continuation has a SynchronizationContext
     // that we should continue on.
-    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT = 16,
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT = 4,
     // If this bit is set the continuation has a TaskScheduler
     // that we should continue on.
-    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_TASK_SCHEDULER = 32,
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_TASK_SCHEDULER = 8,
 };
 
 struct CORINFO_ASYNC_INFO
@@ -1735,10 +1728,6 @@ struct CORINFO_ASYNC_INFO
     CORINFO_FIELD_HANDLE continuationStateFldHnd;
     // 'Flags' field
     CORINFO_FIELD_HANDLE continuationFlagsFldHnd;
-    // 'Data' field
-    CORINFO_FIELD_HANDLE continuationDataFldHnd;
-    // 'GCData' field
-    CORINFO_FIELD_HANDLE continuationGCDataFldHnd;
     // Whether or not the continuation needs to be allocated through the
     // helper that also takes a method handle
     bool continuationsNeedMethodHandle;
@@ -1751,6 +1740,15 @@ struct CORINFO_ASYNC_INFO
     CORINFO_METHOD_HANDLE restoreContextsMethHnd;
 };
 
+// Offsets for continuation data layout
+struct CORINFO_CONTINUATION_DATA_OFFSETS
+{
+    uint32_t Result;
+    uint32_t Exception;
+    uint32_t ContinuationContext;
+    uint32_t KeepAlive;
+};
+
 // Flags passed from JIT to runtime.
 enum CORINFO_GET_TAILCALL_HELPERS_FLAGS
 {
@@ -1960,6 +1958,8 @@ struct CORINFO_FPSTRUCT_LOWERING
 #define OFFSETOF__CORINFO_Span__reference                 0
 #define OFFSETOF__CORINFO_Span__length                    TARGET_POINTER_SIZE
 
+#define OFFSETOF__CORINFO_Continuation__data              (SIZEOF__CORINFO_Object + TARGET_POINTER_SIZE /* Next */ + TARGET_POINTER_SIZE /* Resume */ + TARGET_POINTER_SIZE /* Flags + maybe padding */)
+
 
 /* data to optimize delegate construction */
 struct DelegateCtorArgs
@@ -3342,6 +3342,12 @@ class ICorDynamicInfo : public ICorStaticInfo
 
     virtual CORINFO_METHOD_HANDLE getAsyncResumptionStub() = 0;
 
+    virtual CORINFO_CLASS_HANDLE getContinuationType(
+        size_t dataSize,
+        bool* objRefs,
+        const CORINFO_CONTINUATION_DATA_OFFSETS& dataOffsets
+        ) = 0;
+
     // Optionally, convert calli to regular method call. This is for PInvoke argument marshalling.
     virtual bool convertPInvokeCalliToCall(
             CORINFO_RESOLVED_TOKEN *    pResolvedToken,

src/coreclr/inc/dacvars.h

@@ -175,6 +175,8 @@ DEFINE_DACVAR(UNKNOWN_POINTER_TYPE, dac__g_TypedReferenceMT, ::g_TypedReferenceM
 DEFINE_DACVAR(UNKNOWN_POINTER_TYPE, dac__g_pWeakReferenceClass, ::g_pWeakReferenceClass)
 DEFINE_DACVAR(UNKNOWN_POINTER_TYPE, dac__g_pWeakReferenceOfTClass, ::g_pWeakReferenceOfTClass)
 
+DEFINE_DACVAR_VOLATILE(UNKNOWN_POINTER_TYPE, dac__g_pContinuationClassIfSubTypeCreated, ::g_pContinuationClassIfSubTypeCreated)
+
 #ifdef FEATURE_COMINTEROP
 DEFINE_DACVAR(UNKNOWN_POINTER_TYPE, dac__g_pBaseCOMObject, ::g_pBaseCOMObject)
 #endif

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -655,6 +655,11 @@ bool getTailCallHelpers(
           CORINFO_GET_TAILCALL_HELPERS_FLAGS flags,
           CORINFO_TAILCALL_HELPERS* pResult) override;
 
+CORINFO_CLASS_HANDLE getContinuationType(
+          size_t dataSize,
+          bool* objRefs,
+          const CORINFO_CONTINUATION_DATA_OFFSETS& dataOffsets) override;
+
 CORINFO_METHOD_HANDLE getAsyncResumptionStub() override;
 
 bool convertPInvokeCalliToCall(

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* 3d2bdd20-eced-4a07-b9fb-227ce7f55fcd */
-    0x3d2bdd20,
-    0xeced,
-    0x4a07,
-    {0xb9, 0xfb, 0x22, 0x7c, 0xe7, 0xf5, 0x5f, 0xcd}
+constexpr GUID JITEEVersionIdentifier = { /* 156f1a7f-c06a-40fc-a984-987c8b4c2ba7 */
+    0x156f1a7f,
+    0xc06a,
+    0x40fc,
+    {0xa9, 0x84, 0x98, 0x7c, 0x8b, 0x4c, 0x2b, 0xa7}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -159,6 +159,7 @@ DEF_CLR_API(getFieldThreadLocalStoreID)
 DEF_CLR_API(GetDelegateCtor)
 DEF_CLR_API(MethodCompileComplete)
 DEF_CLR_API(getTailCallHelpers)
+DEF_CLR_API(getContinuationType)
 DEF_CLR_API(getAsyncResumptionStub)
 DEF_CLR_API(convertPInvokeCalliToCall)
 DEF_CLR_API(notifyInstructionSetUsage)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -1536,6 +1536,17 @@ bool WrapICorJitInfo::getTailCallHelpers(
     return temp;
 }
 
+CORINFO_CLASS_HANDLE WrapICorJitInfo::getContinuationType(
+          size_t dataSize,
+          bool* objRefs,
+          const CORINFO_CONTINUATION_DATA_OFFSETS& dataOffsets)
+{
+    API_ENTER(getContinuationType);
+    CORINFO_CLASS_HANDLE temp = wrapHnd->getContinuationType(dataSize, objRefs, dataOffsets);
+    API_LEAVE(getContinuationType);
+    return temp;
+}
+
 CORINFO_METHOD_HANDLE WrapICorJitInfo::getAsyncResumptionStub()
 {
     API_ENTER(getAsyncResumptionStub);