Fix SPMI to handle replays of BBINSTR jit method contexts (#41386)

Move the handling of allocMethodBlockCounts to the method context.
While this data is both an input and an output, we're more interested
in the input side.

We could also treat it as an output, if we say wanted to verify that
the replay jit wrote the same sequence of IL offsets into the BlockCount
records, but that doesn't seem crucial as changes here would also lead
to code diffs.

Fix the code that checks the AddressMap for reloc hints so it doesn't
fail if the exact lookup fails, but instead falls back to a search.
This is needed because the base of the replay count buffer is recorded
as the map key, but the instrumentation probes in jitted code refer
to offsets from this value.

Fixes #37270.

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.cpp b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.cpp
index 75d8ab5..086d4c4 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.cpp
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.cpp
@@ -895,9 +895,17 @@ void* CompileResult::repAddressMap(void* replayAddress)
 {
     if (AddressMap == nullptr)
         return nullptr;
-    Agnostic_AddressMap value;
-    value = AddressMap->Get((DWORDLONG)replayAddress);
-    return (void*)value.Address;
+
+    int index = AddressMap->GetIndex((DWORDLONG)replayAddress);
+
+    if (index != -1)
+    {
+        Agnostic_AddressMap value;
+        value = AddressMap->Get((DWORDLONG)replayAddress);
+        return (void*)value.Address;
+    }
+
+    return nullptr;
 }
 void* CompileResult::searchAddressMap(void* newAddress)
 {
@@ -962,39 +970,6 @@ void CompileResult::dmpAllocUnwindInfo(DWORD key, const Agnostic_AllocUnwindInfo
            value.pUnwindBlock_index, value.funcKind);
 }
 
-void CompileResult::recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result)
-{
-    if (AllocMethodBlockCounts == nullptr)
-        AllocMethodBlockCounts = new LightWeightMap<DWORD, Agnostic_AllocMethodBlockCounts>();
-
-    Agnostic_AllocMethodBlockCounts value;
-
-    value.count  = (DWORD)count;
-    value.result = (DWORD)result;
-    value.pBlockCounts_index =
-        AllocMethodBlockCounts->AddBuffer((unsigned char*)*pBlockCounts, count * sizeof(ICorJitInfo::BlockCounts));
-
-    AllocMethodBlockCounts->Add((DWORD)0, value);
-}
-void CompileResult::dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value)
-{
-    printf("AllocMethodBlockCounts key %u, value cnt-%u ind-%u res-%08X", key, value.count, value.pBlockCounts_index,
-           value.result);
-}
-HRESULT CompileResult::repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts)
-{
-    Agnostic_AllocMethodBlockCounts value;
-    value = AllocMethodBlockCounts->Get((DWORD)0);
-
-    if (count != value.count)
-        __debugbreak();
-
-    HRESULT result = (HRESULT)value.result;
-    *pBlockCounts = (ICorJitInfo::BlockCounts*)AllocMethodBlockCounts->GetBuffer(value.pBlockCounts_index);
-    recAddressMap((void*)0x4242, (void*)*pBlockCounts, count * (sizeof(ICorJitInfo::BlockCounts)));
-    return result;
-}
-
 void CompileResult::recRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle)
 {
     repRecordCallSite(instrOffset, callSig, methodHandle);

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.h b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.h
index 09ec839..125d9fe 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.h
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/compileresult.h
@@ -166,12 +166,6 @@ public:
         DWORD HandlerLength;
         DWORD ClassToken; // one view of symetric union
     };
-    struct Agnostic_AllocMethodBlockCounts
-    {
-        DWORD count;
-        DWORD pBlockCounts_index;
-        DWORD result;
-    };
     struct Agnostic_CORINFO_SIG_INFO2
     {
         DWORD     callConv;
@@ -328,10 +322,6 @@ public:
                             CorJitFuncKind funcKind);
     void dmpAllocUnwindInfo(DWORD key, const Agnostic_AllocUnwindInfo& value);
 
-    void recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result);
-    void dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value);
-    HRESULT repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts);
-
     void recRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle);
     void dmpRecordCallSite(DWORD key, const Agnostic_RecordCallSite& value);
     void repRecordCallSite(ULONG instrOffset, CORINFO_SIG_INFO* callSig, CORINFO_METHOD_HANDLE methodHandle);

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/crlwmlist.h b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/crlwmlist.h
index b42077d..12b473a 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/crlwmlist.h
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/crlwmlist.h
@@ -19,7 +19,6 @@
 #endif
 
 LWM(AddressMap, DWORDLONG, CompileResult::Agnostic_AddressMap)
-LWM(AllocMethodBlockCounts, DWORD, CompileResult::Agnostic_AllocMethodBlockCounts)
 LWM(AllocGCInfo, DWORD, CompileResult::Agnostic_AllocGCInfo)
 LWM(AllocMem, DWORD, CompileResult::Agnostic_AllocMemDetails)
 DENSELWM(AllocUnwindInfo, CompileResult::Agnostic_AllocUnwindInfo)

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lightweightmap.h b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lightweightmap.h
index 2a5594f..adce2c1 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lightweightmap.h
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lightweightmap.h
@@ -78,6 +78,33 @@ public:
         return newOffset + sizeof(unsigned int);
     }
 
+    const unsigned char* CreateBuffer(unsigned int len)
+    {
+        if (len == 0)
+        {
+            return nullptr;
+        }
+
+        if (locked)
+        {
+            LogError("Added item that extended the buffer after it was locked by a call to GetBuffer()");
+            __debugbreak();
+        }
+
+        unsigned int   newbuffsize = bufferLength + sizeof(unsigned int) + len;
+        unsigned char* newbuffer   = new unsigned char[newbuffsize];
+        unsigned int   newOffset   = bufferLength;
+        if (bufferLength > 0)
+            memcpy(newbuffer, buffer, bufferLength);
+        memset(newbuffer + bufferLength + sizeof(unsigned int), 0, len);
+        *((unsigned int*)(newbuffer + bufferLength)) = len;
+        bufferLength += sizeof(unsigned int) + len;
+        if (buffer != nullptr)
+            delete[] buffer;
+        buffer = newbuffer;
+        return buffer + newOffset + sizeof(unsigned int);
+    }
+
     unsigned char* GetBuffer(unsigned int offset)
     {
         if (offset == (unsigned int)-1)

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lwmlist.h b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lwmlist.h
index 1922e7d..b973119 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lwmlist.h
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/lwmlist.h
@@ -18,6 +18,7 @@
 #define DENSELWM(map, value) LWM(map, this_is_an_error, value)
 #endif
 
+LWM(AllocMethodBlockCounts, DWORD, Agnostic_AllocMethodBlockCounts)
 LWM(AppendClassName, Agnostic_AppendClassName, DWORD)
 LWM(AreTypesEquivalent, DLDL, DWORD)
 LWM(AsCorInfoType, DWORDLONG, DWORD)

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.cpp b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.cpp
index d45c5e0..4a134fe 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.cpp
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.cpp
@@ -5133,6 +5133,50 @@ DWORD MethodContext::repGetFieldThreadLocalStoreID(CORINFO_FIELD_HANDLE field, v
     return (DWORD)value.B;
 }
 
+
+void MethodContext::recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result)
+{
+    if (AllocMethodBlockCounts == nullptr)
+        AllocMethodBlockCounts = new LightWeightMap<DWORD, Agnostic_AllocMethodBlockCounts>();
+
+    Agnostic_AllocMethodBlockCounts value;
+
+    value.address = (DWORDLONG)*pBlockCounts;
+    value.count  = (DWORD)count;
+    value.result = (DWORD)result;
+
+    AllocMethodBlockCounts->Add((DWORD)0, value);
+}
+void MethodContext::dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value)
+{
+    printf("AllocMethodBlockCounts key %u, value addr-%016llX cnt-%u res-%08X", key, value.address, value.count, value.result);
+}
+HRESULT MethodContext::repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts)
+{
+    Agnostic_AllocMethodBlockCounts value;
+    value = AllocMethodBlockCounts->Get((DWORD)0);
+
+    if (count != value.count)
+    {
+        LogWarning("AllocMethodBlockCount mismatch: record %d, replay %d", value.count, count);
+    }
+
+    HRESULT result = (HRESULT)value.result;
+
+    // Allocate a scratch buffer, linked to method context via AllocMethodBlockCounts, so it gets
+    // cleaned up when the method context does.
+    //
+    // We won't bother recording this via AddBuffer because currently SPMI will never look at it.
+    // But we need a writeable buffer because the jit will store IL offsets inside.
+    //
+    // Todo, perhaps: record the buffer as a compile result instead, and defer copying until
+    // jit completion so we can snapshot the offsets the jit writes.
+    //
+    *pBlockCounts = (ICorJitInfo::BlockCounts*)AllocMethodBlockCounts->CreateBuffer(count * sizeof(ICorJitInfo::BlockCounts));
+    cr->recAddressMap((void*)value.address, (void*)*pBlockCounts, count * (sizeof(ICorJitInfo::BlockCounts)));
+    return result;
+}
+
 void MethodContext::recGetMethodBlockCounts(CORINFO_METHOD_HANDLE        ftnHnd,
                                             UINT32 *                     pCount,
                                             ICorJitInfo::BlockCounts**   pBlockCounts,

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.h b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.h
index 36c0fbb..0061803 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.h
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shared/methodcontext.h
@@ -422,6 +422,12 @@ public:
         DWORDLONG method;
         DWORDLONG delegateCls;
     };
+    struct Agnostic_AllocMethodBlockCounts
+    {
+        DWORDLONG address;
+        DWORD count;
+        DWORD result;
+    };
     struct Agnostic_GetMethodBlockCounts
     {
         DWORD count;
@@ -1168,6 +1174,10 @@ public:
     void dmpGetFieldThreadLocalStoreID(DWORDLONG key, DLD value);
     DWORD repGetFieldThreadLocalStoreID(CORINFO_FIELD_HANDLE field, void** ppIndirection);
 
+    void recAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts, HRESULT result);
+    void dmpAllocMethodBlockCounts(DWORD key, const Agnostic_AllocMethodBlockCounts& value);
+    HRESULT repAllocMethodBlockCounts(ULONG count, ICorJitInfo::BlockCounts** pBlockCounts);
+
     void recGetMethodBlockCounts(CORINFO_METHOD_HANDLE        ftnHnd,
                                  UINT32 *                     pCount,
                                  ICorJitInfo::BlockCounts**   pBlockCounts,
@@ -1338,6 +1348,7 @@ private:
 // *************************************************************************************
 enum mcPackets
 {
+    Packet_AllocMethodBlockCounts     = 131,
     Packet_AppendClassName            = 149, // Added 8/6/2014 - needed for SIMD
     Packet_AreTypesEquivalent         = 1,
     Packet_AsCorInfoType              = 2,
@@ -1493,7 +1504,6 @@ enum mcPackets
     Packet_ShouldEnforceCallvirtRestriction              = 112, // Retired 2/18/2020
 
     PacketCR_AddressMap                        = 113,
-    PacketCR_AllocMethodBlockCounts            = 131,
     PacketCR_AllocGCInfo                       = 114,
     PacketCR_AllocMem                          = 115,
     PacketCR_AllocUnwindInfo                   = 132,

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp b/src/coreclr/src/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp
index f5b6929..bb1bba8 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi-shim-collector/icorjitinfo.cpp
@@ -2030,7 +2030,7 @@ HRESULT interceptor_ICJI::allocMethodBlockCounts(UINT32          count, // The n
 {
     mc->cr->AddCall("allocMethodBlockCounts");
     HRESULT result = original_ICorJitInfo->allocMethodBlockCounts(count, pBlockCounts);
-    mc->cr->recAllocMethodBlockCounts(count, pBlockCounts, result);
+    mc->recAllocMethodBlockCounts(count, pBlockCounts, result);
     return result;
 }
 

diff --git a/src/coreclr/src/ToolBox/superpmi/superpmi/icorjitinfo.cpp b/src/coreclr/src/ToolBox/superpmi/superpmi/icorjitinfo.cpp
index f0e2114..6b8f454 100644 (file)
--- a/src/coreclr/src/ToolBox/superpmi/superpmi/icorjitinfo.cpp
+++ b/src/coreclr/src/ToolBox/superpmi/superpmi/icorjitinfo.cpp
@@ -1784,7 +1784,7 @@ HRESULT MyICJI::allocMethodBlockCounts(UINT32          count, // The number of b
                                        BlockCounts**   pBlockCounts)
 {
     jitInstance->mc->cr->AddCall("allocMethodBlockCounts");
-    return jitInstance->mc->cr->repAllocMethodBlockCounts(count, pBlockCounts);
+    return jitInstance->mc->repAllocMethodBlockCounts(count, pBlockCounts);
 }
 
 // get profile information to be used for optimizing the current method.  The format