Implement host tracked memory manager mode (#6356)

* Add host tracked memory manager mode

* Skipping flush is no longer needed

* Formatting + revert unrelated change

* LightningJit: Ensure that dest register is saved for load ops that do partial updates

* avoid allocations when doing address space lookup

Add missing improvement

* IsRmwMemory -> IsPartialRegisterUpdateMemory

* Ensure we iterate all private allocations in range

* PR feedback and potential fixes

* Simplified bridges a lot

* Skip calling SignalMappingChanged if Guest is true

* Late map bridge too

* Force address masking for prefetch instructions

* Reprotection for bridges

* Move partition list validation to separate debug method

* Move host tracked related classes to HostTracked folder

* New HostTracked namespace

* Move host tracked modes to the end of enum to avoid PPTC invalidation

---------

Co-authored-by: riperiperi <rhy3756547@hotmail.com>

8 files changed, 112 insertions, 20 deletions

diff --git a/src/Ryujinx.Cpu/LightningJit/Arm32/Target/Arm64/InstEmitMemory.cs b/src/Ryujinx.Cpu/LightningJit/Arm32/Target/Arm64/InstEmitMemory.cs
index 6ab4b949..d8caee6e 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm32/Target/Arm64/InstEmitMemory.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm32/Target/Arm64/InstEmitMemory.cs
@@ -1126,11 +1126,23 @@ namespace Ryujinx.Cpu.LightningJit.Arm32.Target.Arm64
             Operand destination64 = new(destination.Kind, OperandType.I64, destination.Value);
             Operand basePointer = new(regAlloc.FixedPageTableRegister, RegisterType.Integer, OperandType.I64);
 
-            if (mmType == MemoryManagerType.HostMapped || mmType == MemoryManagerType.HostMappedUnsafe)
+            // We don't need to mask the address for the safe mode, since it is already naturally limited to 32-bit
+            // and can never reach out of the guest address space.
+
+            if (mmType.IsHostTracked())
             {
-                // We don't need to mask the address for the safe mode, since it is already naturally limited to 32-bit
-                // and can never reach out of the guest address space.
+                int tempRegister = regAlloc.AllocateTempGprRegister();
+
+                Operand pte = new(tempRegister, RegisterType.Integer, OperandType.I64);
+
+                asm.Lsr(pte, guestAddress, new Operand(OperandKind.Constant, OperandType.I32, 12));
+                asm.LdrRr(pte, basePointer, pte, ArmExtensionType.Uxtx, true);
+                asm.Add(destination64, pte, guestAddress);
 
+                regAlloc.FreeTempGprRegister(tempRegister);
+            }
+            else if (mmType.IsHostMapped())
+            {
                 asm.Add(destination64, basePointer, guestAddress);
             }
             else
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/InstName.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/InstName.cs
index 36564064..3391a2c1 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/InstName.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/InstName.cs
@@ -1131,5 +1131,37 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
 
             return false;
         }
+
+        public static bool IsPartialRegisterUpdateMemory(this InstName name)
+        {
+            switch (name)
+            {
+                case InstName.Ld1AdvsimdSnglAsNoPostIndex:
+                case InstName.Ld1AdvsimdSnglAsPostIndex:
+                case InstName.Ld2AdvsimdSnglAsNoPostIndex:
+                case InstName.Ld2AdvsimdSnglAsPostIndex:
+                case InstName.Ld3AdvsimdSnglAsNoPostIndex:
+                case InstName.Ld3AdvsimdSnglAsPostIndex:
+                case InstName.Ld4AdvsimdSnglAsNoPostIndex:
+                case InstName.Ld4AdvsimdSnglAsPostIndex:
+                    return true;
+            }
+
+            return false;
+        }
+
+        public static bool IsPrefetchMemory(this InstName name)
+        {
+            switch (name)
+            {
+                case InstName.PrfmImm:
+                case InstName.PrfmLit:
+                case InstName.PrfmReg:
+                case InstName.Prfum:
+                    return true;
+            }
+
+            return false;
+        }
     }
 }
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterAllocator.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterAllocator.cs
index c9a93209..1c6eab0d 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterAllocator.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterAllocator.cs
@@ -1,15 +1,12 @@
+using ARMeilleure.Memory;
 using Ryujinx.Cpu.LightningJit.CodeGen.Arm64;
 using System;
-using System.Diagnostics;
 using System.Numerics;
 
 namespace Ryujinx.Cpu.LightningJit.Arm64
 {
     class RegisterAllocator
     {
-        public const int MaxTemps = 1;
-        public const int MaxTempsInclFixed = MaxTemps + 2;
-
         private uint _gprMask;
         private readonly uint _fpSimdMask;
         private readonly uint _pStateMask;
@@ -25,7 +22,7 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
         public uint AllFpSimdMask => _fpSimdMask;
         public uint AllPStateMask => _pStateMask;
 
-        public RegisterAllocator(uint gprMask, uint fpSimdMask, uint pStateMask, bool hasHostCall)
+        public RegisterAllocator(MemoryManagerType mmType, uint gprMask, uint fpSimdMask, uint pStateMask, bool hasHostCall)
         {
             _gprMask = gprMask;
             _fpSimdMask = fpSimdMask;
@@ -56,7 +53,7 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
 
             BuildRegisterMap(_registerMap);
 
-            Span<int> tempRegisters = stackalloc int[MaxTemps];
+            Span<int> tempRegisters = stackalloc int[CalculateMaxTemps(mmType)];
 
             for (int index = 0; index < tempRegisters.Length; index++)
             {
@@ -150,5 +147,15 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
         {
             mask &= ~(1u << index);
         }
+
+        public static int CalculateMaxTemps(MemoryManagerType mmType)
+        {
+            return mmType.IsHostMapped() ? 1 : 2;
+        }
+
+        public static int CalculateMaxTempsInclFixed(MemoryManagerType mmType)
+        {
+            return CalculateMaxTemps(mmType) + 2;
+        }
     }
 }
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterUtils.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterUtils.cs
index eb3fc229..191e03e7 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterUtils.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/RegisterUtils.cs
@@ -247,7 +247,7 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
                     }
                 }
 
-                if (!flags.HasFlag(InstFlags.ReadRt))
+                if (!flags.HasFlag(InstFlags.ReadRt) || name.IsPartialRegisterUpdateMemory())
                 {
                     if (flags.HasFlag(InstFlags.Rt))
                     {
@@ -281,7 +281,7 @@ namespace Ryujinx.Cpu.LightningJit.Arm64
                     gprMask |= MaskFromIndex(ExtractRd(flags, encoding));
                 }
 
-                if (!flags.HasFlag(InstFlags.ReadRt))
+                if (!flags.HasFlag(InstFlags.ReadRt) || name.IsPartialRegisterUpdateMemory())
                 {
                     if (flags.HasFlag(InstFlags.Rt))
                     {
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Compiler.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Compiler.cs
index 7ef3bf49..7a6d761e 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Compiler.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Compiler.cs
@@ -316,7 +316,7 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             uint pStateUseMask = multiBlock.GlobalUseMask.PStateMask;
 
             CodeWriter writer = new();
-            RegisterAllocator regAlloc = new(gprUseMask, fpSimdUseMask, pStateUseMask, multiBlock.HasHostCall);
+            RegisterAllocator regAlloc = new(memoryManager.Type, gprUseMask, fpSimdUseMask, pStateUseMask, multiBlock.HasHostCall);
             RegisterSaveRestore rsr = new(
                 regAlloc.AllGprMask & AbiConstants.GprCalleeSavedRegsMask,
                 regAlloc.AllFpSimdMask & AbiConstants.FpSimdCalleeSavedRegsMask,
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Decoder.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Decoder.cs
index 00a1758f..d5e1eb19 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Decoder.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/Decoder.cs
@@ -274,7 +274,8 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
 
                 uint tempGprUseMask = gprUseMask | instGprReadMask | instGprWriteMask;
 
-                if (CalculateAvailableTemps(tempGprUseMask) < CalculateRequiredGprTemps(tempGprUseMask) || totalInsts++ >= MaxInstructionsPerFunction)
+                if (CalculateAvailableTemps(tempGprUseMask) < CalculateRequiredGprTemps(memoryManager.Type, tempGprUseMask) ||
+                    totalInsts++ >= MaxInstructionsPerFunction)
                 {
                     isTruncated = true;
                     address -= 4UL;
@@ -378,9 +379,9 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             return false;
         }
 
-        private static int CalculateRequiredGprTemps(uint gprUseMask)
+        private static int CalculateRequiredGprTemps(MemoryManagerType mmType, uint gprUseMask)
         {
-            return BitOperations.PopCount(gprUseMask & RegisterUtils.ReservedRegsMask) + RegisterAllocator.MaxTempsInclFixed;
+            return BitOperations.PopCount(gprUseMask & RegisterUtils.ReservedRegsMask) + RegisterAllocator.CalculateMaxTempsInclFixed(mmType);
         }
 
         private static int CalculateAvailableTemps(uint gprUseMask)
diff --git a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/InstEmitMemory.cs b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/InstEmitMemory.cs
index e03d9373..790a7de9 100644
--- a/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/InstEmitMemory.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Arm64/Target/Arm64/InstEmitMemory.cs
@@ -55,6 +55,16 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             ulong pc,
             uint encoding)
         {
+            if (name.IsPrefetchMemory() && mmType == MemoryManagerType.HostTrackedUnsafe)
+            {
+                // Prefetch to invalid addresses do not cause faults, so for memory manager
+                // types where we need to access the page table before doing the prefetch,
+                // we should make sure we won't try to access an out of bounds page table region.
+                // To do this, we force the masked memory manager variant to be used.
+
+                mmType = MemoryManagerType.HostTracked;
+            }
+
             switch (addressForm)
             {
                 case AddressForm.OffsetReg:
@@ -511,18 +521,48 @@ namespace Ryujinx.Cpu.LightningJit.Arm64.Target.Arm64
             WriteAddressTranslation(asBits, mmType, regAlloc, ref asm, destination, guestAddress);
         }
 
-        private static void WriteAddressTranslation(int asBits, MemoryManagerType mmType, RegisterAllocator regAlloc, ref Assembler asm, Operand destination, ulong guestAddress)
+        private static void WriteAddressTranslation(
+            int asBits,
+            MemoryManagerType mmType,
+            RegisterAllocator regAlloc,
+            ref Assembler asm,
+            Operand destination,
+            ulong guestAddress)
         {
             asm.Mov(destination, guestAddress);
 
             WriteAddressTranslation(asBits, mmType, regAlloc, ref asm, destination, destination);
         }
 
-        private static void WriteAddressTranslation(int asBits, MemoryManagerType mmType, RegisterAllocator regAlloc, ref Assembler asm, Operand destination, Operand guestAddress)
+        private static void WriteAddressTranslation(
+            int asBits,
+            MemoryManagerType mmType,
+            RegisterAllocator regAlloc,
+            ref Assembler asm,
+            Operand destination,
+            Operand guestAddress)
         {
             Operand basePointer = new(regAlloc.FixedPageTableRegister, RegisterType.Integer, OperandType.I64);
 
-            if (mmType == MemoryManagerType.HostMapped || mmType == MemoryManagerType.HostMappedUnsafe)
+            if (mmType.IsHostTracked())
+            {
+                int tempRegister = regAlloc.AllocateTempGprRegister();
+
+                Operand pte = new(tempRegister, RegisterType.Integer, OperandType.I64);
+
+                asm.Lsr(pte, guestAddress, new Operand(OperandKind.Constant, OperandType.I32, 12));
+
+                if (mmType == MemoryManagerType.HostTracked)
+                {
+                    asm.And(pte, pte, new Operand(OperandKind.Constant, OperandType.I64, ulong.MaxValue >> (64 - (asBits - 12))));
+                }
+
+                asm.LdrRr(pte, basePointer, pte, ArmExtensionType.Uxtx, true);
+                asm.Add(destination, pte, guestAddress);
+
+                regAlloc.FreeTempGprRegister(tempRegister);
+            }
+            else if (mmType.IsHostMapped())
             {
                 if (mmType == MemoryManagerType.HostMapped)
                 {
diff --git a/src/Ryujinx.Cpu/LightningJit/Translator.cs b/src/Ryujinx.Cpu/LightningJit/Translator.cs
index c883c1d6..d6241025 100644
--- a/src/Ryujinx.Cpu/LightningJit/Translator.cs
+++ b/src/Ryujinx.Cpu/LightningJit/Translator.cs
@@ -68,9 +68,9 @@ namespace Ryujinx.Cpu.LightningJit
 
             FunctionTable.Fill = (ulong)Stubs.SlowDispatchStub;
 
-            if (memory.Type.IsHostMapped())
+            if (memory.Type.IsHostMappedOrTracked())
             {
-                NativeSignalHandler.InitializeSignalHandler(MemoryBlock.GetPageSize());
+                NativeSignalHandler.InitializeSignalHandler();
             }
         }