CPU Threading and NUMA Optimization¶

pub const ThreadPoolConfig = struct {
    num_threads: u32,
    numa_policy: NumaPolicy = .balanced,
    work_stealing: bool = true,
    affinity_enabled: bool = true,
    stack_size: usize = 2 * 1024 * 1024,  // 2 MB per thread

    pub fn detect() ThreadPoolConfig {
        const cpu_count = @as(u32, @intCast(std.Thread.getCpuCount() catch 4));
        return ThreadPoolConfig{
            .num_threads = @max(1, cpu_count - 1),
        };
    }
};

pub const Worker = struct {
    id: u32,
    thread: std.Thread,
    queue: WorkStealingQueue,
    pool: *ThreadPool,
    should_stop: std.atomic.Atomic(bool),

    fn workerLoop(self: *Worker) void {
        while (!self.should_stop.load(.Acquire)) {
            // 1. Try own queue
            var work_item = self.queue.pop();

            // 2. If empty, steal from a peer
            if (work_item == null) {
                work_item = self.stealWork();
            }

            // 3. Execute or yield
            if (work_item) |item| {
                item.execute();
            } else {
                std.Thread.yield() catch {};
            }
        }
    }
};

pub const ThreadPool = struct {
    workers: []Worker,
    config: ThreadPoolConfig,
    topology: CpuTopology,
    allocator: std.mem.Allocator,
    next_worker: std.atomic.Atomic(u32),

    pub fn submit(self: *ThreadPool, work_item: *WorkItem) bool {
        // Round-robin assignment
        const id = self.next_worker.fetchAdd(1, .AcqRel)
                   % @as(u32, @intCast(self.workers.len));
        return self.workers[id].queue.push(work_item);
    }
};

pub const WorkStealingQueue = struct {
    items: []?*WorkItem,
    head: std.atomic.Atomic(u32),
    tail: std.atomic.Atomic(u32),
    capacity: u32,

    /// Owner pushes to tail
    pub fn push(self: *WorkStealingQueue, item: *WorkItem) bool {
        const tail = self.tail.load(.Acquire);
        const next = (tail + 1) % self.capacity;
        if (next == self.head.load(.Acquire)) return false;  // full
        self.items[tail] = item;
        self.tail.store(next, .Release);
        return true;
    }

    /// Owner pops from tail (LIFO)
    pub fn pop(self: *WorkStealingQueue) ?*WorkItem {
        const tail = self.tail.load(.Acquire);
        if (tail == self.head.load(.Acquire)) return null;  // empty
        const prev = if (tail == 0) self.capacity - 1 else tail - 1;
        const item = self.items[prev];
        self.items[prev] = null;
        self.tail.store(prev, .Release);
        return item;
    }

    /// Thief steals from head (FIFO) using CAS
    pub fn steal(self: *WorkStealingQueue) ?*WorkItem {
        const head = self.head.load(.Acquire);
        if (head == self.tail.load(.Acquire)) return null;  // empty
        const item = self.items[head] orelse return null;
        const next = (head + 1) % self.capacity;
        if (self.head.compareAndSwap(head, next, .AcqRel, .Monotonic) == null) {
            return item;  // CAS succeeded
        }
        return null;  // contention, try again later
    }
};

pub const NumaPolicy = enum {
    local,       // Prefer local NUMA node for all allocations
    interleave,  // Round-robin pages across all NUMA nodes
    balanced,    // Automatic balancing based on topology
    disabled,    // No NUMA awareness (default for single-socket)
};

pub const CpuTopology = struct {
    num_cores: u32,
    num_threads: u32,
    num_numa_nodes: u32,
    cache_line_size: u32,
    l1_cache_size: u64,
    l2_cache_size: u64,
    l3_cache_size: u64,

    pub fn detect() CpuTopology {
        const cpu_count = @as(u32, @intCast(std.Thread.getCpuCount() catch 4));
        return CpuTopology{
            .num_cores = cpu_count,
            .num_threads = cpu_count,
            .num_numa_nodes = @max(1, cpu_count / 8),
            .cache_line_size = 64,
            .l1_cache_size = 32 * 1024,
            .l2_cache_size = 256 * 1024,
            .l3_cache_size = 8 * 1024 * 1024,
        };
    }
};

const MatMulContext = struct {
    a: *const Tensor,
    b: *const Tensor,
    c: *Tensor,
    start_row: u32,
    end_row: u32,

    fn execute(work_item: *WorkItem) void {
        const ctx = @as(*MatMulContext, @ptrCast(@alignCast(work_item.data.?)));
        for (ctx.start_row..ctx.end_row) |i| {
            for (0..ctx.c.shape[1]) |j| {
                var sum: f32 = 0.0;
                for (0..ctx.a.shape[1]) |k| {
                    sum += ctx.a.get(.{i, k}) * ctx.b.get(.{k, j});
                }
                ctx.c.set(.{i, j}, sum);
            }
        }
    }
};

pub const ThreadingUtils = struct {
    pub fn allocAligned(allocator: std.mem.Allocator,
                        comptime T: type, count: usize) ![]T {
        const cache_line_size = getCacheLineSize();  // 64
        const alignment = @max(@alignOf(T), cache_line_size);
        const total_size = count * @sizeOf(T) + alignment;
        const raw_mem = try allocator.alloc(u8, total_size);
        const addr = @intFromPtr(raw_mem.ptr);
        const aligned_addr = (addr + alignment - 1) & ~(alignment - 1);
        return @as([*]T, @ptrFromInt(aligned_addr))[0..count];
    }
};

pub const ThreadPoolStats = struct {
    total_tasks_executed: std.atomic.Atomic(u64),
    total_work_stolen: std.atomic.Atomic(u64),
    average_queue_depth: f64,
    cpu_utilization: f64,
};