lib/accy/src/kernel/oracle/machine.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const testing = std.testing;
3 const accy_root = @import("../../root.zig");
4
5 const core = @import("../model/core/root.zig");
6 const eval = @import("../../eval/root.zig");
7
8 const builder = core.builder;
9 const schedule_mod = core.schedule;
10
11 pub const enabled = true;
12 pub const Launch = schedule_mod.Launch;
13 pub const Argument = eval.KernelArgument;
14 pub const Machine = eval.Machine;
15 pub const Scalar = eval.Scalar;
16 pub const ExecutionDiagnostic = eval.Diagnostic;
17 pub const shared_alloc_sentinel = eval.shared_alloc_sentinel;
18 pub const tf32RoundF32 = eval.tf32RoundF32;
19
20 pub fn argumentBuffer(comptime T: type, values: []T) Argument {
21 return .{ .memref = std.mem.sliceAsBytes(values) };
22 }
23
24 pub fn argumentBool(value: bool) Argument {
25 return .{ .scalar = .{ .bool = value } };
26 }
27
28 pub fn argumentI32(value: i32) Argument {
29 return .{ .scalar = .{ .i32 = value } };
30 }
31
32 pub fn argumentU32(value: u32) Argument {
33 return .{ .scalar = .{ .u32 = value } };
34 }
35
36 pub fn argumentI64(value: i64) Argument {
37 return .{ .scalar = .{ .i64 = value } };
38 }
39
40 pub fn argumentF32(value: f32) Argument {
41 return .{ .scalar = .{ .f32 = value } };
42 }
43
44 pub fn argumentF64(value: f64) Argument {
45 return .{ .scalar = .{ .f64 = value } };
46 }
47
48 pub fn run(allocator: std.mem.Allocator, kernel: *builder.Kernel, args: []const Argument, launch: Launch) !void {
49 try kernel.verify();
50 try eval.executeKernelFunction(allocator, kernel.func().op, args, launch);
51 }
52
53 pub fn runWithDiagnostic(
54 allocator: std.mem.Allocator,
55 kernel: *builder.Kernel,
56 args: []const Argument,
57 launch: Launch,
58 diagnostic: *ExecutionDiagnostic,
59 ) !void {
60 diagnostic.len = 0;
61 try kernel.verifyWithDiagnostic(diagnostic);
62 try eval.executeKernelFunctionWithDiagnostic(allocator, kernel.func().op, args, launch, diagnostic);
63 }
64
65 pub fn run1D(allocator: std.mem.Allocator, kernel: *builder.Kernel, args: []const Argument, extent: usize) !void {
66 try run(allocator, kernel, args, Launch.oneDim(extent));
67 }
68
69 pub fn run1DWithDiagnostic(
70 allocator: std.mem.Allocator,
71 kernel: *builder.Kernel,
72 args: []const Argument,
73 extent: usize,
74 diagnostic: *ExecutionDiagnostic,
75 ) !void {
76 try runWithDiagnostic(allocator, kernel, args, Launch.oneDim(extent), diagnostic);
77 }
78
79 fn expectApproxSlice(actual: []const f32, expected: []const f32) !void {
80 try testing.expectEqual(expected.len, actual.len);
81 for (actual, expected) |a, e| {
82 try testing.expectApproxEqAbs(e, a, 0.0001);
83 }
84 }
85
86 test "run1D executes Choir-backed copy kernels on CPU" {
87 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "copy_f32", &.{
88 builder.dynamicBuffer(.f32),
89 builder.dynamicBuffer(.f32),
90 }, .{});
91 errdefer b.deinit();
92
93 const src = b.argument(0);
94 const dst = b.argument(1);
95 const i = try b.globalId(.x);
96 const x = try b.load(src, i);
97 try b.store(x, dst, i);
98 try b.return_();
99
100 var kernel = try b.finish();
101 defer kernel.deinit();
102
103 var input = [_]f32{ 1.5, -2.0, 3.25, 4.5 };
104 var output = [_]f32{ 0, 0, 0, 0 };
105 try run1D(testing.allocator, &kernel, &.{
106 argumentBuffer(f32, input[0..]),
107 argumentBuffer(f32, output[0..]),
108 }, input.len);
109
110 try expectApproxSlice(output[0..], input[0..]);
111 }
112
113 test "run1D executes structured control flow kernels on CPU" {
114 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "clamp_positive", &.{
115 builder.dynamicBuffer(.f32),
116 builder.dynamicBuffer(.f32),
117 }, .{});
118 errdefer b.deinit();
119
120 const src = b.argument(0);
121 const dst = b.argument(1);
122 const i = try b.globalId(.x);
123 const x = try b.load(src, i);
124 const zero = try b.constantFloat(.f32, 0.0);
125 const pred = try b.compare(.gt, x, zero);
126 var if_op = try b.if_(pred, &.{x.valueType()});
127
128 {
129 var scope = try b.enterBlock(if_op.thenBlock());
130 defer scope.leave();
131 try b.yield_(&.{x});
132 }
133
134 {
135 var scope = try b.enterBlock(if_op.elseBlock().?);
136 defer scope.leave();
137 try b.yield_(&.{zero});
138 }
139
140 try b.store(if_op.result(0).?, dst, i);
141 try b.return_();
142
143 var kernel = try b.finish();
144 defer kernel.deinit();
145
146 var input = [_]f32{ -3.0, 2.5, 0.0, 9.0 };
147 var output = [_]f32{ 7.0, 7.0, 7.0, 7.0 };
148 try run1D(testing.allocator, &kernel, &.{
149 argumentBuffer(f32, input[0..]),
150 argumentBuffer(f32, output[0..]),
151 }, input.len);
152
153 try expectApproxSlice(output[0..], &.{ 0.0, 2.5, 0.0, 9.0 });
154 }
155
156 test "run1D binds scalar kernel arguments on CPU" {
157 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "fill_scalar", &.{
158 builder.dynamicBuffer(.f32),
159 builder.scalar(.f32),
160 }, .{});
161 errdefer b.deinit();
162
163 const dst = b.argument(0);
164 const value = b.argument(1);
165 const i = try b.globalId(.x);
166 try b.store(value, dst, i);
167 try b.return_();
168
169 var kernel = try b.finish();
170 defer kernel.deinit();
171
172 var output = [_]f32{ 0, 0, 0 };
173 try run1D(testing.allocator, &kernel, &.{
174 argumentBuffer(f32, output[0..]),
175 argumentF32(3.5),
176 }, output.len);
177
178 try expectApproxSlice(output[0..], &.{ 3.5, 3.5, 3.5 });
179 }
180
181 test "run1D executes u32 scalar and buffer kernels on CPU" {
182 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "wrap_neg_u32", &.{
183 builder.dynamicBuffer(.u32),
184 builder.dynamicBuffer(.u32),
185 builder.scalar(.u32),
186 }, .{});
187 errdefer b.deinit();
188
189 const src = b.argument(0);
190 const dst = b.argument(1);
191 const delta = b.argument(2);
192 const i = try b.globalId(.x);
193 const loaded = try b.load(src, i);
194 const sum = try b.add(loaded, delta);
195 const wrapped = try b.neg(sum);
196 try b.store(wrapped, dst, i);
197 try b.return_();
198
199 var kernel = try b.finish();
200 defer kernel.deinit();
201
202 var input = [_]u32{ 0, std.math.maxInt(u32), 41 };
203 var output = [_]u32{ 7, 7, 7 };
204 try run1D(testing.allocator, &kernel, &.{
205 argumentBuffer(u32, input[0..]),
206 argumentBuffer(u32, output[0..]),
207 argumentU32(1),
208 }, input.len);
209
210 try testing.expectEqualSlices(u32, &.{
211 std.math.maxInt(u32),
212 0,
213 std.math.maxInt(u32) - 41,
214 }, output[0..]);
215 }
216
217 test "run executes loop kernels on CPU" {
218 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "fill_loop", &.{
219 builder.dynamicBuffer(.i32),
220 }, .{});
221 errdefer b.deinit();
222
223 const dst = b.argument(0);
224 const lower = try b.constantIndex(0);
225 const upper = try b.constantIndex(4);
226 const step = try b.constantIndex(1);
227 var loop = try b.for_(lower, upper, step, &.{}, &.{});
228
229 {
230 var scope = try b.enterBlock(loop.bodyBlock());
231 defer scope.leave();
232 const one = try b.constantInt(.i32, 1);
233 try b.store(one, dst, loop.inductionVar());
234 try b.yield_(&.{});
235 }
236
237 try b.return_();
238
239 var kernel = try b.finish();
240 defer kernel.deinit();
241
242 var output = [_]i32{ 0, 0, 0, 0 };
243 try run(testing.allocator, &kernel, &.{
244 argumentBuffer(i32, output[0..]),
245 }, .{});
246
247 try testing.expectEqualSlices(i32, &.{ 1, 1, 1, 1 }, output[0..]);
248 }
249
250 test "run1D executes atomic add kernels with old-value results on CPU" {
251 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "atomic_add_f32", &.{
252 builder.dynamicBuffer(.f32),
253 builder.dynamicBuffer(.f32),
254 }, .{});
255 errdefer b.deinit();
256
257 const bins = b.argument(0);
258 const olds = b.argument(1);
259 const i = try b.globalId(.x);
260 const bin = try b.constantIndex(0);
261 const delta = try b.constantFloat(.f32, 2.5);
262 const old = try b.atomicRmw(.add, delta, bins, bin);
263 try b.store(old, olds, i);
264 try b.return_();
265
266 var kernel = try b.finish();
267 defer kernel.deinit();
268
269 var bin_values = [_]f32{0};
270 var old_values = [_]f32{ -1, -1, -1, -1 };
271 try run1D(testing.allocator, &kernel, &.{
272 argumentBuffer(f32, bin_values[0..]),
273 argumentBuffer(f32, old_values[0..]),
274 }, old_values.len);
275
276 try testing.expectApproxEqAbs(@as(f32, 10.0), bin_values[0], 0.0001);
277 try expectApproxSlice(old_values[0..], &.{ 0.0, 2.5, 5.0, 7.5 });
278 }
279
280 test "run1D executes integer atomic max and exchange kernels on CPU" {
281 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "atomic_max_exchange_i32", &.{
282 builder.dynamicBuffer(.i32),
283 builder.dynamicBuffer(.i32),
284 builder.dynamicBuffer(.i32),
285 }, .{});
286 errdefer b.deinit();
287
288 const peak = b.argument(0);
289 const last = b.argument(1);
290 const values = b.argument(2);
291 const i = try b.globalId(.x);
292 const slot = try b.constantIndex(0);
293 const value = try b.load(values, i);
294 _ = try b.atomicRmw(.max, value, peak, slot);
295 _ = try b.atomicRmw(.exchange, value, last, slot);
296 try b.return_();
297
298 var kernel = try b.finish();
299 defer kernel.deinit();
300
301 var peak_values = [_]i32{std.math.minInt(i32)};
302 var last_values = [_]i32{0};
303 var inputs = [_]i32{ 4, 9, -3, 7 };
304 try run1D(testing.allocator, &kernel, &.{
305 argumentBuffer(i32, peak_values[0..]),
306 argumentBuffer(i32, last_values[0..]),
307 argumentBuffer(i32, inputs[0..]),
308 }, inputs.len);
309
310 try testing.expectEqual(@as(i32, 9), peak_values[0]);
311 try testing.expectEqual(@as(i32, 7), last_values[0]);
312 }
313
314 test "run1D executes integer atomic cas kernels on CPU" {
315 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "atomic_cas_i32", &.{
316 builder.dynamicBuffer(.i32),
317 builder.dynamicBuffer(.i32),
318 }, .{});
319 errdefer b.deinit();
320
321 const slot_values = b.argument(0);
322 const old_values = b.argument(1);
323 const i = try b.globalId(.x);
324 const slot = try b.constantIndex(0);
325 const expected = try b.constantInt(.i32, 0);
326 const desired = try b.constantInt(.i32, 11);
327 const old = try b.atomicCas(expected, desired, slot_values, slot);
328 try b.store(old, old_values, i);
329 try b.return_();
330
331 var kernel = try b.finish();
332 defer kernel.deinit();
333
334 var slot_data = [_]i32{0};
335 var old_data = [_]i32{ -1, -1, -1, -1 };
336 try run1D(testing.allocator, &kernel, &.{
337 argumentBuffer(i32, slot_data[0..]),
338 argumentBuffer(i32, old_data[0..]),
339 }, old_data.len);
340
341 try testing.expectEqual(@as(i32, 11), slot_data[0]);
342 try testing.expectEqualSlices(i32, &.{ 0, 11, 11, 11 }, old_data[0..]);
343 }
344
345 test "run executes warp mma.sync fragments against a scalar product" {
346 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "mma_sync_m16n8k8", &.{
347 builder.dynamicBuffer(.f32),
348 builder.dynamicBuffer(.f32),
349 builder.dynamicBuffer(.f32),
350 builder.dynamicBuffer(.f32),
351 }, .{});
352 errdefer b.deinit();
353
354 const a_buf = b.argument(0);
355 const b_buf = b.argument(1);
356 const c_buf = b.argument(2);
357 const d_buf = b.argument(3);
358
359 const lane = try b.laneId();
360 const four = try b.constantIndex(4);
361 const eight = try b.constantIndex(8);
362 const group = try b.div(lane, four);
363 const tid = try b.sub(lane, try b.mul(group, four));
364 const group_hi = try b.add(group, eight);
365 const tid_hi = try b.add(tid, four);
366
367 const row_lo = try b.mul(group, eight);
368 const row_hi = try b.mul(group_hi, eight);
369 const a = [4]builder.Value{
370 try b.load(a_buf, try b.add(row_lo, tid)),
371 try b.load(a_buf, try b.add(row_hi, tid)),
372 try b.load(a_buf, try b.add(row_lo, tid_hi)),
373 try b.load(a_buf, try b.add(row_hi, tid_hi)),
374 };
375 const bv = [2]builder.Value{
376 try b.load(b_buf, try b.add(try b.mul(tid, eight), group)),
377 try b.load(b_buf, try b.add(try b.mul(tid_hi, eight), group)),
378 };
379 const two = try b.constantIndex(2);
380 const one = try b.constantIndex(1);
381 const col_lo = try b.mul(tid, two);
382 const col_hi = try b.add(col_lo, one);
383 const c = [4]builder.Value{
384 try b.load(c_buf, try b.add(row_lo, col_lo)),
385 try b.load(c_buf, try b.add(row_lo, col_hi)),
386 try b.load(c_buf, try b.add(row_hi, col_lo)),
387 try b.load(c_buf, try b.add(row_hi, col_hi)),
388 };
389
390 const d = try b.mmaSync(.{ .m = 16, .n = 8, .k = 8 }, a, bv, c);
391 try b.store(d[0], d_buf, try b.add(row_lo, col_lo));
392 try b.store(d[1], d_buf, try b.add(row_lo, col_hi));
393 try b.store(d[2], d_buf, try b.add(row_hi, col_lo));
394 try b.store(d[3], d_buf, try b.add(row_hi, col_hi));
395 try b.return_();
396
397 var kernel = try b.finish();
398 defer kernel.deinit();
399
400 var a_values: [16 * 8]f32 = undefined;
401 var b_values: [8 * 8]f32 = undefined;
402 var c_values: [16 * 8]f32 = undefined;
403 var d_values: [16 * 8]f32 = @splat(0);
404 for (&a_values, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 13)) * 0.25 - 1.0;
405 for (&b_values, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 7)) * 0.5 - 1.5;
406 for (&c_values, 0..) |*value, index| value.* = @as(f32, @floatFromInt(index % 5)) - 2.0;
407
408 try run(testing.allocator, &kernel, &.{
409 argumentBuffer(f32, a_values[0..]),
410 argumentBuffer(f32, b_values[0..]),
411 argumentBuffer(f32, c_values[0..]),
412 argumentBuffer(f32, d_values[0..]),
413 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 32, 1, 1 } });
414
415 for (0..16) |row| {
416 for (0..8) |col| {
417 var expected = c_values[row * 8 + col];
418 for (0..8) |k| {
419 expected += a_values[row * 8 + k] * b_values[k * 8 + col];
420 }
421 try testing.expectApproxEqAbs(expected, d_values[row * 8 + col], 0.001);
422 }
423 }
424 }
425
426 test "run1D executes shared atomic accumulation on CPU" {
427 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "shared_atomic_add_i32", &.{
428 builder.dynamicBuffer(.i32),
429 }, .{});
430 errdefer b.deinit();
431
432 const out = b.argument(0);
433 const i = try b.globalId(.x);
434 const bins = try b.sharedBuffer(.i32, 4);
435 const zero = try b.constantInt(.i32, 0);
436 try b.store(zero, bins, i);
437 try b.barrier(.block);
438 const slot = try b.constantIndex(0);
439 const one = try b.constantInt(.i32, 1);
440 const old = try b.atomicRmw(.add, one, bins, slot);
441 try b.store(old, out, i);
442 try b.return_();
443
444 var kernel = try b.finish();
445 defer kernel.deinit();
446
447 var olds = [_]i32{ -1, -1, -1, -1 };
448 try run(testing.allocator, &kernel, &.{
449 argumentBuffer(i32, olds[0..]),
450 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 4, 1, 1 } });
451
452 try testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3 }, olds[0..]);
453 }
454
455 test "run1D stages async shared copies before reading them back" {
456 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "cp_async_shared_f32", &.{
457 builder.dynamicBuffer(.f32),
458 builder.dynamicBuffer(.f32),
459 }, .{});
460 errdefer b.deinit();
461
462 const input = b.argument(0);
463 const out = b.argument(1);
464 const i = try b.globalId(.x);
465 const staged = try b.sharedBuffer(.f32, 8);
466 const four = try b.constantIndex(4);
467 const quad_base = try b.mul(i, four);
468 try b.asyncCopyShared(staged, quad_base, input, quad_base, 16);
469 try b.asyncCopyCommit();
470 try b.asyncCopyWait(0);
471 try b.barrier(.block);
472 const seven = try b.constantIndex(7);
473 const flipped = try b.sub(seven, i);
474 const value = try b.load(staged, flipped);
475 try b.store(value, out, i);
476 try b.return_();
477
478 var kernel = try b.finish();
479 defer kernel.deinit();
480
481 var input_values = [_]f32{ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0 };
482 var out_values = @as([8]f32, @splat(0.0));
483 try run(testing.allocator, &kernel, &.{
484 argumentBuffer(f32, input_values[0..]),
485 argumentBuffer(f32, out_values[0..]),
486 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 2, 1, 1 } });
487
488 try testing.expectEqualSlices(f32, &.{ 8.0, 7.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 }, out_values[0..]);
489 }
490
491 test "run1D executes shared integer atomic cas kernels on CPU" {
492 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "shared_atomic_cas_i32", &.{
493 builder.dynamicBuffer(.i32),
494 }, .{});
495 errdefer b.deinit();
496
497 const out = b.argument(0);
498 const i = try b.globalId(.x);
499 const bins = try b.sharedBuffer(.i32, 4);
500 const zero = try b.constantInt(.i32, 0);
501 try b.store(zero, bins, i);
502 try b.barrier(.block);
503 const slot = try b.constantIndex(0);
504 const desired = try b.constantInt(.i32, 7);
505 const old = try b.atomicCas(zero, desired, bins, slot);
506 try b.store(old, out, i);
507 try b.return_();
508
509 var kernel = try b.finish();
510 defer kernel.deinit();
511
512 var olds = [_]i32{ -1, -1, -1, -1 };
513 try run(testing.allocator, &kernel, &.{
514 argumentBuffer(i32, olds[0..]),
515 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 4, 1, 1 } });
516
517 try testing.expectEqualSlices(i32, &.{ 0, 7, 7, 7 }, olds[0..]);
518 }
519
520 test "run1D resets unseeded shared memory to the sentinel per block" {
521 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "shared_atomic_unseeded_i32", &.{
522 builder.dynamicBuffer(.i32),
523 }, .{});
524 errdefer b.deinit();
525
526 const out = b.argument(0);
527 const i = try b.globalId(.x);
528 const bins = try b.sharedBuffer(.i32, 4);
529 const slot = try b.constantIndex(0);
530 const one = try b.constantInt(.i32, 1);
531 const old = try b.atomicRmw(.add, one, bins, slot);
532 try b.store(old, out, i);
533 try b.return_();
534
535 var kernel = try b.finish();
536 defer kernel.deinit();
537
538 const fill: u8 = shared_alloc_sentinel;
539 const word: u32 = @as(u32, fill) << 24 | @as(u32, fill) << 16 | @as(u32, fill) << 8 | @as(u32, fill);
540 const sentinel_value: i32 = @bitCast(word);
541
542 var per_block = [_]i32{ -1, -1, -1, -1 };
543 try run1D(testing.allocator, &kernel, &.{
544 argumentBuffer(i32, per_block[0..]),
545 }, per_block.len);
546
547 try testing.expectEqualSlices(
548 i32,
549 &.{ sentinel_value, sentinel_value, sentinel_value, sentinel_value },
550 per_block[0..],
551 );
552 }
553
554 test "run1D executes while loops with divergent lane trip counts on CPU" {
555 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "while_triangular_i32", &.{
556 builder.dynamicBuffer(.i32),
557 }, .{});
558 errdefer b.deinit();
559
560 const dst = b.argument(0);
561 const i = try b.globalId(.x);
562 const start = try b.cast(i, .i32);
563 const zero = try b.constantInt(.i32, 0);
564 const one = try b.constantInt(.i32, 1);
565
566 var scope = try b.whileScope(&.{ start, zero }, &.{ start.valueType(), zero.valueType() });
567 errdefer scope.abort();
568 const remaining = scope.beforeArg(0).?;
569 const total = scope.beforeArg(1).?;
570 const proceed = try b.compare(.gt, remaining, zero);
571 try scope.condition(proceed, &.{ remaining, total });
572 const after_remaining = scope.afterArg(0).?;
573 const after_total = scope.afterArg(1).?;
574 const next_total = try b.add(after_total, after_remaining);
575 const next_remaining = try b.sub(after_remaining, one);
576 try scope.leave(&.{ next_remaining, next_total });
577
578 try b.store(scope.result(1).?, dst, i);
579 try b.return_();
580
581 var kernel = try b.finish();
582 defer kernel.deinit();
583
584 var out = @as([8]i32, @splat(-1));
585 try run1D(testing.allocator, &kernel, &.{
586 argumentBuffer(i32, out[0..]),
587 }, out.len);
588
589 try testing.expectEqualSlices(i32, &.{ 0, 1, 3, 6, 10, 15, 21, 28 }, out[0..]);
590 }
591
592 fn buildClampedFillKernel(allocator: std.mem.Allocator) !builder.Kernel {
593 var b = try builder.Builder.init(allocator, builder.Builder.Limits.testing, "clamped_fill_f32", &.{
594 builder.dynamicBuffer(.f32),
595 builder.dynamicBuffer(.f32),
596 }, .{});
597 errdefer b.deinit();
598
599 const src = b.argument(0);
600 const dst = b.argument(1);
601 const i = try b.globalId(.x);
602 const x = try b.load(src, i);
603 const zero = try b.constantFloat(.f32, 0.0);
604 const pred = try b.compare(.gt, x, zero);
605 var if_op = try b.if_(pred, &.{x.valueType()});
606
607 {
608 var scope = try b.enterBlock(if_op.thenBlock());
609 defer scope.leave();
610 try b.yield_(&.{x});
611 }
612
613 {
614 var scope = try b.enterBlock(if_op.elseBlock().?);
615 defer scope.leave();
616 try b.yield_(&.{zero});
617 }
618
619 const clamped = if_op.result(0).?;
620 const lower = try b.constantIndex(0);
621 const upper = try b.constantIndex(4);
622 const step = try b.constantIndex(1);
623 var loop = try b.for_(lower, upper, step, &.{}, &.{});
624
625 {
626 var scope = try b.enterBlock(loop.bodyBlock());
627 defer scope.leave();
628 try b.store(clamped, dst, i);
629 try b.yield_(&.{});
630 }
631
632 try b.return_();
633 return b.finish();
634 }
635
636 const clamped_fill_input = [_]f32{ 1.5, -2.0, 3.25, -4.5, 0.5, -1.0, 2.0, 6.0 };
637 const clamped_fill_expected = [_]f32{ 1.5, 0.0, 3.25, 0.0, 0.5, 0.0, 2.0, 6.0 };
638
639 test "Machine steady launches make no allocator calls" {
640 var kernel = try buildClampedFillKernel(testing.allocator);
641 defer kernel.deinit();
642 try kernel.verify();
643
644 var failing = std.testing.FailingAllocator.init(testing.allocator, .{});
645 var machine = Machine.init(failing.allocator());
646 defer machine.deinit();
647
648 var input = clamped_fill_input;
649 var output = @as([8]f32, @splat(0));
650 const launch = Launch{ .grid = .{ 2, 1, 1 }, .block = .{ 4, 1, 1 } };
651 try machine.run(kernel.func().op, &.{
652 argumentBuffer(f32, input[0..]),
653 argumentBuffer(f32, output[0..]),
654 }, launch);
655 try expectApproxSlice(output[0..], clamped_fill_expected[0..]);
656
657 failing.fail_index = failing.alloc_index;
658 failing.resize_fail_index = failing.resize_index;
659 for (0..8) |_| {
660 @memset(output[0..], -1.0);
661 try machine.run(kernel.func().op, &.{
662 argumentBuffer(f32, input[0..]),
663 argumentBuffer(f32, output[0..]),
664 }, launch);
665 try expectApproxSlice(output[0..], clamped_fill_expected[0..]);
666 }
667 try testing.expect(!failing.has_induced_failure);
668 }
669
670 test "Machine grown grids reuse warmed block storage" {
671 var kernel = try buildClampedFillKernel(testing.allocator);
672 defer kernel.deinit();
673 try kernel.verify();
674
675 var failing = std.testing.FailingAllocator.init(testing.allocator, .{});
676 var machine = Machine.init(failing.allocator());
677 defer machine.deinit();
678
679 var input = clamped_fill_input;
680 var output = @as([8]f32, @splat(0));
681 try machine.run(kernel.func().op, &.{
682 argumentBuffer(f32, input[0..4]),
683 argumentBuffer(f32, output[0..4]),
684 }, Launch{ .grid = .{ 1, 1, 1 }, .block = .{ 4, 1, 1 } });
685 try expectApproxSlice(output[0..4], clamped_fill_expected[0..4]);
686
687 failing.fail_index = failing.alloc_index;
688 failing.resize_fail_index = failing.resize_index;
689 @memset(output[0..], -1.0);
690 try machine.run(kernel.func().op, &.{
691 argumentBuffer(f32, input[0..]),
692 argumentBuffer(f32, output[0..]),
693 }, Launch{ .grid = .{ 2, 1, 1 }, .block = .{ 4, 1, 1 } });
694 try expectApproxSlice(output[0..], clamped_fill_expected[0..]);
695 try testing.expect(!failing.has_induced_failure);
696 }
697
698 fn buildUnseededSharedKernel(allocator: std.mem.Allocator) !builder.Kernel {
699 var b = try builder.Builder.init(allocator, builder.Builder.Limits.testing, "machine_shared_unseeded_i32", &.{
700 builder.dynamicBuffer(.i32),
701 }, .{});
702 errdefer b.deinit();
703
704 const out = b.argument(0);
705 const i = try b.globalId(.x);
706 const bins = try b.sharedBuffer(.i32, 4);
707 const slot = try b.constantIndex(0);
708 const one = try b.constantInt(.i32, 1);
709 const old = try b.atomicRmw(.add, one, bins, slot);
710 try b.store(old, out, i);
711 try b.return_();
712 return b.finish();
713 }
714
715 test "Machine steady shared launches reuse bytes and reset the sentinel" {
716 var kernel = try buildUnseededSharedKernel(testing.allocator);
717 defer kernel.deinit();
718 try kernel.verify();
719
720 const fill: u8 = shared_alloc_sentinel;
721 const word: u32 = @as(u32, fill) << 24 | @as(u32, fill) << 16 | @as(u32, fill) << 8 | @as(u32, fill);
722 const sentinel_value: i32 = @bitCast(word);
723
724 var failing = std.testing.FailingAllocator.init(testing.allocator, .{});
725 var machine = Machine.init(failing.allocator());
726 defer machine.deinit();
727
728 var out = [_]i32{ -1, -1, -1, -1 };
729 try machine.run(kernel.func().op, &.{
730 argumentBuffer(i32, out[0..]),
731 }, Launch.oneDim(out.len));
732 try testing.expectEqualSlices(i32, &.{ sentinel_value, sentinel_value, sentinel_value, sentinel_value }, out[0..]);
733
734 failing.fail_index = failing.alloc_index;
735 failing.resize_fail_index = failing.resize_index;
736 for (0..8) |_| {
737 @memset(out[0..], -1);
738 try machine.run(kernel.func().op, &.{
739 argumentBuffer(i32, out[0..]),
740 }, Launch.oneDim(out.len));
741 try testing.expectEqualSlices(i32, &.{ sentinel_value, sentinel_value, sentinel_value, sentinel_value }, out[0..]);
742 }
743 try testing.expect(!failing.has_induced_failure);
744 }
745
746 test "Machine failed acquisition leaves the machine reusable" {
747 var kernel = try buildClampedFillKernel(testing.allocator);
748 defer kernel.deinit();
749 try kernel.verify();
750
751 var failing = std.testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
752 var machine = Machine.init(failing.allocator());
753 defer machine.deinit();
754
755 var input = clamped_fill_input;
756 var output = @as([8]f32, @splat(0));
757 const launch = Launch{ .grid = .{ 2, 1, 1 }, .block = .{ 4, 1, 1 } };
758 try testing.expectError(error.OutOfMemory, machine.run(kernel.func().op, &.{
759 argumentBuffer(f32, input[0..]),
760 argumentBuffer(f32, output[0..]),
761 }, launch));
762 try testing.expect(failing.has_induced_failure);
763
764 failing.fail_index = std.math.maxInt(usize);
765 try machine.run(kernel.func().op, &.{
766 argumentBuffer(f32, input[0..]),
767 argumentBuffer(f32, output[0..]),
768 }, launch);
769 try expectApproxSlice(output[0..], clamped_fill_expected[0..]);
770 }
771
772 fn runMachineLifecycle(allocator: std.mem.Allocator, kernel: *builder.Kernel) !void {
773 var machine = Machine.init(allocator);
774 defer machine.deinit();
775
776 var input = clamped_fill_input;
777 var output = @as([8]f32, @splat(0));
778 try machine.run(kernel.func().op, &.{
779 argumentBuffer(f32, input[0..4]),
780 argumentBuffer(f32, output[0..4]),
781 }, Launch{ .grid = .{ 1, 1, 1 }, .block = .{ 4, 1, 1 } });
782 for (output[0..4], clamped_fill_expected[0..4]) |actual, expected| {
783 if (actual != expected) return error.TestUnexpectedResult;
784 }
785
786 try machine.run(kernel.func().op, &.{
787 argumentBuffer(f32, input[0..]),
788 argumentBuffer(f32, output[0..]),
789 }, Launch{ .grid = .{ 1, 1, 1 }, .block = .{ 8, 1, 1 } });
790 for (output[0..], clamped_fill_expected[0..]) |actual, expected| {
791 if (actual != expected) return error.TestUnexpectedResult;
792 }
793 }
794
795 test "Machine lifecycle releases every acquisition failure and retries" {
796 var kernel = try buildClampedFillKernel(testing.allocator);
797 defer kernel.deinit();
798 try kernel.verify();
799
800 try @import("../../fixture/root.zig").checkAllAllocationFailures(runMachineLifecycle, .{&kernel});
801 try runMachineLifecycle(testing.allocator, &kernel);
802 }
803
804 test "Machine runs zero-thread launches without touching output" {
805 var kernel = try buildClampedFillKernel(testing.allocator);
806 defer kernel.deinit();
807 try kernel.verify();
808
809 var machine = Machine.init(testing.allocator);
810 defer machine.deinit();
811
812 var input = clamped_fill_input;
813 var output = @as([8]f32, @splat(-7.0));
814 try machine.run(kernel.func().op, &.{
815 argumentBuffer(f32, input[0..]),
816 argumentBuffer(f32, output[0..]),
817 }, Launch{ .grid = .{ 2, 1, 1 }, .block = .{ 0, 1, 1 } });
818 try expectApproxSlice(output[0..], &(@as([8]f32, @splat(-7.0))));
819 }
820
821 test "run executes popcount on CPU" {
822 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "popcount_i32", &.{
823 builder.dynamicBuffer(.i32),
824 builder.dynamicBuffer(.i32),
825 }, .{});
826 errdefer b.deinit();
827
828 const out = b.argument(0);
829 const keys = b.argument(1);
830 const i = try b.globalId(.x);
831 const key = try b.load(keys, i);
832 const bits = try b.popcount(key);
833 try b.store(bits, out, i);
834 try b.return_();
835
836 var kernel = try b.finish();
837 defer kernel.deinit();
838
839 var keys_data = [_]i32{ 0, 1, 3, 7, 255, -1, 0x55555555, @bitCast(@as(u32, 0x80000000)) };
840 var counts = @as([8]i32, @splat(-1));
841 try run(testing.allocator, &kernel, &.{
842 argumentBuffer(i32, counts[0..]),
843 argumentBuffer(i32, keys_data[0..]),
844 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 8, 1, 1 } });
845
846 try testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3, 8, 32, 16, 1 }, counts[0..]);
847 }
848
849 test "run executes ballot rank on CPU" {
850 var b = try builder.Builder.init(testing.allocator, builder.Builder.Limits.testing, "ballot_rank_i32", &.{
851 builder.dynamicBuffer(.i32),
852 builder.dynamicBuffer(.i32),
853 }, .{});
854 errdefer b.deinit();
855
856 const out = b.argument(0);
857 const keys = b.argument(1);
858 const i = try b.globalId(.x);
859 const key = try b.load(keys, i);
860 const one = try b.constantInt(.i32, 1);
861 const low_bit = try b.and_(key, one);
862 const is_odd = try b.compare(.eq, low_bit, one);
863 const mask = try b.ballotSync(is_odd);
864 const lane = try b.laneId();
865 const lane_i32 = try b.cast(lane, .i32);
866 const lane_bit = try b.shl(one, lane_i32);
867 const lower = try b.sub(lane_bit, one);
868 const below = try b.and_(mask, lower);
869 const rank = try b.popcount(below);
870 try b.store(rank, out, i);
871 try b.return_();
872
873 var kernel = try b.finish();
874 defer kernel.deinit();
875
876 var keys_data: [32]i32 = undefined;
877 for (&keys_data, 0..) |*entry_value, index| entry_value.* = @intCast(index);
878 var ranks = @as([32]i32, @splat(-1));
879 try run(testing.allocator, &kernel, &.{
880 argumentBuffer(i32, ranks[0..]),
881 argumentBuffer(i32, keys_data[0..]),
882 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 32, 1, 1 } });
883
884 for (ranks, 0..) |observed_rank, index| {
885 const odd_below: i32 = @intCast(index / 2);
886 try testing.expectEqual(odd_below, observed_rank);
887 }
888 }