SaveState: Keep JIT blocks in module code a load doesn't change

Loading a state started the JIT over, so every block the game runs had to
be compiled again, which stutters. Now the loaded modules' code is
compared with the state's as RAM is loaded (a chunk that already matches
isn't copied either), and blocks entirely inside unchanged code stay.
Everything else is invalidated as before, and without a JIT, or with no
unchanged code, the JIT still starts over.

The JIT throws its blocks away if the loaded state needs code they weren't
compiled for (rounding mode switches, or no default prefixes at block
start). The other way round it keeps them, along with the flags they were
compiled under: before, a loaded hasSetRounding of false, or the reset of
lastSetRounding, had the next compile rebuild everything anyway.

In the three seconds after loading a state, Ridge Racer compiles 28 blocks
instead of 4823, Force Unleashed 31 instead of 7727. Frames rendered after
the load are identical either way.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Core/HLE/sceKernelModule.cpp

@@ -534,6 +534,19 @@ static void __KernelModuleInit() {
 	eventModuleLoadDone = CoreTiming::RegisterEvent("LoadModuleDone", &__KernelLoadModuleDone);
 }
 
+std::vector<std::pair<u32, u32>> __KernelModuleCodeRanges() {
+	std::vector<std::pair<u32, u32>> ranges;
+	u32 error;
+	for (SceUID moduleId : loadedModules) {
+		const PSPModule *module = kernelObjects.Get<PSPModule>(moduleId, error);
+		// textEnd is the last instruction.
+		if (module && !module->isFake && module->textStart != 0 && module->textEnd >= module->textStart) {
+			ranges.emplace_back(module->textStart, module->textEnd + 4 - module->textStart);
+		}
+	}
+	return ranges;
+}
+
 void __KernelModuleDoState(PointerWrap &p) {
 	auto s = p.Section("sceKernelModule", 1, 5);
 	if (!s)

Core/HLE/sceKernelModule.h

@@ -229,6 +229,8 @@ public:
 
 KernelObject *__KernelModuleObject();
 void __KernelModuleDoState(PointerWrap &p);
+// The code of the loaded modules, as (address, size).
+std::vector<std::pair<u32, u32>> __KernelModuleCodeRanges();
 void __KernelModuleShutdown();
 
 class MIPSState;

Core/MIPS/IR/IRFrontend.cpp

@@ -44,11 +44,14 @@ IRFrontend::IRFrontend(bool startDefaultPrefix) {
 	g_breakpoints.ClearSkipFirst();
 }
 
-void IRFrontend::DoState(PointerWrap &p) {
+bool IRFrontend::DoState(PointerWrap &p) {
 	auto s = p.Section("Jit", 1, 2);
 	if (!s)
-		return;
+		return false;
 
+	const bool oldStartDefaultPrefix = js.startDefaultPrefix;
+	const bool oldHasSetRounding = js.hasSetRounding;
+	const u8 oldLastSetRounding = js.lastSetRounding;
 	Do(p, js.startDefaultPrefix);
 	if (s >= 2) {
 		Do(p, js.hasSetRounding);
@@ -62,6 +65,23 @@ void IRFrontend::DoState(PointerWrap &p) {
 	// The debugger sets this so that "go" on a breakpoint will actually... go.
 	// But if they reset, we can end up hitting it by mistake, since it's based on PC and ticks.
 	g_breakpoints.ClearSkipFirst();
+
+	if (p.mode != PointerWrap::MODE_READ) {
+		return false;
+	}
+	// Blocks only switch rounding modes if hasSetRounding was set when they were compiled, and
+	// assume default prefixes at their start if startDefaultPrefix was.
+	if ((js.hasSetRounding && !oldHasSetRounding) || (!js.startDefaultPrefix && oldStartDefaultPrefix)) {
+		return true;
+	}
+	// The other way round, the blocks already compiled just do more than they need to. Keep the
+	// flags that way, or a loaded hasSetRounding of false would have UpdateFCR31() throw them all
+	// away for a rounding mode they handle fine. lastSetRounding says they were compiled with the
+	// rounding checks, so the next compile doesn't throw them away to add them either.
+	js.hasSetRounding = js.hasSetRounding || oldHasSetRounding;
+	js.startDefaultPrefix = js.startDefaultPrefix && oldStartDefaultPrefix;
+	js.lastSetRounding = oldLastSetRounding;
+	return false;
 }
 
 bool IRFrontend::UpdateFCR31(u32 fcr31) {

Core/MIPS/IR/IRFrontend.h

@@ -85,7 +85,8 @@ public:
 	void Comp_Vbfy(MIPSOpcode op) override;
 
 	int Replace_fabsf() override;
-	void DoState(PointerWrap &p);
+	// Returns true if a loaded state means blocks compiled before it can't run any more.
+	bool DoState(PointerWrap &p);
 	bool UpdateFCR31(u32 fcr31);  // returns true if the block cache needs clearing
 	bool CheckRounding(u32 blockAddress);  // returns true if we need a do-over
 

Core/MIPS/IR/IRJit.cpp

@@ -99,7 +99,9 @@ IRJit::~IRJit() {
 }
 
 void IRJit::DoState(PointerWrap &p) {
-	frontend_.DoState(p);
+	if (frontend_.DoState(p)) {
+		ClearCache();
+	}
 }
 
 void IRJit::UpdateFCR31() {
@@ -173,18 +175,55 @@ void IRJit::InvalidateCacheAt(u32 em_address, int length) {
 	invalidatedBlocks_ += (int)numbers.size();
 
 	for (int block_num : numbers) {
-		auto block = blocks_.GetBlock(block_num);
 		// TODO: We are invalidating a lot of blocks that are already invalid (yu gi oh).
-		// INFO_LOG(Log::JIT, "Block at %08x invalidated: valid: %d", block->GetOriginalStart(), block->IsValid());
-		// If we're a native JIT (IR->JIT, not just IR interpreter), we write native offsets into the blocks.
-		int cookie = compileToNative_ ? block->GetNativeOffset() : block->GetIRArenaOffset();
-		// Before Destroy, which forgets the block's address. Blocks linked to it jump straight in.
-		InvalidateNativeBlock(&blocks_, block_num);
-		blocks_.RemoveBlockFromPageLookup(block_num);
-		block->Destroy(cookie);
+		InvalidateBlock(block_num);
 	}
 }
 
+void IRJit::InvalidateBlock(int block_num) {
+	IRBlock *block = blocks_.GetBlock(block_num);
+	// If we're a native JIT (IR->JIT, not just IR interpreter), we write native offsets into the blocks.
+	int cookie = compileToNative_ ? block->GetNativeOffset() : block->GetIRArenaOffset();
+	// Before Destroy, which forgets the block's address. Blocks linked to it jump straight in.
+	InvalidateNativeBlock(&blocks_, block_num);
+	blocks_.RemoveBlockFromPageLookup(block_num);
+	block->Destroy(cookie);
+}
+
+void IRJit::InvalidateCacheExcept(const std::vector<std::pair<u32, u32>> &keep) {
+	if (keep.empty()) {
+		ClearCache();
+		return;
+	}
+
+	int kept = 0;
+	int invalidated = 0;
+	for (int block_num = 0; block_num < blocks_.GetNumBlocks(); ++block_num) {
+		const IRBlock *block = blocks_.GetBlock(block_num);
+		if (!block->IsValid()) {
+			continue;
+		}
+		u32 start, size;
+		block->GetRange(&start, &size);
+		start &= 0x3FFFFFFF;
+		bool inside = false;
+		for (const auto &[keepStart, keepSize] : keep) {
+			const u32 maskedStart = keepStart & 0x3FFFFFFF;
+			if (start >= maskedStart && start + size <= maskedStart + keepSize) {
+				inside = true;
+				break;
+			}
+		}
+		if (inside) {
+			kept++;
+		} else {
+			InvalidateBlock(block_num);
+			invalidated++;
+		}
+	}
+	INFO_LOG(Log::JIT, "IRJit: Kept %d blocks, invalidated %d", kept, invalidated);
+}
+
 void IRJit::Compile(u32 em_address) {
 	_dbg_assert_(compilerEnabled_);
 

Core/MIPS/IR/IRJit.h

@@ -209,6 +209,7 @@ public:
 
 	void ClearCache() override;
 	void InvalidateCacheAt(u32 em_address, int length = 4) override;
+	void InvalidateCacheExcept(const std::vector<std::pair<u32, u32>> &keep) override;
 	void UpdateFCR31() override;
 	void GetStats(JitStats *stats) const override;
 
@@ -238,6 +239,7 @@ protected:
 	virtual bool CompileNativeBlock(IRBlockCache *irBlockCache, int block_num) { return true; }
 	virtual void FinalizeNativeBlock(IRBlockCache *irBlockCache, int block_num) {}
 	virtual void InvalidateNativeBlock(IRBlockCache *irBlockCache, int block_num) {}
+	void InvalidateBlock(int block_num);
 
 	bool compileToNative_;
 

Core/MIPS/JitCommon/JitCommon.h

@@ -148,6 +148,8 @@ namespace MIPSComp {
 		virtual const u8 *GetCrashHandler() const = 0;
 		virtual JitBlockCacheDebugInterface *GetBlockCacheDebugInterface() = 0;
 		virtual void InvalidateCacheAt(u32 em_address, int length = 4) = 0;
+		// Invalidates every block not entirely inside one of the (address, size) ranges.
+		virtual void InvalidateCacheExcept(const std::vector<std::pair<u32, u32>> &keep) = 0;
 		virtual void DoState(PointerWrap &p) = 0;
 		virtual void RunLoopUntil(u64 globalticks) = 0;
 		virtual void Compile(u32 em_address) = 0;

Core/MIPS/MIPS.cpp

@@ -326,6 +326,18 @@ void MIPSState::Reset() {
 }
 
 void MIPSState::Init() {
+	InitRegisters();
+	if (PSP_CoreParameter().cpuCore == CPUCore::JIT || PSP_CoreParameter().cpuCore == CPUCore::JIT_IR) {
+		MIPSComp::jit = MIPSComp::CreateNativeJit(this);
+	} else if (PSP_CoreParameter().cpuCore == CPUCore::IR_INTERPRETER) {
+		MIPSComp::jit = new MIPSComp::IRJit(this, false);
+	} else {
+		MIPSComp::jit = nullptr;
+	}
+	UpdateBlockShadowHook();
+}
+
+void MIPSState::InitRegisters() {
 	memset(r, 0, sizeof(r));
 	memset(f, 0, sizeof(f));
 	memset(v, 0, sizeof(v));
@@ -359,15 +371,6 @@ void MIPSState::Init() {
 	downcount = 0;
 
 	memset(vcmpResult, 0, sizeof(vcmpResult));
-
-	if (PSP_CoreParameter().cpuCore == CPUCore::JIT || PSP_CoreParameter().cpuCore == CPUCore::JIT_IR) {
-		MIPSComp::jit = MIPSComp::CreateNativeJit(this);
-	} else if (PSP_CoreParameter().cpuCore == CPUCore::IR_INTERPRETER) {
-		MIPSComp::jit = new MIPSComp::IRJit(this, false);
-	} else {
-		MIPSComp::jit = nullptr;
-	}
-	UpdateBlockShadowHook();
 }
 
 const float vfpuPrefixConstants[8] = { 0.f, 1.f, 2.f, 0.5f, 3.f, 1.f / 3.f, 0.25f, 1.f / 6.f };
@@ -427,9 +430,17 @@ void MIPSState::DoState(PointerWrap &p) {
 	if (!s)
 		return;
 
-	// Reset the jit if we're loading.
-	if (p.mode == p.MODE_READ)
-		Reset();
+	if (p.mode == p.MODE_READ) {
+		if (MIPSComp::jit && !keepJitBlocks_.empty()) {
+			// Blocks in code the load didn't change can stay. Queued invalidations stay too.
+			InitRegisters();
+			MIPSComp::jit->InvalidateCacheExcept(keepJitBlocks_);
+		} else {
+			// Start the JIT over.
+			Reset();
+		}
+		keepJitBlocks_.clear();
+	}
 	// Assume we're not saving state during a CPU core reset, so no lock.
 	if (MIPSComp::jit)
 		MIPSComp::jit->DoState(p);

Core/MIPS/MIPS.h

@@ -22,6 +22,7 @@
 #include <cstddef>
 #include <cstring>
 #include <type_traits>
+#include <utility>
 #include <vector>
 
 #include "Common/CommonTypes.h"
@@ -180,11 +181,18 @@ public:
 	~MIPSState();
 
 	void Init();
+	// Init without the JIT.
+	void InitRegisters();
 	void Shutdown();
 	void Reset();
 	void UpdateCore(CPUCore desired);
 
 	void DoState(PointerWrap &p);
+	// For the next DoState load: keep the JIT blocks entirely inside these (address, size) ranges,
+	// code the load leaves as it is, instead of starting the JIT over.
+	void KeepJitBlocksOnNextLoad(std::vector<std::pair<u32, u32>> ranges) {
+		keepJitBlocks_ = std::move(ranges);
+	}
 
 	// MUST start with r and be followed by f, v, and t!
 	u32 r[32];
@@ -295,6 +303,7 @@ public:
 	// Doesn't need save stating.
 	std::vector<PendingCacheOperation> pendingInvalidates_;
 	bool invalidateAll_ = false;
+	std::vector<std::pair<u32, u32>> keepJitBlocks_;
 };
 
 // Access control ('private') break standard layout. Let's not allow it, to avoid excessive

Core/MemMap.cpp

@@ -449,7 +449,60 @@ bool Reinit() {
 	return success;
 }
 
-static void DoMemoryVoid(PointerWrap &p, uint32_t start, uint32_t size) {
+// Loads size bytes at start, comparing the parts of it in compare->ranges with what's there as it
+// goes. A chunk that already matches doesn't need copying, and the matching parts are reported in
+// compare->unchanged. The rest is a plain copy.
+static void CopyComparing(uint8_t *dest, const uint8_t *src, u32 start, u32 size, RamCompare *compare) {
+	// In the size of the JIT's block lookup pages. One changed byte costs its chunk's blocks.
+	constexpr u32 CHUNK_SIZE = 1024;
+
+	// As offsets into the copy, in order, merged.
+	std::vector<std::pair<u32, u32>> ranges;
+	for (const auto &[address, length] : compare->ranges) {
+		const u32 rangeStart = std::max(address & 0x3FFFFFFF, start);
+		const u32 rangeEnd = std::min((address & 0x3FFFFFFF) + length, start + size);
+		if (rangeStart < rangeEnd) {
+			ranges.emplace_back(rangeStart - start, rangeEnd - start);
+		}
+	}
+	std::sort(ranges.begin(), ranges.end());
+
+	u32 pos = 0;
+	for (const auto &[rangeStart, rangeEnd] : ranges) {
+		const u32 from = std::max(rangeStart, pos);
+		if (from >= rangeEnd) {
+			continue;
+		}
+		if (from > pos) {
+			ParallelMemcpy(&g_threadManager, dest + pos, src + pos, from - pos);
+		}
+		u32 runStart = 0;
+		bool inRun = false;
+		for (u32 chunk = from; chunk < rangeEnd; chunk += CHUNK_SIZE) {
+			const u32 chunkSize = std::min(CHUNK_SIZE, rangeEnd - chunk);
+			const bool same = memcmp(dest + chunk, src + chunk, chunkSize) == 0;
+			if (same && !inRun) {
+				runStart = chunk;
+				inRun = true;
+			} else if (!same) {
+				memcpy(dest + chunk, src + chunk, chunkSize);
+				if (inRun) {
+					compare->unchanged.emplace_back(start + runStart, chunk - runStart);
+					inRun = false;
+				}
+			}
+		}
+		if (inRun) {
+			compare->unchanged.emplace_back(start + runStart, rangeEnd - runStart);
+		}
+		pos = rangeEnd;
+	}
+	if (pos < size) {
+		ParallelMemcpy(&g_threadManager, dest + pos, src + pos, size - pos);
+	}
+}
+
+static void DoMemoryVoid(PointerWrap &p, uint32_t start, uint32_t size, RamCompare *compare = nullptr) {
 	uint8_t *d = GetPointerWriteOrException(start);
 	uint8_t *&storage = *p.ptr;
 
@@ -463,7 +516,11 @@ static void DoMemoryVoid(PointerWrap &p, uint32_t start, uint32_t size) {
 
 	switch (p.mode) {
 	case PointerWrap::MODE_READ:
-		ParallelMemcpy(&g_threadManager, d, storage, size);
+		if (compare && !compare->ranges.empty()) {
+			CopyComparing(d, storage, start, size, compare);
+		} else {
+			ParallelMemcpy(&g_threadManager, d, storage, size);
+		}
 		break;
 	case PointerWrap::MODE_WRITE:
 		ParallelMemcpy(&g_threadManager, storage, d, size);
@@ -483,7 +540,7 @@ static void DoMemoryVoid(PointerWrap &p, uint32_t start, uint32_t size) {
 	storage += size;
 }
 
-void DoState(PointerWrap &p) {
+void DoState(PointerWrap &p, RamCompare *compare) {
 	auto s = p.Section("Memory", 1, 3);
 	if (!s)
 		return;
@@ -529,7 +586,7 @@ void DoState(PointerWrap &p) {
 		}
 	}
 
-	DoMemoryVoid(p, PSP_GetKernelMemoryBase(), g_MemorySize);
+	DoMemoryVoid(p, PSP_GetKernelMemoryBase(), g_MemorySize, compare);
 	p.DoMarker("RAM");
 
 	DoMemoryVoid(p, PSP_GetVidMemBase(), VRAM_SIZE);

Core/MemMap.h

@@ -21,6 +21,8 @@
 
 #include <cstring>
 #include <cstdint>
+#include <utility>
+#include <vector>
 #ifndef offsetof
 #include <stddef.h>
 #endif
@@ -169,10 +171,19 @@ enum class MemMapSetupFlags {
 };
 ENUM_CLASS_BITOPS(MemMapSetupFlags);
 
+// For a savestate load: parts of RAM to compare with what's loaded over them, as they're loaded.
+// Addresses can be in any mirror. The JIT keeps its blocks in code the load doesn't change.
+struct RamCompare {
+	// In: (address, size) ranges.
+	std::vector<std::pair<u32, u32>> ranges;
+	// Out: the parts of them the load left as they were.
+	std::vector<std::pair<u32, u32>> unchanged;
+};
+
 // Init and Shutdown
 bool Init(MemMapSetupFlags flags);
 void Shutdown();
-void DoState(PointerWrap &p);
+void DoState(PointerWrap &p, RamCompare *compare = nullptr);
 
 // False when shutdown has already been called.
 bool IsActive();

Core/SaveState.cpp

@@ -53,6 +53,7 @@
 #include "Core/HLE/ReplaceTables.h"
 #include "Core/HLE/sceDisplay.h"
 #include "Core/HLE/sceKernel.h"
+#include "Core/HLE/sceKernelModule.h"
 #include "Core/HLE/sceUtility.h"
 #include "Core/HLE/sceSas.h"
 #include "Core/HLE/sceIo.h"
@@ -60,6 +61,7 @@
 #include "GPU/GPU.h"
 #include "GPU/GPUCommon.h"
 #include "Core/MIPS/JitCommon/JitBlockCache.h"
+#include "Core/MIPS/JitCommon/JitCommon.h"
 #include "Core/RetroAchievements.h"
 #include "HW/MemoryStick.h"
 
@@ -233,13 +235,20 @@ int g_screenshotFailures;
 
 		// The state may hold different code. KernelModuleDoState hooks the functions again.
 		std::vector<std::pair<u32, u32>> hookedOps;
+		// The loaded modules' code is compared with the state's as it's loaded. Where it's the same,
+		// the JIT keeps its blocks, so they don't all have to be compiled again.
+		Memory::RamCompare codeCompare;
 		if (p.mode == p.MODE_READ) {
 			hookedOps = Replacement_GetHookedOps();
 			RestoreReplacedInstructions(0, 0xFFFFFFFF);
+			if (MIPSComp::jit) {
+				codeCompare.ranges = __KernelModuleCodeRanges();
+			}
 		}
-		Memory::DoState(p);
+		Memory::DoState(p, &codeCompare);
 		if (p.mode == p.MODE_READ) {
 			Replacement_RepairStaleHookOps(hookedOps);
+			currentMIPS->KeepJitBlocksOnNextLoad(std::move(codeCompare.unchanged));
 		}
 
 		if (s >= 3) {