JitArm64: Don't store immediate values in register cache

Like the previous commit did for Jit64, JitArm64 can now handle the
combination of a value simultaneously being in a host register and being
a known immediate.

Unlike with Jit64, I've put the codegen-affecting changes in this commit
and the move away from the RegType enum in a follow-up commit. This is
in part because the design of JitArm64 made it easy to implement the
codegen-affecting changes without combining it with a big bang
refactorization, and in part because we need to keep RegType around for
keeping track of different float formats in Arm64FPRCache, complicating
the refactorization a bit.

3 files changed, 63 insertions, 52 deletions

diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
index 975f5f95f7..37bfc4c24a 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
@@ -1362,7 +1362,10 @@ bool JitArm64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
         m_constant_propagation.Apply(constant_propagation_result);
 
         if (constant_propagation_result.gpr >= 0)
+        {
+          // Mark the GPR as dirty in the register cache
           gpr.SetImmediate(constant_propagation_result.gpr, constant_propagation_result.gpr_value);
+        }
 
         if (constant_propagation_result.instruction_fully_executed)
         {
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.cpp b/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.cpp
index d0c9ac71bd..161a1e01e3 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.cpp
@@ -115,7 +115,7 @@ void Arm64RegCache::FlushMostStaleRegister()
     const u32 last_used = reg.GetLastUsed();
 
     if (last_used > most_stale_amount && reg.GetType() != RegType::NotLoaded &&
-        reg.GetType() != RegType::Discarded && reg.GetType() != RegType::Immediate)
+        reg.GetType() != RegType::Discarded)
     {
       most_stale_preg = i;
       most_stale_amount = last_used;
@@ -145,6 +145,19 @@ void Arm64GPRCache::Start(PPCAnalyst::BlockRegStats& stats)
 {
 }
 
+// Returns if a register is set as an immediate. Only valid for guest GPRs.
+bool Arm64GPRCache::IsImm(size_t preg) const
+{
+  return m_jit->GetConstantPropagation().HasGPR(preg);
+}
+
+// Gets the immediate that a register is set to. Only valid for guest GPRs.
+u32 Arm64GPRCache::GetImm(size_t preg) const
+{
+  ASSERT(m_jit->GetConstantPropagation().HasGPR(preg));
+  return m_jit->GetConstantPropagation().GetGPR(preg);
+}
+
 bool Arm64GPRCache::IsCallerSaved(ARM64Reg reg) const
 {
   return ARM64XEmitter::CALLER_SAVED_GPRS[DecodeReg(reg)];
@@ -186,6 +199,7 @@ void Arm64GPRCache::FlushRegister(size_t index, FlushMode mode, ARM64Reg tmp_reg
   GuestRegInfo guest_reg = GetGuestByIndex(index);
   OpArg& reg = guest_reg.reg;
   size_t bitsize = guest_reg.bitsize;
+  const bool is_gpr = index >= GUEST_GPR_OFFSET && index < GUEST_GPR_OFFSET + GUEST_GPR_COUNT;
 
   if (reg.GetType() == RegType::Register)
   {
@@ -199,11 +213,12 @@ void Arm64GPRCache::FlushRegister(size_t index, FlushMode mode, ARM64Reg tmp_reg
       reg.Flush();
     }
   }
-  else if (reg.GetType() == RegType::Immediate)
+  else if (is_gpr && IsImm(index - GUEST_GPR_OFFSET))
   {
     if (reg.IsDirty())
     {
-      if (!reg.GetImm())
+      const u32 imm = GetImm(index - GUEST_GPR_OFFSET);
+      if (imm == 0)
       {
         m_emit->STR(IndexType::Unsigned, bitsize == 64 ? ARM64Reg::ZR : ARM64Reg::WZR, PPC_REG,
                     u32(guest_reg.ppc_offset));
@@ -225,7 +240,7 @@ void Arm64GPRCache::FlushRegister(size_t index, FlushMode mode, ARM64Reg tmp_reg
 
         const ARM64Reg encoded_tmp_reg = bitsize != 64 ? tmp_reg : EncodeRegTo64(tmp_reg);
 
-        m_emit->MOVI2R(encoded_tmp_reg, reg.GetImm());
+        m_emit->MOVI2R(encoded_tmp_reg, imm);
         m_emit->STR(IndexType::Unsigned, encoded_tmp_reg, PPC_REG, u32(guest_reg.ppc_offset));
 
         if (allocated_tmp_reg)
@@ -244,10 +259,10 @@ void Arm64GPRCache::FlushRegisters(BitSet32 regs, FlushMode mode, ARM64Reg tmp_r
   for (auto iter = regs.begin(); iter != regs.end(); ++iter)
   {
     const int i = *iter;
-
     ASSERT_MSG(DYNA_REC,
                ignore_discarded_registers != IgnoreDiscardedRegisters::No ||
-                   m_guest_registers[GUEST_GPR_OFFSET + i].GetType() != RegType::Discarded,
+                   m_guest_registers[GUEST_GPR_OFFSET + i].GetType() != RegType::Discarded ||
+                   IsImm(i),
                "Attempted to flush discarded register");
 
     if (i + 1 < int(GUEST_GPR_COUNT) && regs[i + 1])
@@ -255,10 +270,10 @@ void Arm64GPRCache::FlushRegisters(BitSet32 regs, FlushMode mode, ARM64Reg tmp_r
       // We've got two guest registers in a row to store
       OpArg& reg1 = m_guest_registers[GUEST_GPR_OFFSET + i];
       OpArg& reg2 = m_guest_registers[GUEST_GPR_OFFSET + i + 1];
-      const bool reg1_imm = reg1.GetType() == RegType::Immediate;
-      const bool reg2_imm = reg2.GetType() == RegType::Immediate;
-      const bool reg1_zero = reg1_imm && reg1.GetImm() == 0;
-      const bool reg2_zero = reg2_imm && reg2.GetImm() == 0;
+      const bool reg1_imm = IsImm(i);
+      const bool reg2_imm = IsImm(i + 1);
+      const bool reg1_zero = reg1_imm && GetImm(i) == 0;
+      const bool reg2_zero = reg2_imm && GetImm(i + 1) == 0;
       const bool flush_all = mode == FlushMode::All;
       if (reg1.IsDirty() && reg2.IsDirty() &&
           (reg1.GetType() == RegType::Register || (reg1_imm && (reg1_zero || flush_all))) &&
@@ -334,6 +349,7 @@ ARM64Reg Arm64GPRCache::BindForRead(size_t index)
   GuestRegInfo guest_reg = GetGuestByIndex(index);
   OpArg& reg = guest_reg.reg;
   size_t bitsize = guest_reg.bitsize;
+  const bool is_gpr = index >= GUEST_GPR_OFFSET && index < GUEST_GPR_OFFSET + GUEST_GPR_COUNT;
 
   IncrementAllUsed();
   reg.ResetLastUsed();
@@ -342,17 +358,15 @@ ARM64Reg Arm64GPRCache::BindForRead(size_t index)
   {
   case RegType::Register:  // already in a reg
     return reg.GetReg();
-  case RegType::Immediate:  // Is an immediate
+  case RegType::Discarded:  // Is an immediate or discarded
   {
+    ASSERT_MSG(DYNA_REC, is_gpr && IsImm(index - GUEST_GPR_OFFSET),
+               "Attempted to read discarded register");
     ARM64Reg host_reg = bitsize != 64 ? GetReg() : EncodeRegTo64(GetReg());
-    m_emit->MOVI2R(host_reg, reg.GetImm());
+    m_emit->MOVI2R(host_reg, GetImm(index - GUEST_GPR_OFFSET));
     reg.Load(host_reg);
     return host_reg;
   }
-  break;
-  case RegType::Discarded:
-    ASSERT_MSG(DYNA_REC, false, "Attempted to read discarded register");
-    break;
   case RegType::NotLoaded:  // Register isn't loaded at /all/
   {
     // This is a bit annoying. We try to keep these preloaded as much as possible
@@ -379,7 +393,7 @@ void Arm64GPRCache::SetImmediateInternal(size_t index, u32 imm, bool dirty)
   OpArg& reg = guest_reg.reg;
   if (reg.GetType() == RegType::Register)
     UnlockRegister(EncodeRegTo32(reg.GetReg()));
-  reg.LoadToImm(imm);
+  reg.Discard();
   reg.SetDirty(dirty);
   m_jit->GetConstantPropagation().SetGPR(index - GUEST_GPR_OFFSET, imm);
 }
@@ -394,28 +408,32 @@ void Arm64GPRCache::BindForWrite(size_t index, bool will_read, bool will_write)
   reg.ResetLastUsed();
 
   const RegType reg_type = reg.GetType();
-  if (reg_type == RegType::NotLoaded || reg_type == RegType::Discarded)
+  if (reg_type != RegType::Register)
   {
-    const ARM64Reg host_reg = bitsize != 64 ? GetReg() : EncodeRegTo64(GetReg());
-    reg.Load(host_reg);
-    reg.SetDirty(will_write);
-    if (will_read)
+    if (is_gpr && IsImm(index - GUEST_GPR_OFFSET))
     {
-      ASSERT_MSG(DYNA_REC, reg_type != RegType::Discarded, "Attempted to load a discarded value");
-      m_emit->LDR(IndexType::Unsigned, host_reg, PPC_REG, u32(guest_reg.ppc_offset));
+      const ARM64Reg host_reg = bitsize != 64 ? GetReg() : EncodeRegTo64(GetReg());
+      if (will_read || !will_write)
+      {
+        // TODO: Emitting this instruction when (!will_read && !will_write) would be unnecessary if
+        // we had some way to indicate to Flush that the immediate value should be written to
+        // ppcState even though there is a host register allocated
+        m_emit->MOVI2R(host_reg, GetImm(index - GUEST_GPR_OFFSET));
+      }
+      reg.Load(host_reg);
     }
-  }
-  else if (reg_type == RegType::Immediate)
-  {
-    const ARM64Reg host_reg = bitsize != 64 ? GetReg() : EncodeRegTo64(GetReg());
-    if (will_read || !will_write)
+    else
     {
-      // TODO: Emitting this instruction when (!will_read && !will_write) would be unnecessary if we
-      // had some way to indicate to Flush that the immediate value should be written to ppcState
-      // even though there is a host register allocated
-      m_emit->MOVI2R(host_reg, reg.GetImm());
+      const ARM64Reg host_reg = bitsize != 64 ? GetReg() : EncodeRegTo64(GetReg());
+      reg.Load(host_reg);
+      reg.SetDirty(will_write);
+      if (will_read)
+      {
+        ASSERT_MSG(DYNA_REC, reg_type != RegType::Discarded, "Attempted to load a discarded value");
+        m_emit->LDR(IndexType::Unsigned, host_reg, PPC_REG, u32(guest_reg.ppc_offset));
+      }
+      return;
     }
-    reg.Load(host_reg);
   }
 
   if (will_write)
@@ -521,7 +539,7 @@ void Arm64FPRCache::Flush(FlushMode mode, ARM64Reg tmp_reg,
       ASSERT_MSG(DYNA_REC, ignore_discarded_registers != IgnoreDiscardedRegisters::No,
                  "Attempted to flush discarded register");
     }
-    else if (reg_type != RegType::NotLoaded && reg_type != RegType::Immediate)
+    else if (reg_type != RegType::NotLoaded)
     {
       FlushRegister(i, mode, tmp_reg);
     }
@@ -785,7 +803,7 @@ void Arm64FPRCache::FlushByHost(ARM64Reg host_reg, ARM64Reg tmp_reg)
     const RegType reg_type = reg.GetType();
 
     if (reg_type != RegType::NotLoaded && reg_type != RegType::Discarded &&
-        reg_type != RegType::Immediate && reg.GetReg() == host_reg)
+        reg.GetReg() == host_reg)
     {
       FlushRegister(i, FlushMode::All, tmp_reg);
       return;
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.h b/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.h
index c973e3bc28..82c0d94142 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.h
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_RegCache.h
@@ -61,9 +61,8 @@ static_assert(PPCSTATE_OFF(xer_so_ov) < 4096, "STRB can't store xer_so_ov!");
 enum class RegType
 {
   NotLoaded,
-  Discarded,   // Reg is not loaded because we know it won't be read before the next write
+  Discarded,   // Reg is in ConstantPropagation, or isn't loaded at all
   Register,    // Reg type is register
-  Immediate,   // Reg is really a IMM
   LowerPair,   // Only the lower pair of a paired register
   Duplicated,  // The lower reg is the same as the upper one (physical upper doesn't actually have
                // the duplicated value)
@@ -94,24 +93,17 @@ public:
 
   RegType GetType() const { return m_type; }
   Arm64Gen::ARM64Reg GetReg() const { return m_reg; }
-  u32 GetImm() const { return m_value; }
   void Load(Arm64Gen::ARM64Reg reg, RegType type = RegType::Register)
   {
     m_type = type;
     m_reg = reg;
   }
-  void LoadToImm(u32 imm)
-  {
-    m_type = RegType::Immediate;
-    m_value = imm;
-
-    m_reg = Arm64Gen::ARM64Reg::INVALID_REG;
-  }
   void Discard()
   {
     // Invalidate any previous information
     m_type = RegType::Discarded;
     m_reg = Arm64Gen::ARM64Reg::INVALID_REG;
+    m_dirty = true;
 
     // Arbitrarily large value that won't roll over on a lot of increments
     m_last_used = 0xFFFF;
@@ -121,6 +113,7 @@ public:
     // Invalidate any previous information
     m_type = RegType::NotLoaded;
     m_reg = Arm64Gen::ARM64Reg::INVALID_REG;
+    m_dirty = false;
 
     // Arbitrarily large value that won't roll over on a lot of increments
     m_last_used = 0xFFFF;
@@ -137,9 +130,6 @@ private:
   RegType m_type = RegType::NotLoaded;                         // store type
   Arm64Gen::ARM64Reg m_reg = Arm64Gen::ARM64Reg::INVALID_REG;  // host register we are in
 
-  // For REG_IMM
-  u32 m_value = 0;  // IMM value
-
   u32 m_last_used = 0;
 
   bool m_dirty = false;
@@ -339,11 +329,11 @@ public:
     SetImmediateInternal(GUEST_GPR_OFFSET + preg, imm, dirty);
   }
 
-  // Returns if a register is set as an immediate. Only valid for guest GPRs.
-  bool IsImm(size_t preg) const { return GetGuestGPROpArg(preg).GetType() == RegType::Immediate; }
+  // Returns whether a register is set as an immediate. Only valid for guest GPRs.
+  bool IsImm(size_t preg) const;
 
   // Gets the immediate that a register is set to. Only valid for guest GPRs.
-  u32 GetImm(size_t preg) const { return GetGuestGPROpArg(preg).GetImm(); }
+  u32 GetImm(size_t preg) const;
 
   bool IsImm(size_t preg, u32 imm) const { return IsImm(preg) && GetImm(preg) == imm; }