CPU_Manager: Unload/Reload threads on preemption on SingleCore

4 files changed, 64 insertions, 7 deletions

diff --git a/src/core/cpu_manager.cpp b/src/core/cpu_manager.cpp
index e72f89808..95842aad1 100644
--- a/src/core/cpu_manager.cpp
+++ b/src/core/cpu_manager.cpp
@@ -225,7 +225,7 @@ void CpuManager::SingleCoreRunGuestLoop() {
         }
         physical_core.ClearExclusive();
         PreemptSingleCore();
-        auto& scheduler = physical_core.Scheduler();
+        auto& scheduler = kernel.Scheduler(current_core);
         scheduler.TryDoContextSwitch();
     }
 }
@@ -260,11 +260,15 @@ void CpuManager::SingleCoreRunSuspendThread() {
 void CpuManager::PreemptSingleCore() {
     preemption_count = 0;
     std::size_t old_core = current_core;
-    current_core = (current_core + 1) % Core::Hardware::NUM_CPU_CORES;
+    current_core.store((current_core + 1) % Core::Hardware::NUM_CPU_CORES);
     auto& scheduler = system.Kernel().Scheduler(old_core);
-    Kernel::Thread* current_thread = system.Kernel().Scheduler(old_core).GetCurrentThread();
-    Kernel::Thread* next_thread = system.Kernel().Scheduler(current_core).GetCurrentThread();
-    Common::Fiber::YieldTo(current_thread->GetHostContext(), next_thread->GetHostContext());
+    Kernel::Thread* current_thread = scheduler.GetCurrentThread();
+    scheduler.Unload();
+    auto& next_scheduler = system.Kernel().Scheduler(current_core);
+    Common::Fiber::YieldTo(current_thread->GetHostContext(), next_scheduler.ControlContext());
+    /// May have changed scheduler
+    auto& current_scheduler = system.Kernel().Scheduler(current_core);
+    current_scheduler.Reload();
 }
 
 void CpuManager::SingleCorePause(bool paused) {
diff --git a/src/core/cpu_manager.h b/src/core/cpu_manager.h
index 1e81481ec..ff1935d5c 100644
--- a/src/core/cpu_manager.h
+++ b/src/core/cpu_manager.h
@@ -5,6 +5,7 @@
 #pragma once
 
 #include <array>
+#include <atomic>
 #include <functional>
 #include <memory>
 #include <thread>
@@ -45,7 +46,7 @@ public:
     void* GetStartFuncParamater();
 
     std::size_t CurrentCore() const {
-        return current_core;
+        return current_core.load();
     }
 
 private:
@@ -88,7 +89,7 @@ private:
     std::array<CoreData, Core::Hardware::NUM_CPU_CORES> core_data{};
 
     bool is_multicore{};
-    std::size_t current_core{};
+    std::atomic<std::size_t> current_core{};
     std::size_t preemption_count{};
     static constexpr std::size_t max_cycle_runs = 5;
 
diff --git a/src/core/hle/kernel/scheduler.cpp b/src/core/hle/kernel/scheduler.cpp
index d68d86cdf..00322d997 100644
--- a/src/core/hle/kernel/scheduler.cpp
+++ b/src/core/hle/kernel/scheduler.cpp
@@ -602,6 +602,48 @@ void Scheduler::OnThreadStart() {
     SwitchContextStep2();
 }
 
+void Scheduler::Unload() {
+    Thread* thread = current_thread.get();
+    if (thread) {
+        thread->last_running_ticks = system.CoreTiming().GetCPUTicks();
+        thread->SetIsRunning(false);
+        if (!thread->IsHLEThread()) {
+            auto& cpu_core = system.ArmInterface(core_id);
+            cpu_core.SaveContext(thread->GetContext32());
+            cpu_core.SaveContext(thread->GetContext64());
+            // Save the TPIDR_EL0 system register in case it was modified.
+            thread->SetTPIDR_EL0(cpu_core.GetTPIDR_EL0());
+            cpu_core.ClearExclusiveState();
+        }
+        thread->context_guard.unlock();
+    }
+}
+
+void Scheduler::Reload() {
+    Thread* thread = current_thread.get();
+    if (thread) {
+        ASSERT_MSG(thread->GetSchedulingStatus() == ThreadSchedStatus::Runnable,
+                   "Thread must be runnable.");
+
+        // Cancel any outstanding wakeup events for this thread
+        thread->SetIsRunning(true);
+        thread->last_running_ticks = system.CoreTiming().GetCPUTicks();
+
+        auto* const thread_owner_process = thread->GetOwnerProcess();
+        if (thread_owner_process != nullptr) {
+            system.Kernel().MakeCurrentProcess(thread_owner_process);
+        }
+        if (!thread->IsHLEThread()) {
+            auto& cpu_core = system.ArmInterface(core_id);
+            cpu_core.LoadContext(thread->GetContext32());
+            cpu_core.LoadContext(thread->GetContext64());
+            cpu_core.SetTlsAddress(thread->GetTLSAddress());
+            cpu_core.SetTPIDR_EL0(thread->GetTPIDR_EL0());
+            cpu_core.ClearExclusiveState();
+        }
+    }
+}
+
 void Scheduler::SwitchContextStep2() {
     Thread* previous_thread = current_thread_prev.get();
     Thread* new_thread = selected_thread.get();
diff --git a/src/core/hle/kernel/scheduler.h b/src/core/hle/kernel/scheduler.h
index 5e062bf59..f63cc5085 100644
--- a/src/core/hle/kernel/scheduler.h
+++ b/src/core/hle/kernel/scheduler.h
@@ -210,6 +210,12 @@ public:
     /// Reschedules to the next available thread (call after current thread is suspended)
     void TryDoContextSwitch();
 
+    /// The next two are for SingleCore Only.
+    /// Unload current thread before preempting core.
+    void Unload();
+    /// Reload current thread after core preemption.
+    void Reload();
+
     /// Gets the current running thread
     Thread* GetCurrentThread() const;
 
@@ -230,6 +236,10 @@ public:
 
     void OnThreadStart();
 
+    std::shared_ptr<Common::Fiber> ControlContext() {
+        return switch_fiber;
+    }
+
 private:
     friend class GlobalScheduler;