From cf4da5c4d1cfa3fd63002a5d33e563eb3a850bd3 Mon Sep 17 00:00:00 2001
From: Eladash <elad3356p@gmail.com>
Date: Wed, 14 Sep 2022 11:47:25 +0300
Subject: [PATCH] CPU preemption control: bugfixes

---
 rpcs3/Emu/CPU/CPUThread.cpp |  5 +++--
 rpcs3/Emu/Cell/lv2/lv2.cpp  | 40 +++++++++++++++++++++++++++++++------
 rpcs3/Emu/RSX/RSXThread.cpp |  8 ++++----
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/rpcs3/Emu/CPU/CPUThread.cpp b/rpcs3/Emu/CPU/CPUThread.cpp
index ed71fe1d5f..c01b4b9eb2 100644
--- a/rpcs3/Emu/CPU/CPUThread.cpp
+++ b/rpcs3/Emu/CPU/CPUThread.cpp
@@ -654,7 +654,7 @@ bool cpu_thread::check_state() noexcept
 	{
 		// Process all flags in a single atomic op
 		bs_t<cpu_flag> state1;
-		const auto state0 = state.fetch_op([&](bs_t<cpu_flag>& flags)
+		auto state0 = state.fetch_op([&](bs_t<cpu_flag>& flags)
 		{
 			bool store = false;
 
@@ -784,7 +784,8 @@ bool cpu_thread::check_state() noexcept
 			if (cpu_flag::wait - state0)
 			{
 				// Yield itself
-				s_dummy_atomic.wait(0, 1u << 30, atomic_wait_timeout{80'000});
+				escape = false;
+				state0 += cpu_flag::yield;
 			}
 
 			if (const u128 bits = s_cpu_bits)
diff --git a/rpcs3/Emu/Cell/lv2/lv2.cpp b/rpcs3/Emu/Cell/lv2/lv2.cpp
index d37a8f9caf..ca51d4a82a 100644
--- a/rpcs3/Emu/Cell/lv2/lv2.cpp
+++ b/rpcs3/Emu/Cell/lv2/lv2.cpp
@@ -1668,14 +1668,42 @@ void lv2_obj::schedule_all(u64 current_time)
 
 	if (const u64 freq = s_yield_frequency)
 	{
-		if (auto cpu = cpu_thread::get_current())
-		{
-			const u64 tsc = utils::get_tsc();
-			const u64 last_tsc = s_last_yield_tsc;
+		const u64 tsc = utils::get_tsc();
+		const u64 last_tsc = s_last_yield_tsc;
 
-			if (tsc >= last_tsc && tsc <= s_max_allowed_yield_tsc && tsc - last_tsc >= freq)
+		if (tsc >= last_tsc && tsc <= s_max_allowed_yield_tsc && tsc - last_tsc >= freq)
+		{
+			auto target = +g_ppu;
+			cpu_thread* cpu = nullptr;
+
+			for (usz x = g_cfg.core.ppu_threads;; target = target->next_ppu, x--)
+			{
+				if (!target || !x)
+				{
+					if (g_ppu && cpu_flag::preempt - g_ppu->state)
+					{
+						// Don't be picky, pick up any running PPU thread even it has a wait flag 
+						cpu = g_ppu;
+					}
+					// TODO: If this case is common enough it may be valuable to iterate over all CPU threads to find a perfect candidate (one without a wait or suspend flag)
+					else if (auto current = cpu_thread::get_current(); current && cpu_flag::suspend - current->state)
+					{
+						// May be an SPU or RSX thread, use them as a last resort
+						cpu = current;
+					}
+
+					break;
+				}
+
+				if (target->state.none_of(cpu_flag::preempt + cpu_flag::wait))
+				{
+					cpu = target;
+					break;
+				}
+			}
+
+			if (cpu && cpu_flag::preempt - cpu->state && !cpu->state.test_and_set(cpu_flag::preempt))
 			{
-				cpu->state += cpu_flag::preempt;
 				s_last_yield_tsc = tsc;
 				g_lv2_preempts_taken.release(g_lv2_preempts_taken.load() + 1); // Has a minor race but performance is more important
 				rsx::set_rsx_yield_flag();
diff --git a/rpcs3/Emu/RSX/RSXThread.cpp b/rpcs3/Emu/RSX/RSXThread.cpp
index ba97252390..f9ef46efcd 100644
--- a/rpcs3/Emu/RSX/RSXThread.cpp
+++ b/rpcs3/Emu/RSX/RSXThread.cpp
@@ -3558,7 +3558,7 @@ namespace rsx
 				prev_preempt_count = frame_times[i].preempt_count;
 			}
 
-			preempt_count = frame_times.back().preempt_count;
+			preempt_count = std::min<u32>(frame_times.back().preempt_count, max_preempt_count);
 
 			u32 fails = 0;
 			u32 hard_fails = 0;
@@ -3630,9 +3630,9 @@ namespace rsx
 					{
 						prevent_preempt_increase_tickets--;
 					}
-					else if (preempt_count < max_preempt_count)
+					else
 					{
-						preempt_count += 4;
+						preempt_count = std::min<u32>(preempt_count + 4, max_preempt_count);
 					}
 				}
 				else
@@ -3650,7 +3650,7 @@ namespace rsx
 
 			if (hard_measures_taken)
 			{
-				preempt_fail_old_preempt_count = std::max<u32>(preempt_fail_old_preempt_count, frame_times.back().preempt_count);
+				preempt_fail_old_preempt_count = std::max<u32>(preempt_fail_old_preempt_count, std::min<u32>(frame_times.back().preempt_count, max_preempt_count));
 			}
 			else if (preempt_fail_old_preempt_count)
 			{