kern: SvcGetThreadCoreMask, SvcSetThreadCoreMask

libraries/libmesosphere/include/mesosphere/kern_k_thread.hpp

@@ -296,7 +296,11 @@ namespace ams::kern {
 
             constexpr KThreadContext &GetContext() { return this->thread_context; }
             constexpr const KThreadContext &GetContext() const { return this->thread_context; }
+
             constexpr const KAffinityMask &GetAffinityMask() const { return this->affinity_mask; }
+            Result GetCoreMask(int32_t *out_ideal_core, u64 *out_affinity_mask);
+            Result SetCoreMask(int32_t ideal_core, u64 affinity_mask);
+
             constexpr ThreadState GetState() const { return static_cast<ThreadState>(this->thread_state & ThreadState_Mask); }
             constexpr ThreadState GetRawState() const { return this->thread_state; }
             NOINLINE void SetState(ThreadState state);

libraries/libmesosphere/source/kern_k_thread.cpp

@@ -409,7 +409,7 @@ namespace ams::kern {
 
             /* Restore our ideals. */
             this->ideal_core_id = this->original_ideal_core_id;
-            this->original_affinity_mask = this->affinity_mask;
+            this->affinity_mask = this->original_affinity_mask;
 
             if (this->affinity_mask.GetAffinityMask() != old_mask.GetAffinityMask()) {
                 const s32 active_core = this->GetActiveCore();
@@ -426,6 +426,74 @@ namespace ams::kern {
         }
     }
 
+    Result KThread::GetCoreMask(int32_t *out_ideal_core, u64 *out_affinity_mask) {
+        MESOSPHERE_ASSERT_THIS();
+        {
+            KScopedSchedulerLock sl;
+            MESOSPHERE_ASSERT(this->num_core_migration_disables >= 0);
+
+            /* Select between core mask and original core mask. */
+            if (this->num_core_migration_disables == 0) {
+                *out_ideal_core    = this->ideal_core_id;
+                *out_affinity_mask = this->affinity_mask.GetAffinityMask();
+            } else {
+                *out_ideal_core    = this->original_ideal_core_id;
+                *out_affinity_mask = this->original_affinity_mask.GetAffinityMask();
+            }
+        }
+
+        return ResultSuccess();
+    }
+
+    Result KThread::SetCoreMask(int32_t ideal_core, u64 affinity_mask) {
+        MESOSPHERE_ASSERT_THIS();
+        MESOSPHERE_ASSERT(this->parent != nullptr);
+        MESOSPHERE_ASSERT(affinity_mask != 0);
+        {
+            KScopedSchedulerLock sl;
+            MESOSPHERE_ASSERT(this->num_core_migration_disables >= 0);
+
+            /* If the core id is no-update magic, preserve the ideal core id. */
+            if (ideal_core == ams::svc::IdealCoreNoUpdate) {
+                if (this->num_core_migration_disables == 0) {
+                    ideal_core = this->ideal_core_id;
+                } else {
+                    ideal_core = this->original_ideal_core_id;
+                }
+
+                R_UNLESS(((1ul << ideal_core) & affinity_mask) != 0, svc::ResultInvalidCombination());
+            }
+
+            /* If we haven't disabled migration, perform an affinity change. */
+            if (this->num_core_migration_disables == 0) {
+                const KAffinityMask old_mask = this->affinity_mask;
+
+                /* Set our new ideals. */
+                this->ideal_core_id = ideal_core;
+                this->affinity_mask.SetAffinityMask(affinity_mask);
+
+                if (this->affinity_mask.GetAffinityMask() != old_mask.GetAffinityMask()) {
+                    const s32 active_core = this->GetActiveCore();
+
+                    if (active_core >= 0) {
+                        if (!this->affinity_mask.GetAffinity(active_core)) {
+                            this->SetActiveCore(this->ideal_core_id);
+                        } else {
+                            this->SetActiveCore(BITSIZEOF(unsigned long long) - 1 - __builtin_clzll(this->affinity_mask.GetAffinityMask()));
+                        }
+                    }
+                    KScheduler::OnThreadAffinityMaskChanged(this, old_mask, active_core);
+                }
+            } else {
+                /* Otherwise, we edit the original affinity for restoration later. */
+                this->original_ideal_core_id = ideal_core;
+                this->original_affinity_mask.SetAffinityMask(affinity_mask);
+            }
+        }
+
+        return ResultSuccess();
+    }
+
     void KThread::SetBasePriority(s32 priority) {
         MESOSPHERE_ASSERT_THIS();
         MESOSPHERE_ASSERT(ams::svc::HighestThreadPriority <= priority && priority <= ams::svc::LowestThreadPriority);

libraries/libmesosphere/source/svc/kern_svc_thread.cpp

@@ -145,6 +145,46 @@ namespace ams::kern::svc {
             return ResultSuccess();
         }
 
+        Result GetThreadCoreMask(int32_t *out_core_id, uint64_t *out_affinity_mask, ams::svc::Handle thread_handle) {
+            /* Get the thread from its handle. */
+            KScopedAutoObject thread = GetCurrentProcess().GetHandleTable().GetObject<KThread>(thread_handle);
+            R_UNLESS(thread.IsNotNull(), svc::ResultInvalidHandle());
+
+            /* Get the core mask. */
+            R_TRY(thread->GetCoreMask(out_core_id, out_affinity_mask));
+
+            return ResultSuccess();
+        }
+
+        Result SetThreadCoreMask(ams::svc::Handle thread_handle, int32_t core_id, uint64_t affinity_mask) {
+            /* Determine the core id/affinity mask. */
+            if (core_id == ams::svc::IdealCoreUseProcessValue) {
+                core_id       = GetCurrentProcess().GetIdealCoreId();
+                affinity_mask = (1ul << core_id);
+            } else {
+                /* Validate the affinity mask. */
+                const u64 process_core_mask = GetCurrentProcess().GetCoreMask();
+                R_UNLESS((affinity_mask | process_core_mask) == process_core_mask, svc::ResultInvalidCoreId());
+                R_UNLESS(affinity_mask != 0,                                       svc::ResultInvalidCombination());
+
+                /* Validate the core id. */
+                if (IsValidCoreId(core_id)) {
+                    R_UNLESS(((1ul << core_id) & affinity_mask) != 0, svc::ResultInvalidCombination());
+                } else {
+                    R_UNLESS(core_id == ams::svc::IdealCoreNoUpdate || core_id == ams::svc::IdealCoreDontCare, svc::ResultInvalidCoreId());
+                }
+            }
+
+            /* Get the thread from its handle. */
+            KScopedAutoObject thread = GetCurrentProcess().GetHandleTable().GetObject<KThread>(thread_handle);
+            R_UNLESS(thread.IsNotNull(), svc::ResultInvalidHandle());
+
+            /* Set the core mask. */
+            R_TRY(thread->SetCoreMask(core_id, affinity_mask));
+
+            return ResultSuccess();
+        }
+
         Result GetThreadId(uint64_t *out_thread_id, ams::svc::Handle thread_handle) {
             /* Get the thread from its handle. */
             KScopedAutoObject thread = GetCurrentProcess().GetHandleTable().GetObject<KThread>(thread_handle);
@@ -184,11 +224,11 @@ namespace ams::kern::svc {
     }
 
     Result GetThreadCoreMask64(int32_t *out_core_id, uint64_t *out_affinity_mask, ams::svc::Handle thread_handle) {
-        MESOSPHERE_PANIC("Stubbed SvcGetThreadCoreMask64 was called.");
+        return GetThreadCoreMask(out_core_id, out_affinity_mask, thread_handle);
     }
 
     Result SetThreadCoreMask64(ams::svc::Handle thread_handle, int32_t core_id, uint64_t affinity_mask) {
-        MESOSPHERE_PANIC("Stubbed SvcSetThreadCoreMask64 was called.");
+        return SetThreadCoreMask(thread_handle, core_id, affinity_mask);
     }
 
     Result GetThreadId64(uint64_t *out_thread_id, ams::svc::Handle thread_handle) {
@@ -238,11 +278,11 @@ namespace ams::kern::svc {
     }
 
     Result GetThreadCoreMask64From32(int32_t *out_core_id, uint64_t *out_affinity_mask, ams::svc::Handle thread_handle) {
-        MESOSPHERE_PANIC("Stubbed SvcGetThreadCoreMask64From32 was called.");
+        return GetThreadCoreMask(out_core_id, out_affinity_mask, thread_handle);
     }
 
     Result SetThreadCoreMask64From32(ams::svc::Handle thread_handle, int32_t core_id, uint64_t affinity_mask) {
-        MESOSPHERE_PANIC("Stubbed SvcSetThreadCoreMask64From32 was called.");
+        return SetThreadCoreMask(thread_handle, core_id, affinity_mask);
     }
 
     Result GetThreadId64From32(uint64_t *out_thread_id, ams::svc::Handle thread_handle) {