JIT Sparse Function Table (#250)

More up to date build of the JIT Sparse PR for continued development.
JIT Sparse Function Table was originally developed by riperiperi for the
original Ryujinx project, and decreased the amount of layers in the
Function Table structure, to decrease lookup times at the cost of
slightly higher RAM usage.
This PR rebalances the JIT Sparse Function Table to be a bit more RAM
intensive, but faster in workloads where the JIT Function Table is a
bottleneck. Faster RAM will see a bigger impact and slower RAM (DDR3 and
potentially slow DDR4) will see a slight performance decrease.
This PR also implements a base for a PPTC profile system that could
allow for PPTC with ExeFS mods enabled in the future.
This PR also potentially fixes a strange issue where Avalonia would time
out in some rare instances, e.g. when running ExeFS mods with TotK and a
strange controller configuration.

---------

Co-authored-by: Evan Husted <gr33m11@gmail.com>

src/Ryujinx.Cpu/AddressTable.cs

@@ -0,0 +1,482 @@
+using ARMeilleure.Memory;
+using Ryujinx.Common;
+using Ryujinx.Cpu.Signal;
+using Ryujinx.Memory;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.InteropServices;
+using System.Threading;
+using static Ryujinx.Cpu.MemoryEhMeilleure;
+
+namespace ARMeilleure.Common
+{
+    /// <summary>
+    /// Represents a table of guest address to a value.
+    /// </summary>
+    /// <typeparam name="TEntry">Type of the value</typeparam>
+    public unsafe class AddressTable<TEntry> : IAddressTable<TEntry> where TEntry : unmanaged
+    {
+        /// <summary>
+        /// Represents a page of the address table.
+        /// </summary>
+        private readonly struct AddressTablePage
+        {
+            /// <summary>
+            /// True if the allocation belongs to a sparse block, false otherwise.
+            /// </summary>
+            public readonly bool IsSparse;
+
+            /// <summary>
+            /// Base address for the page.
+            /// </summary>
+            public readonly IntPtr Address;
+
+            public AddressTablePage(bool isSparse, IntPtr address)
+            {
+                IsSparse = isSparse;
+                Address = address;
+            }
+        }
+
+        /// <summary>
+        /// A sparsely mapped block of memory with a signal handler to map pages as they're accessed.
+        /// </summary>
+        private readonly struct TableSparseBlock : IDisposable
+        {
+            public readonly SparseMemoryBlock Block;
+            private readonly TrackingEventDelegate _trackingEvent;
+
+            public TableSparseBlock(ulong size, Action<IntPtr> ensureMapped, PageInitDelegate pageInit)
+            {
+                var block = new SparseMemoryBlock(size, pageInit, null);
+
+                _trackingEvent = (ulong address, ulong size, bool write) =>
+                {
+                    ulong pointer = (ulong)block.Block.Pointer + address;
+                    ensureMapped((IntPtr)pointer);
+                    return pointer;
+                };
+
+                bool added = NativeSignalHandler.AddTrackedRegion(
+                    (nuint)block.Block.Pointer,
+                    (nuint)(block.Block.Pointer + (IntPtr)block.Block.Size),
+                    Marshal.GetFunctionPointerForDelegate(_trackingEvent));
+
+                if (!added)
+                {
+                    throw new InvalidOperationException("Number of allowed tracked regions exceeded.");
+                }
+
+                Block = block;
+            }
+
+            public void Dispose()
+            {
+                NativeSignalHandler.RemoveTrackedRegion((nuint)Block.Block.Pointer);
+
+                Block.Dispose();
+            }
+        }
+
+        private bool _disposed;
+        private TEntry** _table;
+        private readonly List<AddressTablePage> _pages;
+        private TEntry _fill;
+
+        private readonly MemoryBlock _sparseFill;
+        private readonly SparseMemoryBlock _fillBottomLevel;
+        private readonly TEntry* _fillBottomLevelPtr;
+
+        private readonly List<TableSparseBlock> _sparseReserved;
+        private readonly ReaderWriterLockSlim _sparseLock;
+
+        private ulong _sparseBlockSize;
+        private ulong _sparseReservedOffset;
+
+        public bool Sparse { get; }
+
+        /// <inheritdoc/>
+        public ulong Mask { get; }
+
+        /// <inheritdoc/>
+        public AddressTableLevel[] Levels { get; }
+
+        /// <inheritdoc/>
+        public TEntry Fill
+        {
+            get
+            {
+                return _fill;
+            }
+            set
+            {
+                UpdateFill(value);
+            }
+        }
+
+        /// <inheritdoc/>
+        public IntPtr Base
+        {
+            get
+            {
+                ObjectDisposedException.ThrowIf(_disposed, this);
+
+                lock (_pages)
+                {
+                    return (IntPtr)GetRootPage();
+                }
+            }
+        }
+
+        /// <summary>
+        /// Constructs a new instance of the <see cref="AddressTable{TEntry}"/> class with the specified list of
+        /// <see cref="Level"/>.
+        /// </summary>
+        /// <param name="levels">Levels for the address table</param>
+        /// <param name="sparse">True if the bottom page should be sparsely mapped</param>
+        /// <exception cref="ArgumentNullException"><paramref name="levels"/> is null</exception>
+        /// <exception cref="ArgumentException">Length of <paramref name="levels"/> is less than 2</exception>
+        public AddressTable(AddressTableLevel[] levels, bool sparse)
+        {
+            ArgumentNullException.ThrowIfNull(levels);
+
+            _pages = new List<AddressTablePage>(capacity: 16);
+
+            Levels = levels;
+            Mask = 0;
+
+            foreach (var level in Levels)
+            {
+                Mask |= level.Mask;
+            }
+
+            Sparse = sparse;
+
+            if (sparse)
+            {
+                // If the address table is sparse, allocate a fill block
+
+                _sparseFill = new MemoryBlock(268435456ul, MemoryAllocationFlags.Mirrorable); //low Power TC uses size: 65536ul
+
+                ulong bottomLevelSize = (1ul << levels.Last().Length) * (ulong)sizeof(TEntry);
+
+                _fillBottomLevel = new SparseMemoryBlock(bottomLevelSize, null, _sparseFill);
+                _fillBottomLevelPtr = (TEntry*)_fillBottomLevel.Block.Pointer;
+
+                _sparseReserved = new List<TableSparseBlock>();
+                _sparseLock = new ReaderWriterLockSlim();
+
+                _sparseBlockSize = bottomLevelSize;
+            }
+        }
+
+        /// <summary>
+        /// Create an <see cref="AddressTable{TEntry}"/> instance for an ARM function table.
+        /// Selects the best table structure for A32/A64, taking into account the selected memory manager type.
+        /// </summary>
+        /// <param name="for64Bits">True if the guest is A64, false otherwise</param>
+        /// <param name="type">Memory manager type</param>
+        /// <returns>An <see cref="AddressTable{TEntry}"/> for ARM function lookup</returns>
+        public static AddressTable<TEntry> CreateForArm(bool for64Bits, MemoryManagerType type)
+        {
+            // Assume software memory means that we don't want to use any signal handlers.
+            bool sparse = type != MemoryManagerType.SoftwareMmu && type != MemoryManagerType.SoftwarePageTable;
+
+            return new AddressTable<TEntry>(AddressTablePresets.GetArmPreset(for64Bits, sparse), sparse);
+        }
+
+        /// <summary>
+        /// Update the fill value for the bottom level of the table.
+        /// </summary>
+        /// <param name="fillValue">New fill value</param>
+        private void UpdateFill(TEntry fillValue)
+        {
+            if (_sparseFill != null)
+            {
+                Span<byte> span = _sparseFill.GetSpan(0, (int)_sparseFill.Size);
+                MemoryMarshal.Cast<byte, TEntry>(span).Fill(fillValue);
+            }
+
+            _fill = fillValue;
+        }
+
+        /// <summary>
+        /// Signal that the given code range exists.
+        /// </summary>
+        /// <param name="address"></param>
+        /// <param name="size"></param>
+        public void SignalCodeRange(ulong address, ulong size)
+        {
+            AddressTableLevel bottom = Levels.Last();
+            ulong bottomLevelEntries = 1ul << bottom.Length;
+
+            ulong entryIndex = address >> bottom.Index;
+            ulong entries = size >> bottom.Index;
+            entries += entryIndex - BitUtils.AlignDown(entryIndex, bottomLevelEntries);
+
+            _sparseBlockSize = Math.Max(_sparseBlockSize, BitUtils.AlignUp(entries, bottomLevelEntries) * (ulong)sizeof(TEntry));
+        }
+
+        /// <inheritdoc/>
+        public bool IsValid(ulong address)
+        {
+            return (address & ~Mask) == 0;
+        }
+
+        /// <inheritdoc/>
+        public ref TEntry GetValue(ulong address)
+        {
+            ObjectDisposedException.ThrowIf(_disposed, this);
+
+            if (!IsValid(address))
+            {
+                throw new ArgumentException($"Address 0x{address:X} is not mapped onto the table.", nameof(address));
+            }
+
+            lock (_pages)
+            {
+                TEntry* page = GetPage(address);
+
+                int index = Levels[^1].GetValue(address);
+
+                EnsureMapped((IntPtr)(page + index));
+
+                return ref page[index];
+            }
+        }
+
+        /// <summary>
+        /// Gets the leaf page for the specified guest <paramref name="address"/>.
+        /// </summary>
+        /// <param name="address">Guest address</param>
+        /// <returns>Leaf page for the specified guest <paramref name="address"/></returns>
+        private TEntry* GetPage(ulong address)
+        {
+            TEntry** page = GetRootPage();
+
+            for (int i = 0; i < Levels.Length - 1; i++)
+            {
+                ref AddressTableLevel level = ref Levels[i];
+                ref TEntry* nextPage = ref page[level.GetValue(address)];
+
+                if (nextPage == null || nextPage == _fillBottomLevelPtr)
+                {
+                    ref AddressTableLevel nextLevel = ref Levels[i + 1];
+
+                    if (i == Levels.Length - 2)
+                    {
+                        nextPage = (TEntry*)Allocate(1 << nextLevel.Length, Fill, leaf: true);
+                    }
+                    else
+                    {
+                        nextPage = (TEntry*)Allocate(1 << nextLevel.Length, GetFillValue(i), leaf: false);
+                    }
+                }
+
+                page = (TEntry**)nextPage;
+            }
+
+            return (TEntry*)page;
+        }
+
+        /// <summary>
+        /// Ensure the given pointer is mapped in any overlapping sparse reservations.
+        /// </summary>
+        /// <param name="ptr">Pointer to be mapped</param>
+        private void EnsureMapped(IntPtr ptr)
+        {
+            if (Sparse)
+            {
+                // Check sparse allocations to see if the pointer is in any of them.
+                // Ensure the page is committed if there's a match.
+
+                _sparseLock.EnterReadLock();
+
+                try
+                {
+                    foreach (TableSparseBlock reserved in _sparseReserved)
+                    {
+                        SparseMemoryBlock sparse = reserved.Block;
+
+                        if (ptr >= sparse.Block.Pointer && ptr < sparse.Block.Pointer + (IntPtr)sparse.Block.Size)
+                        {
+                            sparse.EnsureMapped((ulong)(ptr - sparse.Block.Pointer));
+
+                            break;
+                        }
+                    }
+                }
+                finally
+                {
+                    _sparseLock.ExitReadLock();
+                }
+            }
+        }
+
+        /// <summary>
+        /// Get the fill value for a non-leaf level of the table.
+        /// </summary>
+        /// <param name="level">Level to get the fill value for</param>
+        /// <returns>The fill value</returns>
+        private IntPtr GetFillValue(int level)
+        {
+            if (_fillBottomLevel != null && level == Levels.Length - 2)
+            {
+                return (IntPtr)_fillBottomLevelPtr;
+            }
+            else
+            {
+                return IntPtr.Zero;
+            }
+        }
+
+        /// <summary>
+        /// Lazily initialize and get the root page of the <see cref="AddressTable{TEntry}"/>.
+        /// </summary>
+        /// <returns>Root page of the <see cref="AddressTable{TEntry}"/></returns>
+        private TEntry** GetRootPage()
+        {
+            if (_table == null)
+            {
+                if (Levels.Length == 1)
+                    _table = (TEntry**)Allocate(1 << Levels[0].Length, Fill, leaf: true);
+                else
+                    _table = (TEntry**)Allocate(1 << Levels[0].Length, GetFillValue(0), leaf: false);
+            }
+
+            return _table;
+        }
+
+        /// <summary>
+        /// Initialize a leaf page with the fill value.
+        /// </summary>
+        /// <param name="page">Page to initialize</param>
+        private void InitLeafPage(Span<byte> page)
+        {
+            MemoryMarshal.Cast<byte, TEntry>(page).Fill(_fill);
+        }
+
+        /// <summary>
+        /// Reserve a new sparse block, and add it to the list.
+        /// </summary>
+        /// <returns>The new sparse block that was added</returns>
+        private TableSparseBlock ReserveNewSparseBlock()
+        {
+            var block = new TableSparseBlock(_sparseBlockSize, EnsureMapped, InitLeafPage);
+
+            _sparseReserved.Add(block);
+            _sparseReservedOffset = 0;
+
+            return block;
+        }
+
+        /// <summary>
+        /// Allocates a block of memory of the specified type and length.
+        /// </summary>
+        /// <typeparam name="T">Type of elements</typeparam>
+        /// <param name="length">Number of elements</param>
+        /// <param name="fill">Fill value</param>
+        /// <param name="leaf"><see langword="true"/> if leaf; otherwise <see langword="false"/></param>
+        /// <returns>Allocated block</returns>
+        private IntPtr Allocate<T>(int length, T fill, bool leaf) where T : unmanaged
+        {
+            var size = sizeof(T) * length;
+
+            AddressTablePage page;
+
+            if (Sparse && leaf)
+            {
+                _sparseLock.EnterWriteLock();
+
+                SparseMemoryBlock block;
+
+                if (_sparseReserved.Count == 0)
+                {
+                    block = ReserveNewSparseBlock().Block;
+                }
+                else
+                {
+                    block = _sparseReserved.Last().Block;
+
+                    if (_sparseReservedOffset == block.Block.Size)
+                    {
+                        block = ReserveNewSparseBlock().Block;
+                    }
+                }
+
+                page = new AddressTablePage(true, block.Block.Pointer + (IntPtr)_sparseReservedOffset);
+
+                _sparseReservedOffset += (ulong)size;
+
+                _sparseLock.ExitWriteLock();
+            }
+            else
+            {
+                var address = (IntPtr)NativeAllocator.Instance.Allocate((uint)size);
+                page = new AddressTablePage(false, address);
+
+                var span = new Span<T>((void*)page.Address, length);
+                span.Fill(fill);
+            }
+
+            _pages.Add(page);
+
+            //TranslatorEventSource.Log.AddressTableAllocated(size, leaf);
+
+            return page.Address;
+        }
+
+        /// <summary>
+        /// Releases all resources used by the <see cref="AddressTable{TEntry}"/> instance.
+        /// </summary>
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        /// <summary>
+        /// Releases all unmanaged and optionally managed resources used by the <see cref="AddressTable{TEntry}"/>
+        /// instance.
+        /// </summary>
+        /// <param name="disposing"><see langword="true"/> to dispose managed resources also; otherwise just unmanaged resouces</param>
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!_disposed)
+            {
+                foreach (var page in _pages)
+                {
+                    if (!page.IsSparse)
+                    {
+                        Marshal.FreeHGlobal(page.Address);
+                    }
+                }
+
+                if (Sparse)
+                {
+                    foreach (TableSparseBlock block in _sparseReserved)
+                    {
+                        block.Dispose();
+                    }
+
+                    _sparseReserved.Clear();
+
+                    _fillBottomLevel.Dispose();
+                    _sparseFill.Dispose();
+                    _sparseLock.Dispose();
+                }
+
+                _disposed = true;
+            }
+        }
+
+        /// <summary>
+        /// Frees resources used by the <see cref="AddressTable{TEntry}"/> instance.
+        /// </summary>
+        ~AddressTable()
+        {
+            Dispose(false);
+        }
+    }
+}

src/Ryujinx.Cpu/AppleHv/HvCpuContext.cs

@@ -32,7 +32,7 @@ namespace Ryujinx.Cpu.AppleHv
         {
         }
 
-        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled)
+        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled, string cacheSelector)
         {
             return new DummyDiskCacheLoadState();
         }

src/Ryujinx.Cpu/ICpuContext.cs

@@ -48,7 +48,7 @@ namespace Ryujinx.Cpu
         /// <param name="displayVersion">Version of the application</param>
         /// <param name="enabled">True if the cache should be loaded from disk if it exists, false otherwise</param>
         /// <returns>Disk cache load progress reporter and manager</returns>
-        IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled);
+        IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled, string cacheSelector);
 
         /// <summary>
         /// Indicates that code has been loaded into guest memory, and that it might be executed in the future.

src/Ryujinx.Cpu/Jit/JitCpuContext.cs

@@ -1,3 +1,4 @@
+using ARMeilleure.Common;
 using ARMeilleure.Memory;
 using ARMeilleure.Translation;
 using Ryujinx.Cpu.Signal;
@@ -9,11 +10,13 @@ namespace Ryujinx.Cpu.Jit
     {
         private readonly ITickSource _tickSource;
         private readonly Translator _translator;
+        private readonly AddressTable<ulong> _functionTable;
 
         public JitCpuContext(ITickSource tickSource, IMemoryManager memory, bool for64Bit)
         {
             _tickSource = tickSource;
-            _translator = new Translator(new JitMemoryAllocator(forJit: true), memory, for64Bit);
+            _functionTable = AddressTable<ulong>.CreateForArm(for64Bit, memory.Type);
+            _translator = new Translator(new JitMemoryAllocator(forJit: true), memory, _functionTable);
 
             if (memory.Type.IsHostMappedOrTracked())
             {
@@ -47,14 +50,15 @@ namespace Ryujinx.Cpu.Jit
         }
 
         /// <inheritdoc/>
-        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled)
+        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled, string cacheSelector)
         {
-            return new JitDiskCacheLoadState(_translator.LoadDiskCache(titleIdText, displayVersion, enabled));
+            return new JitDiskCacheLoadState(_translator.LoadDiskCache(titleIdText, displayVersion, enabled, cacheSelector));
         }
 
         /// <inheritdoc/>
         public void PrepareCodeRange(ulong address, ulong size)
         {
+            _functionTable.SignalCodeRange(address, size);
             _translator.PrepareCodeRange(address, size);
         }
 

src/Ryujinx.Cpu/LightningJit/Arm32/Target/Arm64/InstEmitFlow.cs

@@ -140,6 +140,10 @@ namespace Ryujinx.Cpu.LightningJit.Arm32.Target.Arm64
             bool isTail = false)
         {
             int tempRegister;
+            int tempGuestAddress = -1;
+
+            bool inlineLookup = guestAddress.Kind != OperandKind.Constant && 
+                                funcTable is { Sparse: true };
 
             if (guestAddress.Kind == OperandKind.Constant)
             {
@@ -153,9 +157,16 @@ namespace Ryujinx.Cpu.LightningJit.Arm32.Target.Arm64
             else
             {
                 asm.StrRiUn(guestAddress, Register(regAlloc.FixedContextRegister), NativeContextOffsets.DispatchAddressOffset);
+
+                if (inlineLookup && guestAddress.Value == 0)
+                {
+                    // X0 will be overwritten. Move the address to a temp register.
+                    tempGuestAddress = regAlloc.AllocateTempGprRegister();
+                    asm.Mov(Register(tempGuestAddress), guestAddress);
+                }
             }
 
-            tempRegister = regAlloc.FixedContextRegister == 1 ? 2 : 1;
+            tempRegister = NextFreeRegister(1, tempGuestAddress);
 
             if (!isTail)
             {
@@ -176,6 +187,40 @@ namespace Ryujinx.Cpu.LightningJit.Arm32.Target.Arm64
                 asm.Mov(rn, funcPtrLoc & ~0xfffUL);
                 asm.LdrRiUn(rn, rn, (int)(funcPtrLoc & 0xfffUL));
             }
+            else if (inlineLookup)
+            {
+                // Inline table lookup. Only enabled when the sparse function table is enabled with 2 levels.
+
+                Operand indexReg = Register(NextFreeRegister(tempRegister + 1, tempGuestAddress));
+
+                if (tempGuestAddress != -1)
+                {
+                    guestAddress = Register(tempGuestAddress);
+                }
+
+                ulong tableBase = (ulong)funcTable.Base;
+
+                // Index into the table.
+                asm.Mov(rn, tableBase);
+
+                for (int i = 0; i < funcTable.Levels.Length; i++)
+                {
+                    var level = funcTable.Levels[i];
+                    asm.Ubfx(indexReg, guestAddress, level.Index, level.Length);
+                    asm.Lsl(indexReg, indexReg, Const(3));
+
+                    // Index into the page.
+                    asm.Add(rn, rn, indexReg);
+
+                    // Load the page address.
+                    asm.LdrRiUn(rn, rn, 0);
+                }
+
+                if (tempGuestAddress != -1)
+                {
+                    regAlloc.FreeTempGprRegister(tempGuestAddress);
+                }
+            }
             else
             {
                 asm.Mov(rn, (ulong)funcPtr);
@@ -252,5 +297,20 @@ namespace Ryujinx.Cpu.LightningJit.Arm32.Target.Arm64
         {
             return new Operand(register, RegisterType.Integer, type);
         }
+
+        private static Operand Const(long value, OperandType type = OperandType.I64)
+        {
+            return new Operand(type, (ulong)value);
+        }
+
+        private static int NextFreeRegister(int start, int avoid)
+        {
+            if (start == avoid)
+            {
+                start++;
+            }
+
+            return start;
+        }
     }
 }

src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/InstEmitSystem.cs

@@ -305,6 +305,10 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             bool isTail = false)
         {
             int tempRegister;
+            int tempGuestAddress = -1;
+
+            bool inlineLookup = guestAddress.Kind != OperandKind.Constant && 
+                                funcTable is { Sparse: true };
 
             if (guestAddress.Kind == OperandKind.Constant)
             {
@@ -318,9 +322,16 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             else
             {
                 asm.StrRiUn(guestAddress, Register(regAlloc.FixedContextRegister), NativeContextOffsets.DispatchAddressOffset);
+
+                if (inlineLookup && guestAddress.Value == 0)
+                {
+                    // X0 will be overwritten. Move the address to a temp register.
+                    tempGuestAddress = regAlloc.AllocateTempGprRegister();
+                    asm.Mov(Register(tempGuestAddress), guestAddress);
+                }
             }
 
-            tempRegister = regAlloc.FixedContextRegister == 1 ? 2 : 1;
+            tempRegister = NextFreeRegister(1, tempGuestAddress);
 
             if (!isTail)
             {
@@ -341,6 +352,40 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
                 asm.Mov(rn, funcPtrLoc & ~0xfffUL);
                 asm.LdrRiUn(rn, rn, (int)(funcPtrLoc & 0xfffUL));
             }
+            else if (inlineLookup)
+            {
+                // Inline table lookup. Only enabled when the sparse function table is enabled with 2 levels.
+
+                Operand indexReg = Register(NextFreeRegister(tempRegister + 1, tempGuestAddress));
+
+                if (tempGuestAddress != -1)
+                {
+                    guestAddress = Register(tempGuestAddress);
+                }
+
+                ulong tableBase = (ulong)funcTable.Base;
+
+                // Index into the table.
+                asm.Mov(rn, tableBase);
+
+                for (int i = 0; i < funcTable.Levels.Length; i++)
+                {
+                    var level = funcTable.Levels[i];
+                    asm.Ubfx(indexReg, guestAddress, level.Index, level.Length);
+                    asm.Lsl(indexReg, indexReg, Const(3));
+
+                    // Index into the page.
+                    asm.Add(rn, rn, indexReg);
+
+                    // Load the page address.
+                    asm.LdrRiUn(rn, rn, 0);
+                }
+
+                if (tempGuestAddress != -1)
+                {
+                    regAlloc.FreeTempGprRegister(tempGuestAddress);
+                }
+            }
             else
             {
                 asm.Mov(rn, (ulong)funcPtr);
@@ -613,5 +658,20 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
         {
             return new Operand(register, RegisterType.Integer, type);
         }
+
+        private static Operand Const(long value, OperandType type = OperandType.I64)
+        {
+            return new Operand(type, (ulong)value);
+        }
+
+        private static int NextFreeRegister(int start, int avoid)
+        {
+            if (start == avoid)
+            {
+                start++;
+            }
+
+            return start;
+        }
     }
 }

src/Ryujinx.Cpu/LightningJit/LightningJitCpuContext.cs

@@ -1,3 +1,4 @@
+using ARMeilleure.Common;
 using ARMeilleure.Memory;
 using Ryujinx.Cpu.Jit;
 using Ryujinx.Cpu.LightningJit.State;
@@ -8,11 +9,16 @@ namespace Ryujinx.Cpu.LightningJit
     {
         private readonly ITickSource _tickSource;
         private readonly Translator _translator;
+        private readonly AddressTable<ulong> _functionTable;
 
         public LightningJitCpuContext(ITickSource tickSource, IMemoryManager memory, bool for64Bit)
         {
             _tickSource = tickSource;
-            _translator = new Translator(memory, for64Bit);
+
+            _functionTable = AddressTable<ulong>.CreateForArm(for64Bit, memory.Type);
+
+            _translator = new Translator(memory, _functionTable);
+
             memory.UnmapEvent += UnmapHandler;
         }
 
@@ -40,7 +46,7 @@ namespace Ryujinx.Cpu.LightningJit
         }
 
         /// <inheritdoc/>
-        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled)
+        public IDiskCacheLoadState LoadDiskCache(string titleIdText, string displayVersion, bool enabled, string cacheSelector)
         {
             return new DummyDiskCacheLoadState();
         }
@@ -48,6 +54,7 @@ namespace Ryujinx.Cpu.LightningJit
         /// <inheritdoc/>
         public void PrepareCodeRange(ulong address, ulong size)
         {
+            _functionTable.SignalCodeRange(address, size);
         }
 
         public void Dispose()

src/Ryujinx.Cpu/LightningJit/Translator.cs

@@ -19,25 +19,6 @@ namespace Ryujinx.Cpu.LightningJit
         // Should be enabled on platforms that enforce W^X.
         private static bool IsNoWxPlatform => false;
 
-        private static readonly AddressTable<ulong>.Level[] _levels64Bit =
-            new AddressTable<ulong>.Level[]
-            {
-                new(31, 17),
-                new(23,  8),
-                new(15,  8),
-                new( 7,  8),
-                new( 2,  5),
-            };
-
-        private static readonly AddressTable<ulong>.Level[] _levels32Bit =
-            new AddressTable<ulong>.Level[]
-            {
-                new(23, 9),
-                new(15, 8),
-                new( 7, 8),
-                new( 1, 6),
-            };
-
         private readonly ConcurrentQueue<KeyValuePair<ulong, TranslatedFunction>> _oldFuncs;
         private readonly NoWxCache _noWxCache;
         private bool _disposed;
@@ -47,7 +28,7 @@ namespace Ryujinx.Cpu.LightningJit
         internal TranslatorStubs Stubs { get; }
         internal IMemoryManager Memory { get; }
 
-        public Translator(IMemoryManager memory, bool for64Bits)
+        public Translator(IMemoryManager memory, AddressTable<ulong> functionTable)
         {
             Memory = memory;
 
@@ -63,7 +44,7 @@ namespace Ryujinx.Cpu.LightningJit
             }
 
             Functions = new TranslatorCache<TranslatedFunction>();
-            FunctionTable = new AddressTable<ulong>(for64Bits ? _levels64Bit : _levels32Bit);
+            FunctionTable = functionTable;
             Stubs = new TranslatorStubs(FunctionTable, _noWxCache);
 
             FunctionTable.Fill = (ulong)Stubs.SlowDispatchStub;

src/Ryujinx.Cpu/LightningJit/TranslatorStubs.cs

@@ -23,7 +23,7 @@ namespace Ryujinx.Cpu.LightningJit
 
         private bool _disposed;
 
-        private readonly AddressTable<ulong> _functionTable;
+        private readonly IAddressTable<ulong> _functionTable;
         private readonly NoWxCache _noWxCache;
         private readonly GetFunctionAddressDelegate _getFunctionAddressRef;
         private readonly nint _getFunctionAddress;
@@ -79,7 +79,7 @@ namespace Ryujinx.Cpu.LightningJit
         /// <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(AddressTable<ulong> functionTable, NoWxCache noWxCache)
+        public TranslatorStubs(IAddressTable<ulong> functionTable, NoWxCache noWxCache)
         {
             ArgumentNullException.ThrowIfNull(functionTable);