tiny.simd.IndexRangePartition

1pub const IndexRangePartition = struct {2    range: IndexRange,3    task_size: u32,4    task_count: u32,56    pub fn single(7        task_size: usize,8    ) error{ EmptyTask, TaskTooLarge }!IndexRangePartition {9        if (task_size == 0) return error.EmptyTask;10        if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;11        return .{12            .range = .{ .begin = 0, .end = task_size },13            .task_size = @intCast(task_size),14            .task_count = 1,15        };16    }1718    pub fn init(19        range: IndexRange,20        task_size: usize,21    ) error{ InvalidRange, EmptyTask, RangeTooLarge, TaskTooLarge }!IndexRangePartition {22        if (range.begin >= range.end) return error.InvalidRange;23        if (task_size == 0) return error.EmptyTask;24        if (range.count() > std.math.maxInt(u32)) return error.RangeTooLarge;25        if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;26        const task_count = std.math.divCeil(27            u32,28            @intCast(range.count()),29            @intCast(task_size),30        ) catch unreachable;31        std.debug.assert(task_count != 0);32        return .{33            .range = range,34            .task_size = @intCast(task_size),35            .task_count = task_count,36        };37    }3839    pub fn taskSize(self: IndexRangePartition) usize {40        return self.task_size;41    }4243    pub fn taskCount(self: IndexRangePartition) usize {44        return self.task_count;45    }4647    pub fn rangeAt(48        self: IndexRangePartition,49        task_index: usize,50    ) error{TaskOutOfBounds}!IndexRange {51        if (task_index >= self.task_count) return error.TaskOutOfBounds;52        const begin = self.range.begin + task_index * self.task_size;53        return .{54            .begin = begin,55            .end = begin + @min(self.task_size, self.range.end - begin),56        };57    }5859    pub fn visitAll(60        self: IndexRangePartition,61        context: anytype,62        comptime visit: fn (@TypeOf(context), IndexRange) void,63    ) void {64        for (0..self.task_count) |task_index| {65            visit(context, self.rangeAt(task_index) catch unreachable);66        }67    }6869    pub fn visitFirst(70        self: IndexRangePartition,71        context: anytype,72        comptime visit: fn (@TypeOf(context), IndexRange) void,73    ) void {74        visit(context, self.rangeAt(0) catch unreachable);75    }7677    pub fn visitRemaining(78        self: IndexRangePartition,79        context: anytype,80        comptime visit: fn (@TypeOf(context), IndexRange) void,81    ) void {82        for (1..self.task_count) |task_index| {83            visit(context, self.rangeAt(task_index) catch unreachable);84        }85    }86};

1pub const IndexRangePartition = thread.IndexRangePartition;