tiny.simd.ThreadPool

1pub const ThreadPool = struct {2    pub const State = enum(u8) {3        initialization,4        steady,5        teardown,6    };78    pub const Capacity = struct {9        requested_threads: usize,10        threads: u8,11        workers: u8,1213        pub fn derive(requested_threads: usize) Capacity {14            const supported = topology.haveThreadingSupport() and15                sys.thread.threadsSupported();16            const thread_count: usize = if (supported)17                @min(requested_threads, max_threads)18            else19                0;20            const worker_count: usize = thread_count + 1;21            return .{22                .requested_threads = requested_threads,23                .threads = @intCast(thread_count),24                .workers = @intCast(worker_count),25            };26        }2728        pub fn wasClamped(self: Capacity) bool {29            return self.requested_threads != self.threads;30        }31    };3233    pub const StartError = sys.thread.SpawnError || error{34        AlreadyDeinitialized,35        PoolMoved,36    };3738    pub const RunError = StartError || error{39        Busy,40        InvalidRange,41        TooManyTasks,42    };4344    state: State = .initialization,45    capacity: Capacity,46    mapping: PoolWorkerMapping,47    address: usize = 0,48    wait_mode: std.atomic.Value(u8) =49        std.atomic.Value(u8).init(@backingInt(PoolWaitMode.block)),50    config_bits: std.atomic.Value(u16) =51        std.atomic.Value(u16).init((Config{ .wait_type = .block }).encode()),52    epoch: std.atomic.Value(u32) = std.atomic.Value(u32).init(0),53    work_available: std.atomic.Value(bool) =54        std.atomic.Value(bool).init(false),55    exiting: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),56    busy: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),57    synchronization: sys.thread.Mutex = .{},58    done: sys.thread.Condition = .{},59    started_threads: usize = 0,60    handles: [max_threads]sys.thread.JoinHandle = undefined,61    workers: [max_workers]Worker align(64) = @splat(.{}),62    task_context: *anyopaque = undefined,63    task_callback: Callback = undefined,64    current_begin: u64 = 0,65    current_tasks: usize = 0,66    stats: PoolStats = .{},67    caller_stats: [max_callers]CallerStats = @splat(.{}),68    block_tuner: AutoTuner = .{},69    spin_tuner: AutoTuner = .{},7071    pub fn init(72        requested_threads: usize,73        mapping: PoolWorkerMapping,74    ) ThreadPool {75        var self = ThreadPool{76            .capacity = Capacity.derive(requested_threads),77            .mapping = mapping,78        };79        var candidates: [max_configs]Config = undefined;80        self.block_tuner.setCandidates(81            Config.candidates(.block, &candidates),82        ) catch unreachable;83        self.spin_tuner.setCandidates(84            Config.candidates(.spin, &candidates),85        ) catch unreachable;86        self.config_bits.store(87            self.block_tuner.nextConfig().encode(),88            .monotonic,89        );90        self.configureWorkers();91        return self;92    }9394    pub fn maxThreads() usize {95        if (!topology.haveThreadingSupport() or96            !sys.thread.threadsSupported())97        {98            return 0;99        }100        var affinity = topology.LogicalProcessorSet{};101        if (topology.getThreadAffinity(&affinity)) {102            const count = affinity.count();103            if (count != 0) return @min(count - 1, max_threads);104        }105        return @min(106            topology.totalLogicalProcessors() -| 1,107            max_threads,108        );109    }110111    pub fn numThreadsFromCores(112        allocator: std.mem.Allocator,113    ) std.mem.Allocator.Error!usize {114        if (!topology.haveThreadingSupport() or115            !sys.thread.threadsSupported())116        {117            return 0;118        }119        var detected = try topology.init(allocator);120        defer detected.deinit(allocator);121        if (detected.packages.len == 0) return maxThreads();122        return @min(detected.packages[0].cores.len -| 1, max_threads);123    }124125    pub fn start(self: *ThreadPool) StartError!void {126        if (self.state == .teardown) return error.AlreadyDeinitialized;127        if (self.state == .steady) {128            try self.requireStableAddress();129            return;130        }131        self.address = @intFromPtr(self);132        if (self.capacity.threads == 0) {133            self.state = .steady;134            return;135        }136        var spawned: usize = 0;137        while (spawned < self.capacity.threads) : (spawned += 1) {138            self.handles[spawned] = sys.thread.spawn(workerMain, .{139                WorkerContext{140                    .pool = self,141                    .worker_index = spawned + 1,142                },143            }) catch |err| {144                self.stopSpawned(spawned);145                return err;146            };147            nameWorker(self.handles[spawned], spawned);148        }149        self.synchronization.lock();150        while (self.started_threads != self.capacity.threads) {151            self.done.wait(&self.synchronization);152        }153        self.synchronization.unlock();154        self.state = .steady;155    }156157    pub fn deinit(self: *ThreadPool) void {158        if (self.state == .teardown) return;159        if (self.state == .initialization) {160            self.state = .teardown;161            self.address = 0;162            return;163        }164        self.requireStableAddress() catch unreachable;165        std.debug.assert(!self.busy.load(.acquire));166        if (self.capacity.threads != 0) {167            self.signalExit();168            for (self.handles[0..self.capacity.threads]) |handle| {169                handle.join();170            }171        }172        self.state = .teardown;173        self.address = 0;174    }175176    pub fn numWorkers(self: *const ThreadPool) usize {177        return self.capacity.workers;178    }179180    pub fn wasClamped(self: *const ThreadPool) bool {181        return self.capacity.wasClamped();182    }183184    pub fn setWaitMode(185        self: *ThreadPool,186        mode: PoolWaitMode,187    ) (error{Busy} || StartError)!void {188        if (self.state == .teardown) return error.AlreadyDeinitialized;189        if (self.busy.cmpxchgStrong(190            false,191            true,192            .acquire,193            .monotonic,194        ) != null) return error.Busy;195        defer self.busy.store(false, .release);196        self.wait_mode.store(@backingInt(mode), .release);197        const mode_tuner = self.tuner(mode);198        const selected = mode_tuner.best() orelse mode_tuner.nextConfig();199        self.config_bits.store(selected.encode(), .release);200        if (self.state == .steady and self.capacity.threads != 0) {201            try self.requireStableAddress();202            const epoch = self.wakeWorkers(false, true);203            self.waitForWorkers(epoch);204        }205    }206207    pub fn waitMode(self: *const ThreadPool) PoolWaitMode {208        return @fromBackingInt(@intCast(self.wait_mode.load(.acquire)));209    }210211    pub fn config(self: *const ThreadPool) Config {212        return Config.decode(self.config_bits.load(.acquire));213    }214215    pub fn autoTuneComplete(self: *const ThreadPool) bool {216        return self.tunerConst(self.waitMode()).best() != null;217    }218219    pub fn autoTuneCosts(self: *ThreadPool) []autotune.CostDistribution {220        return self.tuner(self.waitMode()).costs();221    }222223    pub fn statsSnapshot(self: *ThreadPool) PoolStats {224        self.synchronization.lock();225        defer self.synchronization.unlock();226        return self.stats;227    }228229    pub fn callerStats(230        self: *ThreadPool,231        caller: Caller,232    ) ?CallerStats {233        if (caller.index >= max_callers) return null;234        self.synchronization.lock();235        defer self.synchronization.unlock();236        return self.caller_stats[caller.index];237    }238239    pub fn resetStats(self: *ThreadPool) error{Busy}!void {240        if (self.busy.cmpxchgStrong(241            false,242            true,243            .acquire,244            .monotonic,245        ) != null) return error.Busy;246        defer self.busy.store(false, .release);247        self.synchronization.lock();248        defer self.synchronization.unlock();249        self.stats = .{};250        self.caller_stats = @splat(.{});251    }252253    pub fn globalWorkerIndex(254        self: *const ThreadPool,255        local_worker_index: usize,256    ) error{WorkerOutOfBounds}!usize {257        if (local_worker_index >= self.numWorkers()) {258            return error.WorkerOutOfBounds;259        }260        return self.mapping.globalIndex(local_worker_index);261    }262263    pub fn run(264        self: *ThreadPool,265        begin: u64,266        end: u64,267        context: anytype,268        comptime body: fn (@TypeOf(context), u64, usize) void,269    ) RunError!void {270        return self.runWithCaller(begin, end, .{}, context, body);271    }272273    pub fn runWithCaller(274        self: *ThreadPool,275        begin: u64,276        end: u64,277        caller: Caller,278        context: anytype,279        comptime body: fn (@TypeOf(context), u64, usize) void,280    ) RunError!void {281        if (begin > end) return error.InvalidRange;282        const task_count_u64 = end - begin;283        if (task_count_u64 > std.math.maxInt(usize)) {284            return error.TooManyTasks;285        }286        const task_count: usize = @intCast(task_count_u64);287        const Context = @TypeOf(context);288        const Adapter = CallbackAdapter(Context, body);289        var context_storage = context;290        const started_ns = nowNanoseconds();291        if (task_count <= 1 or self.numWorkers() == 1) {292            var task = begin;293            while (task < end) : (task += 1) body(context, task, 0);294            self.recordRun(295                caller,296                task_count,297                false,298                0,299                elapsedNanoseconds(started_ns),300            );301            return;302        }303        if (self.busy.cmpxchgStrong(304            false,305            true,306            .acquire,307            .monotonic,308        ) != null) return error.Busy;309        defer self.busy.store(false, .release);310        try self.start();311        try self.requireStableAddress();312        self.task_context = @ptrCast(&context_storage);313        self.task_callback = Adapter.call;314        self.current_begin = begin;315        self.current_tasks = task_count;316        self.divideRange(begin, end);317        for (self.workers[0..self.numWorkers()]) |*worker| {318            worker.last_tasks = 0;319            worker.last_stolen = 0;320        }321        const epoch = self.wakeWorkers(322            true,323            self.config().wait_type == .block,324        );325        self.runWorker(0);326        self.waitForWorkers(epoch);327        var stolen_tasks: usize = 0;328        for (self.workers[0..self.numWorkers()]) |worker| {329            stolen_tasks += worker.last_stolen;330        }331        const elapsed_ns = elapsedNanoseconds(started_ns);332        self.recordRun(333            caller,334            task_count,335            true,336            stolen_tasks,337            elapsed_ns,338        );339        self.notifyAutotune(elapsed_ns);340    }341342    fn configureWorkers(self: *ThreadPool) void {343        const worker_count: u32 = self.capacity.workers;344        for (self.workers[0..worker_count], 0..) |*worker, worker_index| {345            worker.* = .{};346            worker.victim_count = @intCast(@min(max_victims, worker_count));347            const coprime = ShuffledIota.findAnotherCoprime(348                worker_count,349                @intCast((worker_index + 1) * 257 + worker_index * 13),350            );351            const shuffled = ShuffledIota.init(coprime);352            worker.victims[0] = @intCast(worker_index);353            for (1..worker.victim_count) |victim_index| {354                worker.victims[victim_index] = @intCast(shuffled.next(355                    worker.victims[victim_index - 1],356                    worker_count,357                ));358            }359        }360    }361362    fn divideRange(self: *ThreadPool, begin: u64, end: u64) void {363        for (self.workers[0..self.numWorkers()], 0..) |*worker, index| {364            const range = workerRange(365                0,366                end - begin,367                self.numWorkers(),368                index,369            ) catch unreachable;370            worker.begin.store(@intCast(range.begin), .monotonic);371            worker.end = @intCast(range.end);372        }373    }374375    fn runWorker(self: *ThreadPool, worker_index: usize) void {376        var task_count: usize = 0;377        var stolen_count: usize = 0;378        if (self.current_tasks <= self.numWorkers()) {379            const offset = self.workers[0].begin.load(.monotonic) +380                worker_index;381            const end = self.workers[self.numWorkers() - 1].end;382            if (offset < end) {383                self.task_callback(384                    self.task_context,385                    self.current_begin + offset,386                    worker_index,387                );388                task_count = 1;389            }390        } else {391            const worker = &self.workers[worker_index];392            for (worker.victims[0..worker.victim_count]) |victim_index| {393                const victim = &self.workers[victim_index];394                while (true) {395                    const offset = victim.begin.fetchAdd(1, .monotonic);396                    if (offset >= victim.end) break;397                    self.task_callback(398                        self.task_context,399                        self.current_begin + offset,400                        worker_index,401                    );402                    task_count += 1;403                    stolen_count += @intFromBool(victim_index != worker_index);404                }405            }406        }407        self.workers[worker_index].last_tasks = task_count;408        self.workers[worker_index].last_stolen = stolen_count;409    }410411    fn wakeWorkers(412        self: *ThreadPool,413        work_available: bool,414        wake_blocked: bool,415    ) u32 {416        self.work_available.store(work_available, .monotonic);417        const epoch = self.epoch.load(.monotonic) +% 1;418        self.epoch.store(epoch, .release);419        for (self.workers[1..self.numWorkers()]) |*worker| {420            worker.wait_epoch.store(epoch, .release);421        }422        if (wake_blocked and self.capacity.threads != 0) {423            wait.wakeAll(&self.workers[1].wait_epoch);424        }425        return epoch;426    }427428    fn waitForWorkers(self: *ThreadPool, epoch: u32) void {429        spin.callWithSpin(430            self.config().spin_type,431            BarrierContext{ .pool = self, .epoch = epoch },432            waitAtBarrier,433        );434    }435436    fn waitForWork(437        self: *ThreadPool,438        worker_index: usize,439        observed_epoch: u32,440    ) ?u32 {441        while (true) {442            if (self.exiting.load(.acquire)) return null;443            const config_value = self.config();444            switch (config_value.wait_type) {445                .block => {446                    const next_epoch = wait.blockUntilDifferent(447                        observed_epoch,448                        &self.workers[1].wait_epoch,449                    );450                    if (self.exiting.load(.acquire)) return null;451                    return next_epoch;452                },453                .spin_shared, .spin_separate => {454                    const watched = if (config_value.wait_type ==455                        .spin_separate)456                        &self.workers[worker_index].wait_epoch457                    else458                        &self.epoch;459                    var next_epoch: u32 = observed_epoch;460                    spin.callWithSpin(461                        config_value.spin_type,462                        SpinWaitContext{463                            .previous = observed_epoch,464                            .watched = watched,465                            .next = &next_epoch,466                        },467                        waitUntilDifferent,468                    );469                    if (self.exiting.load(.acquire)) return null;470                    return next_epoch;471                },472            }473        }474    }475476    fn workerStarted(self: *ThreadPool) void {477        self.synchronization.lock();478        self.started_threads += 1;479        self.done.broadcast();480        self.synchronization.unlock();481    }482483    fn workerReached(self: *ThreadPool, worker_index: usize, epoch: u32) void {484        self.workers[worker_index].barrier_epoch.store(epoch, .release);485    }486487    fn stopSpawned(self: *ThreadPool, spawned: usize) void {488        self.signalExit();489        for (self.handles[0..spawned]) |handle| handle.join();490        self.exiting.store(false, .monotonic);491        self.epoch.store(0, .monotonic);492        self.started_threads = 0;493        self.address = 0;494        self.work_available.store(false, .monotonic);495        self.configureWorkers();496    }497498    fn signalExit(self: *ThreadPool) void {499        self.exiting.store(true, .release);500        const epoch = self.epoch.load(.monotonic) +% 1;501        self.epoch.store(epoch, .release);502        for (self.workers[1..self.numWorkers()]) |*worker| {503            worker.wait_epoch.store(epoch, .release);504        }505        if (self.capacity.threads != 0) {506            wait.wakeAll(&self.workers[1].wait_epoch);507        }508    }509510    fn requireStableAddress(self: *const ThreadPool) error{PoolMoved}!void {511        if (self.address != @intFromPtr(self)) return error.PoolMoved;512    }513514    fn recordRun(515        self: *ThreadPool,516        caller: Caller,517        tasks: usize,518        threaded: bool,519        stolen_tasks: usize,520        elapsed_ns: u64,521    ) void {522        self.synchronization.lock();523        defer self.synchronization.unlock();524        self.stats.runs +|= 1;525        self.stats.serial_runs +|= @intFromBool(!threaded);526        self.stats.threaded_runs +|= @intFromBool(threaded);527        self.stats.tasks +|= tasks;528        self.stats.stolen_tasks +|= stolen_tasks;529        self.stats.elapsed_ns +|= elapsed_ns;530        const caller_index = if (caller.index < max_callers)531            caller.index532        else533            0;534        const stats = &self.caller_stats[caller_index];535        stats.runs +|= 1;536        stats.tasks +|= tasks;537        stats.workers +|= if (threaded) self.numWorkers() else 1;538        stats.elapsed_ns +|= elapsed_ns;539    }540541    fn notifyAutotune(self: *ThreadPool, elapsed_ns: u64) void {542        const mode = self.waitMode();543        const selected = blk: {544            const tuner_value = self.tuner(mode);545            if (tuner_value.best()) |best| break :blk best.*;546            tuner_value.notifyCost(@max(elapsed_ns, 1));547            break :blk if (tuner_value.best()) |best|548                best.*549            else550                tuner_value.nextConfig().*;551        };552        self.config_bits.store(selected.encode(), .release);553    }554555    fn tuner(self: *ThreadPool, mode: PoolWaitMode) *AutoTuner {556        return switch (mode) {557            .block => &self.block_tuner,558            .spin => &self.spin_tuner,559        };560    }561562    fn tunerConst(self: *const ThreadPool, mode: PoolWaitMode) *const AutoTuner {563        return switch (mode) {564            .block => &self.block_tuner,565            .spin => &self.spin_tuner,566        };567    }568};

1pub const ThreadPool = thread.ThreadPool;