Support dynamically creating tailored continuation layouts (#120411)

Add ability for the VM to dynamically create continuation layout types
and for the JIT to request such types to be created.

Author: jakobbotsch

.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/AsyncHelpers.CoreCLR.cs

@@ -1698,29 +1698,34 @@ struct CORINFO_EE_INFO
     CORINFO_OS  osType;
 };
 
+// Keep in sync with ContinuationFlags enum in BCL sources
 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,
+    // Note: the following 'Has' members determine the members present at
+    // the beginning of the continuation's data chunk. Each field is
+    // pointer sized when present, apart from the result that has variable
+    // size.
+
+    // Whether or not the continuation starts with an OSR IL offset.
+    CORINFO_CONTINUATION_HAS_OSR_ILOFFSET = 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.
+    CORINFO_CONTINUATION_HAS_EXCEPTION = 2,
+    // If this bit is set the continuation has space for a continuation
+    // context.
+    CORINFO_CONTINUATION_HAS_CONTINUATION_CONTEXT = 4,
+    // If this bit is set the continuation has space to store a result
+    // returned by the callee.
+    CORINFO_CONTINUATION_HAS_RESULT = 8,
     // If this bit is set the continuation should continue on the thread
     // pool.
-    CORINFO_CONTINUATION_CONTINUE_ON_THREAD_POOL = 8,
-    // If this bit is set the continuation has a SynchronizationContext
-    // that we should continue on.
-    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT = 16,
-    // 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_THREAD_POOL = 16,
+    // If this bit is set the continuation context is a
+    // SynchronizationContext that we should continue on.
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT = 32,
+    // If this bit is set the continuation context is a TaskScheduler that
+    // we should continue on.
+    CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_TASK_SCHEDULER = 64,
 };
 
 struct CORINFO_ASYNC_INFO
@@ -1735,13 +1740,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;
     // Method handle for AsyncHelpers.CaptureExecutionContext
     CORINFO_METHOD_HANDLE captureExecutionContextMethHnd;
     // Method handle for AsyncHelpers.RestoreExecutionContext
@@ -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 */ + 8 /* Flags + State */)
+
 
 /* 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,
+        size_t objRefsSize
+        ) = 0;
+
     // Optionally, convert calli to regular method call. This is for PInvoke argument marshalling.
     virtual bool convertPInvokeCalliToCall(
             CORINFO_RESOLVED_TOKEN *    pResolvedToken,

src/coreclr/inc/corinfo.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/dacvars.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,
+          size_t objRefsSize) override;
+
 CORINFO_METHOD_HANDLE getAsyncResumptionStub() override;
 
 bool convertPInvokeCalliToCall(

src/coreclr/inc/icorjitinfoimpl_generated.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 = { /* 68e93e9d-dd28-49ca-9ebc-a01a54532bb3 */
+    0x68e93e9d,
+    0xdd28,
+    0x49ca,
+    {0x9e, 0xbc, 0xa0, 0x1a, 0x54, 0x53, 0x2b, 0xb3}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/inc/jiteeversionguid.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_names_generated.h

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