tiny.choir.threading

1const std = @import("std");2const sys = @import("sys");34pub const Options = struct {5    max_threads: usize = 1,6    worker_allocator: ?std.mem.Allocator = null,78    pub fn workerCount(self: Options, item_count: usize) usize {9        if (item_count == 0) return 1;10        const configured = if (self.max_threads == 0)11            sys.thread.cpuCount()12        else13            self.max_threads;14        return @max(@as(usize, 1), @min(configured, item_count));15    }1617    pub fn workerAllocator(self: Options, fallback: std.mem.Allocator) std.mem.Allocator {18        return self.worker_allocator orelse fallback;19    }2021    pub fn requestsParallelism(self: Options) bool {22        return self.max_threads == 0 or self.max_threads > 1;23    }24};2526pub const ExecutionState = struct {27    active: std.atomic.Value(usize),2829    pub fn init() ExecutionState {30        return .{ .active = std.atomic.Value(usize).init(0) };31    }3233    pub fn enter(self: *ExecutionState) ExecutionGuard {34        _ = self.active.fetchAdd(1, .acq_rel);35        return .{ .state = self };36    }3738    pub fn isActive(self: *const ExecutionState) bool {39        return self.active.load(.acquire) != 0;40    }4142    fn exit(self: *ExecutionState) void {43        const prior = self.active.fetchSub(1, .acq_rel);44        std.debug.assert(prior != 0);45    }46};4748pub const ExecutionGuard = struct {49    state: ?*ExecutionState,5051    pub fn deinit(self: *ExecutionGuard) void {52        if (self.state) |state| {53            state.exit();54            self.state = null;55        }56    }57};5859pub const ParallelError = std.mem.Allocator.Error || sys.thread.SpawnError;6061pub fn parallelForEachIndex(62    allocator: std.mem.Allocator,63    options: Options,64    item_count: usize,65    context: anytype,66    comptime run: fn (@TypeOf(context), usize) void,67) ParallelError!void {68    const worker_count = options.workerCount(item_count);69    if (worker_count == 1) {70        for (0..item_count) |index| {71            run(context, index);72        }73        return;74    }7576    const Context = @TypeOf(context);77    const Batch = struct {78        task_context: Context,79        next_index: *std.atomic.Value(usize),80        item_count: usize,81    };82    const Worker = struct {83        fn work(batch: *Batch) void {84            while (true) {85                const index = batch.next_index.fetchAdd(1, .monotonic);86                if (index >= batch.item_count) return;87                run(batch.task_context, index);88            }89        }90    };9192    var next_index = std.atomic.Value(usize).init(0);93    var batch = Batch{94        .task_context = context,95        .next_index = &next_index,96        .item_count = item_count,97    };9899    const spawned_count = worker_count - 1;100    const threads = try allocator.alloc(sys.thread.JoinHandle, spawned_count);101    defer allocator.free(threads);102103    var spawned: usize = 0;104    errdefer {105        for (threads[0..spawned]) |thread| {106            thread.join();107        }108    }109110    for (threads) |*thread| {111        thread.* = try sys.thread.spawn(Worker.work, .{&batch});112        spawned += 1;113    }114    Worker.work(&batch);115    for (threads) |thread| {116        thread.join();117    }118}119120test "ThreadingOptions counts workers with item and cpu limits" {121    const testing = std.testing;122123    try testing.expectEqual(@as(usize, 1), (Options{}).workerCount(0));124    try testing.expectEqual(@as(usize, 1), (Options{}).workerCount(8));125    try testing.expectEqual(@as(usize, 2), (Options{ .max_threads = 2 }).workerCount(8));126    try testing.expectEqual(@as(usize, 2), (Options{ .max_threads = 4 }).workerCount(2));127    try testing.expect((Options{ .max_threads = 0 }).workerCount(8) >= 1);128}129130test "ThreadingExecutionState tracks nested execution guards" {131    const testing = std.testing;132133    var state = ExecutionState.init();134    try testing.expect(!state.isActive());135136    var outer = state.enter();137    defer outer.deinit();138    try testing.expect(state.isActive());139140    {141        var inner = state.enter();142        defer inner.deinit();143        try testing.expect(state.isActive());144    }145146    try testing.expect(state.isActive());147    outer.deinit();148    try testing.expect(!state.isActive());149}150151test "parallelForEachIndex visits every item" {152    const testing = std.testing;153    if (!sys.thread.threadsSupported()) return error.SkipZigTest;154155    var hits: [8]std.atomic.Value(usize) = undefined;156    for (&hits) |*hit| hit.* = std.atomic.Value(usize).init(0);157158    const Task = struct {159        values: []std.atomic.Value(usize),160    };161    const Runner = struct {162        fn run(task: *Task, index: usize) void {163            _ = task.values[index].fetchAdd(1, .acq_rel);164        }165    };166167    var task = Task{ .values = &hits };168    try parallelForEachIndex(testing.allocator, .{ .max_threads = 2 }, hits.len, &task, Runner.run);169170    for (&hits) |*hit| {171        try testing.expectEqual(@as(usize, 1), hit.load(.acquire));172    }173}

1pub const threading = ir.threading;