cpu: preserve full thread affinity on Windows and Linux (#626)

Windows: pool workers pin with SetThreadGroupAffinity (the processor group was dropped: a 3990X's 63 workers
collapsed to 32 group-local masks); the host thread selects a CPU Set and restores the previous selection.
Linux: the complete, dynamically sized affinity mask is kept and restored (CPU IDs past 63).
pool_affinity_test (Windows) and pool_affinity_linux_test. Squashed from the PR's 8 commits (net diff).

CMakeLists.txt

@@ -1062,6 +1062,15 @@ if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_
 endif()
 
 if(STRATA_BUILD_TESTS)
+  if(WIN32)
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    add_executable(pool_affinity_test src/kernels/cpu/pool_affinity_linux_test.cpp)
+    target_link_libraries(pool_affinity_test PRIVATE strata_kernels_cpu)
+    add_test(NAME pool_affinity_test COMMAND pool_affinity_test)
+  endif()
   # Plan v0.3 P6: one expert applied to several tokens, bitwise vs the single-token kernel (synthetic, CPU only).
   # plan v0.3 P5: the device dequantizers vs the artifact's CPU dequantizers on the real shards (not a ctest: needs
   # the 38 GB model; run it by hand or from tools/gpu_session.py)

include/strata/core/session.hpp

@@ -16,6 +16,7 @@
 
 #include "strata/core/hit_hook.hpp"
 #include "strata/core/layer.hpp"
+#include "strata/kernels/cpu/pool.hpp"
 
 #include <cuda_runtime.h>
 
@@ -369,7 +370,7 @@ struct SessionLoopScratch {
     float* y_miss = nullptr;        ///< pinned host staging for the pool's answer, `parts_bytes` long
     size_t parts_bytes = 0;
     cudaEvent_t probe = nullptr;
-    long long pinned_core = -1;     ///< the affinity to restore, or -1 if the host was never pinned
+    strata::kernels::cpu::ThreadAffinity pinned_core{};  ///< affinity to restore, or invalid if the host was not pinned
     bool pinned = false;
 
     /// Allocates the buffers and pins the host thread.  Call ONCE, at session setup.

include/strata/kernels/cpu/pool.hpp

@@ -76,7 +76,7 @@ struct CpuTopology {
     int p_cores = 0;                ///< Physical performance cores
     int p_threads = 0;              ///< Total logical threads on performance cores
     int e_cores = 0;                ///< Efficient cores
-    std::vector<int> worker_cores;  ///< Ordered logical core IDs for workers (excluding host if skip_first)
+    std::vector<int> worker_cores;  ///< Ordered CPU IDs; on Windows, group * 64 + processor within the group
     int host_core = -1;             ///< Logical core reserved for host thread
 };
 
@@ -100,9 +100,23 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity = PoolAff
 /// - exactly 5/6 of L9's 44.14 on 6 - and at **26.9 GB/s inside the host loop**, where the unpinned spinning
 /// host is free to land on a worker's core or its SMT sibling.  That 1.35x is not the kernel.
 ///
-/// Returns the PREVIOUS affinity mask, or -1 if the platform refused; pass it to `restore_thread_affinity`.
-long long pin_current_thread(int core);
-void restore_thread_affinity(long long previous);
+/// The previous host placement. `valid` is false when querying or setting placement failed.
+/// Windows uses a reversible CPU Set selection, leaving hard affinity (including Windows 11's implicit
+/// all-group eligibility) untouched. An empty selection restores inheritance from the process defaults.
+/// Linux stores the complete, dynamically sized native CPU mask, including CPU IDs above 63.
+struct ThreadAffinity {
+#if defined(_WIN32)
+    std::vector<unsigned long> cpu_sets;
+#else
+    std::vector<unsigned long> mask;
+#endif
+    bool valid = false;
+};
+
+/// Selects one encoded processor for the caller; Windows CPU Sets respect existing hard-affinity limits.
+/// Restore on the same calling thread. Pool-owned workers use hard group affinity instead.
+ThreadAffinity pin_current_thread(int core);
+void restore_thread_affinity(const ThreadAffinity& previous);
 
 class ExpertPool {
 public:

src/core/session.cpp

@@ -547,7 +547,7 @@ bool SessionLoopScratch::init(size_t parts_bytes_in, std::string& err) {
     const std::vector<int> cores = strata::kernels::cpu::physical_cores(false);
     if (!cores.empty()) {
         pinned_core = strata::kernels::cpu::pin_current_thread(cores[0]);
-        pinned = true;
+        pinned = pinned_core.valid;
     }
     return true;
 }
@@ -558,7 +558,7 @@ void SessionLoopScratch::free() {
     if (pinned) {
         strata::kernels::cpu::restore_thread_affinity(pinned_core);
         pinned = false;
-        pinned_core = -1;
+        pinned_core = {};
     }
     if (probe != nullptr) { cudaEventDestroy(probe); probe = nullptr; }
     if (y_miss != nullptr) { cudaFreeHost(y_miss); y_miss = nullptr; }

src/kernels/cpu/pool.cpp

@@ -14,9 +14,11 @@
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
+#include "pool_affinity_win.hpp"
 #else
 #include <pthread.h>
 #include <sched.h>
+#include "pool_affinity_linux.hpp"
 #endif
 
 namespace strata::kernels::cpu {
@@ -36,7 +38,13 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     DWORD len = 0;
     GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &len);
     if (len == 0) {
-        for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) topo.worker_cores.push_back((int) i);
+        // Preserve group identity even when detailed core topology is unavailable.
+        const WORD groups = GetActiveProcessorGroupCount();
+        for (WORD group = 0; group < groups; ++group) {
+            const DWORD count = GetActiveProcessorCount(group);
+            if (count == 0 || count == (DWORD) -1 || count > 64) continue;
+            for (DWORD i = 0; i < count; ++i) topo.worker_cores.push_back((int) group * 64 + (int) i);
+        }
         if (skip_first && !topo.worker_cores.empty()) {
             topo.host_core = topo.worker_cores.front();
             topo.worker_cores.erase(topo.worker_cores.begin());
@@ -60,10 +68,12 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
                 CoreDesc cd;
                 cd.efficiency = e->Processor.EfficiencyClass;
                 cd.has_smt = (e->Processor.Flags & LTP_PC_SMT) != 0;
-                const GROUP_AFFINITY& g = e->Processor.GroupMask[0];
-                for (int bit = 0; bit < 64; ++bit) {
-                    if (g.Mask & (1ull << bit)) {
-                        cd.lps.push_back((int) (g.Group * 64 + bit));
+                for (WORD group = 0; group < e->Processor.GroupCount; ++group) {
+                    const GROUP_AFFINITY& g = e->Processor.GroupMask[group];
+                    for (int bit = 0; bit < 64; ++bit) {
+                        if (g.Mask & (KAFFINITY(1) << bit)) {
+                            cd.lps.push_back((int) (g.Group * 64 + bit));
+                        }
                     }
                 }
                 if (!cd.lps.empty()) {
@@ -168,12 +178,8 @@ CpuTopology detect_cpu_topology(bool skip_first, PoolAffinity affinity) {
     };
 
     std::vector<int> allowed;
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    if (sched_getaffinity(0, sizeof set, &set) == 0) {
-        for (int i = 0; i < CPU_SETSIZE; ++i)
-            if (CPU_ISSET(i, &set)) allowed.push_back(i);
-    } else {
+    std::vector<unsigned long> allowed_mask;
+    if (detail::get_thread_affinity(allowed_mask, &allowed) != 0) {
         for (unsigned i = 0; i < std::thread::hardware_concurrency(); ++i) allowed.push_back((int) i);
     }
 
@@ -273,49 +279,60 @@ std::vector<int> physical_cores(bool skip_first, PoolAffinity affinity) {
 
 namespace {
 
-void pin_this_thread(int core) {
-    if (core < 0) return;
+bool pin_this_thread(int core, [[maybe_unused]] int worker = -1) {
+    if (core < 0) return false;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
+    return detail::set_thread_group_affinity(core, worker);
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    CPU_SET(core, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::pin_thread_to_cpu(core);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: affinity for worker %d (CPU %d) failed: %d; previous affinity kept\n",
+                     worker, core, error);
+    return error == 0;
 #endif
 }
 
 }  // namespace
 
-long long pin_current_thread(int core) {
-    if (core < 0) return -1;
+ThreadAffinity pin_current_thread(int core) {
+    if (core < 0) return {};
 #if defined(_WIN32)
-    // `SetThreadAffinityMask` RETURNS the previous mask, or 0 on failure - so 0 doubles as the error, which is
-    // why the caller must not treat it as a restorable value.
-    const DWORD_PTR prev = SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) 1 << (core & 63));
-    return prev == 0 ? -1 : (long long) prev;
+    ThreadAffinity previous;
+    ULONG target = 0;
+    if (!detail::get_thread_cpu_sets(previous.cpu_sets) || !detail::cpu_set_for_core(core, target) ||
+        !SetThreadSelectedCpuSets(GetCurrentThread(), &target, 1)) {
+        std::fprintf(stderr, "strata cpu pool: host CPU Set selection for processor %d failed: %lu; previous placement kept\n",
+                     core, (unsigned long) GetLastError());
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #else
-    cpu_set_t prev;
-    CPU_ZERO(&prev);
-    if (pthread_getaffinity_np(pthread_self(), sizeof prev, &prev) != 0) return -1;
-    unsigned long mask = 0;
-    for (int i = 0; i < CPU_SETSIZE && i < 64; ++i)
-        if (CPU_ISSET(i, &prev)) mask |= 1ul << i;
-    pin_this_thread(core);
-    return (long long) mask;
+    ThreadAffinity previous;
+    int error = detail::get_thread_affinity(previous.mask);
+    if (error == 0) error = detail::pin_thread_to_cpu(core);
+    if (error != 0) {
+        std::fprintf(stderr, "strata cpu pool: host affinity for CPU %d failed: %d; previous affinity kept\n", core, error);
+        return {};
+    }
+    previous.valid = true;
+    return previous;
 #endif
 }
 
-void restore_thread_affinity(long long previous) {
-    if (previous <= 0) return;
+void restore_thread_affinity(const ThreadAffinity& previous) {
+    if (!previous.valid) return;
 #if defined(_WIN32)
-    SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR) previous);
+    // Clearing an originally empty selection restores process-default/all-group eligibility without
+    // turning the caller's implicit Windows 11 affinity into an explicit single-group hard mask.
+    if (!SetThreadSelectedCpuSets(GetCurrentThread(), previous.cpu_sets.empty() ? nullptr : previous.cpu_sets.data(),
+                                 (ULONG) previous.cpu_sets.size()))
+        std::fprintf(stderr, "strata cpu pool: host CPU Set restoration failed: %lu\n",
+                     (unsigned long) GetLastError());
 #else
-    cpu_set_t set;
-    CPU_ZERO(&set);
-    for (int i = 0; i < 64; ++i)
-        if ((previous >> i) & 1) CPU_SET(i, &set);
-    pthread_setaffinity_np(pthread_self(), sizeof set, &set);
+    const int error = detail::set_thread_affinity(previous.mask);
+    if (error != 0)
+        std::fprintf(stderr, "strata cpu pool: host affinity restoration failed: %d\n", error);
 #endif
 }
 
@@ -375,7 +392,7 @@ ExpertPool::ExpertPool(int n_workers, bool pin, bool host_works, PoolAffinity af
     for (int i = 0; i < n_; ++i) {
         const int core = pin ? (i < (int) topo_.worker_cores.size() ? topo_.worker_cores[(size_t) i] : -1) : -1;
         threads_.emplace_back([this, i, core] {
-            pin_this_thread(core);
+            pin_this_thread(core, i);
             worker(i);
         });
     }

src/kernels/cpu/pool_affinity_linux.hpp

@@ -0,0 +1,69 @@
+#pragma once
+
+#if defined(__linux__)
+#include <pthread.h>
+#include <sched.h>
+
+#include <cerrno>
+#include <climits>
+#include <cstring>
+#include <vector>
+
+namespace strata::kernels::cpu::detail {
+
+// CPU_ALLOC and the sized macros also handle kernels whose CPU mask exceeds CPU_SETSIZE.
+// Keep the native mask opaque: copy its bytes for storage and use CPU_*_S to inspect it.
+class LinuxCpuSet {
+public:
+    explicit LinuxCpuSet(int cpus) : bytes(CPU_ALLOC_SIZE(cpus)), set(CPU_ALLOC(cpus)) {
+        if (set) CPU_ZERO_S(bytes, set);
+    }
+    ~LinuxCpuSet() { CPU_FREE(set); }
+    LinuxCpuSet(const LinuxCpuSet&) = delete;
+    LinuxCpuSet& operator=(const LinuxCpuSet&) = delete;
+
+    const size_t bytes;
+    cpu_set_t* const set;
+};
+
+// pthread affinity functions return an error number directly; they do not set errno.
+// A too-small query buffer returns EINVAL, even when the number of online CPUs is small.
+inline int get_thread_affinity(std::vector<unsigned long>& mask, std::vector<int>* cpus = nullptr) {
+    for (int count = CPU_SETSIZE;;) {
+        LinuxCpuSet current(count);
+        if (!current.set) return ENOMEM;
+        const int error = pthread_getaffinity_np(pthread_self(), current.bytes, current.set);
+        if (error == 0) {
+            mask.resize(current.bytes / sizeof(unsigned long));
+            std::memcpy(mask.data(), current.set, current.bytes);
+            if (cpus) {
+                cpus->clear();
+                for (int cpu = 0; cpu < count; ++cpu)
+                    if (CPU_ISSET_S(cpu, current.bytes, current.set)) cpus->push_back(cpu);
+            }
+            return 0;
+        }
+        if (error != EINVAL || count > INT_MAX / 2) return error;
+        count *= 2;
+    }
+}
+
+inline int set_thread_affinity(const std::vector<unsigned long>& mask) {
+    constexpr size_t word_bits = sizeof(unsigned long) * CHAR_BIT;
+    if (mask.empty() || mask.size() > INT_MAX / word_bits) return EINVAL;
+    LinuxCpuSet target(static_cast<int>(mask.size() * word_bits));
+    if (!target.set) return ENOMEM;
+    std::memcpy(target.set, mask.data(), target.bytes);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+inline int pin_thread_to_cpu(int core) {
+    if (core < 0 || core == INT_MAX) return EINVAL;
+    LinuxCpuSet target(core < CPU_SETSIZE ? CPU_SETSIZE : core + 1);
+    if (!target.set) return ENOMEM;
+    CPU_SET_S(core, target.bytes, target.set);
+    return pthread_setaffinity_np(pthread_self(), target.bytes, target.set);
+}
+
+}  // namespace strata::kernels::cpu::detail
+#endif

src/kernels/cpu/pool_affinity_linux_test.cpp

@@ -0,0 +1,290 @@
+#ifndef _GNU_SOURCE
+#define _GNU_SOURCE
+#endif
+
+#include "strata/kernels/cpu/pool.hpp"
+
+#include <algorithm>
+#include <cerrno>
+#include <climits>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include <thread>
+#include <vector>
+
+#if defined(__linux__)
+#include "pool_affinity_linux.hpp"
+
+#include <pthread.h>
+#include <sched.h>
+#include <unistd.h>
+
+namespace cpu = strata::kernels::cpu;
+
+class DynamicCpuSet {
+public:
+    explicit DynamicCpuSet(size_t count)
+        : count_(count), bytes_(CPU_ALLOC_SIZE(count)), set_(CPU_ALLOC(count)) {
+        if (set_) CPU_ZERO_S(bytes_, set_);
+    }
+
+    ~DynamicCpuSet() {
+        if (set_) CPU_FREE(set_);
+    }
+
+    DynamicCpuSet(const DynamicCpuSet&) = delete;
+    DynamicCpuSet& operator=(const DynamicCpuSet&) = delete;
+
+    cpu_set_t* get() { return set_; }
+    const cpu_set_t* get() const { return set_; }
+    size_t bytes() const { return bytes_; }
+    size_t count() const { return count_; }
+    explicit operator bool() const { return set_ != nullptr; }
+
+private:
+    size_t count_;
+    size_t bytes_;
+    cpu_set_t* set_;
+};
+
+static size_t initial_cpu_count() {
+    const long configured = sysconf(_SC_NPROCESSORS_CONF);
+    return (std::max)(size_t(128), configured > 0 ? (size_t) configured : size_t(0));
+}
+
+// Query with CPU_ALLOC rather than cpu_set_t: this must observe CPUs beyond CPU_SETSIZE too.
+static bool get_affinity(std::vector<int>& ids) {
+    size_t count = initial_cpu_count();
+    for (int attempt = 0; attempt < 8; ++attempt) {
+        DynamicCpuSet set(count);
+        if (!set) return false;
+        const int error = pthread_getaffinity_np(pthread_self(), set.bytes(), set.get());
+        if (error == 0) {
+            ids.clear();
+            for (size_t i = 0; i < set.count(); ++i)
+                if (CPU_ISSET_S(i, set.bytes(), set.get())) ids.push_back((int) i);
+            return !ids.empty();
+        }
+        if (error != EINVAL || count > (size_t(1) << 20)) {
+            std::fprintf(stderr, "pthread_getaffinity_np failed: %s\n", std::strerror(error));
+            return false;
+        }
+        count *= 2;
+    }
+    std::fprintf(stderr, "pthread_getaffinity_np still needs a larger CPU set\n");
+    return false;
+}
+
+static bool set_affinity(const std::vector<int>& ids) {
+    if (ids.empty()) return false;
+    const int highest = *std::max_element(ids.begin(), ids.end());
+    if (highest < 0) return false;
+    const size_t count = (std::max)(initial_cpu_count(), (size_t) highest + 1);
+    DynamicCpuSet set(count);
+    if (!set) return false;
+    for (int id : ids) {
+        if (id < 0 || (size_t) id >= count) return false;
+        CPU_SET_S((size_t) id, set.bytes(), set.get());
+    }
+    const int error = pthread_setaffinity_np(pthread_self(), set.bytes(), set.get());
+    if (error != 0) {
+        std::fprintf(stderr, "pthread_setaffinity_np failed: %s\n", std::strerror(error));
+        return false;
+    }
+    return true;
+}
+
+class OriginalAffinity {
+public:
+    bool capture() { return get_affinity(ids_); }
+    const std::vector<int>& ids() const { return ids_; }
+    bool restore() { return ids_.empty() || set_affinity(ids_); }
+    ~OriginalAffinity() { if (!ids_.empty()) set_affinity(ids_); }
+
+private:
+    std::vector<int> ids_;
+};
+
+static bool expect_affinity(const std::vector<int>& expected, const char* label) {
+    std::vector<int> actual;
+    if (!get_affinity(actual)) {
+        std::fprintf(stderr, "Linux affinity test: %s: cannot query current affinity\n", label);
+        return false;
+    }
+    if (actual != expected) {
+        std::fprintf(stderr, "Linux affinity test: %s: expected %zu CPUs, observed %zu\n",
+                     label, expected.size(), actual.size());
+        return false;
+    }
+    return true;
+}
+
+static bool round_trip_full_mask(const std::vector<int>& allowed) {
+    const int target = allowed.front();
+    const auto previous = cpu::pin_current_thread(target);
+    const bool pinned = previous.valid && expect_affinity({target}, "pin from full allowed mask");
+    cpu::restore_thread_affinity(previous);
+    return pinned && expect_affinity(allowed, "restore full allowed mask");
+}
+
+static bool topology_respects_current_mask() {
+    std::vector<int> allowed;
+    if (!get_affinity(allowed)) return false;
+    const auto topology = cpu::detect_cpu_topology(false);
+    if (topology.worker_cores.empty()) return false;
+    for (int core : topology.worker_cores) {
+        if (!std::binary_search(allowed.begin(), allowed.end(), core)) {
+            std::fprintf(stderr, "Linux affinity test: topology selected disallowed CPU %d\n", core);
+            return false;
+        }
+    }
+    return true;
+}
+
+static bool round_trip_sparse_nested(const std::vector<int>& allowed) {
+    if (allowed.size() < 2) {
+        std::puts("SKIP sparse/nested affinity: fewer than two allowed CPUs");
+        return true;
+    }
+    const std::vector<int> sparse{allowed.front(), allowed.back()};
+    if (!set_affinity(sparse) || !expect_affinity(sparse, "establish sparse mask") ||
+        !topology_respects_current_mask()) return false;
+
+    const auto outer = cpu::pin_current_thread(sparse.front());
+    const bool outer_pinned = outer.valid && expect_affinity({sparse.front()}, "outer pin");
+    if (!outer_pinned) {
+        cpu::restore_thread_affinity(outer);
+        return false;
+    }
+
+    const auto inner = cpu::pin_current_thread(sparse.back());
+    const bool inner_pinned = inner.valid && expect_affinity({sparse.back()}, "inner pin");
+    cpu::restore_thread_affinity(inner);
+    const bool outer_restored = expect_affinity({sparse.front()}, "restore inner pin to outer pin");
+    cpu::restore_thread_affinity(outer);
+    const bool sparse_restored = expect_affinity(sparse, "restore outer pin to sparse mask");
+    const bool full_restored = set_affinity(allowed) && expect_affinity(allowed, "restore full mask after sparse test");
+    return inner_pinned && outer_restored && sparse_restored && full_restored;
+}
+
+static bool round_trip_high_cpu(const std::vector<int>& allowed, bool require_high) {
+    const auto it = std::find_if(allowed.begin(), allowed.end(), [](int id) { return id >= 64; });
+    if (it == allowed.end()) {
+        if (require_high) {
+            std::fprintf(stderr, "Linux affinity test: --require-high-cpu requested, but no allowed CPU ID is >= 64\n");
+            return false;
+        }
+        std::puts("SKIP high-CPU affinity path: no allowed CPU ID is >= 64 (use --require-high-cpu to require it)");
+        return true;
+    }
+    const int target = allowed.front() == *it ? (allowed.size() > 1 ? allowed.back() : *it) : allowed.front();
+    if (target == *it) {
+        if (require_high) {
+            std::fprintf(stderr, "Linux affinity test: --require-high-cpu needs at least two allowed CPUs to test restoration\n");
+            return false;
+        }
+        std::puts("SKIP high-CPU affinity round trip: only one CPU is allowed, so no meaningful pin transition is possible");
+        return true;
+    }
+
+    // Restrict the caller to one high-numbered CPU, then pin to a different originally allowed CPU.
+    // Restoring the high-only mask catches truncation even when the original process mask is sparse.
+    const std::vector<int> high_only{*it};
+    if (!set_affinity(high_only) || !expect_affinity(high_only, "establish high-only mask") ||
+        !topology_respects_current_mask()) return false;
+    const auto previous = cpu::pin_current_thread(target);
+    const bool pinned = previous.valid && expect_affinity({target}, "pin away from high-only mask");
+    cpu::restore_thread_affinity(previous);
+    const bool restored = pinned && expect_affinity(high_only, "restore high-only mask");
+    if (restored) std::printf("High-CPU affinity round trip: saved CPU %d, pin target %d\n", *it, target);
+    return restored;
+}
+
+static bool negative_cpu_is_noop() {
+    std::vector<int> before;
+    if (!get_affinity(before)) return false;
+    const auto invalid = cpu::pin_current_thread(-1);
+    std::vector<int> after;
+    const bool unchanged = get_affinity(after) && before == after;
+    cpu::restore_thread_affinity(invalid);
+    return !invalid.valid && unchanged && expect_affinity(before, "negative CPU no-op");
+}
+
+static bool worker_thread_can_pin(int target) {
+    bool ok = false;
+    std::thread worker([&] {
+        std::vector<int> before;
+        if (!get_affinity(before)) return;
+        std::vector<unsigned long> saved;
+        if (cpu::detail::get_thread_affinity(saved) != 0) return;
+        const int pin_error = cpu::detail::pin_thread_to_cpu(target);
+        const bool pinned = pin_error == 0 && expect_affinity({target}, "worker thread helper pin") &&
+                            sched_getcpu() == target;
+        const int restore_error = cpu::detail::set_thread_affinity(saved);
+        ok = pinned && restore_error == 0 && expect_affinity(before, "worker thread helper restore");
+    });
+    worker.join();
+    return ok;
+}
+
+static bool failed_pin_preserves_affinity() {
+    std::vector<int> before;
+    if (!get_affinity(before)) return false;
+    // An ID outside the caller's current mask may still be online and legal, so it is not a
+    // reliable failure case. INT_MAX is rejected before allocating or changing any CPU set.
+    const int error = cpu::detail::pin_thread_to_cpu(INT_MAX);
+    std::vector<int> after;
+    const bool helper_unchanged = get_affinity(after) && before == after;
+    if (error == 0) {
+        std::fprintf(stderr, "Linux affinity test: pin_thread_to_cpu(INT_MAX) unexpectedly succeeded\n");
+        return false;
+    }
+    const auto invalid = cpu::pin_current_thread(INT_MAX);
+    cpu::restore_thread_affinity(invalid);
+    std::vector<int> after_public;
+    const bool public_unchanged = get_affinity(after_public) && before == after_public;
+    if (!helper_unchanged || invalid.valid || !public_unchanged)
+        std::fprintf(stderr, "Linux affinity test: failed pin changed the observed affinity or returned a valid snapshot\n");
+    return helper_unchanged && !invalid.valid && public_unchanged;
+}
+
+int main(int argc, char** argv) {
+    bool require_high = false;
+    for (int i = 1; i < argc; ++i) {
+        if (std::strcmp(argv[i], "--require-high-cpu") == 0) require_high = true;
+        else {
+            std::fprintf(stderr, "Unknown argument: %s\n", argv[i]);
+            return 2;
+        }
+    }
+
+    OriginalAffinity original;
+    if (!original.capture()) {
+        std::fprintf(stderr, "Linux affinity test: cannot capture original thread affinity\n");
+        return 1;
+    }
+    const auto allowed = original.ids();
+
+    bool ok = round_trip_full_mask(allowed) && worker_thread_can_pin(allowed.back());
+    if (ok) ok = round_trip_sparse_nested(allowed);
+    if (ok) ok = round_trip_high_cpu(allowed, require_high);
+    if (ok) ok = negative_cpu_is_noop();
+    if (ok) ok = failed_pin_preserves_affinity();
+    const bool restored = original.restore() && expect_affinity(allowed, "final original affinity cleanup");
+    if (!ok || !restored) return 1;
+
+    std::printf("Linux affinity test: full, sparse/nested, worker helper, invalid-target, failed-pin, and cleanup checks passed (%zu allowed CPUs)\n",
+                allowed.size());
+    return 0;
+}
+
+#else
+
+int main() {
+    std::puts("SKIP Linux affinity test: Linux is required");
+    return 0;
+}
+
+#endif

src/kernels/cpu/pool_affinity_test.cpp

@@ -0,0 +1,148 @@
+#include "strata/kernels/cpu/pool.hpp"
+#include "pool_affinity_win.hpp"
+
+#include <algorithm>
+#include <cstdio>
+#include <thread>
+#include <vector>
+
+namespace cpu = strata::kernels::cpu;
+
+static bool check(bool ok, const char* message) {
+    if (!ok) std::fprintf(stderr, "Windows affinity test: %s (error %lu)\n", message, (unsigned long) GetLastError());
+    return ok;
+}
+
+static bool select(const std::vector<ULONG>& ids) {
+    return SetThreadSelectedCpuSets(GetCurrentThread(), ids.empty() ? nullptr : ids.data(), (ULONG) ids.size()) != 0;
+}
+
+static bool selected_is(std::vector<ULONG> expected) {
+    std::vector<ULONG> actual;
+    if (!cpu::detail::get_thread_cpu_sets(actual)) return false;
+    std::sort(actual.begin(), actual.end());
+    std::sort(expected.begin(), expected.end());
+    return actual == expected;
+}
+
+// Establish reachability before the round trip: CPU Sets are soft affinity and cannot override an
+// existing job/hard-affinity restriction. Afterwards, require the same observable reachability.
+static bool observe(int core) {
+    for (int i = 0; i < 200; ++i) {
+        Sleep(1);
+        PROCESSOR_NUMBER current{};
+        GetCurrentProcessorNumberEx(&current);
+        if ((int) current.Group * 64 + current.Number == core) return true;
+    }
+    return false;
+}
+
+static bool host_round_trips(const std::vector<int>& cores, const std::vector<ULONG>& ids) {
+    if (!check(select({}), "clear initial thread CPU Sets")) return false;
+    std::vector<bool> reachable;
+    for (size_t i = 0; i < cores.size(); ++i) {
+        if (!check(select({ids[i]}), "select baseline CPU Set")) return false;
+        reachable.push_back(observe(cores[i]));
+    }
+    if (!check(select({}), "clear baseline CPU Sets")) return false;
+
+    for (size_t i = 0; i < cores.size(); ++i) {
+        const auto previous = cpu::pin_current_thread(cores[i]);
+        const bool pinned = previous.valid && previous.cpu_sets.empty() && selected_is({ids[i]});
+        cpu::restore_thread_affinity(previous);
+        if (!check(pinned && selected_is({}), "restore empty thread assignment")) return false;
+    }
+
+    // GetThreadGroupAffinity alone misses a lost implicit all-group state. Selecting each CPU Set
+    // again must still permit execution there, without ever setting a hard group affinity.
+    for (size_t i = 0; i < cores.size(); ++i) {
+        if (!check(select({ids[i]}), "select CPU Set after restore")) return false;
+        if (reachable[i]) {
+            if (!check(observe(cores[i]), "lost CPU/group eligibility after host restore")) return false;
+            std::printf("Host restore: still executes on group %d processor %d\n", cores[i] / 64, cores[i] & 63);
+        } else {
+            std::printf("SKIP scheduling observation for group %d: not reachable before test\n", cores[i] / 64);
+        }
+    }
+
+    // A nonempty selection may span groups. Nested host selections must restore the entire set.
+    if (!check(select(ids), "set existing multi-CPU selection")) return false;
+    const auto outer = cpu::pin_current_thread(cores.front());
+    const auto inner = cpu::pin_current_thread(cores.back());
+    const bool inner_ok = inner.valid && selected_is({ids.back()});
+    cpu::restore_thread_affinity(inner);
+    const bool outer_ok = outer.valid && selected_is({ids.front()});
+    cpu::restore_thread_affinity(outer);
+    if (!check(inner_ok && outer_ok && selected_is(ids), "restore nested / existing CPU Set selections")) return false;
+    const auto invalid = cpu::pin_current_thread(-1);
+    cpu::restore_thread_affinity(invalid);
+    if (!check(!invalid.valid && selected_is(ids), "invalid target must not mutate selection")) return false;
+    return check(select({}), "clear final host selection");
+}
+
+static bool inherited_default(int core, ULONG id, int temporary_core, ULONG temporary_id) {
+    ULONG count = 0;
+    if (!GetProcessDefaultCpuSets(GetCurrentProcess(), nullptr, 0, &count) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    std::vector<ULONG> saved(count);
+    if (count && !GetProcessDefaultCpuSets(GetCurrentProcess(), saved.data(), count, &count)) return false;
+    saved.resize(count);
+    if (!check(SetProcessDefaultCpuSets(GetCurrentProcess(), &id, 1) != 0, "set process default")) return false;
+    bool ok = false;
+    std::thread host([&] {
+        if (!select({})) return;
+        const bool reachable = observe(core);
+        const auto previous = cpu::pin_current_thread(temporary_core);
+        const bool overridden = previous.valid && selected_is({temporary_id});
+        cpu::restore_thread_affinity(previous);
+        ULONG inherited = 0, required = 0;
+        ok = overridden && previous.cpu_sets.empty() && selected_is({}) &&
+             GetProcessDefaultCpuSets(GetCurrentProcess(), &inherited, 1, &required) &&
+             required == 1 && inherited == id && (!reachable || observe(core));
+    });
+    host.join();
+    const bool restored = SetProcessDefaultCpuSets(GetCurrentProcess(), saved.empty() ? nullptr : saved.data(),
+                                                  (ULONG) saved.size()) != 0;
+    return check(ok && restored, "restore inherited process default");
+}
+
+int main() {
+    const auto all = cpu::physical_cores(false);
+    std::vector<int> cores;
+    std::vector<ULONG> ids;
+    for (int core : all) {
+        if (std::any_of(cores.begin(), cores.end(), [core](int other) { return other / 64 == core / 64; })) continue;
+        ULONG id = 0;
+        if (!cpu::detail::cpu_set_for_core(core, id)) continue;
+        cores.push_back(core);
+        ids.push_back(id);
+    }
+    if (!check(!cores.empty(), "no usable CPU Sets")) return 1;
+    if (!check(cores.size() == GetActiveProcessorGroupCount(), "CPU Sets do not cover all active groups")) return 1;
+
+    // Hard pins belong on pool-owned threads; never contaminate the caller under test with one.
+    for (size_t i = 0; i < cores.size(); ++i) {
+        bool ok = false;
+        std::thread worker([&] {
+            if (!cpu::detail::set_thread_group_affinity(cores[i], (int) i)) return;
+            GROUP_AFFINITY pinned{};
+            if (!GetThreadGroupAffinity(GetCurrentThread(), &pinned)) return;
+            ok = pinned.Group == cores[i] / 64 && pinned.Mask == (KAFFINITY(1) << (cores[i] & 63));
+            // Host placement must also preserve an existing restrictive hard affinity.
+            const auto previous = cpu::pin_current_thread(cores.back());
+            cpu::restore_thread_affinity(previous);
+            GROUP_AFFINITY after{};
+            ok = ok && previous.valid && GetThreadGroupAffinity(GetCurrentThread(), &after) &&
+                 after.Group == pinned.Group && after.Mask == pinned.Mask;
+        });
+        worker.join();
+        if (!check(ok, "worker hard pin / caller hard-affinity preservation")) return 1;
+    }
+
+    bool host_ok = false;
+    std::thread host([&] { host_ok = host_round_trips(cores, ids); });
+    host.join();
+    if (!host_ok || !inherited_default(cores.back(), ids.back(), cores.front(), ids.front())) return 1;
+    std::printf("Windows affinity test: worker pins, host CPU Sets and process-default restoration passed (%zu groups)\n", cores.size());
+    return 0;
+}

src/kernels/cpu/pool_affinity_win.hpp

@@ -0,0 +1,64 @@
+#pragma once
+
+#if defined(_WIN32)
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN
+#endif
+#include <windows.h>
+
+#include <cstdio>
+#include <cstddef>
+#include <vector>
+
+namespace strata::kernels::cpu::detail {
+
+// Hard affinity is only for pool-owned threads, which terminate with the pool. A caller's implicit
+// all-group affinity cannot be restored from PreviousGroupAffinity, so the host uses CPU Sets below.
+inline bool set_thread_group_affinity(int core, int worker) {
+    if (core < 0) return false;
+    GROUP_AFFINITY target{};
+    target.Group = (WORD) (core / 64);
+    target.Mask = KAFFINITY(1) << (core & 63);
+    if (SetThreadGroupAffinity(GetCurrentThread(), &target, nullptr)) return true;
+    std::fprintf(stderr, "strata cpu pool: SetThreadGroupAffinity for worker %d (group %u, mask 0x%llx) failed: %lu; previous affinity kept\n",
+                 worker, (unsigned) target.Group, (unsigned long long) target.Mask, (unsigned long) GetLastError());
+    return false;
+}
+
+inline bool get_thread_cpu_sets(std::vector<ULONG>& ids) {
+    ULONG count = 0;
+    if (!GetThreadSelectedCpuSets(GetCurrentThread(), nullptr, 0, &count) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    ids.resize(count);
+    if (count == 0) return true;
+    if (!GetThreadSelectedCpuSets(GetCurrentThread(), ids.data(), count, &count)) return false;
+    ids.resize(count);
+    return true;
+}
+
+// CPU Set IDs are opaque: find the entry by both group and group-relative processor number.
+inline bool cpu_set_for_core(int core, ULONG& id) {
+    ULONG bytes = 0;
+    if (!GetSystemCpuSetInformation(nullptr, 0, &bytes, GetCurrentProcess(), 0) &&
+        GetLastError() != ERROR_INSUFFICIENT_BUFFER) return false;
+    std::vector<unsigned char> buffer(bytes);
+    if (bytes && !GetSystemCpuSetInformation(reinterpret_cast<PSYSTEM_CPU_SET_INFORMATION>(buffer.data()),
+                                            bytes, &bytes, GetCurrentProcess(), 0)) return false;
+    constexpr size_t header_size = offsetof(SYSTEM_CPU_SET_INFORMATION, CpuSet);
+    for (size_t offset = 0; offset + header_size <= bytes;) {
+        const auto* info = reinterpret_cast<const SYSTEM_CPU_SET_INFORMATION*>(buffer.data() + offset);
+        if (info->Size < header_size || info->Size > bytes - offset) break;
+        if (info->Type == CpuSetInformation && info->Size >= sizeof(SYSTEM_CPU_SET_INFORMATION) && info->CpuSet.Group == core / 64 &&
+            info->CpuSet.LogicalProcessorIndex == (core & 63) &&
+            (!info->CpuSet.Allocated || info->CpuSet.AllocatedToTargetProcess)) {
+            id = info->CpuSet.Id;
+            return true;
+        }
+        offset += info->Size;
+    }
+    SetLastError(ERROR_NOT_FOUND);
+    return false;
+}
+
+}  // namespace strata::kernels::cpu::detail
+#endif