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>

src/Ryujinx.Cpu/LightningJit/Cache/CacheMemoryAllocator.cs

@@ -24,11 +24,30 @@ namespace Ryujinx.Cpu.LightningJit.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)
@@ -101,12 +120,13 @@ namespace Ryujinx.Cpu.LightningJit.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);
@@ -114,7 +134,8 @@ namespace Ryujinx.Cpu.LightningJit.Cache
                 }
             }
 
-            if (index > 0)
+            // Don't merge blocks from different allocations
+            if (index > 0 && block.Offset % _regionSize != 0)
             {
                 MemoryBlock prev = _blocks[index - 1];
 

src/Ryujinx.Cpu/LightningJit/Cache/JitCache.cs

@@ -4,83 +4,50 @@ using Ryujinx.Common.Logging;
 using Ryujinx.Memory;
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Runtime.InteropServices;
-using System.Runtime.Versioning;
 using System.Threading;
 
 namespace Ryujinx.Cpu.LightningJit.Cache
 {
-    static partial class JitCache
+    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.
-        // TODO: JIT Cache size should be application dependent, not global.
-        private const int CacheSize = 1024 * (1024 * 1024); // Megabytes * Size of Megabytes (since its in bytes).
+        private const uint CacheSize = 256 * 1024 * 1024; // Megabytes * Size of Megabytes (since its in bytes).
 
-        private static JitCacheInvalidation _jitCacheInvalidator;
+        private readonly JitCacheInvalidation _jitCacheInvalidator;
 
-        private static CacheMemoryAllocator _cacheAllocator;
+        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 static readonly List<ReservedRegion> _jitRegions = [];
-        private static int _activeRegionIndex = 0;
+        private readonly Lock _lock = new();
 
-        [SupportedOSPlatform("windows")]
-        [LibraryImport("kernel32.dll", SetLastError = true)]
-        public static partial nint FlushInstructionCache(nint hProcess, nint lpAddress, nuint dwSize);
+        private readonly List<ReservedRegion> _jitRegions = [];
 
-        [SupportedOSPlatform("macos")]
-        [LibraryImport("libSystem.dylib", EntryPoint = "sys_icache_invalidate")]
-        internal static partial void SysICacheInvalidate(nint start, nuint len);
-
-        [SupportedOSPlatform("linux")]
-        [LibraryImport("libgcc_s.so.1", EntryPoint = "__clear_cache")]
-        internal static partial void ClearCache(nint begin, nint end);
-
-        public static void Initialize(IJitMemoryAllocator allocator)
+        public JitCache(IJitMemoryAllocator allocator)
         {
-            if (_initialized)
-            {
-                return;
-            }
-
             lock (_lock)
             {
-                if (_initialized)
-                {
-                    return;
-                }
-
-                ReservedRegion firstRegion = new(allocator, CacheSize);
-                _jitRegions.Add(firstRegion);
-                _activeRegionIndex = 0;
+                _jitRegions.Add(new(allocator, CacheSize));
+                
+                _cacheAllocator = new CacheMemoryAllocator((int)CacheSize);
 
                 if (!OperatingSystem.IsWindows() && !OperatingSystem.IsMacOS())
                 {
                     _jitCacheInvalidator = new JitCacheInvalidation(allocator);
                 }
-
-                _cacheAllocator = new CacheMemoryAllocator(CacheSize);
-
-                _initialized = true;
             }
         }
 
-        public unsafe static nint Map(ReadOnlySpan<byte> code)
+        public unsafe nint Map(ReadOnlySpan<byte> 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)
                 {
@@ -94,18 +61,11 @@ namespace Ryujinx.Cpu.LightningJit.Cache
                 }
                 else
                 {
-                    ReprotectAsWritable(targetRegion, funcOffset, code.Length);
+                    ReprotectAsWritable(funcOffset, code.Length);
                     code.CopyTo(new Span<byte>((void*)funcPtr, code.Length));
-                    ReprotectAsExecutable(targetRegion, funcOffset, code.Length);
+                    ReprotectAsExecutable(funcOffset, code.Length);
 
-                    if (OperatingSystem.IsWindows() && RuntimeInformation.ProcessArchitecture == Architecture.Arm64)
-                    {
-                        FlushInstructionCache(Process.GetCurrentProcess().Handle, funcPtr, (nuint)code.Length);
-                    }
-                    else
-                    {
-                        _jitCacheInvalidator?.Invalidate(funcPtr, (ulong)code.Length);
-                    }
+                    _jitCacheInvalidator?.Invalidate(funcPtr, (ulong)code.Length);
                 }
 
                 Add(funcOffset, code.Length);
@@ -114,21 +74,20 @@ namespace Ryujinx.Cpu.LightningJit.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++)
                 {
+                    ReservedRegion region = _jitRegions[i];
                     if (pointer.ToInt64() < region.Pointer.ToInt64() ||
-                        pointer.ToInt64() >= (region.Pointer + CacheSize).ToInt64())
+                        pointer.ToInt64() >= (region.Pointer + (nint)CacheSize).ToInt64())
                     {
                         continue;
                     }
 
-                    int funcOffset = (int)(pointer.ToInt64() - region.Pointer.ToInt64());
+                    int funcOffset = (int)(pointer.ToInt64() - region.Pointer.ToInt64() + i * CacheSize);
 
                     if (TryFind(funcOffset, out CacheEntry entry, out int entryIndex) && entry.Offset == funcOffset)
                     {
@@ -141,59 +100,63 @@ namespace Ryujinx.Cpu.LightningJit.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);
-
-            for (int i = _activeRegionIndex; i < _jitRegions.Count; i++)
-            {
-                int allocOffset = _cacheAllocator.Allocate(codeSize);
+            
+            int allocOffset = _cacheAllocator.Allocate(codeSize); 
         
-                if (allocOffset >= 0)
-                {
-                    _jitRegions[i].ExpandIfNeeded((ulong)allocOffset + (ulong)codeSize);
-                    _activeRegionIndex = i;
-                    return allocOffset;
-                }
+            if (allocOffset >= 0)
+            {
+                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;
-            
-            int newRegionNumber = _activeRegionIndex;
 
-            Logger.Warning?.Print(LogClass.Cpu, $"JIT Cache Region {exhaustedRegion} exhausted, creating new Cache Region {newRegionNumber} ({((long)(newRegionNumber + 1) * CacheSize).Bytes()} Total Allocation).");
-        
-            _cacheAllocator = new CacheMemoryAllocator(CacheSize);
-
-            int allocOffsetNew = _cacheAllocator.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)
@@ -201,7 +164,7 @@ namespace Ryujinx.Cpu.LightningJit.Cache
             return checked(codeSize + (CodeAlignment - 1)) & ~(CodeAlignment - 1);
         }
 
-        private static void Add(int offset, int size)
+        private void Add(int offset, int size)
         {
             CacheEntry entry = new(offset, size);
 
@@ -215,7 +178,7 @@ namespace Ryujinx.Cpu.LightningJit.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)
             {
@@ -238,5 +201,13 @@ namespace Ryujinx.Cpu.LightningJit.Cache
             entryIndex = 0;
             return false;
         }
+
+        public void Dispose()
+        {
+            foreach (ReservedRegion jitRegion in _jitRegions)
+            {
+                jitRegion.Dispose();
+            }
+        }
     }
 }

src/Ryujinx.Cpu/LightningJit/Translator.cs

@@ -19,6 +19,7 @@ namespace Ryujinx.Cpu.LightningJit
         private static bool IsNoWxPlatform => false;
 
         private readonly ConcurrentQueue<KeyValuePair<ulong, TranslatedFunction>> _oldFuncs;
+        private readonly JitCache _jitCache;
         private readonly NoWxCache _noWxCache;
         private bool _disposed;
 
@@ -39,12 +40,12 @@ namespace Ryujinx.Cpu.LightningJit
             }
             else
             {
-                JitCache.Initialize(new JitMemoryAllocator(forJit: true));
+                _jitCache = new(new JitMemoryAllocator(forJit: true));
             }
 
             Functions = new TranslatorCache<TranslatedFunction>();
             FunctionTable = functionTable;
-            Stubs = new TranslatorStubs(FunctionTable, _noWxCache);
+            Stubs = new TranslatorStubs(_jitCache, FunctionTable, _noWxCache);
 
             FunctionTable.Fill = (ulong)Stubs.SlowDispatchStub;
 
@@ -100,7 +101,7 @@ namespace Ryujinx.Cpu.LightningJit
 
                 if (oldFunc != func)
                 {
-                    JitCache.Unmap(func.FuncPointer);
+                    _jitCache.Unmap(func.FuncPointer);
                     func = oldFunc;
                 }
 
@@ -121,7 +122,7 @@ namespace Ryujinx.Cpu.LightningJit
         private TranslatedFunction Translate(ulong address, ExecutionMode mode)
         {
             CompiledFunction func = Compile(address, mode);
-            nint funcPointer = JitCache.Map(func.Code);
+            nint funcPointer = _jitCache.Map(func.Code);
 
             return new TranslatedFunction(funcPointer, (ulong)func.GuestCodeLength);
         }
@@ -164,14 +165,14 @@ namespace Ryujinx.Cpu.LightningJit
 
             foreach (TranslatedFunction func in functions)
             {
-                JitCache.Unmap(func.FuncPointer);
+                _jitCache.Unmap(func.FuncPointer);
             }
 
             Functions.Clear();
 
             while (_oldFuncs.TryDequeue(out KeyValuePair<ulong, TranslatedFunction> kv))
             {
-                JitCache.Unmap(kv.Value.FuncPointer);
+                _jitCache.Unmap(kv.Value.FuncPointer);
             }
         }
 
@@ -188,6 +189,7 @@ namespace Ryujinx.Cpu.LightningJit
                     else
                     {
                         ClearJitCache();
+                        _jitCache.Dispose();
                     }
 
                     Stubs.Dispose();

src/Ryujinx.Cpu/LightningJit/TranslatorStubs.cs

@@ -19,6 +19,8 @@ namespace Ryujinx.Cpu.LightningJit
     {
         private delegate ulong GetFunctionAddressDelegate(nint framePointer, ulong address);
 
+        private readonly JitCache _jitCache;
+
         private readonly Lazy<nint> _slowDispatchStub;
 
         private bool _disposed;
@@ -76,13 +78,16 @@ namespace Ryujinx.Cpu.LightningJit
         /// Initializes a new instance of the <see cref="TranslatorStubs"/> class with the specified
         /// <see cref="Translator"/> instance.
         /// </summary>
+        /// <param name="jitCache">Jit cache to map functions in</param>
         /// <param name="functionTable">Function table used to store pointers to the functions that the guest code will call</param>
         /// <param name="noWxCache">Cache used on platforms that enforce W^X, otherwise should be null</param>
         /// <exception cref="ArgumentNullException"><paramref name="translator"/> is null</exception>
-        public TranslatorStubs(IAddressTable<ulong> functionTable, NoWxCache noWxCache)
+        public TranslatorStubs(JitCache jitCache, IAddressTable<ulong> functionTable, NoWxCache noWxCache)
         {
             ArgumentNullException.ThrowIfNull(functionTable);
 
+            _jitCache = jitCache;
+
             _functionTable = functionTable;
             _noWxCache = noWxCache;
             _getFunctionAddressRef = NativeInterface.GetFunctionAddress;
@@ -113,12 +118,12 @@ namespace Ryujinx.Cpu.LightningJit
                 {
                     if (_dispatchStub.IsValueCreated)
                     {
-                        JitCache.Unmap(_dispatchStub.Value);
+                        _jitCache.Unmap(_dispatchStub.Value);
                     }
 
                     if (_dispatchLoop.IsValueCreated)
                     {
-                        JitCache.Unmap(Marshal.GetFunctionPointerForDelegate(_dispatchLoop.Value));
+                        _jitCache.Unmap(Marshal.GetFunctionPointerForDelegate(_dispatchLoop.Value));
                     }
                 }
 
@@ -363,7 +368,7 @@ namespace Ryujinx.Cpu.LightningJit
             }
             else
             {
-                return JitCache.Map(code);
+                return _jitCache.Map(code);
             }
         }