lib/simd/src/thread/range.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const IndexRange = struct {
4 begin: usize = 0,
5 end: usize = 0,
6
7 pub fn init(begin: usize, end: usize) error{InvalidRange}!IndexRange {
8 if (begin >= end) return error.InvalidRange;
9 return .{ .begin = begin, .end = end };
10 }
11
12 pub fn count(self: IndexRange) usize {
13 std.debug.assert(self.begin <= self.end);
14 return self.end - self.begin;
15 }
16
17 pub fn contains(self: IndexRange, index: usize) bool {
18 return self.begin <= index and index < self.end;
19 }
20
21 pub fn containsRange(self: IndexRange, other: IndexRange) bool {
22 return other.begin >= self.begin and other.end <= self.end;
23 }
24 };
25
26 pub fn makeIndexRange(
27 begin: usize,
28 end: usize,
29 max_size: usize,
30 ) error{ InvalidRange, EmptyTask }!IndexRange {
31 if (begin >= end) return error.InvalidRange;
32 if (max_size == 0) return error.EmptyTask;
33 return .{
34 .begin = begin,
35 .end = begin + @min(max_size, end - begin),
36 };
37 }
38
39 pub const IndexRangePartition = struct {
40 range: IndexRange,
41 task_size: u32,
42 task_count: u32,
43
44 pub fn single(
45 task_size: usize,
46 ) error{ EmptyTask, TaskTooLarge }!IndexRangePartition {
47 if (task_size == 0) return error.EmptyTask;
48 if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;
49 return .{
50 .range = .{ .begin = 0, .end = task_size },
51 .task_size = @intCast(task_size),
52 .task_count = 1,
53 };
54 }
55
56 pub fn init(
57 range: IndexRange,
58 task_size: usize,
59 ) error{ InvalidRange, EmptyTask, RangeTooLarge, TaskTooLarge }!IndexRangePartition {
60 if (range.begin >= range.end) return error.InvalidRange;
61 if (task_size == 0) return error.EmptyTask;
62 if (range.count() > std.math.maxInt(u32)) return error.RangeTooLarge;
63 if (task_size > std.math.maxInt(u32)) return error.TaskTooLarge;
64 const task_count = std.math.divCeil(
65 u32,
66 @intCast(range.count()),
67 @intCast(task_size),
68 ) catch unreachable;
69 std.debug.assert(task_count != 0);
70 return .{
71 .range = range,
72 .task_size = @intCast(task_size),
73 .task_count = task_count,
74 };
75 }
76
77 pub fn taskSize(self: IndexRangePartition) usize {
78 return self.task_size;
79 }
80
81 pub fn taskCount(self: IndexRangePartition) usize {
82 return self.task_count;
83 }
84
85 pub fn rangeAt(
86 self: IndexRangePartition,
87 task_index: usize,
88 ) error{TaskOutOfBounds}!IndexRange {
89 if (task_index >= self.task_count) return error.TaskOutOfBounds;
90 const begin = self.range.begin + task_index * self.task_size;
91 return .{
92 .begin = begin,
93 .end = begin + @min(self.task_size, self.range.end - begin),
94 };
95 }
96
97 pub fn visitAll(
98 self: IndexRangePartition,
99 context: anytype,
100 comptime visit: fn (@TypeOf(context), IndexRange) void,
101 ) void {
102 for (0..self.task_count) |task_index| {
103 visit(context, self.rangeAt(task_index) catch unreachable);
104 }
105 }
106
107 pub fn visitFirst(
108 self: IndexRangePartition,
109 context: anytype,
110 comptime visit: fn (@TypeOf(context), IndexRange) void,
111 ) void {
112 visit(context, self.rangeAt(0) catch unreachable);
113 }
114
115 pub fn visitRemaining(
116 self: IndexRangePartition,
117 context: anytype,
118 comptime visit: fn (@TypeOf(context), IndexRange) void,
119 ) void {
120 for (1..self.task_count) |task_index| {
121 visit(context, self.rangeAt(task_index) catch unreachable);
122 }
123 }
124 };
125
126 const RangeVisits = struct {
127 ranges: [8]IndexRange = @splat(.{}),
128 count: usize = 0,
129
130 fn append(self: *RangeVisits, range: IndexRange) void {
131 self.ranges[self.count] = range;
132 self.count += 1;
133 }
134 };
135
136 test "Highway index ranges preserve half-open containment" {
137 const range = try IndexRange.init(7, 19);
138 try std.testing.expectEqual(@as(usize, 12), range.count());
139 try std.testing.expect(range.contains(7));
140 try std.testing.expect(range.contains(18));
141 try std.testing.expect(!range.contains(19));
142 try std.testing.expect(range.containsRange(try IndexRange.init(9, 13)));
143 try std.testing.expect(!range.containsRange(try IndexRange.init(6, 13)));
144 try std.testing.expectEqual(
145 IndexRange{ .begin = 17, .end = 19 },
146 try makeIndexRange(17, 19, std.math.maxInt(usize)),
147 );
148 }
149
150 test "Highway index range partitions cover the remainder once" {
151 const partition = try IndexRangePartition.init(
152 try IndexRange.init(11, 40),
153 6,
154 );
155 try std.testing.expectEqual(@as(usize, 6), partition.taskSize());
156 try std.testing.expectEqual(@as(usize, 5), partition.taskCount());
157 var cursor: usize = 11;
158 for (0..partition.taskCount()) |task_index| {
159 const range = try partition.rangeAt(task_index);
160 try std.testing.expectEqual(cursor, range.begin);
161 try std.testing.expect(range.count() >= 1);
162 try std.testing.expect(range.count() <= partition.taskSize());
163 cursor = range.end;
164 }
165 try std.testing.expectEqual(@as(usize, 40), cursor);
166 try std.testing.expectError(
167 error.TaskOutOfBounds,
168 partition.rangeAt(partition.taskCount()),
169 );
170 var all = RangeVisits{};
171 partition.visitAll(&all, RangeVisits.append);
172 try std.testing.expectEqual(partition.taskCount(), all.count);
173 for (all.ranges[0..all.count], 0..) |range, task| {
174 try std.testing.expectEqual(try partition.rangeAt(task), range);
175 }
176 var first = RangeVisits{};
177 partition.visitFirst(&first, RangeVisits.append);
178 try std.testing.expectEqual(@as(usize, 1), first.count);
179 try std.testing.expectEqual(try partition.rangeAt(0), first.ranges[0]);
180 var remaining = RangeVisits{};
181 partition.visitRemaining(&remaining, RangeVisits.append);
182 try std.testing.expectEqual(partition.taskCount() - 1, remaining.count);
183 for (remaining.ranges[0..remaining.count], 1..) |range, task| {
184 try std.testing.expectEqual(try partition.rangeAt(task), range);
185 }
186 }
187
188 test "Highway one-task partition covers the requested prefix" {
189 const partition = try IndexRangePartition.single(17);
190 try std.testing.expectEqual(@as(usize, 17), partition.taskSize());
191 try std.testing.expectEqual(@as(usize, 1), partition.taskCount());
192 try std.testing.expectEqual(
193 IndexRange{ .begin = 0, .end = 17 },
194 try partition.rangeAt(0),
195 );
196 try std.testing.expectError(
197 error.EmptyTask,
198 IndexRangePartition.single(0),
199 );
200 }
201
202 test "Highway index range partition rejects unrepresentable limits" {
203 try std.testing.expectError(
204 error.EmptyTask,
205 IndexRangePartition.init(try IndexRange.init(0, 1), 0),
206 );
207 if (@sizeOf(usize) > @sizeOf(u32)) {
208 try std.testing.expectError(
209 error.RangeTooLarge,
210 IndexRangePartition.init(
211 .{ .begin = 0, .end = @as(usize, std.math.maxInt(u32)) + 2 },
212 1,
213 ),
214 );
215 try std.testing.expectError(
216 error.TaskTooLarge,
217 IndexRangePartition.init(
218 try IndexRange.init(0, 1),
219 @as(usize, std.math.maxInt(u32)) + 1,
220 ),
221 );
222 }
223 }