Add debug information for runtime async methods

.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/inc/cordebuginfo.h

@@ -431,4 +431,31 @@ class ICorDebugInfo
         // Source information about the IL instruction in the inlinee
         SourceTypes Source;
     };
+
+    struct AsyncContinuationVarInfo
+    {
+        // IL number of variable (or one of the special IL numbers, like TYPECTXT_ILNUM)
+        uint32_t VarNumber;
+        // Offset in continuation's data where this variable is stored
+        uint32_t Offset;
+    };
+
+    struct AsyncSuspensionPoint
+    {
+        // IL offset in the root method that resulted in the creation of this suspension point.
+        uint32_t RootILOffset;
+        // Index of inline tree node containing the IL offset (0 for root)
+        uint32_t Inlinee;
+        // IL offset that resulted in the creation of the suspension point.
+        uint32_t ILOffset;
+        // Count of AsyncContinuationVarInfo in array of locals starting where
+        // the previous suspension point's locals end.
+        uint32_t NumContinuationVars;
+    };
+
+    struct AsyncInfo
+    {
+        // Number of suspension points in the method.
+        uint32_t NumSuspensionPoints;
+    };
 };

src/coreclr/inc/corinfo.h

@@ -2912,6 +2912,16 @@ class ICorStaticInfo
             uint32_t                          numMappings         // [IN] Number of rich mappings
             ) = 0;
 
+    // Report async debug information to EE.
+    // The arrays are expected to be allocated with allocateArray
+    // and ownership is transferred to the EE with this call.
+    virtual void reportAsyncDebugInfo(
+            ICorDebugInfo::AsyncInfo*             asyncInfo,         // [IN] Async method information
+            ICorDebugInfo::AsyncSuspensionPoint*  suspensionPoints,  // [IN] Array of async suspension points, indexed by state number
+            ICorDebugInfo::AsyncContinuationVarInfo* vars,           // [IN] Array of async continuation variable info
+            uint32_t                              numVars            // [IN] Number of entries in the async vars array
+            ) = 0;
+
     // Report back some metadata about the compilation to the EE -- for
     // example, metrics about the compilation.
     virtual void reportMetadata(

src/coreclr/inc/icorjitinfoimpl_generated.h

@@ -453,6 +453,12 @@ void reportRichMappings(
           ICorDebugInfo::RichOffsetMapping* mappings,
           uint32_t numMappings) override;
 
+void reportAsyncDebugInfo(
+          ICorDebugInfo::AsyncInfo* asyncInfo,
+          ICorDebugInfo::AsyncSuspensionPoint* suspensionPoints,
+          ICorDebugInfo::AsyncContinuationVarInfo* vars,
+          uint32_t numVars) override;
+
 void reportMetadata(
           const char* key,
           const void* value,

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 = { /* e1c4ef4f-62a1-4055-813e-23837fdfdffa */
+    0xe1c4ef4f,
+    0x62a1,
+    0x4055,
+    {0x81, 0x3e, 0x23, 0x83, 0x7f, 0xdf, 0xdf, 0xfa}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/inc/readytorun.h

@@ -19,10 +19,10 @@
 //  src/coreclr/nativeaot/Runtime/inc/ModuleHeaders.h
 // If you update this, ensure you run `git grep MINIMUM_READYTORUN_MAJOR_VERSION`
 // and handle pending work.
-#define READYTORUN_MAJOR_VERSION 16
+#define READYTORUN_MAJOR_VERSION 17
 #define READYTORUN_MINOR_VERSION 0x0000
 
-#define MINIMUM_READYTORUN_MAJOR_VERSION 16
+#define MINIMUM_READYTORUN_MAJOR_VERSION 17
 
 // R2R Version 2.1 adds the InliningInfo section
 // R2R Version 2.2 adds the ProfileDataInfo section
@@ -47,6 +47,7 @@
 // R2R Version 14 changed x86 code generation to use funclets
 // R2R Version 15 removes double to int/uint helper calls
 // R2R Version 16 replaces the compression format for debug boundaries with a new format that is smaller and more efficient to parse
+// R2R Version 17 adds support for producing "fat" debug information (that e.g. can include async debug info)
 
 struct READYTORUN_CORE_HEADER
 {

src/coreclr/jit/ICorJitInfo_names_generated.h

@@ -112,6 +112,7 @@ DEF_CLR_API(setBoundaries)
 DEF_CLR_API(getVars)
 DEF_CLR_API(setVars)
 DEF_CLR_API(reportRichMappings)
+DEF_CLR_API(reportAsyncDebugInfo)
 DEF_CLR_API(reportMetadata)
 DEF_CLR_API(allocateArray)
 DEF_CLR_API(freeArray)

src/coreclr/jit/ICorJitInfo_wrapper_generated.hpp

@@ -1067,6 +1067,17 @@ void WrapICorJitInfo::reportRichMappings(
     API_LEAVE(reportRichMappings);
 }
 
+void WrapICorJitInfo::reportAsyncDebugInfo(
+          ICorDebugInfo::AsyncInfo* asyncInfo,
+          ICorDebugInfo::AsyncSuspensionPoint* suspensionPoints,
+          ICorDebugInfo::AsyncContinuationVarInfo* vars,
+          uint32_t numVars)
+{
+    API_ENTER(reportAsyncDebugInfo);
+    wrapHnd->reportAsyncDebugInfo(asyncInfo, suspensionPoints, vars, numVars);
+    API_LEAVE(reportAsyncDebugInfo);
+}
+
 void WrapICorJitInfo::reportMetadata(
           const char* key,
           const void* value,

src/coreclr/jit/async.cpp

@@ -760,6 +760,8 @@ PhaseStatus AsyncTransformation::Run()
 
     m_comp->fgInvalidateDfsTree();
 
+    ReportDebugInfo();
+
     return PhaseStatus::MODIFIED_EVERYTHING;
 }
 
@@ -818,6 +820,8 @@ void AsyncTransformation::Transform(
     BasicBlock* resumeBB = CreateResumption(block, *remainder, call, callDefInfo, stateNum, layout);
 
     m_resumptionBBs.push_back(resumeBB);
+
+    CreateDebugInfoForSuspensionPoint(call, layout);
 }
 
 //------------------------------------------------------------------------
@@ -2202,6 +2206,78 @@ GenTreeStoreInd* AsyncTransformation::StoreAtOffset(GenTree* base, unsigned offs
     return store;
 }
 
+//------------------------------------------------------------------------
+// AsyncTransformation::CreateDebugInfoForSuspensionPoint:
+//   Create debug info for the specific suspension point we just created.
+//
+// Parameters:
+//   asyncCall - Call node resulting in the suspension point
+//   stateNum  - State number that was assigned to the suspension point
+//   layout    - Layout of continuation
+//
+void AsyncTransformation::CreateDebugInfoForSuspensionPoint(GenTreeCall* asyncCall, const ContinuationLayout& layout)
+{
+    if (!m_comp->opts.compDbgInfo)
+    {
+        return;
+    }
+
+    uint32_t numLocals = 0;
+    for (const LiveLocalInfo& local : layout.Locals)
+    {
+        unsigned ilVarNum = m_comp->compMap2ILvarNum(local.LclNum);
+        if (ilVarNum == (unsigned)ICorDebugInfo::UNKNOWN_ILNUM)
+        {
+            continue;
+        }
+
+        ICorDebugInfo::AsyncContinuationVarInfo varInf;
+        varInf.VarNumber = ilVarNum;
+        varInf.Offset    = local.DataOffset;
+        m_dbgContinuationVars.push_back(varInf);
+        numLocals++;
+    }
+
+    ICorDebugInfo::AsyncSuspensionPoint suspensionPoint;
+    const DebugInfo&                    di = asyncCall->GetAsyncInfo().DebugInfo;
+    suspensionPoint.RootILOffset           = di.GetRoot().GetLocation().GetOffset();
+    suspensionPoint.Inlinee                = di.GetInlineContext()->GetOrdinal();
+    suspensionPoint.ILOffset               = di.GetLocation().GetOffset();
+    suspensionPoint.NumContinuationVars    = numLocals;
+
+    m_dbgSuspensionPoints.push_back(suspensionPoint);
+}
+
+//------------------------------------------------------------------------
+// AsyncTransformation::ReportDebugInfo:
+//   Report debug info back to EE.
+//
+void AsyncTransformation::ReportDebugInfo()
+{
+    if (!m_comp->opts.compDbgInfo)
+    {
+        return;
+    }
+
+    ICorDebugInfo::AsyncInfo asyncInfo;
+    asyncInfo.NumSuspensionPoints = static_cast<uint32_t>(m_dbgSuspensionPoints.size());
+
+    ICorDebugInfo::AsyncSuspensionPoint* hostSuspensionPoints =
+        static_cast<ICorDebugInfo::AsyncSuspensionPoint*>(m_comp->info.compCompHnd->allocateArray(
+            m_dbgSuspensionPoints.size() * sizeof(ICorDebugInfo::AsyncSuspensionPoint)));
+    for (size_t i = 0; i < m_dbgSuspensionPoints.size(); i++)
+        hostSuspensionPoints[i] = m_dbgSuspensionPoints[i];
+
+    ICorDebugInfo::AsyncContinuationVarInfo* hostVars =
+        static_cast<ICorDebugInfo::AsyncContinuationVarInfo*>(m_comp->info.compCompHnd->allocateArray(
+            m_dbgContinuationVars.size() * sizeof(ICorDebugInfo::AsyncContinuationVarInfo)));
+    for (size_t i = 0; i < m_dbgContinuationVars.size(); i++)
+        hostVars[i] = m_dbgContinuationVars[i];
+
+    m_comp->info.compCompHnd->reportAsyncDebugInfo(&asyncInfo, hostSuspensionPoints, hostVars,
+                                                   static_cast<uint32_t>(m_dbgContinuationVars.size()));
+}
+
 //------------------------------------------------------------------------
 // AsyncTransformation::GetDataArrayVar:
 //   Create a new local to hold the data array of the continuation object. This

src/coreclr/jit/async.h

@@ -63,6 +63,10 @@ class AsyncTransformation
     BasicBlock*                   m_lastResumptionBB        = nullptr;
     BasicBlock*                   m_sharedReturnBB          = nullptr;
 
+    // Debug info
+    jitstd::vector<ICorDebugInfo::AsyncContinuationVarInfo> m_dbgContinuationVars;
+    jitstd::vector<ICorDebugInfo::AsyncSuspensionPoint>     m_dbgSuspensionPoints;
+
     bool IsLive(unsigned lclNum);
     void Transform(BasicBlock*               block,
                    GenTreeCall*              call,
@@ -132,6 +136,9 @@ class AsyncTransformation
                                     GenTreeFlags indirFlags = GTF_IND_NONFAULTING);
     GenTreeStoreInd* StoreAtOffset(GenTree* base, unsigned offset, GenTree* value, var_types storeType);
 
+    void CreateDebugInfoForSuspensionPoint(GenTreeCall* asyncCall, const ContinuationLayout& layout);
+    void ReportDebugInfo();
+
     unsigned GetDataArrayVar();
     unsigned GetGCDataArrayVar();
     unsigned GetResultBaseVar();
@@ -147,6 +154,8 @@ class AsyncTransformation
         : m_comp(comp)
         , m_liveLocalsScratch(comp->getAllocator(CMK_Async))
         , m_resumptionBBs(comp->getAllocator(CMK_Async))
+        , m_dbgContinuationVars(comp->getAllocator(CMK_Async))
+        , m_dbgSuspensionPoints(comp->getAllocator(CMK_Async))
     {
     }
 

src/coreclr/jit/compiler.h

@@ -4601,7 +4601,7 @@ class Compiler
                             CORINFO_CALL_INFO* callInfo,
                             IL_OFFSET          rawILOffset);
 
-    void impSetupAndSpillForAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags);
+    void impSetupAndSpillForAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags, const DebugInfo& callInstDI);
 
     void impInsertAsyncContinuationForLdvirtftnCall(GenTreeCall* call);
 

src/coreclr/jit/gentree.h

@@ -4400,6 +4400,9 @@ struct AsyncCallInfo
     bool                        SaveAndRestoreSynchronizationContextField = false;
     bool                        HasSuspensionIndicatorDef                 = false;
     unsigned                    SynchronizationContextLclNum              = BAD_VAR_NUM;
+
+    // Exact debug info of call IL instruction
+    DebugInfo DebugInfo;
 };
 
 // Return type descriptor of a GT_CALL node.

src/coreclr/jit/importercalls.cpp

@@ -386,7 +386,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
 
                 if (sig->isAsyncCall())
                 {
-                    impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags);
+                    impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags, di);
                 }
 
                 impPopCallArgs(sig, call->AsCall());
@@ -691,7 +691,7 @@ var_types Compiler::impImportCall(OPCODE                  opcode,
 
     if (sig->isAsyncCall())
     {
-        impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags);
+        impSetupAndSpillForAsyncCall(call->AsCall(), opcode, prefixFlags, di);
     }
 
     // Now create the argument list.
@@ -6800,10 +6800,15 @@ void Compiler::impCheckForPInvokeCall(
 //    call        - The call
 //    opcode      - The IL opcode for the call
 //    prefixFlags - Flags containing context handling information from IL
+//    callInstDI  - Debug info for the exact call instruction
 //
-void Compiler::impSetupAndSpillForAsyncCall(GenTreeCall* call, OPCODE opcode, unsigned prefixFlags)
+void Compiler::impSetupAndSpillForAsyncCall(GenTreeCall*     call,
+                                            OPCODE           opcode,
+                                            unsigned         prefixFlags,
+                                            const DebugInfo& callInstDI)
 {
     AsyncCallInfo asyncInfo;
+    asyncInfo.DebugInfo = callInstDI;
 
     if ((prefixFlags & PREFIX_IS_TASK_AWAIT) != 0)
     {