Add allocMem support for allocating multiple data chunks (#123378)

Rework `allocMem` to support allocating an arbitrary number of data
chunks. The motivation is to allow ilc to put async resumption info
chunks into sections that support relocations to .text, but additionally
will allow for sharing read only data between multiple functions.

Author: jakobbotsch

src/coreclr/inc/corjit.h

@@ -46,12 +46,10 @@ enum CorJitResult
 // to guide the memory allocation for the code, readonly data, and read-write data
 enum CorJitAllocMemFlag
 {
-    CORJIT_ALLOCMEM_DEFAULT_CODE_ALIGN = 0x00000000, // The code will use the normal alignment
-    CORJIT_ALLOCMEM_FLG_16BYTE_ALIGN   = 0x00000001, // The code will be 16-byte aligned
-    CORJIT_ALLOCMEM_FLG_RODATA_16BYTE_ALIGN = 0x00000002, // The read-only data will be 16-byte aligned
-    CORJIT_ALLOCMEM_FLG_32BYTE_ALIGN   = 0x00000004, // The code will be 32-byte aligned
-    CORJIT_ALLOCMEM_FLG_RODATA_32BYTE_ALIGN = 0x00000008, // The read-only data will be 32-byte aligned
-    CORJIT_ALLOCMEM_FLG_RODATA_64BYTE_ALIGN = 0x00000010, // The read-only data will be 64-byte aligned
+    CORJIT_ALLOCMEM_HOT_CODE = 1,
+    CORJIT_ALLOCMEM_COLD_CODE = 2,
+    CORJIT_ALLOCMEM_READONLY_DATA = 4,
+    CORJIT_ALLOCMEM_HAS_POINTERS_TO_CODE = 8,
 };
 
 inline CorJitAllocMemFlag operator |(CorJitAllocMemFlag a, CorJitAllocMemFlag b)
@@ -77,22 +75,28 @@ enum CheckedWriteBarrierKinds {
     CWBKind_AddrOfLocal,     // Store through the address of a local (arguably a bug that this happens at all).
 };
 
+struct AllocMemChunk
+{
+    // Alignment of the chunk. Must be a power of two with the following restrictions:
+    // - For the hot code chunk the max supported alignment is 32.
+    // - For the cold code chunk the value must always be 1.
+    // - For read-only data chunks the max supported alignment is 64.
+    uint32_t alignment;
+    uint32_t size;
+    CorJitAllocMemFlag flags;
+
+    // out
+    uint8_t* block;
+    uint8_t* blockRW;
+};
+
 struct AllocMemArgs
 {
-    // Input arguments
-    uint32_t hotCodeSize;
-    uint32_t coldCodeSize;
-    uint32_t roDataSize;
+    // Chunks to allocate. Supports one hot code chunk, one cold code chunk,
+    // and an arbitrary number of data chunks.
+    AllocMemChunk* chunks;
+    unsigned chunksCount;
     uint32_t xcptnsCount;
-    CorJitAllocMemFlag flag;
-
-    // Output arguments
-    void* hotCodeBlock;
-    void* hotCodeBlockRW;
-    void* coldCodeBlock;
-    void* coldCodeBlockRW;
-    void* roDataBlock;
-    void* roDataBlockRW;
 };
 
 #include "corjithost.h"

src/coreclr/inc/jiteeversionguid.h

@@ -37,11 +37,11 @@
 
 #include <minipal/guid.h>
 
-constexpr GUID JITEEVersionIdentifier = { /* 976a4d6d-d1b2-4096-a3c9-46ddcae71196 */
-    0x976a4d6d,
-    0xd1b2,
-    0x4096,
-    {0xa3, 0xc9, 0x46, 0xdd, 0xca, 0xe7, 0x11, 0x96}
+constexpr GUID JITEEVersionIdentifier = { /* db46fd97-a8e8-4bda-9cec-d7feb061154c */
+    0xdb46fd97,
+    0xa8e8,
+    0x4bda,
+    {0x9c, 0xec, 0xd7, 0xfe, 0xb0, 0x61, 0x15, 0x4c}
   };
 
 #endif // JIT_EE_VERSIONING_GUID_H

src/coreclr/interpreter/eeinterp.cpp

@@ -128,20 +128,23 @@ CorJitResult CILInterp::compileMethod(ICorJitInfo*         compHnd,
             uint32_t sizeOfCode = sizeof(InterpMethod*) + IRCodeSize * sizeof(int32_t);
             uint8_t unwindInfo[8] = {0, 0, 0, 0, 0, 0, 0, 0};
 
+            AllocMemChunk codeChunk {};
+            codeChunk.alignment = 1;
+            codeChunk.size = sizeOfCode;
+            codeChunk.flags = CORJIT_ALLOCMEM_HOT_CODE;
+
             AllocMemArgs args {};
-            args.hotCodeSize = sizeOfCode;
-            args.coldCodeSize = 0;
-            args.roDataSize = 0;
+            args.chunks = &codeChunk;
+            args.chunksCount = 1;
             args.xcptnsCount = 0;
-            args.flag = CORJIT_ALLOCMEM_DEFAULT_CODE_ALIGN;
             compHnd->allocMem(&args);
 
             // We store first the InterpMethod pointer as the code header, followed by the actual code
-            *(InterpMethod**)args.hotCodeBlockRW = pMethod;
-            memcpy ((uint8_t*)args.hotCodeBlockRW + sizeof(InterpMethod*), pIRCode, IRCodeSize * sizeof(int32_t));
+            *(InterpMethod**)codeChunk.blockRW = pMethod;
+            memcpy ((uint8_t*)codeChunk.blockRW + sizeof(InterpMethod*), pIRCode, IRCodeSize * sizeof(int32_t));
 
-            compiler.UpdateWithFinalMethodByteCodeAddress((InterpByteCodeStart*)args.hotCodeBlock);
-            *entryAddress = (uint8_t*)args.hotCodeBlock;
+            compiler.UpdateWithFinalMethodByteCodeAddress((InterpByteCodeStart*)codeChunk.block);
+            *entryAddress = (uint8_t*)codeChunk.block;
             *nativeSizeOfCode = sizeOfCode;
 
             // We can't do this until we've called allocMem

src/coreclr/jit/codegen.h

@@ -195,8 +195,6 @@ class CodeGen final : public CodeGenInterface
     unsigned  codeSize;
     void*     coldCodePtr;
     void*     coldCodePtrRW;
-    void*     consPtr;
-    void*     consPtrRW;
 
     // Last instr we have displayed for dspInstrs
     unsigned genCurDispOffset;

src/coreclr/jit/codegencommon.cpp

@@ -2131,7 +2131,7 @@ void CodeGen::genEmitMachineCode()
     codeSize =
         GetEmitter()->emitEndCodeGen(compiler, trackedStackPtrsContig, GetInterruptible(), IsFullPtrRegMapRequired(),
                                      compiler->compHndBBtabCount, &prologSize, &epilogSize, codePtr, &codePtrRW,
-                                     &coldCodePtr, &coldCodePtrRW, &consPtr, &consPtrRW DEBUGARG(&instrCount));
+                                     &coldCodePtr, &coldCodePtrRW DEBUGARG(&instrCount));
 
 #ifdef DEBUG
     assert(compiler->compCodeGenDone == false);

src/coreclr/jit/compiler.h

@@ -8542,7 +8542,11 @@ class Compiler
 
     // ICorJitInfo wrappers
 
-    void eeAllocMem(AllocMemArgs* args, const UNATIVE_OFFSET roDataSectionAlignment);
+    void eeAllocMem(AllocMemChunk& codeChunk,
+                    AllocMemChunk* coldCodeChunk,
+                    AllocMemChunk* dataChunks,
+                    unsigned       numDataChunks,
+                    unsigned       numExceptions);
 
     void eeReserveUnwindInfo(bool isFunclet, bool isColdCode, ULONG unwindSize);
 

src/coreclr/jit/ee_il_dll.cpp

@@ -1129,63 +1129,117 @@ void Compiler::eeDispLineInfos()
  * (e.g., host AMD64, target ARM64), then VM will get confused anyway.
  */
 
-void Compiler::eeAllocMem(AllocMemArgs* args, const UNATIVE_OFFSET roDataSectionAlignment)
+void Compiler::eeAllocMem(AllocMemChunk& codeChunk,
+                          AllocMemChunk* coldCodeChunk,
+                          AllocMemChunk* dataChunks,
+                          unsigned       numDataChunks,
+                          unsigned       numExceptions)
 {
+    ArrayStack<AllocMemChunk> chunks(getAllocator(CMK_Codegen));
+
+    chunks.Push(codeChunk);
+
+    int coldCodeChunkIndex = -1;
+    if (coldCodeChunk != nullptr)
+    {
+        coldCodeChunkIndex = chunks.Height();
+        chunks.Push(*coldCodeChunk);
+    }
+
 #ifdef DEBUG
 
     // Fake splitting implementation: place hot/cold code in contiguous section.
-    UNATIVE_OFFSET coldCodeOffset = 0;
-    if (JitConfig.JitFakeProcedureSplitting() && (args->coldCodeSize > 0))
+    if (JitConfig.JitFakeProcedureSplitting() && (coldCodeChunk != nullptr))
     {
-        coldCodeOffset = args->hotCodeSize;
-        assert(coldCodeOffset > 0);
-        args->hotCodeSize += args->coldCodeSize;
-        args->coldCodeSize = 0;
+        // Keep offset into hot code in the block/blockRW pointers
+        coldCodeChunk->block   = (uint8_t*)(uintptr_t)chunks.BottomRef(0).size;
+        coldCodeChunk->blockRW = (uint8_t*)(uintptr_t)chunks.BottomRef(0).size;
+        chunks.BottomRef(0).size += coldCodeChunk->size;
+        // Remove cold chunk
+        chunks.Pop();
+        coldCodeChunkIndex = -1;
     }
 
 #endif // DEBUG
 
+    int firstDataChunk = chunks.Height();
+
 #if defined(TARGET_ARM64) || defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
 
     // For arm64/LoongArch64/RISCV64, we want to allocate JIT data always adjacent to code similar to what native
     // compiler does.
     // This way allows us to use a single `ldr` to access such data like float constant/jmp table.
     // For LoongArch64 using `pcaddi + ld` to access such data.
 
-    UNATIVE_OFFSET roDataAlignmentDelta = 0;
-    if (args->roDataSize > 0)
+    for (unsigned i = 0; i < numDataChunks; i++)
     {
-        roDataAlignmentDelta = AlignmentPad(args->hotCodeSize, roDataSectionAlignment);
+        // Increase size of the hot code chunk and store offset in data chunk
+        AllocMemChunk& codeChunk = chunks.BottomRef(0);
+
+        codeChunk.size        = AlignUp(codeChunk.size, dataChunks[i].alignment);
+        dataChunks[i].block   = (uint8_t*)(uintptr_t)codeChunk.size;
+        dataChunks[i].blockRW = (uint8_t*)(uintptr_t)codeChunk.size;
+        codeChunk.size += dataChunks[i].size;
+
+        codeChunk.alignment = max(codeChunk.alignment, dataChunks[i].alignment);
     }
 
-    const UNATIVE_OFFSET roDataOffset = args->hotCodeSize + roDataAlignmentDelta;
-    args->hotCodeSize                 = roDataOffset + args->roDataSize;
-    args->roDataSize                  = 0;
+#else
+
+    for (unsigned i = 0; i < numDataChunks; i++)
+    {
+        chunks.Push(dataChunks[i]);
+    }
 
 #endif // defined(TARGET_ARM64) || defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
 
-    info.compCompHnd->allocMem(args);
+    AllocMemArgs args;
+    args.chunks      = chunks.Data();
+    args.chunksCount = (unsigned)chunks.Height();
+    args.xcptnsCount = numExceptions;
+
+    info.compCompHnd->allocMem(&args);
+
+    codeChunk.block   = chunks.BottomRef(0).block;
+    codeChunk.blockRW = chunks.BottomRef(0).blockRW;
+
+    if (coldCodeChunkIndex != -1)
+    {
+        coldCodeChunk->block   = chunks.BottomRef(coldCodeChunkIndex).block;
+        coldCodeChunk->blockRW = chunks.BottomRef(coldCodeChunkIndex).blockRW;
+    }
 
 #ifdef DEBUG
 
-    if (JitConfig.JitFakeProcedureSplitting() && (coldCodeOffset > 0))
+    if (JitConfig.JitFakeProcedureSplitting() && (coldCodeChunk != nullptr))
     {
         // Fix up cold code pointers. Cold section is adjacent to hot section.
-        assert(args->coldCodeBlock == nullptr);
-        assert(args->coldCodeBlockRW == nullptr);
-        args->coldCodeBlock   = ((BYTE*)args->hotCodeBlock) + coldCodeOffset;
-        args->coldCodeBlockRW = ((BYTE*)args->hotCodeBlockRW) + coldCodeOffset;
+        assert(coldCodeChunk != nullptr);
+        coldCodeChunk->block   = codeChunk.block + (uintptr_t)coldCodeChunk->block;
+        coldCodeChunk->blockRW = codeChunk.blockRW + (uintptr_t)coldCodeChunk->blockRW;
     }
 
 #endif // DEBUG
 
+    int curDataChunk = firstDataChunk;
+
 #if defined(TARGET_ARM64) || defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
 
     // Fix up data section pointers.
-    assert(args->roDataBlock == nullptr);
-    assert(args->roDataBlockRW == nullptr);
-    args->roDataBlock   = ((BYTE*)args->hotCodeBlock) + roDataOffset;
-    args->roDataBlockRW = ((BYTE*)args->hotCodeBlockRW) + roDataOffset;
+    for (unsigned i = 0; i < numDataChunks; i++)
+    {
+        dataChunks[i].block   = codeChunk.block + (size_t)dataChunks[i].block;
+        dataChunks[i].blockRW = codeChunk.blockRW + (size_t)dataChunks[i].blockRW;
+    }
+
+#else
+
+    for (unsigned i = 0; i < numDataChunks; i++)
+    {
+        dataChunks[i].block   = chunks.BottomRef(curDataChunk).block;
+        dataChunks[i].blockRW = chunks.BottomRef(curDataChunk).blockRW;
+        curDataChunk++;
+    }
 
 #endif // defined(TARGET_ARM64) || defined(TARGET_LOONGARCH64) || defined(TARGET_RISCV64)
 }