From 2333fc270149a091cd5fc1f8a1495c8ac32b7dad Mon Sep 17 00:00:00 2001
From: Rairii <2650838+Wack0@users.noreply.github.com>
Date: Fri, 10 Nov 2023 12:51:51 +0000
Subject: [PATCH] MMU: Use VSID in segment register as additional TLB lookup
 key

---
 Source/Core/Core/PowerPC/MMU.cpp   | 21 +++++++++++----------
 Source/Core/Core/PowerPC/PowerPC.h |  1 +
 Source/Core/Core/State.cpp         |  2 +-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/Source/Core/Core/PowerPC/MMU.cpp b/Source/Core/Core/PowerPC/MMU.cpp
index aef3dd714e..9df7d2edc4 100644
--- a/Source/Core/Core/PowerPC/MMU.cpp
+++ b/Source/Core/Core/PowerPC/MMU.cpp
@@ -1330,13 +1330,13 @@ enum class TLBLookupResult
 };
 
 static TLBLookupResult LookupTLBPageAddress(PowerPC::PowerPCState& ppc_state,
-                                            const XCheckTLBFlag flag, const u32 vpa, u32* paddr,
-                                            bool* wi)
+                                            const XCheckTLBFlag flag, const u32 vpa, const u32 vsid,
+                                            u32* paddr, bool* wi)
 {
   const u32 tag = vpa >> HW_PAGE_INDEX_SHIFT;
   TLBEntry& tlbe = ppc_state.tlb[IsOpcodeFlag(flag)][tag & HW_PAGE_INDEX_MASK];
 
-  if (tlbe.tag[0] == tag)
+  if (tlbe.tag[0] == tag && tlbe.vsid[0] == vsid)
   {
     UPTE_Hi pte2(tlbe.pte[0]);
 
@@ -1359,7 +1359,7 @@ static TLBLookupResult LookupTLBPageAddress(PowerPC::PowerPCState& ppc_state,
 
     return TLBLookupResult::Found;
   }
-  if (tlbe.tag[1] == tag)
+  if (tlbe.tag[1] == tag && tlbe.vsid[1] == vsid)
   {
     UPTE_Hi pte2(tlbe.pte[1]);
 
@@ -1386,7 +1386,7 @@ static TLBLookupResult LookupTLBPageAddress(PowerPC::PowerPCState& ppc_state,
 }
 
 static void UpdateTLBEntry(PowerPC::PowerPCState& ppc_state, const XCheckTLBFlag flag, UPTE_Hi pte2,
-                           const u32 address)
+                           const u32 address, const u32 vsid)
 {
   if (IsNoExceptionFlag(flag))
     return;
@@ -1398,6 +1398,7 @@ static void UpdateTLBEntry(PowerPC::PowerPCState& ppc_state, const XCheckTLBFlag
   tlbe.paddr[index] = pte2.RPN << HW_PAGE_INDEX_SHIFT;
   tlbe.pte[index] = pte2.Hex;
   tlbe.tag[index] = tag;
+  tlbe.vsid[index] = vsid;
 }
 
 void MMU::InvalidateTLBEntry(u32 address)
@@ -1412,20 +1413,21 @@ void MMU::InvalidateTLBEntry(u32 address)
 template <const XCheckTLBFlag flag>
 MMU::TranslateAddressResult MMU::TranslatePageAddress(const EffectiveAddress address, bool* wi)
 {
+  const auto sr = UReg_SR{m_ppc_state.sr[address.SR]};
+  const u32 VSID = sr.VSID;  // 24 bit
+
   // TLB cache
   // This catches 99%+ of lookups in practice, so the actual page table entry code below doesn't
   // benefit much from optimization.
   u32 translated_address = 0;
   const TLBLookupResult res =
-      LookupTLBPageAddress(m_ppc_state, flag, address.Hex, &translated_address, wi);
+      LookupTLBPageAddress(m_ppc_state, flag, address.Hex, VSID, &translated_address, wi);
   if (res == TLBLookupResult::Found)
   {
     return TranslateAddressResult{TranslateAddressResultEnum::PAGE_TABLE_TRANSLATED,
                                   translated_address};
   }
 
-  const auto sr = UReg_SR{m_ppc_state.sr[address.SR]};
-
   if (sr.T != 0)
     return TranslateAddressResult{TranslateAddressResultEnum::DIRECT_STORE_SEGMENT, 0};
 
@@ -1439,7 +1441,6 @@ MMU::TranslateAddressResult MMU::TranslatePageAddress(const EffectiveAddress add
 
   const u32 offset = address.offset;          // 12 bit
   const u32 page_index = address.page_index;  // 16 bit
-  const u32 VSID = sr.VSID;                   // 24 bit
   const u32 api = address.API;                //  6 bit (part of page_index)
 
   // hash function no 1 "xor" .360
@@ -1494,7 +1495,7 @@ MMU::TranslateAddressResult MMU::TranslatePageAddress(const EffectiveAddress add
 
         // We already updated the TLB entry if this was caused by a C bit.
         if (res != TLBLookupResult::UpdateC)
-          UpdateTLBEntry(m_ppc_state, flag, pte2, address.Hex);
+          UpdateTLBEntry(m_ppc_state, flag, pte2, address.Hex, VSID);
 
         *wi = (pte2.WIMG & 0b1100) != 0;
 
diff --git a/Source/Core/Core/PowerPC/PowerPC.h b/Source/Core/Core/PowerPC/PowerPC.h
index 7612018a93..211d0f37e5 100644
--- a/Source/Core/Core/PowerPC/PowerPC.h
+++ b/Source/Core/Core/PowerPC/PowerPC.h
@@ -62,6 +62,7 @@ struct TLBEntry
 
   WayArray tag{INVALID_TAG, INVALID_TAG};
   WayArray paddr{};
+  WayArray vsid{};
   WayArray pte{};
   u32 recent = 0;
 
diff --git a/Source/Core/Core/State.cpp b/Source/Core/Core/State.cpp
index db855923e6..c8523128dd 100644
--- a/Source/Core/Core/State.cpp
+++ b/Source/Core/Core/State.cpp
@@ -93,7 +93,7 @@ static size_t s_state_writes_in_queue;
 static std::condition_variable s_state_write_queue_is_empty;
 
 // Don't forget to increase this after doing changes on the savestate system
-constexpr u32 STATE_VERSION = 163;  // Last changed in PR 12217
+constexpr u32 STATE_VERSION = 164;  // Last changed in PR 12282
 
 // Increase this if the StateExtendedHeader definition changes
 constexpr u32 EXTENDED_HEADER_VERSION = 1;  // Last changed in PR 12217