From 9192128c50dfd2a904ec84aeb0182e481abc8140 Mon Sep 17 00:00:00 2001
From: JosJuice <josjuice@gmail.com>
Date: Sat, 19 Aug 2023 19:45:35 +0200
Subject: [PATCH] Jit: Handle imm msr in EmitStoreMembase

This makes Jit64::mtmsr more optimized in particular.
---
 Source/Core/Core/PowerPC/Jit64/Jit.cpp           | 16 ++++++++++++----
 .../Core/PowerPC/Jit64/Jit_SystemRegisters.cpp   |  2 +-
 Source/Core/Core/PowerPC/JitArm64/Jit.cpp        | 10 ++++++++++
 Source/Core/Core/PowerPC/JitArm64/Jit.h          |  1 +
 .../JitArm64/JitArm64_SystemRegisters.cpp        |  7 ++++++-
 5 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/Source/Core/Core/PowerPC/Jit64/Jit.cpp b/Source/Core/Core/PowerPC/Jit64/Jit.cpp
index 5ef45f76d0..d2991be698 100644
--- a/Source/Core/Core/PowerPC/Jit64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/Jit64/Jit.cpp
@@ -498,10 +498,18 @@ void Jit64::EmitUpdateMembase()
 void Jit64::EmitStoreMembase(const OpArg& msr, X64Reg scratch_reg)
 {
   auto& memory = m_system.GetMemory();
-  MOV(64, R(RMEM), ImmPtr(memory.GetLogicalBase()));
-  MOV(64, R(scratch_reg), ImmPtr(memory.GetPhysicalBase()));
-  TEST(32, msr, Imm32(1 << (31 - 27)));
-  CMOVcc(64, RMEM, R(scratch_reg), CC_Z);
+  if (msr.IsImm())
+  {
+    MOV(64, R(RMEM),
+        ImmPtr(UReg_MSR(msr.Imm32()).DR ? memory.GetLogicalBase() : memory.GetPhysicalBase()));
+  }
+  else
+  {
+    MOV(64, R(RMEM), ImmPtr(memory.GetLogicalBase()));
+    MOV(64, R(scratch_reg), ImmPtr(memory.GetPhysicalBase()));
+    TEST(32, msr, Imm32(1 << (31 - 27)));
+    CMOVcc(64, RMEM, R(scratch_reg), CC_Z);
+  }
   MOV(64, PPCSTATE(mem_ptr), R(RMEM));
 }
 
diff --git a/Source/Core/Core/PowerPC/Jit64/Jit_SystemRegisters.cpp b/Source/Core/Core/PowerPC/Jit64/Jit_SystemRegisters.cpp
index 5992fe9f28..86b632b3b7 100644
--- a/Source/Core/Core/PowerPC/Jit64/Jit_SystemRegisters.cpp
+++ b/Source/Core/Core/PowerPC/Jit64/Jit_SystemRegisters.cpp
@@ -439,7 +439,7 @@ void Jit64::mtmsr(UGeckoInstruction inst)
     RegCache::Realize(Rs);
     MOV(32, PPCSTATE(msr), Rs);
 
-    EmitStoreMembase(PPCSTATE(msr), RSCRATCH2);
+    EmitStoreMembase(Rs, RSCRATCH2);
   }
 
   gpr.Flush();
diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
index caf40987e4..1d4b07e170 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
@@ -361,6 +361,16 @@ void JitArm64::EmitUpdateMembase()
   LDR(IndexType::Unsigned, MEM_REG, PPC_REG, PPCSTATE_OFF(mem_ptr));
 }
 
+void JitArm64::EmitStoreMembase(u32 msr)
+{
+  auto& memory = m_system.GetMemory();
+  MOVP2R(MEM_REG,
+         UReg_MSR(msr).DR ?
+             (jo.fastmem_arena ? memory.GetLogicalBase() : memory.GetLogicalPageMappingsBase()) :
+             (jo.fastmem_arena ? memory.GetPhysicalBase() : memory.GetPhysicalPageMappingsBase()));
+  STR(IndexType::Unsigned, MEM_REG, PPC_REG, PPCSTATE_OFF(mem_ptr));
+}
+
 void JitArm64::EmitStoreMembase(const ARM64Reg& msr)
 {
   auto& memory = m_system.GetMemory();
diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.h b/Source/Core/Core/PowerPC/JitArm64/Jit.h
index df71e70a93..f97f6c509d 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.h
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.h
@@ -312,6 +312,7 @@ protected:
   void EndTimeProfile(JitBlock* b);
 
   void EmitUpdateMembase();
+  void EmitStoreMembase(u32 msr);
   void EmitStoreMembase(const Arm64Gen::ARM64Reg& msr);
 
   // Exits
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_SystemRegisters.cpp b/Source/Core/Core/PowerPC/JitArm64/JitArm64_SystemRegisters.cpp
index 3df2fea5fa..7205cde681 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_SystemRegisters.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_SystemRegisters.cpp
@@ -91,10 +91,15 @@ void JitArm64::mtmsr(UGeckoInstruction inst)
   JITDISABLE(bJITSystemRegistersOff);
   FALLBACK_IF(jo.fp_exceptions);
 
+  const bool imm_value = gpr.IsImm(inst.RS);
+  if (imm_value)
+    EmitStoreMembase(gpr.GetImm(inst.RS));
+
   gpr.BindToRegister(inst.RS, true);
   STR(IndexType::Unsigned, gpr.R(inst.RS), PPC_REG, PPCSTATE_OFF(msr));
 
-  EmitStoreMembase(gpr.R(inst.RS));
+  if (!imm_value)
+    EmitStoreMembase(gpr.R(inst.RS));
 
   gpr.Flush(FlushMode::All, ARM64Reg::INVALID_REG);
   fpr.Flush(FlushMode::All, ARM64Reg::INVALID_REG);