lib/accy/src/profiling/wos/workload.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3 const accy = @import("accy");
4
5 const options_mod = @import("options.zig");
6
7 const kernel = accy.kernel;
8 const random = accy.kernel.library.random;
9
10 const tau: f32 = 2.0 * std.math.pi;
11 const boundary_radius_floor: f32 = 0.000001;
12 const origin_radius_limit: f32 = 0.9;
13
14 pub const Variant = enum {
15 fold,
16 while_loop,
17
18 pub fn label(self: Variant) []const u8 {
19 return switch (self) {
20 .fold => "fold",
21 .while_loop => "while",
22 };
23 }
24
25 pub fn entryName(self: Variant) []const u8 {
26 return switch (self) {
27 .fold => "accy_wos_bench_fold",
28 .while_loop => "accy_wos_bench_while",
29 };
30 }
31 };
32
33 const WalkValues = struct {
34 x: kernel.Value,
35 y: kernel.Value,
36 steps: kernel.Value,
37 };
38
39 const MoveValues = struct {
40 x: kernel.Value,
41 y: kernel.Value,
42 };
43
44 pub fn buildWalkGraph(
45 allocator: std.mem.Allocator,
46 variant: Variant,
47 epsilon: f32,
48 cap: u32,
49 ) !kernel.Graph {
50 var builder = try kernel.Builder.init(allocator, kernel.Builder.Limits.standard, variant.entryName(), &.{
51 kernel.dynamicBuffer(.f32),
52 kernel.dynamicBuffer(.f32),
53 kernel.dynamicBuffer(.f32),
54 kernel.dynamicBuffer(.i32),
55 kernel.scalar(.i32),
56 kernel.scalar(.i32),
57 });
58 errdefer builder.deinit();
59
60 const xs = builder.argument(0);
61 const ys = builder.argument(1);
62 const dst = builder.argument(2);
63 const steps_out = builder.argument(3);
64 const seed_lo = builder.argument(4);
65 const seed_hi = builder.argument(5);
66
67 const gid = try builder.globalId(.x);
68 const walk_id = try builder.cast(gid, .i32);
69 const x0 = try builder.load(xs, gid);
70 const y0 = try builder.load(ys, gid);
71 const zero_i = try builder.constantInt(.i32, 0);
72 const key_words = try random.block.philoxWords(
73 &builder,
74 random.philox_default_rounds,
75 walk_id,
76 zero_i,
77 seed_lo,
78 seed_hi,
79 );
80
81 const walked = switch (variant) {
82 .fold => try emitFoldWalk(&builder, x0, y0, key_words[0], key_words[1], epsilon, cap),
83 .while_loop => try emitWhileWalk(&builder, x0, y0, key_words[0], key_words[1], epsilon, cap),
84 };
85
86 const boundary = try emitBoundaryValue(&builder, walked.x, walked.y);
87 try builder.store(boundary, dst, gid);
88 try builder.store(walked.steps, steps_out, gid);
89 try builder.return_();
90
91 return builder.finish();
92 }
93
94 fn emitWhileWalk(
95 builder: *kernel.Builder,
96 x0: kernel.Value,
97 y0: kernel.Value,
98 key_lo: kernel.Value,
99 key_hi: kernel.Value,
100 epsilon: f32,
101 cap: u32,
102 ) !WalkValues {
103 const zero_i = try builder.constantInt(.i32, 0);
104 const one_i = try builder.constantInt(.i32, 1);
105 const cap_value = try builder.constantInt(.i32, @intCast(cap));
106 const epsilon_value = try builder.constantFloat(.f32, epsilon);
107
108 const Context = struct {
109 key_lo: kernel.Value,
110 key_hi: kernel.Value,
111 one_i: kernel.Value,
112 cap: kernel.Value,
113 epsilon: kernel.Value,
114 };
115
116 return builder.whileLoop(WalkValues{
117 .x = x0,
118 .y = y0,
119 .steps = zero_i,
120 }, Context{
121 .key_lo = key_lo,
122 .key_hi = key_hi,
123 .one_i = one_i,
124 .cap = cap_value,
125 .epsilon = epsilon_value,
126 }, struct {
127 fn keepGoing(inner: *kernel.Builder, walk: WalkValues, ctx: Context) !kernel.Value {
128 const gate_radius = try emitDistanceToBoundary(inner, walk.x, walk.y);
129 const can_move = try inner.compare(.gt, gate_radius, ctx.epsilon);
130 const below_cap = try inner.compare(.lt, walk.steps, ctx.cap);
131 return inner.and_(can_move, below_cap);
132 }
133 }.keepGoing, struct {
134 fn step(inner: *kernel.Builder, walk: WalkValues, ctx: Context) !WalkValues {
135 const radius = try emitDistanceToBoundary(inner, walk.x, walk.y);
136 const move = try emitMove(inner, walk.x, walk.y, radius, walk.steps, ctx.key_lo, ctx.key_hi);
137 return .{
138 .x = move.x,
139 .y = move.y,
140 .steps = try inner.add(walk.steps, ctx.one_i),
141 };
142 }
143 }.step);
144 }
145
146 fn emitFoldWalk(
147 builder: *kernel.Builder,
148 x0: kernel.Value,
149 y0: kernel.Value,
150 key_lo: kernel.Value,
151 key_hi: kernel.Value,
152 epsilon: f32,
153 cap: u32,
154 ) !WalkValues {
155 const zero_i = try builder.constantInt(.i32, 0);
156 const one_i = try builder.constantInt(.i32, 1);
157 const epsilon_value = try builder.constantFloat(.f32, epsilon);
158
159 var x = x0;
160 var y = y0;
161 var active = try builder.constantBool(true);
162 var steps_taken = zero_i;
163 var step_index: u32 = 0;
164 while (step_index < cap) : (step_index += 1) {
165 const radius = try emitDistanceToBoundary(builder, x, y);
166 const can_move = try builder.compare(.gt, radius, epsilon_value);
167 const moving = try builder.and_(active, can_move);
168 const step_value = try builder.constantInt(.i32, @intCast(step_index));
169 const move = try emitMove(builder, x, y, radius, step_value, key_lo, key_hi);
170 x = try builder.select(moving, move.x, x);
171 y = try builder.select(moving, move.y, y);
172 steps_taken = try builder.add(steps_taken, try builder.select(moving, one_i, zero_i));
173 active = moving;
174 }
175
176 return .{ .x = x, .y = y, .steps = steps_taken };
177 }
178
179 fn emitMove(
180 builder: *kernel.Builder,
181 x: kernel.Value,
182 y: kernel.Value,
183 radius: kernel.Value,
184 step: kernel.Value,
185 key_lo: kernel.Value,
186 key_hi: kernel.Value,
187 ) !MoveValues {
188 const zero_i = try builder.constantInt(.i32, 0);
189 const words = try random.block.philoxWords(
190 builder,
191 random.philox_default_rounds,
192 step,
193 zero_i,
194 key_lo,
195 key_hi,
196 );
197 const unit = try random.block.outputWord(builder, .f32, words[0]);
198 const angle = try builder.mul(unit, try builder.constantFloat(.f32, tau));
199 const dx = try builder.mul(try builder.cos(angle), radius);
200 const dy = try builder.mul(try builder.sin(angle), radius);
201 return .{ .x = try builder.add(x, dx), .y = try builder.add(y, dy) };
202 }
203
204 fn emitDistanceToBoundary(builder: *kernel.Builder, x: kernel.Value, y: kernel.Value) !kernel.Value {
205 const radius2 = try builder.add(try builder.mul(x, x), try builder.mul(y, y));
206 const radius = try builder.sqrt(radius2);
207 const one = try builder.constantFloat(.f32, 1.0);
208 const zero = try builder.constantFloat(.f32, 0.0);
209 return builder.max(try builder.sub(one, radius), zero);
210 }
211
212 fn emitBoundaryValue(builder: *kernel.Builder, x: kernel.Value, y: kernel.Value) !kernel.Value {
213 const radius2 = try builder.add(try builder.mul(x, x), try builder.mul(y, y));
214 const floor = try builder.constantFloat(.f32, boundary_radius_floor);
215 const radius = try builder.sqrt(try builder.max(radius2, floor));
216 const inv = try builder.div(try builder.constantFloat(.f32, 1.0), radius);
217 const bx = try builder.mul(x, inv);
218 const by = try builder.mul(y, inv);
219 const bx2 = try builder.mul(bx, bx);
220 const by2 = try builder.mul(by, by);
221 const cubic = try builder.mul(bx, bx2);
222 const cross = try builder.mul(try builder.mul(bx, by2), try builder.constantFloat(.f32, 3.0));
223 return builder.sub(cubic, cross);
224 }
225
226 pub fn launchGeometry(walks: u32) choir_abi.LaunchGeometry {
227 std.debug.assert(walks % options_mod.threads_per_block == 0);
228 return .{
229 .grid = .{ walks / options_mod.threads_per_block, 1, 1 },
230 .threadgroup = .{ options_mod.threads_per_block, 1, 1 },
231 };
232 }
233
234 pub fn seedArguments() [2]choir_abi.ScalarArgument {
235 return .{
236 .{ .i32 = @bitCast(@as(u32, @truncate(options_mod.walk_seed))) },
237 .{ .i32 = @bitCast(@as(u32, @truncate(options_mod.walk_seed >> 32))) },
238 };
239 }
240
241 pub fn fillWalkOrigins(xs: []f32, ys: []f32) void {
242 std.debug.assert(xs.len == ys.len);
243 var state: u64 = options_mod.walk_seed | 1;
244 for (xs, ys) |*x, *y| {
245 const radius = origin_radius_limit * @sqrt(nextUniform(&state));
246 const angle = tau * nextUniform(&state);
247 x.* = radius * @cos(angle);
248 y.* = radius * @sin(angle);
249 }
250 }
251
252 fn nextUniform(state: *u64) f32 {
253 state.* ^= state.* << 13;
254 state.* ^= state.* >> 7;
255 state.* ^= state.* << 17;
256 return random.uniformFromBits(@truncate(state.* >> 32));
257 }
258
259 pub const HostWalk = struct {
260 boundary: f32,
261 steps: u32,
262 };
263
264 pub fn hostWalk(x0: f32, y0: f32, walk_index: u32, epsilon: f32, cap: u32) HostWalk {
265 const seed_words = random.philoxBlock(
266 random.philox_default_rounds,
267 .{ walk_index, 0, 0, 0 },
268 .{ @truncate(options_mod.walk_seed), @truncate(options_mod.walk_seed >> 32) },
269 );
270 var x = x0;
271 var y = y0;
272 var steps: u32 = 0;
273 while (true) {
274 const radius = @max(0.0, 1.0 - @sqrt(x * x + y * y));
275 if (!(radius > epsilon and steps < cap)) break;
276 const words = random.philoxBlock(
277 random.philox_default_rounds,
278 .{ steps, 0, 0, 0 },
279 .{ seed_words[0], seed_words[1] },
280 );
281 const angle = random.uniformFromBits(words[0]) * tau;
282 x += @cos(angle) * radius;
283 y += @sin(angle) * radius;
284 steps += 1;
285 }
286 const radius = @sqrt(@max(x * x + y * y, boundary_radius_floor));
287 const bx = x / radius;
288 const by = y / radius;
289 return .{ .boundary = bx * bx * bx - 3.0 * bx * by * by, .steps = steps };
290 }
291
292 fn oracleLaunch(walks: u32) kernel.schedule.Launch {
293 const blocks = walks / options_mod.threads_per_block;
294 return .{
295 .grid = .{ @max(blocks, 1), 1, 1 },
296 .block = .{ @min(walks, options_mod.threads_per_block), 1, 1 },
297 };
298 }
299
300 test "fold and while walk kernels agree exactly on the oracle" {
301 const walks = 256;
302 const epsilon: f32 = 1e-2;
303 const cap: u32 = 32;
304
305 var xs: [walks]f32 = undefined;
306 var ys: [walks]f32 = undefined;
307 fillWalkOrigins(xs[0..], ys[0..]);
308
309 var fold_dst = @as([walks]f32, @splat(0.0));
310 var fold_steps = @as([walks]i32, @splat(-1));
311 var while_dst = @as([walks]f32, @splat(0.0));
312 var while_steps = @as([walks]i32, @splat(-1));
313
314 const seeds = seedArguments();
315 inline for (.{
316 .{ Variant.fold, fold_dst[0..], fold_steps[0..] },
317 .{ Variant.while_loop, while_dst[0..], while_steps[0..] },
318 }) |case| {
319 var graph = try buildWalkGraph(std.testing.allocator, case[0], epsilon, cap);
320 defer graph.deinit();
321 try graph.runCpuWithLaunch(std.testing.allocator, &.{
322 kernel.argumentBuffer(f32, xs[0..]),
323 kernel.argumentBuffer(f32, ys[0..]),
324 kernel.argumentBuffer(f32, case[1]),
325 kernel.argumentBuffer(i32, case[2]),
326 kernel.argumentI32(seeds[0].i32),
327 kernel.argumentI32(seeds[1].i32),
328 }, oracleLaunch(walks));
329 }
330
331 try std.testing.expectEqualSlices(i32, fold_steps[0..], while_steps[0..]);
332 for (fold_dst, while_dst) |fold_value, while_value| {
333 try std.testing.expectEqual(fold_value, while_value);
334 }
335 }
336
337 test "oracle walks match the host mirror exactly" {
338 const walks = 64;
339 const epsilon: f32 = 1e-2;
340 const cap: u32 = 32;
341
342 var xs: [walks]f32 = undefined;
343 var ys: [walks]f32 = undefined;
344 fillWalkOrigins(xs[0..], ys[0..]);
345
346 var dst = @as([walks]f32, @splat(0.0));
347 var steps = @as([walks]i32, @splat(-1));
348 const seeds = seedArguments();
349 var graph = try buildWalkGraph(std.testing.allocator, .while_loop, epsilon, cap);
350 defer graph.deinit();
351 try graph.runCpuWithLaunch(std.testing.allocator, &.{
352 kernel.argumentBuffer(f32, xs[0..]),
353 kernel.argumentBuffer(f32, ys[0..]),
354 kernel.argumentBuffer(f32, dst[0..]),
355 kernel.argumentBuffer(i32, steps[0..]),
356 kernel.argumentI32(seeds[0].i32),
357 kernel.argumentI32(seeds[1].i32),
358 }, oracleLaunch(walks));
359
360 for (0..walks) |index| {
361 const mirror = hostWalk(xs[index], ys[index], @intCast(index), epsilon, cap);
362 try std.testing.expectEqual(@as(i32, @intCast(mirror.steps)), steps[index]);
363 try std.testing.expectApproxEqAbs(mirror.boundary, dst[index], 0.00001);
364 }
365 }
366
367 test "walk origins stay inside the sampling disk" {
368 var xs: [128]f32 = undefined;
369 var ys: [128]f32 = undefined;
370 fillWalkOrigins(xs[0..], ys[0..]);
371
372 for (xs, ys) |x, y| {
373 try std.testing.expect(@sqrt(x * x + y * y) <= origin_radius_limit + 0.00001);
374 }
375 }