Jit: Refactor m_back_patch_info/m_fault_to_handler's maps to store a vector per block

(Tested NBA Live 2005)
Profiling codegen performance showed that one of the slower functions was inserting elements to this map.
By inserting instead a vector of all the entries of a block, this disappears from the profiling, and the improvement is noticeable (tested the match intro sequence, though exact numbers are hard to gather).

On x86 (ARM was already doing it), now erasing a block will also delete the corresponding entries from the map.

7 files changed, 151 insertions, 81 deletions

diff --git a/Source/Core/Core/PowerPC/Jit64/Jit.cpp b/Source/Core/Core/PowerPC/Jit64/Jit.cpp
index 91085b6199..d214668783 100644
--- a/Source/Core/Core/PowerPC/Jit64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/Jit64/Jit.cpp
@@ -170,14 +170,24 @@ bool Jit64::BackPatch(SContext* ctx)
   if (!IsInSpace(codePtr))
     return false;  // this will become a regular crash real soon after this
 
-  auto it = m_back_patch_info.find(codePtr);
-  if (it == m_back_patch_info.end())
+  // Find the block
+  auto block_it = m_back_patch_info.lower_bound(codePtr);
+  if (block_it == m_back_patch_info.end())
   {
-    PanicAlertFmt("BackPatch: no register use entry for address {}", fmt::ptr(codePtr));
+    PanicAlertFmt("BackPatch: no block returned for address {}", fmt::ptr(codePtr));
+    return false;
+  }
+  // Then find the entry
+  auto entry_it = std::ranges::find_if(block_it->second, [codePtr](auto& entry) {
+    return codePtr < entry.start + entry.len && codePtr >= entry.start;
+  });
+  if (entry_it == block_it->second.end())
+  {
+    PanicAlertFmt("BackPatch: entry for address {} not found in block", fmt::ptr(codePtr));
     return false;
   }
 
-  TrampolineInfo& info = it->second;
+  TrampolineInfo& entry = *entry_it;
 
   // In the trampoline code, we jump back into the block at the beginning
   // of the next instruction. The next instruction comes immediately
@@ -187,21 +197,19 @@ bool Jit64::BackPatch(SContext* ctx)
   // to insert the backpatch jump.)
 
   js.generatingTrampoline = true;
-  js.trampolineExceptionHandler = info.exception_handler_at_loc;
-  js.compilerPC = info.pc;
+  js.trampolineExceptionHandler = entry.exception_handler_at_loc;
+  js.compilerPC = entry.pc;
 
   // Generate the trampoline.
-  const u8* trampoline = trampolines.GenerateTrampoline(info);
+  const u8* trampoline = trampolines.GenerateTrampoline(entry);
   js.generatingTrampoline = false;
   js.trampolineExceptionHandler = nullptr;
 
-  u8* start = info.start;
-
   // Patch the original memory operation.
-  XEmitter emitter(start, start + info.len);
+  XEmitter emitter(entry.start, entry.start + entry.len);
   emitter.JMP(trampoline);
   // NOPs become dead code
-  const u8* end = info.start + info.len;
+  const u8* end = entry.start + entry.len;
   for (const u8* i = emitter.GetCodePtr(); i < end; ++i)
     emitter.INT3();
 
@@ -209,10 +217,10 @@ bool Jit64::BackPatch(SContext* ctx)
   // before faulting (eg: the store+swap was not an atomic op like MOVBE), let's
   // swap it back so that the swap can happen again (this double swap isn't ideal but
   // only happens the first time we fault).
-  if (info.nonAtomicSwapStoreSrc != Gen::INVALID_REG)
+  if (entry.nonAtomicSwapStoreSrc != Gen::INVALID_REG)
   {
-    u64* ptr = ContextRN(ctx, info.nonAtomicSwapStoreSrc);
-    switch (info.accessSize << 3)
+    u64* ptr = ContextRN(ctx, entry.nonAtomicSwapStoreSrc);
+    switch (entry.accessSize << 3)
     {
     case 8:
       // No need to swap a byte
@@ -234,12 +242,17 @@ bool Jit64::BackPatch(SContext* ctx)
 
   // This is special code to undo the LEA in SafeLoadToReg if it clobbered the address
   // register in the case where reg_value shared the same location as opAddress.
-  if (info.offsetAddedToAddress)
+  if (entry.offsetAddedToAddress)
   {
-    u64* ptr = ContextRN(ctx, info.op_arg.GetSimpleReg());
-    *ptr = static_cast<u32>(*ptr - info.offset);
+    u64* ptr = ContextRN(ctx, entry.op_arg.GetSimpleReg());
+    *ptr = static_cast<u32>(*ptr - entry.offset);
   }
 
+  // Delete the entry
+  block_it->second.erase(entry_it);
+  if (block_it->second.empty())
+    m_back_patch_info.erase(block_it);
+
   ctx->CTX_PC = reinterpret_cast<u64>(trampoline);
 
   return true;
@@ -255,7 +268,7 @@ void Jit64::Init()
 
   jo.optimizeGatherPipe = true;
   jo.accurateSinglePrecision = true;
-  js.fastmemLoadStore = nullptr;
+  js.fastmemLoadStore = false;
   js.compilerPC = 0;
 
   gpr.SetEmitter(this);
@@ -313,7 +326,12 @@ void Jit64::FreeRanges()
   // Check if any code blocks have been freed in the block cache and transfer this information to
   // the local rangesets to allow overwriting them with new code.
   for (const auto& [from, to] : blocks.GetRangesToFreeNear())
+  {
+    // This erase isn't just a matter of memory usage: because lookups first use lower_bound to get
+    // the entries of the currently running block, a destroyed block could interfere if left here.
+    m_back_patch_info.erase(from);
     m_free_ranges_near.insert(from, to);
+  }
   for (const auto& [from, to] : blocks.GetRangesToFreeFar())
     m_free_ranges_far.insert(from, to);
   blocks.ClearRangesToFree();
@@ -857,6 +875,13 @@ void Jit64::Jit(u32 em_address, bool clear_cache_and_retry_on_failure)
       b->far_begin = far_start;
       b->far_end = far_end;
 
+      if (!m_back_patch_info_temp.empty())
+      {
+        // The copy from temp is intentional, to minimize allocations due to exceeded vector
+        // capacity.
+        m_back_patch_info.emplace(near_start, m_back_patch_info_temp);
+      }
+
       blocks.FinalizeBlock(*b, jo.enableBlocklink, code_block, m_code_buffer);
 
 #ifdef JIT_LOG_GENERATED_CODE
@@ -914,6 +939,7 @@ bool Jit64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
   js.curBlock = b;
   js.numLoadStoreInst = 0;
   js.numFloatingPointInst = 0;
+  m_back_patch_info_temp.clear();
 
   // TODO: Test if this or AlignCode16 make a difference from GetCodePtr
   b->normalEntry = AlignCode4();
@@ -995,7 +1021,7 @@ bool Jit64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
     js.instructionsLeft = (code_block.m_num_instructions - 1) - i;
     const GekkoOPInfo* opinfo = op.opinfo;
     js.downcountAmount += opinfo->num_cycles;
-    js.fastmemLoadStore = nullptr;
+    js.fastmemLoadStore = false;
     js.fixupExceptionHandler = false;
 
     if (i == (code_block.m_num_instructions - 1))
@@ -1199,7 +1225,7 @@ bool Jit64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
         }
         else
         {
-          m_back_patch_info.at(js.fastmemLoadStore).exception_handler_at_loc = GetWritableCodePtr();
+          m_back_patch_info_temp.back().exception_handler_at_loc = GetWritableCodePtr();
         }
 
         RCForkGuard gpr_guard = gpr.Fork();
diff --git a/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.cpp b/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.cpp
index d99ec2b155..a8a9300ffd 100644
--- a/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.cpp
+++ b/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.cpp
@@ -336,27 +336,31 @@ void EmuCodeBlock::SafeLoadToReg(X64Reg reg_value, const Gen::OpArg& opAddress,
     MovInfo mov;
     bool offsetAddedToAddress =
         UnsafeLoadToReg(reg_value, opAddress, accessSize, offset, signExtend, &mov);
-    TrampolineInfo& info = m_back_patch_info[mov.address];
-    info.pc = js.compilerPC;
-    info.nonAtomicSwapStoreSrc = mov.nonAtomicSwapStore ? mov.nonAtomicSwapStoreSrc : INVALID_REG;
-    info.start = backpatchStart;
-    info.read = true;
-    info.op_reg = reg_value;
-    info.op_arg = opAddress;
-    info.offsetAddedToAddress = offsetAddedToAddress;
-    info.accessSize = accessSize >> 3;
-    info.offset = offset;
-    info.registersInUse = registersInUse;
-    info.flags = flags;
-    info.signExtend = signExtend;
+
     ptrdiff_t padding = BACKPATCH_SIZE - (GetCodePtr() - backpatchStart);
     if (padding > 0)
     {
       NOP(padding);
     }
-    info.len = static_cast<u16>(GetCodePtr() - info.start);
 
-    js.fastmemLoadStore = mov.address;
+    TrampolineInfo& info = m_back_patch_info_temp.emplace_back();
+    info = {
+        .start = backpatchStart,
+        .len = static_cast<u16>(GetCodePtr() - backpatchStart),
+        .flags = static_cast<u8>(flags),
+        .accessSize = static_cast<u8>(accessSize >> 3),
+        .read = true,
+        .signExtend = signExtend,
+        .offsetAddedToAddress = offsetAddedToAddress,
+        .pc = js.compilerPC,
+        .registersInUse = registersInUse,
+        .nonAtomicSwapStoreSrc = mov.nonAtomicSwapStore ? mov.nonAtomicSwapStoreSrc : INVALID_REG,
+        .offset = offset,
+        .op_reg = reg_value,
+        .op_arg = opAddress,
+    };
+
+    js.fastmemLoadStore = true;
     return;
   }
 
@@ -509,26 +513,30 @@ void EmuCodeBlock::SafeWriteRegToReg(OpArg reg_value, X64Reg reg_addr, int acces
     u8* backpatchStart = GetWritableCodePtr();
     MovInfo mov;
     UnsafeWriteRegToReg(reg_value, reg_addr, accessSize, offset, swap, &mov);
-    TrampolineInfo& info = m_back_patch_info[mov.address];
-    info.pc = js.compilerPC;
-    info.nonAtomicSwapStoreSrc = mov.nonAtomicSwapStore ? mov.nonAtomicSwapStoreSrc : INVALID_REG;
-    info.start = backpatchStart;
-    info.read = false;
-    info.op_arg = reg_value;
-    info.op_reg = reg_addr;
-    info.offsetAddedToAddress = false;
-    info.accessSize = accessSize >> 3;
-    info.offset = offset;
-    info.registersInUse = registersInUse;
-    info.flags = flags;
-    ptrdiff_t padding = BACKPATCH_SIZE - (GetCodePtr() - backpatchStart);
+
+    const ptrdiff_t padding = BACKPATCH_SIZE - (GetCodePtr() - backpatchStart);
     if (padding > 0)
     {
       NOP(padding);
     }
-    info.len = static_cast<u16>(GetCodePtr() - info.start);
 
-    js.fastmemLoadStore = mov.address;
+    TrampolineInfo& info = m_back_patch_info_temp.emplace_back();
+    info = {
+        .start = backpatchStart,
+        .len = static_cast<u16>(GetCodePtr() - backpatchStart),
+        .flags = static_cast<u8>(flags),
+        .accessSize = static_cast<u8>(accessSize >> 3),
+        .read = false,
+        .offsetAddedToAddress = false,
+        .pc = js.compilerPC,
+        .registersInUse = registersInUse,
+        .nonAtomicSwapStoreSrc = mov.nonAtomicSwapStore ? mov.nonAtomicSwapStoreSrc : INVALID_REG,
+        .offset = offset,
+        .op_reg = reg_addr,
+        .op_arg = reg_value,
+    };
+
+    js.fastmemLoadStore = true;
 
     return;
   }
diff --git a/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.h b/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.h
index b53b5da67b..12805de98f 100644
--- a/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.h
+++ b/Source/Core/Core/PowerPC/Jit64Common/EmuCodeBlock.h
@@ -3,7 +3,8 @@
 
 #pragma once
 
-#include <unordered_map>
+#include <map>
+#include <vector>
 
 #include "Common/BitSet.h"
 #include "Common/CommonTypes.h"
@@ -140,6 +141,16 @@ protected:
   u8* m_near_code_end = nullptr;
   bool m_near_code_write_failed = false;
 
-  std::unordered_map<u8*, TrampolineInfo> m_back_patch_info;
-  std::unordered_map<u8*, u8*> m_exception_handler_at_loc;
+  // The keys store the entry point of a block. Lookups find the currently running block through
+  // lower_bound (which is why the comparison is std::greater), and then linearly search the vector
+  // for the specific address.
+  //
+  // Because the total number of entries is in the 100s of thousands, and std::map is inefficient,
+  // this improves performance mostly by allowing to batch insert all the info for each block. (Even
+  // a better ordered map like b-trees would probably benefit.)
+  //
+  // Note that this would have had to be an ordered structure even without this optimization (to
+  // erase the entries when a block is destroyed).
+  std::map<const u8*, std::vector<TrampolineInfo>, std::greater<const u8*>> m_back_patch_info;
+  std::vector<TrampolineInfo> m_back_patch_info_temp;
 };
diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
index aba973346d..7c5d528af2 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
@@ -207,12 +207,9 @@ void JitArm64::FreeRanges()
   // the local rangesets to allow overwriting them with new code.
   for (const auto& [from, to] : blocks.GetRangesToFreeNear())
   {
-    const auto first_fastmem_area = m_fault_to_handler.upper_bound(from);
-    auto last_fastmem_area = first_fastmem_area;
-    const auto end = m_fault_to_handler.end();
-    while (last_fastmem_area != end && last_fastmem_area->first <= to)
-      ++last_fastmem_area;
-    m_fault_to_handler.erase(first_fastmem_area, last_fastmem_area);
+    // This erase isn't just a matter of memory usage: because lookups first use lower_bound to get
+    // the entries of the currently running block, a destroyed block could interfere if left here.
+    m_fault_to_handler.erase(from);
 
     if (from < m_near_code_0.GetCodeEnd())
       m_free_ranges_near_0.insert(from, to);
@@ -1057,6 +1054,13 @@ void JitArm64::Jit(u32 em_address, bool clear_cache_and_retry_on_failure)
       b->far_begin = far_start;
       b->far_end = far_end;
 
+      if (!m_fault_to_handler_temp.empty())
+      {
+        // The copy from temp is intentional, to minimize allocations due to exceeded vector
+        // capacity.
+        m_fault_to_handler.emplace(near_start, m_fault_to_handler_temp);
+      }
+
       blocks.FinalizeBlock(*b, jo.enableBlocklink, code_block, m_code_buffer);
 
 #ifdef JIT_LOG_GENERATED_CODE
@@ -1169,6 +1173,7 @@ bool JitArm64::DoJit(u32 em_address, JitBlock* b, u32 nextPC)
   js.carryFlag = CarryFlag::InPPCState;
   js.numLoadStoreInst = 0;
   js.numFloatingPointInst = 0;
+  m_fault_to_handler_temp.clear();
 
   b->normalEntry = GetWritableCodePtr();
 
diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.h b/Source/Core/Core/PowerPC/JitArm64/Jit.h
index 9c23112e1c..4c1146e541 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.h
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.h
@@ -6,6 +6,7 @@
 #include <cstddef>
 #include <map>
 #include <optional>
+#include <vector>
 
 #include "Common/Arm64Emitter.h"
 #include "Common/RangeSizeSet.h"
@@ -193,7 +194,8 @@ public:
 protected:
   struct FastmemArea
   {
-    const u8* fast_access_code;
+    u8* fast_access_start;
+    u8* fast_access_end;
     const u8* slow_access_code;
   };
 
@@ -389,8 +391,18 @@ protected:
   void SetFPRFIfNeeded(bool single, Arm64Gen::ARM64Reg reg);
   void Force25BitPrecision(Arm64Gen::ARM64Reg output, Arm64Gen::ARM64Reg input);
 
-  // <Fast path fault location, slow path handler location>
-  std::map<const u8*, FastmemArea> m_fault_to_handler{};
+  // The keys store the entry point of a block. Lookups find the currently running block through
+  // lower_bound (which is why the comparison is std::greater), and then linearly search the vector
+  // for the specific address.
+  //
+  // Because the total number of entries is in the 100s of thousands, and std::map is inefficient,
+  // this improves performance mostly by allowing to batch insert all the info for each block. (Even
+  // a better ordered map like b-trees would probably benefit.)
+  //
+  // Note that this would have had to be an ordered structure even without this optimization (to
+  // erase the entries when a block is destroyed).
+  std::map<const u8*, std::vector<FastmemArea>, std::greater<const u8*>> m_fault_to_handler;
+  std::vector<FastmemArea> m_fault_to_handler_temp;
   Arm64GPRCache gpr;
   Arm64FPRCache fpr;
 
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp b/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
index 8368a7cc91..5d4c308879 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
@@ -67,7 +67,7 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, MemAccessMode mode, ARM64Reg RS,
   const bool emit_slow_access = mode != MemAccessMode::AlwaysFastAccess;
 
   bool in_far_code = false;
-  const u8* fast_access_start = GetCodePtr();
+  u8* fast_access_start = GetWritableCodePtr();
   std::optional<FixupBranch> slow_access_fixup;
 
   if (emit_fast_access)
@@ -148,7 +148,7 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, MemAccessMode mode, ARM64Reg RS,
       ByteswapAfterLoad(this, &m_float_emit, RS, RS, flags, true, false);
     }
   }
-  const u8* fast_access_end = GetCodePtr();
+  u8* fast_access_end = GetWritableCodePtr();
 
   if (emit_slow_access)
   {
@@ -161,9 +161,12 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, MemAccessMode mode, ARM64Reg RS,
 
       if (jo.fastmem && !emitting_routine)
       {
-        FastmemArea* fastmem_area = &m_fault_to_handler[fast_access_end];
-        fastmem_area->fast_access_code = fast_access_start;
-        fastmem_area->slow_access_code = GetCodePtr();
+        FastmemArea& fastmem_area = m_fault_to_handler_temp.emplace_back();
+        fastmem_area = {
+            .fast_access_start = fast_access_start,
+            .fast_access_end = fast_access_end,
+            .slow_access_code = GetCodePtr(),
+        };
       }
     }
 
@@ -330,28 +333,34 @@ void JitArm64::FlushPPCStateBeforeSlowAccess(ARM64Reg temp_gpr, ARM64Reg temp_fp
 bool JitArm64::HandleFastmemFault(SContext* ctx)
 {
   const u8* pc = reinterpret_cast<const u8*>(ctx->CTX_PC);
-  auto slow_handler_iter = m_fault_to_handler.upper_bound(pc);
 
-  // no fastmem area found
-  if (slow_handler_iter == m_fault_to_handler.end())
+  // Find the block
+  auto block_it = m_fault_to_handler.lower_bound(pc);
+  if (block_it == m_fault_to_handler.end())
     return false;
 
-  const u8* fastmem_area_start = slow_handler_iter->second.fast_access_code;
-  const u8* fastmem_area_end = slow_handler_iter->first;
-
-  // no overlapping fastmem area found
-  if (pc < fastmem_area_start)
+  // Then find the entry
+  auto entry_it = std::ranges::find_if(block_it->second, [pc](auto& entry) {
+    return pc < entry.fast_access_end && pc >= entry.fast_access_start;
+  });
+  if (entry_it == block_it->second.end())
     return false;
 
+  u8* fastmem_area_start = entry_it->fast_access_start;
+  u8* fastmem_area_end = entry_it->fast_access_end;
+
   const Common::ScopedJITPageWriteAndNoExecute enable_jit_page_writes;
-  ARM64XEmitter emitter(const_cast<u8*>(fastmem_area_start), const_cast<u8*>(fastmem_area_end));
+  ARM64XEmitter emitter(fastmem_area_start, fastmem_area_end);
 
-  emitter.BL(slow_handler_iter->second.slow_access_code);
+  emitter.BL(entry_it->slow_access_code);
 
   while (emitter.GetCodePtr() < fastmem_area_end)
     emitter.NOP();
 
-  m_fault_to_handler.erase(slow_handler_iter);
+  // Delete the entry
+  block_it->second.erase(entry_it);
+  if (block_it->second.empty())
+    m_fault_to_handler.erase(block_it);
 
   emitter.FlushIcache();
 
diff --git a/Source/Core/Core/PowerPC/JitCommon/JitBase.h b/Source/Core/Core/PowerPC/JitCommon/JitBase.h
index 612161b492..3f16d5b00e 100644
--- a/Source/Core/Core/PowerPC/JitCommon/JitBase.h
+++ b/Source/Core/Core/PowerPC/JitCommon/JitBase.h
@@ -6,7 +6,6 @@
 #include <array>
 #include <cstddef>
 #include <iosfwd>
-#include <map>
 #include <string_view>
 #include <unordered_set>
 #include <utility>
@@ -99,8 +98,8 @@ protected:
     u32 downcountAmount;
     u32 numLoadStoreInst;
     u32 numFloatingPointInst;
-    // If this is set, we need to generate an exception handler for the fastmem load.
-    u8* fastmemLoadStore;
+    // If this is true, we need to generate an exception handler for the fastmem load.
+    bool fastmemLoadStore;
     // If this is set, a load or store already prepared a jump to the exception handler for us,
     // so just fixup that branch instead of testing for a DSI again.
     bool fixupExceptionHandler;