From 3a89723d97b8e646cde569030057777813f4371c Mon Sep 17 00:00:00 2001
From: bunnei <bunneidev@gmail.com>
Date: Wed, 29 Dec 2021 21:40:38 -0800
Subject: [PATCH] core: hle: kernel: Implement thread pinning.

- We largely had the mechanics in place for thread pinning, this change hooks these up.
- Validated with tests https://github.com/Atmosphere-NX/Atmosphere/blob/master/tests/TestSvc/source/test_thread_pinning.cpp.
---
 src/core/CMakeLists.txt                       |  2 +
 .../hle/kernel/global_scheduler_context.cpp   |  6 +++
 src/core/hle/kernel/k_interrupt_manager.cpp   | 34 ++++++++++++++
 src/core/hle/kernel/k_interrupt_manager.h     | 17 +++++++
 src/core/hle/kernel/k_process.cpp             | 12 +++--
 src/core/hle/kernel/k_process.h               |  4 +-
 src/core/hle/kernel/k_scheduler.cpp           |  8 ++++
 src/core/hle/kernel/k_thread.cpp              | 44 +++++++++++++++++--
 src/core/hle/kernel/k_thread.h                |  6 ++-
 src/core/hle/kernel/svc.cpp                   | 21 ++++++++-
 10 files changed, 140 insertions(+), 14 deletions(-)
 create mode 100644 src/core/hle/kernel/k_interrupt_manager.cpp
 create mode 100644 src/core/hle/kernel/k_interrupt_manager.h

diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt
index 698c4f912d..b1a7467276 100644
--- a/src/core/CMakeLists.txt
+++ b/src/core/CMakeLists.txt
@@ -187,6 +187,8 @@ add_library(core STATIC
     hle/kernel/k_event.h
     hle/kernel/k_handle_table.cpp
     hle/kernel/k_handle_table.h
+    hle/kernel/k_interrupt_manager.cpp
+    hle/kernel/k_interrupt_manager.h
     hle/kernel/k_light_condition_variable.cpp
     hle/kernel/k_light_condition_variable.h
     hle/kernel/k_light_lock.cpp
diff --git a/src/core/hle/kernel/global_scheduler_context.cpp b/src/core/hle/kernel/global_scheduler_context.cpp
index 4f4e338e3f..baad2c5d67 100644
--- a/src/core/hle/kernel/global_scheduler_context.cpp
+++ b/src/core/hle/kernel/global_scheduler_context.cpp
@@ -9,6 +9,7 @@
 #include "core/hle/kernel/global_scheduler_context.h"
 #include "core/hle/kernel/k_scheduler.h"
 #include "core/hle/kernel/kernel.h"
+#include "core/hle/kernel/physical_core.h"
 
 namespace Kernel {
 
@@ -42,6 +43,11 @@ void GlobalSchedulerContext::PreemptThreads() {
     for (u32 core_id = 0; core_id < Core::Hardware::NUM_CPU_CORES; core_id++) {
         const u32 priority = preemption_priorities[core_id];
         kernel.Scheduler(core_id).RotateScheduledQueue(core_id, priority);
+
+        // Signal an interrupt occurred. For core 3, this is a certainty, as preemption will result
+        // in the rotator thread being scheduled. For cores 0-2, this is to simulate or system
+        // interrupts that may have occurred.
+        kernel.PhysicalCore(core_id).Interrupt();
     }
 }
 
diff --git a/src/core/hle/kernel/k_interrupt_manager.cpp b/src/core/hle/kernel/k_interrupt_manager.cpp
new file mode 100644
index 0000000000..e5dd39751a
--- /dev/null
+++ b/src/core/hle/kernel/k_interrupt_manager.cpp
@@ -0,0 +1,34 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include "core/hle/kernel/k_interrupt_manager.h"
+#include "core/hle/kernel/k_process.h"
+#include "core/hle/kernel/k_scheduler.h"
+#include "core/hle/kernel/k_thread.h"
+#include "core/hle/kernel/kernel.h"
+
+namespace Kernel::KInterruptManager {
+
+void HandleInterrupt(KernelCore& kernel, s32 core_id) {
+    auto* process = kernel.CurrentProcess();
+    if (!process) {
+        return;
+    }
+
+    auto& scheduler = kernel.Scheduler(core_id);
+    auto& current_thread = *scheduler.GetCurrentThread();
+
+    // If the user disable count is set, we may need to pin the current thread.
+    if (current_thread.GetUserDisableCount() && !process->GetPinnedThread(core_id)) {
+        KScopedSchedulerLock sl{kernel};
+
+        // Pin the current thread.
+        process->PinCurrentThread(core_id);
+
+        // Set the interrupt flag for the thread.
+        scheduler.GetCurrentThread()->SetInterruptFlag();
+    }
+}
+
+} // namespace Kernel::KInterruptManager
diff --git a/src/core/hle/kernel/k_interrupt_manager.h b/src/core/hle/kernel/k_interrupt_manager.h
new file mode 100644
index 0000000000..05924801e1
--- /dev/null
+++ b/src/core/hle/kernel/k_interrupt_manager.h
@@ -0,0 +1,17 @@
+// Copyright 2021 yuzu Emulator Project
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include "common/common_types.h"
+
+namespace Kernel {
+
+class KernelCore;
+
+namespace KInterruptManager {
+void HandleInterrupt(KernelCore& kernel, s32 core_id);
+}
+
+} // namespace Kernel
diff --git a/src/core/hle/kernel/k_process.cpp b/src/core/hle/kernel/k_process.cpp
index 73f8bc4fe7..bf98a51e25 100644
--- a/src/core/hle/kernel/k_process.cpp
+++ b/src/core/hle/kernel/k_process.cpp
@@ -220,30 +220,28 @@ bool KProcess::ReleaseUserException(KThread* thread) {
     }
 }
 
-void KProcess::PinCurrentThread() {
+void KProcess::PinCurrentThread(s32 core_id) {
     ASSERT(kernel.GlobalSchedulerContext().IsLocked());
 
     // Get the current thread.
-    const s32 core_id = GetCurrentCoreId(kernel);
-    KThread* cur_thread = GetCurrentThreadPointer(kernel);
+    KThread* cur_thread = kernel.Scheduler(static_cast<std::size_t>(core_id)).GetCurrentThread();
 
     // If the thread isn't terminated, pin it.
     if (!cur_thread->IsTerminationRequested()) {
         // Pin it.
         PinThread(core_id, cur_thread);
-        cur_thread->Pin();
+        cur_thread->Pin(core_id);
 
         // An update is needed.
         KScheduler::SetSchedulerUpdateNeeded(kernel);
     }
 }
 
-void KProcess::UnpinCurrentThread() {
+void KProcess::UnpinCurrentThread(s32 core_id) {
     ASSERT(kernel.GlobalSchedulerContext().IsLocked());
 
     // Get the current thread.
-    const s32 core_id = GetCurrentCoreId(kernel);
-    KThread* cur_thread = GetCurrentThreadPointer(kernel);
+    KThread* cur_thread = kernel.Scheduler(static_cast<std::size_t>(core_id)).GetCurrentThread();
 
     // Unpin it.
     cur_thread->Unpin();
diff --git a/src/core/hle/kernel/k_process.h b/src/core/hle/kernel/k_process.h
index cb93c7e247..e7c8b5838c 100644
--- a/src/core/hle/kernel/k_process.h
+++ b/src/core/hle/kernel/k_process.h
@@ -345,8 +345,8 @@ public:
 
     bool IsSignaled() const override;
 
-    void PinCurrentThread();
-    void UnpinCurrentThread();
+    void PinCurrentThread(s32 core_id);
+    void UnpinCurrentThread(s32 core_id);
     void UnpinThread(KThread* thread);
 
     KLightLock& GetStateLock() {
diff --git a/src/core/hle/kernel/k_scheduler.cpp b/src/core/hle/kernel/k_scheduler.cpp
index 277201de47..31cec990ed 100644
--- a/src/core/hle/kernel/k_scheduler.cpp
+++ b/src/core/hle/kernel/k_scheduler.cpp
@@ -15,6 +15,7 @@
 #include "core/core.h"
 #include "core/core_timing.h"
 #include "core/cpu_manager.h"
+#include "core/hle/kernel/k_interrupt_manager.h"
 #include "core/hle/kernel/k_process.h"
 #include "core/hle/kernel/k_scheduler.h"
 #include "core/hle/kernel/k_scoped_scheduler_lock_and_sleep.h"
@@ -53,6 +54,13 @@ void KScheduler::RescheduleCores(KernelCore& kernel, u64 cores_pending_reschedul
         }
         cores_pending_reschedule &= ~(1ULL << core);
     }
+
+    for (std::size_t core_id = 0; core_id < Core::Hardware::NUM_CPU_CORES; ++core_id) {
+        if (kernel.PhysicalCore(core_id).IsInterrupted()) {
+            KInterruptManager::HandleInterrupt(kernel, static_cast<s32>(core_id));
+        }
+    }
+
     if (must_context_switch) {
         auto core_scheduler = kernel.CurrentScheduler();
         kernel.ExitSVCProfile();
diff --git a/src/core/hle/kernel/k_thread.cpp b/src/core/hle/kernel/k_thread.cpp
index b8c993748f..71e029a3f0 100644
--- a/src/core/hle/kernel/k_thread.cpp
+++ b/src/core/hle/kernel/k_thread.cpp
@@ -3,6 +3,7 @@
 // Refer to the license.txt file included.
 
 #include <algorithm>
+#include <atomic>
 #include <cinttypes>
 #include <optional>
 #include <vector>
@@ -33,6 +34,7 @@
 #include "core/hle/kernel/svc_results.h"
 #include "core/hle/kernel/time_manager.h"
 #include "core/hle/result.h"
+#include "core/memory.h"
 
 #ifdef ARCHITECTURE_x86_64
 #include "core/arm/dynarmic/arm_dynarmic_32.h"
@@ -63,6 +65,13 @@ namespace Kernel {
 
 namespace {
 
+struct ThreadLocalRegion {
+    static constexpr std::size_t MessageBufferSize = 0x100;
+    std::array<u32, MessageBufferSize / sizeof(u32)> message_buffer;
+    std::atomic_uint16_t disable_count;
+    std::atomic_uint16_t interrupt_flag;
+};
+
 class ThreadQueueImplForKThreadSleep final : public KThreadQueueWithoutEndWait {
 public:
     explicit ThreadQueueImplForKThreadSleep(KernelCore& kernel_)
@@ -346,7 +355,7 @@ void KThread::StartTermination() {
     if (parent != nullptr) {
         parent->ReleaseUserException(this);
         if (parent->GetPinnedThread(GetCurrentCoreId(kernel)) == this) {
-            parent->UnpinCurrentThread();
+            parent->UnpinCurrentThread(core_id);
         }
     }
 
@@ -372,7 +381,7 @@ void KThread::StartTermination() {
     this->Close();
 }
 
-void KThread::Pin() {
+void KThread::Pin(s32 current_core) {
     ASSERT(kernel.GlobalSchedulerContext().IsLocked());
 
     // Set ourselves as pinned.
@@ -389,7 +398,6 @@ void KThread::Pin() {
 
         // Bind ourselves to this core.
         const s32 active_core = GetActiveCore();
-        const s32 current_core = GetCurrentCoreId(kernel);
 
         SetActiveCore(current_core);
         physical_ideal_core_id = current_core;
@@ -482,6 +490,36 @@ void KThread::Unpin() {
     }
 }
 
+u16 KThread::GetUserDisableCount() const {
+    if (!IsUserThread()) {
+        // We only emulate TLS for user threads
+        return {};
+    }
+
+    auto& memory = kernel.System().Memory();
+    return memory.Read16(tls_address + offsetof(ThreadLocalRegion, disable_count));
+}
+
+void KThread::SetInterruptFlag() {
+    if (!IsUserThread()) {
+        // We only emulate TLS for user threads
+        return;
+    }
+
+    auto& memory = kernel.System().Memory();
+    memory.Write16(tls_address + offsetof(ThreadLocalRegion, interrupt_flag), 1);
+}
+
+void KThread::ClearInterruptFlag() {
+    if (!IsUserThread()) {
+        // We only emulate TLS for user threads
+        return;
+    }
+
+    auto& memory = kernel.System().Memory();
+    memory.Write16(tls_address + offsetof(ThreadLocalRegion, interrupt_flag), 0);
+}
+
 ResultCode KThread::GetCoreMask(s32* out_ideal_core, u64* out_affinity_mask) {
     KScopedSchedulerLock sl{kernel};
 
diff --git a/src/core/hle/kernel/k_thread.h b/src/core/hle/kernel/k_thread.h
index c8a08bd71f..83dfde69b4 100644
--- a/src/core/hle/kernel/k_thread.h
+++ b/src/core/hle/kernel/k_thread.h
@@ -307,6 +307,10 @@ public:
         return parent != nullptr;
     }
 
+    u16 GetUserDisableCount() const;
+    void SetInterruptFlag();
+    void ClearInterruptFlag();
+
     [[nodiscard]] KThread* GetLockOwner() const {
         return lock_owner;
     }
@@ -490,7 +494,7 @@ public:
         this->GetStackParameters().disable_count--;
     }
 
-    void Pin();
+    void Pin(s32 current_core);
 
     void Unpin();
 
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 63e2dff19f..250ef9042d 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -2027,6 +2027,25 @@ static ResultCode SignalToAddress(Core::System& system, VAddr address, Svc::Sign
                                                                   count);
 }
 
+static void SynchronizePreemptionState(Core::System& system) {
+    auto& kernel = system.Kernel();
+
+    // Lock the scheduler.
+    KScopedSchedulerLock sl{kernel};
+
+    // If the current thread is pinned, unpin it.
+    KProcess* cur_process = system.Kernel().CurrentProcess();
+    const auto core_id = GetCurrentCoreId(kernel);
+
+    if (cur_process->GetPinnedThread(core_id) == GetCurrentThreadPointer(kernel)) {
+        // Clear the current thread's interrupt flag.
+        GetCurrentThread(kernel).ClearInterruptFlag();
+
+        // Unpin the current thread.
+        cur_process->UnpinCurrentThread(core_id);
+    }
+}
+
 static ResultCode SignalToAddress32(Core::System& system, u32 address, Svc::SignalType signal_type,
                                     s32 value, s32 count) {
     return SignalToAddress(system, address, signal_type, value, count);
@@ -2797,7 +2816,7 @@ static const FunctionDef SVC_Table_64[] = {
     {0x33, SvcWrap64<GetThreadContext>, "GetThreadContext"},
     {0x34, SvcWrap64<WaitForAddress>, "WaitForAddress"},
     {0x35, SvcWrap64<SignalToAddress>, "SignalToAddress"},
-    {0x36, nullptr, "SynchronizePreemptionState"},
+    {0x36, SvcWrap64<SynchronizePreemptionState>, "SynchronizePreemptionState"},
     {0x37, nullptr, "Unknown"},
     {0x38, nullptr, "Unknown"},
     {0x39, nullptr, "Unknown"},