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>

3 files changed, 48 insertions, 7 deletions

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)
                 {