instanced jit cache

- jit caches are no longer static.
- JitUnwindWindows now supports multiple jit caches.
- Jit caches should no longer cause crashes due to entry offset collisions.
- cpu tests now run concurrently.

Co-authored-by: LotP1 <68976644+LotP1@users.noreply.github.com>
This commit is contained in:
LotP
2026-07-01 09:54:23 -05:00
committed by KeatonTheBot
co-authored by LotP1
parent 297df259c3
commit a82a721df6
26 changed files with 343 additions and 283 deletions
@@ -4,7 +4,7 @@ using System.Diagnostics.CodeAnalysis;
namespace ARMeilleure.Translation.Cache
{
readonly struct CacheEntry : IComparable<CacheEntry>
public readonly struct CacheEntry : IComparable<CacheEntry>
{
public int Offset { get; }
public int Size { get; }
@@ -23,11 +23,30 @@ namespace ARMeilleure.Translation.Cache
}
}
private readonly int _regionSize;
private int _regionCount;
private readonly List<MemoryBlock> _blocks = [];
public CacheMemoryAllocator(int capacity)
public CacheMemoryAllocator(int regionSize, int initialRegionCount = 1)
{
_blocks.Add(new MemoryBlock(0, capacity));
_regionCount = 0;
_regionSize = regionSize;
for (; initialRegionCount > 0; initialRegionCount--)
{
_blocks.Add(new MemoryBlock(_regionSize * _regionCount, _regionSize));
_regionCount++;
}
}
public void AddNewBlocks(int count)
{
for (; count > 0; count--)
{
_blocks.Add(new MemoryBlock(_regionSize * _regionCount, _regionSize));
_regionCount++;
}
}
public int Allocate(int size)
@@ -66,12 +85,13 @@ namespace ARMeilleure.Translation.Cache
index = ~index;
}
if (index < _blocks.Count)
int endOffs = block.Offset + block.Size;
// Don't merge blocks from different allocations
if (index < _blocks.Count && endOffs % _regionSize != 0)
{
MemoryBlock next = _blocks[index];
int endOffs = block.Offset + block.Size;
if (next.Offset == endOffs)
{
block = new MemoryBlock(block.Offset, block.Size + next.Size);
@@ -79,7 +99,8 @@ namespace ARMeilleure.Translation.Cache
}
}
if (index > 0)
// Don't merge blocks from different allocations
if (index > 0 && block.Offset % _regionSize != 0)
{
MemoryBlock prev = _blocks[index - 1];
+80 -91
View File
@@ -14,62 +14,35 @@ using System.Threading;
namespace ARMeilleure.Translation.Cache
{
static partial class JitCache
public partial class JitCache : IDisposable
{
private static readonly int _pageSize = (int)MemoryBlock.GetPageSize();
private static readonly int _pageMask = _pageSize - 1;
private const int CodeAlignment = 4; // Bytes.
private const int CacheSize = 256 * 1024 * 1024;
private const uint CacheSize = 256 * 1024 * 1024;
private static JitCacheInvalidation _jitCacheInvalidator;
private readonly JitCacheInvalidation _jitCacheInvalidator;
private static List<CacheMemoryAllocator> _cacheAllocators = [];
private readonly CacheMemoryAllocator _cacheAllocator;
private static readonly List<CacheEntry> _cacheEntries = [];
private readonly List<CacheEntry> _cacheEntries = [];
private static readonly Lock _lock = new();
private static bool _initialized;
private readonly Lock _lock = new();
private static readonly List<ReservedRegion> _jitRegions = [];
private static int _activeRegionIndex = 0;
private readonly List<ReservedRegion> _jitRegions = [];
[SupportedOSPlatform("windows")]
[LibraryImport("kernel32.dll", SetLastError = true)]
public static partial nint FlushInstructionCache(nint hProcess, nint lpAddress, nuint dwSize);
private static partial nint FlushInstructionCache(nint hProcess, nint lpAddress, nuint dwSize);
public static void Initialize(IJitMemoryAllocator allocator)
public JitCache(IJitMemoryAllocator allocator)
{
lock (_lock)
{
if (_initialized)
{
if (OperatingSystem.IsWindows())
{
JitUnwindWindows.RemoveFunctionTableHandler(
_jitRegions[0].Pointer);
}
_jitRegions.Add(new(allocator, CacheSize));
for (int i = 0; i < _jitRegions.Count; i++)
{
_jitRegions[i].Dispose();
}
_jitRegions.Clear();
_cacheAllocators.Clear();
}
else
{
_initialized = true;
}
_activeRegionIndex = 0;
ReservedRegion firstRegion = new(allocator, CacheSize);
_jitRegions.Add(firstRegion);
CacheMemoryAllocator firstCacheAllocator = new(CacheSize);
_cacheAllocators.Add(firstCacheAllocator);
_cacheAllocator = new((int)CacheSize);
if (!OperatingSystem.IsWindows() && !OperatingSystem.IsMacOS())
{
@@ -79,23 +52,20 @@ namespace ARMeilleure.Translation.Cache
if (OperatingSystem.IsWindows())
{
JitUnwindWindows.InstallFunctionTableHandler(
firstRegion.Pointer, CacheSize, firstRegion.Pointer + Allocate(_pageSize)
this, _jitRegions[0].Pointer, CacheSize, _jitRegions[0].Pointer + Allocate(_pageSize)
);
}
}
}
public static nint Map(CompiledFunction func)
public nint Map(CompiledFunction func)
{
byte[] code = func.Code;
lock (_lock)
{
Debug.Assert(_initialized);
int funcOffset = Allocate(code.Length);
ReservedRegion targetRegion = _jitRegions[_activeRegionIndex];
nint funcPtr = targetRegion.Pointer + funcOffset;
nint funcPtr = GetFunctionPtr(funcOffset);
if (OperatingSystem.IsMacOS() && RuntimeInformation.ProcessArchitecture == Architecture.Arm64)
{
@@ -109,9 +79,9 @@ namespace ARMeilleure.Translation.Cache
}
else
{
ReprotectAsWritable(targetRegion, funcOffset, code.Length);
ReprotectAsWritable(funcOffset, code.Length);
Marshal.Copy(code, 0, funcPtr, code.Length);
ReprotectAsExecutable(targetRegion, funcOffset, code.Length);
ReprotectAsExecutable(funcOffset, code.Length);
if (OperatingSystem.IsWindows() && RuntimeInformation.ProcessArchitecture == Architecture.Arm64)
{
@@ -129,25 +99,24 @@ namespace ARMeilleure.Translation.Cache
}
}
public static void Unmap(nint pointer)
public void Unmap(nint pointer)
{
lock (_lock)
{
Debug.Assert(_initialized);
foreach (ReservedRegion region in _jitRegions)
for (int i = 0; i < _jitRegions.Count; i++)
{
if (pointer.ToInt64() < region.Pointer.ToInt64() ||
pointer.ToInt64() >= (region.Pointer + CacheSize).ToInt64())
ReservedRegion jitRegion = _jitRegions[i];
if (pointer.ToInt64() < jitRegion.Pointer.ToInt64() ||
pointer.ToInt64() >= (jitRegion.Pointer + (nint)CacheSize).ToInt64())
{
continue;
}
int funcOffset = (int)(pointer.ToInt64() - region.Pointer.ToInt64());
int funcOffset = (int)(pointer.ToInt64() - jitRegion.Pointer.ToInt64() + i * CacheSize);
if (TryFind(funcOffset, out CacheEntry entry, out int entryIndex) && entry.Offset == funcOffset)
{
_cacheAllocators[_activeRegionIndex].Free(funcOffset, AlignCodeSize(entry.Size));
_cacheAllocator.Free(funcOffset, AlignCodeSize(entry.Size));
_cacheEntries.RemoveAt(entryIndex);
}
@@ -156,53 +125,63 @@ namespace ARMeilleure.Translation.Cache
}
}
private static void ReprotectAsWritable(ReservedRegion region, int offset, int size)
private void ReprotectAsWritable(int offset, int size)
{
int endOffs = offset + size;
int regionStart = offset & ~_pageMask;
int regionEnd = (endOffs + _pageMask) & ~_pageMask;
int regionStart = (offset % (int)CacheSize) & ~_pageMask;
int regionEnd = ((endOffs % (int)CacheSize) + _pageMask) & ~_pageMask;
region.Block.MapAsRwx((ulong)regionStart, (ulong)(regionEnd - regionStart));
GetRegion(offset).Block.MapAsRwx((ulong)regionStart, (ulong)(regionEnd - regionStart));
}
private static void ReprotectAsExecutable(ReservedRegion region, int offset, int size)
private void ReprotectAsExecutable(int offset, int size)
{
int endOffs = offset + size;
int regionStart = offset & ~_pageMask;
int regionEnd = (endOffs + _pageMask) & ~_pageMask;
int regionStart = (offset % (int)CacheSize) & ~_pageMask;
int regionEnd = ((endOffs % (int)CacheSize) + _pageMask) & ~_pageMask;
region.Block.MapAsRx((ulong)regionStart, (ulong)(regionEnd - regionStart));
GetRegion(offset).Block.MapAsRx((ulong)regionStart, (ulong)(regionEnd - regionStart));
}
private static int Allocate(int codeSize)
private int Allocate(int codeSize)
{
codeSize = AlignCodeSize(codeSize);
int allocOffset = _cacheAllocators[_activeRegionIndex].Allocate(codeSize);
int allocOffset = _cacheAllocator.Allocate(codeSize);
if (allocOffset >= 0)
{
_jitRegions[_activeRegionIndex].ExpandIfNeeded((ulong)allocOffset + (ulong)codeSize);
GetRegion(allocOffset).ExpandIfNeeded((ulong)(allocOffset % (int)CacheSize) + (ulong)codeSize);
return allocOffset;
}
int exhaustedRegion = _activeRegionIndex;
_cacheAllocator.AddNewBlocks(1);
ReservedRegion newRegion = new(_jitRegions[0].Allocator, CacheSize);
Logger.Warning?.Print(LogClass.Cpu, $"JIT Cache of size {(_jitRegions.Count * CacheSize).Bytes()} exhausted, creating new Cache Region ({((_jitRegions.Count + 1) * CacheSize).Bytes()} Total Allocation).");
_jitRegions.Add(newRegion);
_activeRegionIndex = _jitRegions.Count - 1;
Logger.Warning?.Print(LogClass.Cpu, $"JIT Cache Region {exhaustedRegion} exhausted, creating new Cache Region {_activeRegionIndex} ({((long)(_activeRegionIndex + 1) * CacheSize).Bytes()} Total Allocation).");
_cacheAllocators.Add(new CacheMemoryAllocator(CacheSize));
int allocOffsetNew = _cacheAllocators[_activeRegionIndex].Allocate(codeSize);
if (allocOffsetNew < 0)
allocOffset = _cacheAllocator.Allocate(codeSize);
if (allocOffset < 0)
{
throw new OutOfMemoryException("Failed to allocate in new Cache Region!");
}
newRegion.ExpandIfNeeded((ulong)allocOffsetNew + (ulong)codeSize);
return allocOffsetNew;
GetRegion(allocOffset).ExpandIfNeeded((ulong)(allocOffset % (int)CacheSize) + (ulong)codeSize);
return allocOffset;
}
private nint GetFunctionPtr(int offset)
{
return GetRegion(offset).Pointer + (offset % (int)CacheSize);
}
private ReservedRegion GetRegion(int offset)
{
int index = offset / (int)CacheSize;
return _jitRegions[index];
}
private static int AlignCodeSize(int codeSize)
@@ -210,7 +189,7 @@ namespace ARMeilleure.Translation.Cache
return checked(codeSize + (CodeAlignment - 1)) & ~(CodeAlignment - 1);
}
private static void Add(int offset, int size, UnwindInfo unwindInfo)
private void Add(int offset, int size, UnwindInfo unwindInfo)
{
CacheEntry entry = new(offset, size, unwindInfo);
@@ -224,25 +203,22 @@ namespace ARMeilleure.Translation.Cache
_cacheEntries.Insert(index, entry);
}
public static bool TryFind(int offset, out CacheEntry entry, out int entryIndex)
public bool TryFind(int offset, out CacheEntry entry, out int entryIndex)
{
lock (_lock)
{
foreach (ReservedRegion region in _jitRegions)
int index = _cacheEntries.BinarySearch(new CacheEntry(offset, 0, default));
if (index < 0)
{
int index = _cacheEntries.BinarySearch(new CacheEntry(offset, 0, default));
index = ~index - 1;
}
if (index < 0)
{
index = ~index - 1;
}
if (index >= 0)
{
entry = _cacheEntries[index];
entryIndex = index;
return true;
}
if (index >= 0)
{
entry = _cacheEntries[index];
entryIndex = index;
return true;
}
}
@@ -250,5 +226,18 @@ namespace ARMeilleure.Translation.Cache
entryIndex = 0;
return false;
}
public void Dispose()
{
if (OperatingSystem.IsWindows())
{
JitUnwindWindows.RemoveFunctionTableHandler(_jitRegions[0].Pointer);
}
foreach (ReservedRegion jitRegion in _jitRegions)
{
jitRegion.Dispose();
}
}
}
}
@@ -2,6 +2,8 @@
using ARMeilleure.CodeGen.Unwinding;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.InteropServices;
@@ -27,6 +29,29 @@ namespace ARMeilleure.Translation.Cache
public unsafe fixed ushort UnwindCodes[MaxUnwindCodesArraySize];
}
private unsafe struct InternalFunctionHandler
{
public InternalFunctionHandler(JitCache jitCache, nint workBufferPtr)
{
_jitCache = jitCache;
_runtimeFunction = (RuntimeFunction*)workBufferPtr;
_unwindInfo = (UnwindInfo*)(workBufferPtr + _sizeOfRuntimeFunction);
}
readonly JitCache _jitCache;
readonly RuntimeFunction* _runtimeFunction;
readonly UnwindInfo* _unwindInfo;
public RuntimeFunction* FunctionTableHandler(ulong controlPc, nint context)
{
return JitUnwindWindows.FunctionTableHandler(_jitCache, _runtimeFunction, _unwindInfo, controlPc, context);
}
}
private enum UnwindOp
{
PushNonvol = 0,
@@ -59,27 +84,28 @@ namespace ARMeilleure.Translation.Cache
private static GetRuntimeFunctionCallback _getRuntimeFunctionCallback;
private static int _sizeOfRuntimeFunction;
private static readonly int _sizeOfRuntimeFunction;
private unsafe static RuntimeFunction* _runtimeFunction;
private static readonly ConcurrentDictionary<ulong, InternalFunctionHandler> _functionTableHandlers = new();
private unsafe static UnwindInfo* _unwindInfo;
static JitUnwindWindows()
{
_sizeOfRuntimeFunction = Marshal.SizeOf<RuntimeFunction>();
}
public static void InstallFunctionTableHandler(nint codeCachePointer, uint codeCacheLength, nint workBufferPtr)
public static void InstallFunctionTableHandler(JitCache jitCache, nint codeCachePointer, uint codeCacheLength, nint workBufferPtr)
{
ulong codeCachePtr = (ulong)codeCachePointer.ToInt64();
_sizeOfRuntimeFunction = Marshal.SizeOf<RuntimeFunction>();
bool result;
InternalFunctionHandler handler;
unsafe
{
_runtimeFunction = (RuntimeFunction*)workBufferPtr;
handler = new InternalFunctionHandler(jitCache, workBufferPtr);
_unwindInfo = (UnwindInfo*)(workBufferPtr + _sizeOfRuntimeFunction);
_getRuntimeFunctionCallback = new GetRuntimeFunctionCallback(FunctionTableHandler);
_getRuntimeFunctionCallback = handler.FunctionTableHandler;
result = RtlInstallFunctionTableCallback(
codeCachePtr | 3,
@@ -94,6 +120,8 @@ namespace ARMeilleure.Translation.Cache
{
throw new InvalidOperationException("Failure installing function table callback.");
}
_functionTableHandlers.TryAdd(codeCachePtr, handler);
}
public static void RemoveFunctionTableHandler(nint codeCachePointer)
@@ -111,24 +139,26 @@ namespace ARMeilleure.Translation.Cache
{
throw new InvalidOperationException("Failure removing function table callback.");
}
_functionTableHandlers.Remove(codeCachePtr, out _);
}
private static unsafe RuntimeFunction* FunctionTableHandler(ulong controlPc, nint context)
private static unsafe RuntimeFunction* FunctionTableHandler(JitCache jitCache, RuntimeFunction* runtimeFunction, UnwindInfo* unwindInfo, ulong controlPc, nint context)
{
int offset = (int)((long)controlPc - context.ToInt64());
if (!JitCache.TryFind(offset, out CacheEntry funcEntry, out _))
if (!jitCache.TryFind(offset, out CacheEntry funcEntry, out _))
{
return null; // Not found.
}
CodeGen.Unwinding.UnwindInfo unwindInfo = funcEntry.UnwindInfo;
CodeGen.Unwinding.UnwindInfo funcUnwindInfo = funcEntry.UnwindInfo;
int codeIndex = 0;
for (int index = unwindInfo.PushEntries.Length - 1; index >= 0; index--)
for (int index = funcUnwindInfo.PushEntries.Length - 1; index >= 0; index--)
{
UnwindPushEntry entry = unwindInfo.PushEntries[index];
UnwindPushEntry entry = funcUnwindInfo.PushEntries[index];
switch (entry.PseudoOp)
{
@@ -140,14 +170,14 @@ namespace ARMeilleure.Translation.Cache
if (stackOffset <= 0xFFFF0)
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.SaveXmm128, entry.PrologOffset, entry.RegIndex);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset / 16);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.SaveXmm128, entry.PrologOffset, entry.RegIndex);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset / 16);
}
else
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.SaveXmm128Far, entry.PrologOffset, entry.RegIndex);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 0);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 16);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.SaveXmm128Far, entry.PrologOffset, entry.RegIndex);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 0);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(stackOffset >> 16);
}
break;
@@ -161,18 +191,18 @@ namespace ARMeilleure.Translation.Cache
if (allocSize <= 128)
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocSmall, entry.PrologOffset, (allocSize / 8) - 1);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocSmall, entry.PrologOffset, (allocSize / 8) - 1);
}
else if (allocSize <= 0x7FFF8)
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocLarge, entry.PrologOffset, 0);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize / 8);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocLarge, entry.PrologOffset, 0);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize / 8);
}
else
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocLarge, entry.PrologOffset, 1);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 0);
_unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 16);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.AllocLarge, entry.PrologOffset, 1);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 0);
unwindInfo->UnwindCodes[codeIndex++] = (ushort)(allocSize >> 16);
}
break;
@@ -180,7 +210,7 @@ namespace ARMeilleure.Translation.Cache
case UnwindPseudoOp.PushReg:
{
_unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.PushNonvol, entry.PrologOffset, entry.RegIndex);
unwindInfo->UnwindCodes[codeIndex++] = PackUnwindOp(UnwindOp.PushNonvol, entry.PrologOffset, entry.RegIndex);
break;
}
@@ -192,16 +222,16 @@ namespace ARMeilleure.Translation.Cache
Debug.Assert(codeIndex <= MaxUnwindCodesArraySize);
_unwindInfo->VersionAndFlags = 1; // Flags: The function has no handler.
_unwindInfo->SizeOfProlog = (byte)unwindInfo.PrologSize;
_unwindInfo->CountOfUnwindCodes = (byte)codeIndex;
_unwindInfo->FrameRegister = 0;
unwindInfo->VersionAndFlags = 1; // Flags: The function has no handler.
unwindInfo->SizeOfProlog = (byte)funcUnwindInfo.PrologSize;
unwindInfo->CountOfUnwindCodes = (byte)codeIndex;
unwindInfo->FrameRegister = 0;
_runtimeFunction->BeginAddress = (uint)funcEntry.Offset;
_runtimeFunction->EndAddress = (uint)(funcEntry.Offset + funcEntry.Size);
_runtimeFunction->UnwindData = (uint)_sizeOfRuntimeFunction;
runtimeFunction->BeginAddress = (uint)funcEntry.Offset;
runtimeFunction->EndAddress = (uint)(funcEntry.Offset + funcEntry.Size);
runtimeFunction->UnwindData = (uint)_sizeOfRuntimeFunction;
return _runtimeFunction;
return runtimeFunction;
}
private static ushort PackUnwindOp(UnwindOp op, int prologOffset, int opInfo)