tiny.tldr.parallel

1const std = @import("std");2const allocators = @import("alloc");3const alloc_arena = @import("alloc_arena");4const sys = @import("sys");56const Allocator = std.mem.Allocator;7const Arena = alloc_arena.Arena;89pub const max_workers = 64;1011pub const WorkerArenaPool = struct {12    backing_owner: ?*allocators.LockedAllocator = null,13    arenas: []Arena = &.{},1415    pub fn init(storage: Allocator, worker_count: usize) !WorkerArenaPool {16        std.debug.assert(worker_count <= max_workers);17        const count = @max(worker_count, 1);18        std.debug.assert(count >= 1);19        std.debug.assert(count <= max_workers);2021        const backing_owner = try storage.create(allocators.LockedAllocator);22        errdefer storage.destroy(backing_owner);23        backing_owner.* = allocators.LockedAllocator.init(storage);2425        const arenas = try storage.alloc(Arena, count);26        for (arenas) |*arena| arena.* = Arena.init(backing_owner.allocator());27        return .{ .backing_owner = backing_owner, .arenas = arenas };28    }2930    pub fn deinit(self: *WorkerArenaPool) void {31        std.debug.assert(self.backing_owner != null);32        std.debug.assert(self.arenas.len >= 1);33        const backing_owner = self.backing_owner.?;34        const storage = backing_owner.child;35        for (self.arenas) |*arena| arena.deinit();36        storage.free(self.arenas);37        storage.destroy(backing_owner);38        self.* = .{};39    }4041    pub fn allocator(self: *WorkerArenaPool, worker: usize) Allocator {42        std.debug.assert(self.backing_owner != null);43        std.debug.assert(worker < self.arenas.len);44        return self.arenas[worker].allocator();45    }46};4748pub fn FailureSlots(comptime Failure: type) type {49    return struct {50        items: []Failure = &.{},51        empty: Failure,5253        const Self = @This();5455        pub fn init(storage: Allocator, slot_count: usize, empty: Failure) !Self {56            const count = @max(slot_count, 1);57            const items = try storage.alloc(Failure, count);58            for (items) |*item| item.* = empty;59            return .{ .items = items, .empty = empty };60        }6162        pub fn deinit(self: *Self, storage: Allocator) void {63            storage.free(self.items);64            self.* = .{ .empty = self.empty };65        }6667        pub fn record(self: *Self, worker: usize, failure: Failure) void {68            self.items[worker] = failure;69        }7071        pub fn earliest(72            self: *const Self,73            comptime found: fn (Failure) bool,74            comptime before: fn (Failure, Failure) bool,75        ) ?Failure {76            var best: ?Failure = null;77            for (self.items) |item| {78                if (!found(item)) continue;79                if (best == null or before(item, best.?)) best = item;80            }81            return best;82        }83    };84}8586const ChunkCallback = *const fn (*anyopaque, usize, usize, usize) void;8788const Job = struct {89    total: usize = 0,90    slots: usize = 0,91    context: *anyopaque = undefined,92    callback: ChunkCallback = undefined,93    background: bool = false,94    next_slot: usize = 0,95    pending: usize = 0,96    prev: ?*Job = null,97    next: ?*Job = null,98};99100const Claim = struct {101    job: *Job,102    slot: usize,103};104105const JobQueue = struct {106    head: ?*Job = null,107    tail: ?*Job = null,108109    fn append(self: *JobQueue, job: *Job) void {110        job.prev = self.tail;111        job.next = null;112        if (self.tail) |tail| {113            tail.next = job;114        } else {115            self.head = job;116        }117        self.tail = job;118    }119120    fn remove(self: *JobQueue, job: *Job) void {121        if (job.prev) |prev| {122            prev.next = job.next;123        } else {124            self.head = job.next;125        }126        if (job.next) |next| {127            next.prev = job.prev;128        } else {129            self.tail = job.prev;130        }131        job.prev = null;132        job.next = null;133    }134};135136const WorkerPool = struct {137    mutex: sys.thread.Mutex = .{},138    ready: sys.thread.Condition = .{},139    done: sys.thread.Condition = .{},140    worker_count: usize = 0,141    spawn_failed: bool = false,142    foreground: JobQueue = .{},143    background: JobQueue = .{},144145    fn ensureStarted(self: *WorkerPool, desired_workers: usize) bool {146        if (desired_workers == 0) return false;147        self.mutex.lock();148        defer self.mutex.unlock();149        const desired = @min(desired_workers, max_workers - 1);150        while (self.worker_count < desired and !self.spawn_failed) {151            const handle = sys.thread.spawn(workerLoop, .{self}) catch {152                self.spawn_failed = true;153                break;154            };155            self.worker_count += 1;156            handle.detach();157        }158        return self.worker_count != 0;159    }160161    fn submit(self: *WorkerPool, job: *Job, desired_workers: usize, reserved_slots: usize) bool {162        std.debug.assert(job.slots > reserved_slots);163        if (!self.ensureStarted(desired_workers)) return false;164        self.mutex.lock();165        job.next_slot = reserved_slots;166        job.pending = job.slots;167        self.queueFor(job).append(job);168        self.ready.broadcast();169        self.mutex.unlock();170        return true;171    }172173    fn queueFor(self: *WorkerPool, job: *Job) *JobQueue {174        return if (job.background) &self.background else &self.foreground;175    }176177    fn claimAnyLocked(self: *WorkerPool) ?Claim {178        const job = self.foreground.head orelse self.background.head orelse return null;179        return self.claimFromLocked(job);180    }181182    fn claimFromLocked(self: *WorkerPool, job: *Job) Claim {183        const slot = job.next_slot;184        job.next_slot += 1;185        if (job.next_slot == job.slots) self.queueFor(job).remove(job);186        return .{ .job = job, .slot = slot };187    }188189    fn runClaim(claim: Claim) void {190        const job = claim.job;191        const start = rangeStart(job.total, job.slots, claim.slot);192        const end = rangeStart(job.total, job.slots, claim.slot + 1);193        job.callback(job.context, claim.slot, start, end);194    }195196    fn finishLocked(self: *WorkerPool, job: *Job) void {197        job.pending -= 1;198        if (job.pending == 0) self.done.broadcast();199    }200201    fn workerLoop(self: *WorkerPool) void {202        self.mutex.lock();203        while (true) {204            if (self.claimAnyLocked()) |claim| {205                self.mutex.unlock();206                runClaim(claim);207                self.mutex.lock();208                self.finishLocked(claim.job);209            } else {210                self.ready.wait(&self.mutex);211            }212        }213    }214215    fn wait(self: *WorkerPool, job: *Job) void {216        self.mutex.lock();217        while (job.pending != 0) {218            if (job.next_slot < job.slots) {219                const claim = self.claimFromLocked(job);220                self.mutex.unlock();221                runClaim(claim);222                self.mutex.lock();223                self.finishLocked(job);224            } else {225                self.done.wait(&self.mutex);226            }227        }228        self.mutex.unlock();229    }230231    fn completeReserved(self: *WorkerPool, job: *Job) void {232        runClaim(.{ .job = job, .slot = 0 });233        self.mutex.lock();234        self.finishLocked(job);235        self.mutex.unlock();236        self.wait(job);237    }238};239240var global_worker_pool = WorkerPool{};241242pub fn overlapAvailable() bool {243    return sys.thread.threadsSupported() and sys.thread.cpuCount() > 1;244}245246pub fn chooseWorkers(total: usize, requested: usize) usize {247    if (total <= 1 or !sys.thread.threadsSupported()) return 1;248    const available = sys.thread.cpuCount();249    const ceiling = if (requested == 0) available else @min(available, requested);250    return @max(1, @min(@min(ceiling, total), max_workers));251}252253pub fn rangeStart(total: usize, workers: usize, worker: usize) usize {254    const base = total / workers;255    const remainder = total % workers;256    return worker * base + @min(worker, remainder);257}258259fn ChunkCallbackType(260    comptime Context: type,261    comptime body: fn (Context, usize, usize, usize) void,262) type {263    return struct {264        fn run_opaque(opaque_context: *anyopaque, worker: usize, start: usize, end: usize) void {265            const typed_context: *Context = @ptrCast(@alignCast(opaque_context));266            body(typed_context.*, worker, start, end);267        }268    };269}270271pub fn forChunks(272    total: usize,273    requested_workers: usize,274    context: anytype,275    comptime body: fn (@TypeOf(context), usize, usize, usize) void,276) void {277    const workers = chooseWorkers(total, requested_workers);278    if (workers <= 1) {279        body(context, 0, 0, total);280        return;281    }282283    const Context: type = @TypeOf(context);284    const Callback: type = ChunkCallbackType(Context, body);285286    var context_storage = context;287    var job = Job{288        .total = total,289        .slots = workers,290        .context = @ptrCast(&context_storage),291        .callback = Callback.run_opaque,292    };293    if (global_worker_pool.submit(&job, workers - 1, 1)) {294        global_worker_pool.completeReserved(&job);295        return;296    }297298    var worker: usize = 0;299    while (worker < workers) : (worker += 1) {300        body(301            context,302            worker,303            rangeStart(total, workers, worker),304            rangeStart(total, workers, worker + 1),305        );306    }307}308309fn ItemStateType(comptime Context: type) type {310    return struct {311        user: Context,312        next: std.atomic.Value(usize) = std.atomic.Value(usize).init(0),313        total: usize,314    };315}316317fn ItemCallbackType(318    comptime Context: type,319    comptime State: type,320    comptime body: fn (Context, usize, usize) void,321) type {322    return struct {323        fn run(state: *State, worker: usize, start: usize, end: usize) void {324            _ = start;325            _ = end;326            while (true) {327                const item = state.next.fetchAdd(1, .monotonic);328                if (item >= state.total) break;329                body(state.user, worker, item);330            }331        }332    };333}334335pub fn forItems(336    total: usize,337    requested_workers: usize,338    context: anytype,339    comptime body: fn (@TypeOf(context), usize, usize) void,340) void {341    const workers = chooseWorkers(total, requested_workers);342    if (workers <= 1) {343        var item: usize = 0;344        while (item < total) : (item += 1) body(context, 0, item);345        return;346    }347348    const Context: type = @TypeOf(context);349    const State: type = ItemStateType(Context);350    const Callback: type = ItemCallbackType(Context, State, body);351352    var state = State{ .user = context, .total = total };353    forChunks(workers, workers, &state, Callback.run);354}355356pub fn Background(357    comptime Context: type,358    comptime body: fn (Context, usize) void,359) type {360    return struct {361        job: Job = .{ .background = true },362        context: Context = undefined,363        submitted: bool = false,364365        const Self = @This();366367        pub fn submit(self: *Self, total: usize, requested_workers: usize, context: Context) void {368            std.debug.assert(!self.submitted);369            if (total == 0) return;370            self.context = context;371            self.job.total = total;372            self.job.slots = total;373            self.job.context = @ptrCast(self);374            self.job.callback = runOpaque;375            if (sys.thread.threadsSupported()) {376                const available = sys.thread.cpuCount();377                const ceiling = if (requested_workers == 0)378                    available379                else380                    @min(available, requested_workers);381                const desired = @max(1, @min(@min(ceiling, total), max_workers));382                if (global_worker_pool.submit(&self.job, desired, 0)) {383                    self.submitted = true;384                    return;385                }386            }387            var item: usize = 0;388            while (item < total) : (item += 1) body(context, item);389        }390391        pub fn wait(self: *Self) void {392            if (!self.submitted) return;393            global_worker_pool.wait(&self.job);394            self.submitted = false;395        }396397        fn runOpaque(opaque_context: *anyopaque, slot: usize, start: usize, end: usize) void {398            _ = start;399            _ = end;400            const self: *Self = @ptrCast(@alignCast(opaque_context));401            body(self.context, slot);402        }403    };404}405406fn SortChunkStateType(comptime T: type, comptime Context: type) type {407    return struct {408        items: []T,409        chunk: usize,410        user: Context,411    };412}413414fn SortMergeStateType(comptime T: type, comptime Context: type) type {415    return struct {416        source: []T,417        target: []T,418        width: usize,419        user: Context,420    };421}422423fn SortCallbacksType(424    comptime T: type,425    comptime Context: type,426    comptime SortState: type,427    comptime MergeState: type,428    comptime less: fn (Context, T, T) bool,429) type {430    return struct {431        fn sort_chunk(state: *SortState, worker: usize, index: usize) void {432            _ = worker;433            const low = index * state.chunk;434            if (low >= state.items.len) return;435            const high = @min(low + state.chunk, state.items.len);436            std.sort.pdq(T, state.items[low..high], state.user, less);437        }438439        fn merge_segment(state: *MergeState, worker: usize, index: usize) void {440            _ = worker;441            const low = index * 2 * state.width;442            if (low >= state.source.len) return;443            const middle = @min(low + state.width, state.source.len);444            const high = @min(low + 2 * state.width, state.source.len);445            merge_runs(446                state.user,447                state.source[low..middle],448                state.source[middle..high],449                state.target[low..high],450            );451        }452453        fn merge_runs(user: Context, left: []const T, right: []const T, out: []T) void {454            var left_index: usize = 0;455            var right_index: usize = 0;456            var out_index: usize = 0;457            while (left_index < left.len and right_index < right.len) : (out_index += 1) {458                if (less(user, right[right_index], left[left_index])) {459                    out[out_index] = right[right_index];460                    right_index += 1;461                } else {462                    out[out_index] = left[left_index];463                    left_index += 1;464                }465            }466            if (left_index < left.len) {467                @memcpy(out[out_index..], left[left_index..]);468            } else if (right_index < right.len) {469                @memcpy(out[out_index..], right[right_index..]);470            }471        }472    };473}474475pub fn sortItems(476    comptime T: type,477    items: []T,478    context: anytype,479    comptime less: fn (@TypeOf(context), T, T) bool,480    scratch: []T,481    requested_workers: usize,482) void {483    const workers = chooseWorkers(items.len, requested_workers);484    if (workers <= 1 or items.len < 2) {485        std.sort.pdq(T, items, context, less);486        return;487    }488    std.debug.assert(scratch.len >= items.len);489490    const Context: type = @TypeOf(context);491    const SortState: type = SortChunkStateType(T, Context);492    const MergeState: type = SortMergeStateType(T, Context);493    const Callbacks: type = SortCallbacksType(T, Context, SortState, MergeState, less);494    const total = items.len;495    const chunk = (total + workers - 1) / workers;496497    var sort_state = SortState{ .items = items, .chunk = chunk, .user = context };498    forItems(workers, workers, &sort_state, Callbacks.sort_chunk);499500    var source = items;501    var target = scratch[0..total];502    var width = chunk;503    while (width < total) : (width *= 2) {504        const segments = (total + 2 * width - 1) / (2 * width);505        var merge_state = MergeState{506            .source = source,507            .target = target,508            .width = width,509            .user = context,510        };511        forItems(segments, workers, &merge_state, Callbacks.merge_segment);512        const swapped = source;513        source = target;514        target = swapped;515    }516    if (source.ptr != items.ptr) @memcpy(items, source);517}518519fn lessUsize(_: void, left: usize, right: usize) bool {520    return left < right;521}522523const U8HitsContext = struct {524    hits: []u8,525};526527const U16HitsContext = struct {528    hits: []u16,529};530531const SumContext = struct {532    sum: *usize,533};534535const TestFailure = struct {536    found: bool = false,537    index: usize = 0,538};539540const WorkerArenaPoolAllocationFailures = struct {541    fn run(storage: Allocator) !void {542        var pool = try WorkerArenaPool.init(storage, 4);543        defer pool.deinit();544        _ = try pool.allocator(3).alloc(u8, 128);545    }546};547548fn count_u8_chunks(context: U8HitsContext, _: usize, start: usize, end: usize) void {549    var index = start;550    while (index < end) : (index += 1) context.hits[index] += 1;551}552553fn count_u16_chunks(context: U16HitsContext, _: usize, start: usize, end: usize) void {554    var index = start;555    while (index < end) : (index += 1) context.hits[index] += 1;556}557558fn count_u8_item(context: U8HitsContext, _: usize, item: usize) void {559    context.hits[item] += 1;560}561562fn count_u8_background_item(context: U8HitsContext, item: usize) void {563    context.hits[item] += 1;564}565566fn failure_found(failure: TestFailure) bool {567    return failure.found;568}569570fn failure_before(left: TestFailure, right: TestFailure) bool {571    return left.index < right.index;572}573574const TestBackground = Background(U8HitsContext, count_u8_background_item);575const chained_background_total = 128;576const nested_parallel_inner_total = 64;577578const ChainedBackgroundContext = struct {579    hits: []u8,580    second: *TestBackground,581    second_hits: []u8,582};583584fn start_chained_background(context: ChainedBackgroundContext, item: usize) void {585    context.hits[item] += 1;586    if (item == 0) {587        context.second.submit(chained_background_total, 2, .{ .hits = context.second_hits });588    }589}590591const ChainedBackground = Background(ChainedBackgroundContext, start_chained_background);592593fn count_nested_chunks(context: U8HitsContext, _: usize, start: usize, end: usize) void {594    var chunk = start;595    while (chunk < end) : (chunk += 1) {596        const offset = chunk * nested_parallel_inner_total;597        forItems(nested_parallel_inner_total, 2, U8HitsContext{598            .hits = context.hits[offset..][0..nested_parallel_inner_total],599        }, count_u8_item);600    }601}602603fn sum_chunk(context: SumContext, _: usize, start: usize, end: usize) void {604    context.sum.* += end - start;605}606607test "parallel sort matches serial sort for total orders" {608    const allocator = std.testing.allocator;609    var prng = std.Random.DefaultPrng.init(0x74696e79736f7274);610    const random = prng.random();611612    for ([_]usize{ 0, 1, 2, 63, 4096, 40_001 }) |count| {613        const items = try allocator.alloc(usize, count);614        defer allocator.free(items);615        const expected = try allocator.alloc(usize, count);616        defer allocator.free(expected);617        const scratch = try allocator.alloc(usize, count);618        defer allocator.free(scratch);619620        for (items, 0..) |*item, index| {621            item.* = random.uintLessThan(usize, 1 << 40) * 100_000 + index;622        }623        @memcpy(expected, items);624        std.sort.pdq(usize, expected, {}, lessUsize);625626        sortItems(usize, items, {}, lessUsize, scratch, 8);627        try std.testing.expectEqualSlices(usize, expected, items);628    }629}630631test "worker arena pool routes worker allocations through bounded caller storage" {632    var storage_bytes: [4096]u8 = undefined;633    var storage = std.heap.FixedBufferAllocator.init(&storage_bytes);634    var pool = try WorkerArenaPool.init(storage.allocator(), 4);635    defer pool.deinit();636637    var worker: usize = 0;638    while (worker < 4) : (worker += 1) {639        const bytes = try pool.allocator(worker).alloc(u8, 32);640        @memset(bytes, @intCast(worker + 1));641        try std.testing.expect(storage.ownsSlice(bytes));642        try std.testing.expectEqual(@as(u8, @intCast(worker + 1)), bytes[0]);643    }644    try std.testing.expectError(645        error.OutOfMemory,646        pool.allocator(0).alloc(u8, storage_bytes.len),647    );648}649650test "worker arena pool cleans every caller storage allocation failure" {651    try std.testing.checkAllAllocationFailures(652        std.testing.allocator,653        WorkerArenaPoolAllocationFailures.run,654        .{},655    );656}657658test "forChunks covers every index exactly once" {659    const total = 10_000;660    var hits = @as([total]u8, @splat(0));661    forChunks(total, 8, U8HitsContext{ .hits = &hits }, count_u8_chunks);662    for (hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);663}664665test "forItems covers every item exactly once" {666    const total = 10_000;667    var hits = @as([total]u8, @splat(0));668    forItems(total, 8, U8HitsContext{ .hits = &hits }, count_u8_item);669    for (hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);670}671672test "failure slots merge earliest recorded failure" {673    var slots = try FailureSlots(TestFailure).init(std.testing.allocator, 4, .{});674    defer slots.deinit(std.testing.allocator);675    slots.record(2, .{ .found = true, .index = 40 });676    slots.record(1, .{ .found = true, .index = 11 });677678    const earliest = slots.earliest(failure_found, failure_before) orelse679        return error.MissingFailure;680    try std.testing.expectEqual(@as(usize, 11), earliest.index);681}682683test "forChunks handles repeated pooled calls" {684    const total = 4096;685    const rounds = 8;686    var hits = @as([total]u16, @splat(0));687    var round: usize = 0;688    while (round < rounds) : (round += 1) {689        forChunks(total, 8, U16HitsContext{ .hits = &hits }, count_u16_chunks);690    }691    for (hits) |hit| try std.testing.expectEqual(@as(u16, rounds), hit);692}693694test "background job covers every item exactly once" {695    const total = 512;696    var hits = @as([total]u8, @splat(0));697    var prefetch = TestBackground{};698    prefetch.submit(total, 4, .{ .hits = &hits });699    prefetch.wait();700    for (hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);701}702703test "background job runs alongside foreground chunk jobs" {704    const background_total = 256;705    const foreground_total = 4096;706    var background_hits = @as([background_total]u8, @splat(0));707    var foreground_hits = @as([foreground_total]u8, @splat(0));708709    var prefetch = TestBackground{};710    prefetch.submit(background_total, 2, .{ .hits = &background_hits });711712    var round: usize = 0;713    while (round < 4) : (round += 1) {714        forChunks(foreground_total, 8, U8HitsContext{ .hits = &foreground_hits }, count_u8_chunks);715    }716717    prefetch.wait();718    for (background_hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);719    for (foreground_hits) |hit| try std.testing.expectEqual(@as(u8, 4), hit);720}721722test "background job may chain a following background job" {723    var first_hits = @as([1]u8, @splat(0));724    var second_hits = @as([chained_background_total]u8, @splat(0));725726    var second = TestBackground{};727    var first = ChainedBackground{};728    first.submit(1, 2, .{729        .hits = first_hits[0..1],730        .second = &second,731        .second_hits = &second_hits,732    });733    first.wait();734    second.wait();735736    try std.testing.expectEqual(@as(u8, 1), first_hits[0]);737    for (second_hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);738}739740test "background job with zero items completes without submission" {741    var prefetch = TestBackground{};742    prefetch.submit(0, 4, .{ .hits = &.{} });743    prefetch.wait();744    prefetch.wait();745}746747test "nested parallel sections complete without stalling the pool" {748    const outer_chunks = 4;749    var hits = @as([(outer_chunks * nested_parallel_inner_total)]u8, @splat(0));750    forChunks(outer_chunks, outer_chunks, U8HitsContext{ .hits = &hits }, count_nested_chunks);751    for (hits) |hit| try std.testing.expectEqual(@as(u8, 1), hit);752}753754test "forChunks single worker runs inline over full range" {755    const total = 256;756    var sum: usize = 0;757    forChunks(total, 1, SumContext{ .sum = &sum }, sum_chunk);758    try std.testing.expectEqual(@as(usize, total), sum);759}760761test "rangeStart partitions contiguously and totals correctly" {762    const total = 1003;763    const workers = 7;764    try std.testing.expectEqual(@as(usize, 0), rangeStart(total, workers, 0));765    try std.testing.expectEqual(@as(usize, total), rangeStart(total, workers, workers));766    var worker: usize = 0;767    while (worker < workers) : (worker += 1) {768        try std.testing.expect(769            rangeStart(total, workers, worker) <= rangeStart(total, workers, worker + 1),770        );771    }772}

1pub const parallel = @import("parallel.zig");