support allocating more registers for aarch64 JIT

also some minor fixes for the x64 JIT as well

4 files changed, 29 insertions, 20 deletions

diff --git a/src/ARMJIT_x64/ARMJIT_Branch.cpp b/src/ARMJIT_x64/ARMJIT_Branch.cpp
index 8ca3542..a89cf99 100644
--- a/src/ARMJIT_x64/ARMJIT_Branch.cpp
+++ b/src/ARMJIT_x64/ARMJIT_Branch.cpp
@@ -165,7 +165,7 @@ void Compiler::Comp_JumpTo(Gen::X64Reg addr, bool restoreCPSR)
     bool cpsrDirty = CPSRDirty;
     SaveCPSR();
 
-    PushRegs(restoreCPSR);
+    PushRegs(restoreCPSR, true);
 
     MOV(64, R(ABI_PARAM1), R(RCPU));
     MOV(32, R(ABI_PARAM2), R(addr));
@@ -178,7 +178,7 @@ void Compiler::Comp_JumpTo(Gen::X64Reg addr, bool restoreCPSR)
     else
         CALL((void*)&ARMv4JumpToTrampoline);
 
-    PopRegs(restoreCPSR);
+    PopRegs(restoreCPSR, true);
 
     LoadCPSR();
     // in case this instruction is skipped
diff --git a/src/ARMJIT_x64/ARMJIT_Compiler.cpp b/src/ARMJIT_x64/ARMJIT_Compiler.cpp
index 5ab8c6a..aef8b91 100644
--- a/src/ARMJIT_x64/ARMJIT_Compiler.cpp
+++ b/src/ARMJIT_x64/ARMJIT_Compiler.cpp
@@ -64,7 +64,7 @@ const BitSet32 CallerSavedPushRegs({R10, R11});
 const BitSet32 CallerSavedPushRegs({R9, R10, R11});
 #endif
 
-void Compiler::PushRegs(bool saveHiRegs)
+void Compiler::PushRegs(bool saveHiRegs, bool saveRegsToBeChanged, bool allowUnload)
 {
     BitSet32 loadedRegs(RegCache.LoadedRegs);
 
@@ -83,17 +83,26 @@ void Compiler::PushRegs(bool saveHiRegs)
     }
 
     for (int reg : loadedRegs)
-        if (BitSet32(1 << RegCache.Mapping[reg]) & ABI_ALL_CALLER_SAVED)
-            SaveReg(reg, RegCache.Mapping[reg]);
+    {
+        if (CallerSavedPushRegs[RegCache.Mapping[reg]]
+            && (saveRegsToBeChanged || !((1<<reg) & CurInstr.Info.DstRegs && !((1<<reg) & CurInstr.Info.SrcRegs))))
+        {
+            if ((Thumb || CurInstr.Cond() == 0xE) && !((1 << reg) & (CurInstr.Info.DstRegs|CurInstr.Info.SrcRegs)) && allowUnload)
+                RegCache.UnloadRegister(reg);
+            else
+                SaveReg(reg, RegCache.Mapping[reg]);
+        }
+    }
 }
 
-void Compiler::PopRegs(bool saveHiRegs)
+void Compiler::PopRegs(bool saveHiRegs, bool saveRegsToBeChanged)
 {
     BitSet32 loadedRegs(RegCache.LoadedRegs);
     for (int reg : loadedRegs)
     {
         if ((saveHiRegs && reg >= 8 && reg < 15)
-            || BitSet32(1 << RegCache.Mapping[reg]) & ABI_ALL_CALLER_SAVED)
+            || (CallerSavedPushRegs[RegCache.Mapping[reg]]
+                && (saveRegsToBeChanged || !((1<<reg) & CurInstr.Info.DstRegs && !((1<<reg) & CurInstr.Info.SrcRegs)))))
         {
             LoadReg(reg, RegCache.Mapping[reg]);
         }
@@ -205,14 +214,14 @@ void Compiler::A_Comp_MSR()
             AND(32, R(RSCRATCH2), val);
             OR(32, R(RCPSR), R(RSCRATCH2));
 
-            PushRegs(true);
+            PushRegs(true, true);
 
             MOV(32, R(ABI_PARAM3), R(RCPSR));
             MOV(32, R(ABI_PARAM2), R(RSCRATCH3));
             MOV(64, R(ABI_PARAM1), R(RCPU));
             CALL((void*)&UpdateModeTrampoline);
 
-            PopRegs(true);
+            PopRegs(true, true);
         }
     }
 }
@@ -659,7 +668,7 @@ void Compiler::Comp_SpecialBranchBehaviour(bool taken)
     }
 }
 
-JitBlockEntry Compiler::CompileBlock(ARM* cpu, bool thumb, FetchedInstr instrs[], int instrsCount)
+JitBlockEntry Compiler::CompileBlock(ARM* cpu, bool thumb, FetchedInstr instrs[], int instrsCount, bool hasMemoryInstr)
 {
     if (NearSize - (GetCodePtr() - NearStart) < 1024 * 32) // guess...
     {
diff --git a/src/ARMJIT_x64/ARMJIT_Compiler.h b/src/ARMJIT_x64/ARMJIT_Compiler.h
index ff7095b..eedfc1a 100644
--- a/src/ARMJIT_x64/ARMJIT_Compiler.h
+++ b/src/ARMJIT_x64/ARMJIT_Compiler.h
@@ -79,7 +79,7 @@ public:
 
     void Reset();
 
-    JitBlockEntry CompileBlock(ARM* cpu, bool thumb, FetchedInstr instrs[], int instrsCount);
+    JitBlockEntry CompileBlock(ARM* cpu, bool thumb, FetchedInstr instrs[], int instrsCount, bool hasMemoryInstr);
 
     void LoadReg(int reg, Gen::X64Reg nativeReg);
     void SaveReg(int reg, Gen::X64Reg nativeReg);
@@ -192,8 +192,8 @@ public:
 
     Gen::FixupBranch CheckCondition(u32 cond);
 
-    void PushRegs(bool saveHiRegs);
-    void PopRegs(bool saveHiRegs);
+    void PushRegs(bool saveHiRegs, bool saveRegsToBeChanged, bool allowUnload = true);
+    void PopRegs(bool saveHiRegs, bool saveRegsToBeChanged);
 
     Gen::OpArg MapReg(int reg)
     {
diff --git a/src/ARMJIT_x64/ARMJIT_LoadStore.cpp b/src/ARMJIT_x64/ARMJIT_LoadStore.cpp
index d06c241..88c0a9b 100644
--- a/src/ARMJIT_x64/ARMJIT_LoadStore.cpp
+++ b/src/ARMJIT_x64/ARMJIT_LoadStore.cpp
@@ -266,7 +266,7 @@ void Compiler::Comp_MemAccess(int rd, int rn, const Op2& op2, int size, int flag
     }
     else
     {
-        PushRegs(false);
+        PushRegs(false, false);
 
         void* func = NULL;
         if (addrIsStatic)
@@ -283,7 +283,7 @@ void Compiler::Comp_MemAccess(int rd, int rn, const Op2& op2, int size, int flag
 
             ABI_CallFunction((void (*)())func);
 
-            PopRegs(false);
+            PopRegs(false, false);
 
             if (!(flags & memop_Store))
             {
@@ -370,7 +370,7 @@ void Compiler::Comp_MemAccess(int rd, int rn, const Op2& op2, int size, int flag
                 }
             }
 
-            PopRegs(false);
+            PopRegs(false, false);
             
             if (!(flags & memop_Store))
             {
@@ -508,7 +508,7 @@ s32 Compiler::Comp_MemAccessBlock(int rn, BitSet16 regs, bool store, bool preinc
 
     if (!store)
     {
-        PushRegs(false);
+        PushRegs(false, false, !compileFastPath);
 
         MOV(32, R(ABI_PARAM1), R(RSCRATCH4));
         MOV(32, R(ABI_PARAM3), Imm32(regsCount));
@@ -529,7 +529,7 @@ s32 Compiler::Comp_MemAccessBlock(int rn, BitSet16 regs, bool store, bool preinc
         case 3: CALL((void*)&SlowBlockTransfer7<false, 1>); break;
         }
 
-        PopRegs(false);
+        PopRegs(false, false);
 
         if (allocOffset)
             ADD(64, R(RSP), Imm8(allocOffset));
@@ -606,7 +606,7 @@ s32 Compiler::Comp_MemAccessBlock(int rn, BitSet16 regs, bool store, bool preinc
         if (allocOffset)
             SUB(64, R(RSP), Imm8(allocOffset));
 
-        PushRegs(false);
+        PushRegs(false, false, !compileFastPath);
 
         MOV(32, R(ABI_PARAM1), R(RSCRATCH4));
         if (allocOffset)
@@ -628,7 +628,7 @@ s32 Compiler::Comp_MemAccessBlock(int rn, BitSet16 regs, bool store, bool preinc
 
         ADD(64, R(RSP), stackAlloc <= INT8_MAX ? Imm8(stackAlloc) : Imm32(stackAlloc));
     
-        PopRegs(false);
+        PopRegs(false, false);
     }
 
     if (compileFastPath)