lib/accy/src/kernel/library/scan.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../../artifact/model/root.zig");
6 const shape = @import("../../choir/shape/root.zig");
7 const entry = @import("entry.zig");
8 const extent_mod = @import("extent.zig");
9 const kernel = @import("../root.zig");
10 const tuning = @import("tuning.zig");
11
12 const DType = choir_abi.DType;
13 const indexExtent = extent_mod.indexExtent;
14 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
15
16 pub const PrefixSumMode = enum {
17 inclusive,
18 exclusive,
19
20 pub fn operation(self: PrefixSumMode) entry.ScanOperator {
21 return switch (self) {
22 .inclusive => .prefix_sum,
23 .exclusive => .prefix_sum_exclusive,
24 };
25 }
26
27 pub fn fromOperation(operator: entry.ScanOperator) PrefixSumMode {
28 return switch (operator) {
29 .prefix_sum => .inclusive,
30 .prefix_sum_exclusive => .exclusive,
31 };
32 }
33 };
34
35 pub const PrefixSum = struct {
36 extent: u64,
37 dtype: DType = .f32,
38 mode: PrefixSumMode = .inclusive,
39 threads: u32 = 256,
40 element_axis: []const u8 = "e",
41 };
42
43 pub const prefix_sum_family_version: u32 = 1;
44 pub const prefix_sum_warp_size: u32 = 32;
45 pub const prefix_sum_max_threads: u32 = 1024;
46
47 pub fn prefixSumDTypeSupported(dtype: DType) bool {
48 return switch (dtype) {
49 .f32, .f16, .u32 => true,
50 else => false,
51 };
52 }
53
54 pub fn prefixSumInstanceValid(instance: PrefixSum) bool {
55 if (!prefixSumDTypeSupported(instance.dtype)) return false;
56 if (instance.extent == 0) return false;
57 if (instance.threads == 0 or instance.threads > prefix_sum_max_threads) return false;
58 if (instance.threads % prefix_sum_warp_size != 0) return false;
59 return instance.extent <= instance.threads;
60 }
61
62 pub fn prefixSumThreadsForExtent(extent: u64) ?u32 {
63 if (extent == 0 or extent > prefix_sum_max_threads) return null;
64 const wide: u64 = extent + prefix_sum_warp_size - 1;
65 const rounded: u32 = @intCast((wide / prefix_sum_warp_size) * prefix_sum_warp_size);
66 return @max(rounded, prefix_sum_warp_size);
67 }
68
69 pub const PrefixSumThreadCandidates = struct {
70 count: usize = 0,
71 items: [6]u32 = @as([6]u32, @splat(0)),
72
73 pub fn slice(self: *const PrefixSumThreadCandidates) []const u32 {
74 return self.items[0..self.count];
75 }
76 };
77
78 pub fn prefixSumThreadCandidatesForExtent(extent: u64) PrefixSumThreadCandidates {
79 var result = PrefixSumThreadCandidates{};
80 const base = prefixSumThreadsForExtent(extent) orelse return result;
81 result.items[result.count] = base;
82 result.count += 1;
83 var threads: u32 = prefix_sum_warp_size;
84 while (threads <= prefix_sum_max_threads) : (threads *= 2) {
85 if (threads == base) continue;
86 if (@as(u64, threads) < extent) continue;
87 if (result.count >= result.items.len) break;
88 result.items[result.count] = threads;
89 result.count += 1;
90 }
91 return result;
92 }
93
94 fn prefix_sum_scan_core_seeds_zero(inner: anytype, ctx: anytype) !void {
95 try inner.storeIndex(ctx.zero_value, ctx.warp_sums, ctx.local);
96 }
97
98 fn prefix_sum_scan_core_is_last_lane(inner: anytype, ctx: anytype) !void {
99 try inner.storeIndex(ctx.scanned, ctx.warp_sums, ctx.warp);
100 }
101
102 fn prefix_sum_scan_core_is_first_warp(inner: anytype, ctx: anytype) !void {
103 const warp_sum = try inner.loadIndex(ctx.warp_sums, ctx.lane);
104 const warp_scan = try inner.warpScan(.add, .inclusive, warp_sum);
105 try inner.storeIndex(warp_scan, ctx.warp_sums, ctx.lane);
106 }
107
108 fn prefix_sum_scan_core_in_range(inner: anytype, ctx: anytype) !void {
109 try ctx.args.param(.dst).store(inner, ctx.result, ctx.tid);
110 }
111
112 fn prefixSumScanCore(
113 k: anytype,
114 spec: PrefixSum,
115 args: anytype,
116 extent: kernel.Value,
117 local: kernel.Value,
118 ) !kernel.Value {
119 const tid = try k.globalId(.x);
120 const lane = try k.laneId();
121 const warp = try k.warpId();
122 const zero = try k.constantIndex(0);
123 const one = try k.constantIndex(1);
124
125 const in_range = try k.compare(.lt, tid, extent);
126 const extent_minus_one = try k.sub(extent, one);
127 const clamped_tid = try k.min(tid, extent_minus_one);
128 const loaded = try args.param(.data).load(k, clamped_tid);
129 const accumulator_dtype = comptime prefixSumAccumulatorDType(@TypeOf(loaded).scalar_dtype);
130 const loaded_accumulated = if (comptime @TypeOf(loaded).scalar_dtype == accumulator_dtype)
131 loaded.raw()
132 else
133 (try loaded.cast(k, accumulator_dtype)).raw();
134 const zero_value = try zeroForDType(k, accumulator_dtype);
135 const element = try k.select(in_range, loaded_accumulated, zero_value);
136
137 const scanned = try k.warpScan(.add, .inclusive, element);
138
139 const warp_sums = try k.sharedBuffer(accumulator_dtype, prefix_sum_warp_size);
140 const lane_limit = try k.constantIndex(prefix_sum_warp_size - 1);
141 const warp_count_value = try k.constantIndex(prefix_sum_warp_size);
142
143 const seeds_zero = try k.compare(.lt, local, warp_count_value);
144 try k.guardDo(seeds_zero, .{ .warp_sums = warp_sums, .local = local, .zero_value = zero_value }, prefix_sum_scan_core_seeds_zero);
145 try k.barrier(.block);
146
147 const is_last_lane = try k.compare(.eq, lane, lane_limit);
148 try k.guardDo(is_last_lane, .{ .warp_sums = warp_sums, .warp = warp, .scanned = scanned }, prefix_sum_scan_core_is_last_lane);
149 try k.barrier(.block);
150
151 const is_first_warp = try k.compare(.eq, warp, zero);
152 try k.guardDo(is_first_warp, .{ .warp_sums = warp_sums, .lane = lane }, prefix_sum_scan_core_is_first_warp);
153 try k.barrier(.block);
154
155 const has_base = try k.compare(.gt, warp, zero);
156 const warp_minus_one = try k.sub(warp, one);
157 const base_index = try k.select(has_base, warp_minus_one, zero);
158 const base_loaded = try k.loadIndex(warp_sums, base_index);
159 const base = try k.select(has_base, base_loaded, zero_value);
160 const inclusive = try k.add(scanned, base);
161 const accumulated = switch (spec.mode) {
162 .inclusive => inclusive,
163 .exclusive => try k.sub(inclusive, element),
164 };
165 const dst_dtype = comptime @TypeOf(args.param(.dst)).element_dtype;
166 const result = switch (dst_dtype) {
167 .f32 => accumulated,
168 .u32 => accumulated,
169 .f16 => try k.cast(accumulated, .f16),
170 else => return error.UnsupportedDType,
171 };
172
173 try k.guardDo(in_range, .{ .args = args, .tid = tid, .result = result }, prefix_sum_scan_core_in_range);
174
175 return warp_sums;
176 }
177
178 fn prefixSumAccumulatorDType(dtype: DType) DType {
179 return switch (dtype) {
180 .f16 => .f32,
181 else => dtype,
182 };
183 }
184
185 fn zeroForDType(k: anytype, comptime dtype: DType) !kernel.Value {
186 return switch (dtype) {
187 .f32 => k.constantFloat(.f32, 0.0),
188 .u32 => k.constantInt(.u32, 0),
189 else => error.UnsupportedDType,
190 };
191 }
192
193 fn prefixSumBody(k: anytype, spec: PrefixSum, args: anytype) !void {
194 if (!prefixSumInstanceValid(spec)) return error.UnsupportedPrefixSumInstance;
195 const extent = try k.constantIndex(try indexExtent(spec.extent));
196 const local = try k.threadId(.x);
197 _ = try prefixSumScanCore(k, spec, args, extent, local);
198 }
199
200 fn prefixSumRuntimeBody(k: anytype, spec: PrefixSum, args: anytype) !void {
201 if (!prefixSumInstanceValid(spec)) return error.UnsupportedPrefixSumInstance;
202 const extent = try k.castIndex(args.param(.extent).raw());
203 const local = try k.threadId(.x);
204 _ = try prefixSumScanCore(k, spec, args, extent, local);
205 }
206
207 fn prefixSumFamilySchedule(instance: PrefixSum) kernel.logical.schedule.ThreadBlocks {
208 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
209 }
210
211 fn prefixSumRuntimeFamily(comptime dtype: DType) type {
212 return kernel.logical.Family(.{
213 .name = std.fmt.comptimePrint("accy_kernel_scan_prefix_sum_runtime_{s}", .{dtype.name()}),
214 .parameters = .{
215 .dst = kernel.dynamicBuffer(dtype),
216 .data = kernel.dynamicBuffer(dtype),
217 .extent = kernel.scalar(.i32),
218 },
219 .Instance = PrefixSum,
220 .schedule = prefixSumFamilySchedule,
221 .body = prefixSumRuntimeBody,
222 });
223 }
224
225 pub const PrefixSumRuntimeFamilyF32 = prefixSumRuntimeFamily(.f32);
226 pub const PrefixSumRuntimeFamilyF16 = prefixSumRuntimeFamily(.f16);
227 pub const PrefixSumRuntimeFamilyU32 = prefixSumRuntimeFamily(.u32);
228
229 pub const DeviceScan = struct {
230 extent: u64,
231 dtype: DType = .f32,
232 mode: PrefixSumMode = .inclusive,
233 threads: u32 = 256,
234 element_axis: []const u8 = "e",
235 };
236
237 pub const device_scan_family_version: u32 = 1;
238 pub const device_scan_max_blocks: u32 = prefix_sum_max_threads;
239
240 pub fn deviceScanBlockCount(extent: u64, threads: u32) u64 {
241 return (extent + threads - 1) / threads;
242 }
243
244 pub fn deviceScanDTypeSupported(dtype: DType) bool {
245 return switch (dtype) {
246 .f32, .f16, .u32 => true,
247 else => false,
248 };
249 }
250
251 pub fn deviceScanInstanceValid(instance: DeviceScan) bool {
252 if (!deviceScanDTypeSupported(instance.dtype)) return false;
253 if (instance.extent == 0) return false;
254 if (instance.threads == 0 or instance.threads > prefix_sum_max_threads) return false;
255 if (instance.threads % prefix_sum_warp_size != 0) return false;
256 return extent_mod.blockCountWithinLimit(instance.extent, instance.threads, device_scan_max_blocks);
257 }
258
259 fn deviceScanPrefixSum(spec: DeviceScan) PrefixSum {
260 return .{
261 .extent = spec.extent,
262 .dtype = spec.dtype,
263 .mode = spec.mode,
264 .threads = spec.threads,
265 .element_axis = spec.element_axis,
266 };
267 }
268
269 fn device_scan_block_scan_body_writes_total(inner: anytype, ctx: anytype) !void {
270 try ctx.args.param(.sums).store(inner, ctx.total, ctx.block);
271 }
272
273 fn deviceScanBlockScanBody(k: anytype, spec: DeviceScan, args: anytype) !void {
274 if (!deviceScanInstanceValid(spec)) return error.UnsupportedDeviceScanInstance;
275 const extent = try k.castIndex(args.param(.extent).raw());
276 const local = try k.threadId(.x);
277 const warp_sums = try prefixSumScanCore(k, deviceScanPrefixSum(spec), args, extent, local);
278 const block = try k.blockId(.x);
279 const zero = try k.constantIndex(0);
280 const last_slot = try k.constantIndex(prefix_sum_warp_size - 1);
281 const total = try k.loadIndex(warp_sums, last_slot);
282 const writes_total = try k.compare(.eq, local, zero);
283 try k.guardDo(writes_total, .{ .args = args, .total = total, .block = block }, device_scan_block_scan_body_writes_total);
284 }
285
286 fn device_scan_add_base_body_in_range(inner: anytype, ctx: anytype) !void {
287 const current = try ctx.args.param(.dst).load(inner, ctx.tid);
288 const updated = try inner.add(current.raw(), ctx.base.raw());
289 try ctx.args.param(.dst).store(inner, updated, ctx.tid);
290 }
291
292 fn deviceScanAddBaseBody(k: anytype, spec: DeviceScan, args: anytype) !void {
293 if (!deviceScanInstanceValid(spec)) return error.UnsupportedDeviceScanInstance;
294 const extent = try k.castIndex(args.param(.extent).raw());
295 const tid = try k.globalId(.x);
296 const block = try k.blockId(.x);
297 const base = try args.param(.base).load(k, block);
298 const in_range = try k.compare(.lt, tid, extent);
299 try k.guardDo(in_range, .{ .args = args, .tid = tid, .base = base }, device_scan_add_base_body_in_range);
300 }
301
302 fn deviceScanFamilySchedule(instance: DeviceScan) kernel.logical.schedule.ThreadBlocks {
303 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
304 }
305
306 fn deviceScanBlockScanRuntimeFamily(comptime dtype: DType) type {
307 return kernel.logical.Family(.{
308 .name = std.fmt.comptimePrint("accy_kernel_scan_device_block_scan_runtime_{s}", .{dtype.name()}),
309 .parameters = .{
310 .dst = kernel.dynamicBuffer(deviceScanBlockScanOutputDType(dtype)),
311 .data = kernel.dynamicBuffer(dtype),
312 .sums = kernel.dynamicBuffer(prefixSumAccumulatorDType(dtype)),
313 .extent = kernel.scalar(.i32),
314 },
315 .Instance = DeviceScan,
316 .schedule = deviceScanFamilySchedule,
317 .body = deviceScanBlockScanBody,
318 });
319 }
320
321 fn deviceScanBlockScanOutputDType(comptime dtype: DType) DType {
322 return switch (dtype) {
323 .f16 => .f32,
324 else => dtype,
325 };
326 }
327
328 fn device_scan_add_base_body_f16_in_range(inner: anytype, ctx: anytype) !void {
329 const local = try ctx.args.param(.local).load(inner, ctx.tid);
330 const updated = try inner.add(local.raw(), ctx.base.raw());
331 const result = try inner.cast(updated, .f16);
332 try ctx.args.param(.dst).store(inner, result, ctx.tid);
333 }
334
335 fn deviceScanAddBaseBodyF16(k: anytype, spec: DeviceScan, args: anytype) !void {
336 if (!deviceScanInstanceValid(spec)) return error.UnsupportedDeviceScanInstance;
337 const extent = try k.castIndex(args.param(.extent).raw());
338 const tid = try k.globalId(.x);
339 const block = try k.blockId(.x);
340 const base = try args.param(.base).load(k, block);
341 const in_range = try k.compare(.lt, tid, extent);
342 try k.guardDo(in_range, .{ .args = args, .tid = tid, .base = base }, device_scan_add_base_body_f16_in_range);
343 }
344
345 fn deviceScanAddBaseRuntimeFamily(comptime dtype: DType) type {
346 return kernel.logical.Family(.{
347 .name = std.fmt.comptimePrint("accy_kernel_scan_device_add_base_runtime_{s}", .{dtype.name()}),
348 .parameters = switch (dtype) {
349 .f16 => .{
350 .dst = kernel.dynamicBuffer(dtype),
351 .local = kernel.dynamicBuffer(.f32),
352 .base = kernel.dynamicBuffer(.f32),
353 .extent = kernel.scalar(.i32),
354 },
355 else => .{
356 .dst = kernel.dynamicBuffer(dtype),
357 .base = kernel.dynamicBuffer(prefixSumAccumulatorDType(dtype)),
358 .extent = kernel.scalar(.i32),
359 },
360 },
361 .Instance = DeviceScan,
362 .schedule = deviceScanFamilySchedule,
363 .body = switch (dtype) {
364 .f16 => deviceScanAddBaseBodyF16,
365 else => deviceScanAddBaseBody,
366 },
367 });
368 }
369
370 pub const DeviceScanBlockScanRuntimeFamilyF32 = deviceScanBlockScanRuntimeFamily(.f32);
371 pub const DeviceScanBlockScanRuntimeFamilyF16 = deviceScanBlockScanRuntimeFamily(.f16);
372 pub const DeviceScanBlockScanRuntimeFamilyU32 = deviceScanBlockScanRuntimeFamily(.u32);
373 pub const DeviceScanAddBaseRuntimeFamilyF32 = deviceScanAddBaseRuntimeFamily(.f32);
374 pub const DeviceScanAddBaseRuntimeFamilyF16 = deviceScanAddBaseRuntimeFamily(.f16);
375 pub const DeviceScanAddBaseRuntimeFamilyU32 = deviceScanAddBaseRuntimeFamily(.u32);
376
377 pub const DeviceScanStages = struct {
378 block_count: u32,
379 block_scan: DeviceScan,
380 sums_scan: PrefixSum,
381 add_base: DeviceScan,
382 };
383
384 pub fn deviceScanStages(instance: DeviceScan) !DeviceScanStages {
385 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
386 const blocks: u32 = @intCast(deviceScanBlockCount(instance.extent, instance.threads));
387 const sums_threads = prefixSumThreadsForExtent(blocks) orelse return error.UnsupportedDeviceScanInstance;
388 return .{
389 .block_count = blocks,
390 .block_scan = instance,
391 .sums_scan = .{
392 .extent = blocks,
393 .dtype = prefixSumAccumulatorDType(instance.dtype),
394 .mode = .exclusive,
395 .threads = sums_threads,
396 .element_axis = "b",
397 },
398 .add_base = instance,
399 };
400 }
401
402 pub fn deviceScanBlockScanFamilyTarget(allocator: std.mem.Allocator, instance: DeviceScan) ![]u8 {
403 return std.fmt.allocPrint(
404 allocator,
405 "accy.kernel.scan.device_{s}_block_scan_family_{d}_{s}",
406 .{ prefixSumModePrefix(instance.mode), instance.threads, instance.dtype.name() },
407 );
408 }
409
410 pub fn deviceScanBlockScanFamilyEntryName(allocator: std.mem.Allocator, instance: DeviceScan) ![]u8 {
411 return std.fmt.allocPrint(
412 allocator,
413 "accy_kernel_scan_device_{s}_block_scan_family_{d}_{s}",
414 .{ prefixSumModePrefix(instance.mode), instance.threads, instance.dtype.name() },
415 );
416 }
417
418 pub fn deviceScanAddBaseFamilyTarget(allocator: std.mem.Allocator, instance: DeviceScan) ![]u8 {
419 return std.fmt.allocPrint(
420 allocator,
421 "accy.kernel.scan.device_add_base_family_{d}_{s}",
422 .{ instance.threads, instance.dtype.name() },
423 );
424 }
425
426 pub fn deviceScanAddBaseFamilyEntryName(allocator: std.mem.Allocator, instance: DeviceScan) ![]u8 {
427 return std.fmt.allocPrint(
428 allocator,
429 "accy_kernel_scan_device_add_base_family_{d}_{s}",
430 .{ instance.threads, instance.dtype.name() },
431 );
432 }
433
434 pub fn deviceScanFamilyTarget(allocator: std.mem.Allocator, instance: DeviceScan) ![]u8 {
435 return std.fmt.allocPrint(
436 allocator,
437 "accy.kernel.scan.device_{s}_family_{d}_{s}",
438 .{ prefixSumModePrefix(instance.mode), instance.threads, instance.dtype.name() },
439 );
440 }
441
442 pub fn deviceScanThreadsForExtent(extent: u64) ?u32 {
443 if (extent == 0) return null;
444 const max_extent = @as(u64, prefix_sum_max_threads) * device_scan_max_blocks;
445 if (extent > max_extent) return null;
446 const needed = (extent + device_scan_max_blocks - 1) / device_scan_max_blocks;
447 const wide = needed + prefix_sum_warp_size - 1;
448 const rounded: u32 = @intCast((wide / prefix_sum_warp_size) * prefix_sum_warp_size);
449 return @max(rounded, prefix_sum_warp_size);
450 }
451
452 pub const DeviceScanThreadCandidates = struct {
453 count: usize = 0,
454 items: [6]u32 = @as([6]u32, @splat(0)),
455
456 pub fn slice(self: *const DeviceScanThreadCandidates) []const u32 {
457 return self.items[0..self.count];
458 }
459 };
460
461 pub fn deviceScanThreadCandidatesForExtent(extent: u64) DeviceScanThreadCandidates {
462 var result = DeviceScanThreadCandidates{};
463 const base = deviceScanThreadsForExtent(extent) orelse return result;
464 result.items[result.count] = base;
465 result.count += 1;
466 var threads: u32 = prefix_sum_warp_size;
467 while (threads <= prefix_sum_max_threads) : (threads *= 2) {
468 if (threads == base) continue;
469 if (threads < base) continue;
470 if (result.count >= result.items.len) break;
471 result.items[result.count] = threads;
472 result.count += 1;
473 }
474 return result;
475 }
476
477 pub fn deviceScanFamilySpecialization(backing_allocator: std.mem.Allocator, instance: DeviceScan) !entry.OwnedSpecialization {
478 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
479 var owned = entry.OwnedSpecialization.init(backing_allocator);
480 errdefer owned.deinit();
481 const lifetime_allocator = owned.allocator();
482
483 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
484 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
485
486 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
487 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
488
489 owned.value = .{
490 .dtype = instance.dtype,
491 .operation = .{ .scan = instance.mode.operation() },
492 .inputs = inputs,
493 .outputs = outputs,
494 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.element_axis, instance.extent, instance.threads),
495 };
496 owned.value.launch = owned.value.schedule.?.launch();
497 var family = try deviceScanBlockScanShapeFamily(backing_allocator, instance);
498 errdefer family.deinit();
499 try owned.takeShapeFamily(&family);
500 return owned;
501 }
502
503 pub fn deviceScanInstanceFromSpecialization(specialization: entry.Specialization) ?DeviceScan {
504 if (!specialization.scheduleMatchesLaunch()) return null;
505 const mode: PrefixSumMode = if (specialization.operationIs(.{ .scan = .prefix_sum }))
506 .inclusive
507 else if (specialization.operationIs(.{ .scan = .prefix_sum_exclusive }))
508 .exclusive
509 else
510 return null;
511 const dtype = specialization.dtype orelse return null;
512 if (!deviceScanDTypeSupported(dtype)) return null;
513 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
514 if (specialization.reductions.len != 0) return null;
515 const data = specialization.inputs[0];
516 const output = specialization.outputs[0];
517 if (data.axes.len != 1 or output.axes.len != 1) return null;
518 const extent = data.axes[0].extent;
519 if (output.axes[0].extent != extent) return null;
520 const launch = specialization.launch orelse return null;
521 if (launch.grid[0] <= 1) return null;
522 const instance = DeviceScan{
523 .extent = extent,
524 .dtype = dtype,
525 .mode = mode,
526 .threads = launch.threadgroup[0],
527 .element_axis = data.axes[0].name,
528 };
529 if (!deviceScanInstanceValid(instance)) return null;
530 if (launch.grid[0] != deviceScanBlockCount(extent, instance.threads)) return null;
531 return instance;
532 }
533
534 pub fn deviceScanPipeline(
535 backing_allocator: std.mem.Allocator,
536 instance: DeviceScan,
537 ) !artifact_product.OwnedKernelCallPipeline {
538 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
539 const stages = try deviceScanStages(instance);
540 var owned = artifact_product.OwnedKernelCallPipeline.init(backing_allocator);
541 errdefer owned.deinit();
542 const arena = owned.allocator();
543
544 const block_count_extent = artifact_product.PipelineScalarDerivation{
545 .ceil_div = .{ .argument_index = 0, .divisor = instance.threads },
546 };
547 const accumulator_dtype = prefixSumAccumulatorDType(instance.dtype);
548 const local_prefix_index: ?u32 = if (instance.dtype == .f16) 2 else null;
549 const intermediate_count: usize = if (local_prefix_index != null) 3 else 2;
550 const intermediates = try arena.alloc(artifact_product.PipelineIntermediate, intermediate_count);
551 intermediates[0] = .{ .dtype = accumulator_dtype, .extent = block_count_extent };
552 intermediates[1] = .{ .dtype = accumulator_dtype, .extent = block_count_extent };
553 if (local_prefix_index) |index| {
554 intermediates[index] = .{ .dtype = accumulator_dtype, .extent = .{ .forward = 0 } };
555 }
556 const runtime_scalar_bounds = try arena.alloc(artifact_product.PipelineRuntimeScalarBound, 1);
557 runtime_scalar_bounds[0] = .{
558 .argument_index = 0,
559 .max_u32 = try deviceScanPipelineRuntimeExtentCapacity(instance, stages),
560 };
561
562 const pipeline_stages = try arena.alloc(artifact_product.PipelineStage, 3);
563 pipeline_stages[0] = .{
564 .target = try deviceScanBlockScanFamilyTarget(arena, instance),
565 .version = device_scan_family_version,
566 .buffers = try arena.dupe(artifact_product.PipelineValueRef, if (local_prefix_index) |index| &.{
567 .{ .intermediate = index }, .{ .operand = 0 }, .{ .intermediate = 0 },
568 } else &.{
569 .{ .result = 0 }, .{ .operand = 0 }, .{ .intermediate = 0 },
570 }),
571 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
572 .{ .forward = 0 },
573 }),
574 };
575 pipeline_stages[1] = .{
576 .target = try prefixSumFamilyTarget(arena, stages.sums_scan),
577 .version = prefix_sum_family_version,
578 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
579 .{ .intermediate = 1 }, .{ .intermediate = 0 },
580 }),
581 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
582 block_count_extent,
583 }),
584 };
585 pipeline_stages[2] = .{
586 .target = try deviceScanAddBaseFamilyTarget(arena, instance),
587 .version = device_scan_family_version,
588 .buffers = try arena.dupe(artifact_product.PipelineValueRef, if (local_prefix_index) |index| &.{
589 .{ .result = 0 }, .{ .intermediate = index }, .{ .intermediate = 1 },
590 } else &.{
591 .{ .result = 0 }, .{ .intermediate = 1 },
592 }),
593 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
594 .{ .forward = 0 },
595 }),
596 };
597
598 owned.value = .{
599 .target = try deviceScanFamilyTarget(arena, instance),
600 .version = device_scan_family_version,
601 .operand_count = 1,
602 .result_count = 1,
603 .runtime_scalar_argument_count = 1,
604 .runtime_scalar_bounds = runtime_scalar_bounds,
605 .intermediates = intermediates,
606 .stages = pipeline_stages,
607 };
608 return owned;
609 }
610
611 pub const DeviceScanPipelineArtifacts = struct {
612 block_scan: kernel.OwnedKernelCallArtifact,
613 sums_scan: kernel.OwnedKernelCallArtifact,
614 add_base: kernel.OwnedKernelCallArtifact,
615
616 pub fn entries(self: *const DeviceScanPipelineArtifacts) [3]artifact_product.KernelCallArtifact {
617 return .{ self.block_scan.entry(), self.sums_scan.entry(), self.add_base.entry() };
618 }
619
620 pub fn deinit(self: *DeviceScanPipelineArtifacts) void {
621 self.block_scan.deinit();
622 self.sums_scan.deinit();
623 self.add_base.deinit();
624 self.* = undefined;
625 }
626 };
627
628 pub fn createDeviceScanPipelineArtifacts(
629 allocator: std.mem.Allocator,
630 handle: kernel.BackendHandle,
631 instance: DeviceScan,
632 options: entry.ArtifactOptions,
633 ) !DeviceScanPipelineArtifacts {
634 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
635 const stages = try deviceScanStages(instance);
636 var block_scan = try createDeviceScanBlockScanFamilyArtifact(allocator, handle, instance, options);
637 errdefer block_scan.deinit();
638 var sums_scan = try createPrefixSumFamilyArtifact(allocator, handle, stages.sums_scan, options);
639 errdefer sums_scan.deinit();
640 const add_base = try createDeviceScanAddBaseFamilyArtifact(allocator, handle, instance, options);
641 return .{ .block_scan = block_scan, .sums_scan = sums_scan, .add_base = add_base };
642 }
643
644 pub fn deviceScanRuntimeArguments(instance: DeviceScan) ![1]choir_abi.ScalarArgument {
645 return .{
646 .{ .u32 = try runtimeExtentArgument(instance.extent) },
647 };
648 }
649
650 fn deviceScanPipelineRuntimeExtentCapacity(instance: DeviceScan, stages: DeviceScanStages) !u32 {
651 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
652 return std.math.mul(u32, instance.threads, stages.sums_scan.threads) catch return error.UnsupportedDeviceScanInstance;
653 }
654
655 pub fn deviceScanMaxExtent(instance: DeviceScan) u64 {
656 return @as(u64, instance.threads) * device_scan_max_blocks;
657 }
658
659 pub fn deviceScanShapeProfileDimensions(instance: DeviceScan) [1]artifact_product.KernelCallShapeProfileDimension {
660 return .{
661 .{
662 .name = instance.element_axis,
663 .runtime_scalar_argument_index = 0,
664 .bounds = .{ .min = 1, .max = deviceScanMaxExtent(instance) },
665 },
666 };
667 }
668
669 fn deviceScanLaunch(instance: DeviceScan) !artifact_product.KernelCallLaunch {
670 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
671 return .{ .derived = .{
672 .grid = .{
673 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
674 .{ .fixed = 1 },
675 .{ .fixed = 1 },
676 },
677 .threadgroup = .{ instance.threads, 1, 1 },
678 } };
679 }
680
681 pub fn deviceScanBlockScanShapeFamily(backing_allocator: std.mem.Allocator, instance: DeviceScan) !shape.Family {
682 var builder = try shape.Builder.init(backing_allocator, "device_scan_block_scan");
683 errdefer builder.deinit();
684
685 const elements = try builder.symbol(instance.element_axis);
686 const elements_expr = try builder.symbolExpression(elements);
687 const blocks = try builder.symbol("blocks");
688 const blocks_expr = try builder.symbolExpression(blocks);
689
690 _ = try builder.tensor("data", &.{elements_expr});
691 _ = try builder.tensor("out", &.{elements_expr});
692 _ = try builder.tensor("sums", &.{blocks_expr});
693 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = deviceScanMaxExtent(instance) });
694 try builder.assumeBounds(blocks_expr, .{ .min = 1, .max = device_scan_max_blocks });
695
696 return builder.finish();
697 }
698
699 pub fn deviceScanBlockScanFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: DeviceScan) !u64 {
700 var family = try deviceScanBlockScanShapeFamily(backing_allocator, instance);
701 defer family.deinit();
702 return shape.fingerprint(family);
703 }
704
705 pub fn deviceScanAddBaseShapeFamily(backing_allocator: std.mem.Allocator, instance: DeviceScan) !shape.Family {
706 var builder = try shape.Builder.init(backing_allocator, "device_scan_add_base");
707 errdefer builder.deinit();
708
709 const elements = try builder.symbol(instance.element_axis);
710 const elements_expr = try builder.symbolExpression(elements);
711 const blocks = try builder.symbol("blocks");
712 const blocks_expr = try builder.symbolExpression(blocks);
713
714 _ = try builder.tensor("out", &.{elements_expr});
715 if (instance.dtype == .f16) {
716 _ = try builder.tensor("local", &.{elements_expr});
717 }
718 _ = try builder.tensor("base", &.{blocks_expr});
719 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = deviceScanMaxExtent(instance) });
720 try builder.assumeBounds(blocks_expr, .{ .min = 1, .max = device_scan_max_blocks });
721
722 return builder.finish();
723 }
724
725 pub fn deviceScanAddBaseFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: DeviceScan) !u64 {
726 var family = try deviceScanAddBaseShapeFamily(backing_allocator, instance);
727 defer family.deinit();
728 return shape.fingerprint(family);
729 }
730
731 pub fn createDeviceScanBlockScanFamilyArtifact(
732 allocator: std.mem.Allocator,
733 handle: kernel.BackendHandle,
734 instance: DeviceScan,
735 options: entry.ArtifactOptions,
736 ) !kernel.OwnedKernelCallArtifact {
737 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
738 const target = try deviceScanBlockScanFamilyTarget(allocator, instance);
739 defer allocator.free(target);
740 const entry_name = try deviceScanBlockScanFamilyEntryName(allocator, instance);
741 defer allocator.free(entry_name);
742 const family_fingerprint = options.shape_family_fingerprint orelse try deviceScanBlockScanFamilyFingerprint(allocator, instance);
743 const shape_profile_dimensions = deviceScanShapeProfileDimensions(instance);
744 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
745 .name = "device_scan_block_scan",
746 .fingerprint = family_fingerprint,
747 .dimensions = shape_profile_dimensions[0..],
748 };
749
750 var graph = switch (instance.dtype) {
751 .f32 => try DeviceScanBlockScanRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
752 .f16 => try DeviceScanBlockScanRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
753 .u32 => try DeviceScanBlockScanRuntimeFamilyU32.buildNamed(allocator, options.limits, entry_name, instance),
754 else => return error.UnsupportedDType,
755 };
756 defer graph.deinit();
757 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
758 .target = target,
759 .version = device_scan_family_version,
760 .format = options.format,
761 .kernel_plan = options.kernel_plan,
762 .element_count_argument = options.element_count_argument,
763 .shape_family_fingerprint = family_fingerprint,
764 .shape_profile = shape_profile,
765 .launch = options.launch orelse try deviceScanLaunch(instance),
766 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
767 .static_arguments = options.static_arguments,
768 });
769 }
770
771 pub fn createDeviceScanAddBaseFamilyArtifact(
772 allocator: std.mem.Allocator,
773 handle: kernel.BackendHandle,
774 instance: DeviceScan,
775 options: entry.ArtifactOptions,
776 ) !kernel.OwnedKernelCallArtifact {
777 if (!deviceScanInstanceValid(instance)) return error.UnsupportedDeviceScanInstance;
778 const target = try deviceScanAddBaseFamilyTarget(allocator, instance);
779 defer allocator.free(target);
780 const entry_name = try deviceScanAddBaseFamilyEntryName(allocator, instance);
781 defer allocator.free(entry_name);
782 const family_fingerprint = options.shape_family_fingerprint orelse try deviceScanAddBaseFamilyFingerprint(allocator, instance);
783 const shape_profile_dimensions = deviceScanShapeProfileDimensions(instance);
784 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
785 .name = "device_scan_add_base",
786 .fingerprint = family_fingerprint,
787 .dimensions = shape_profile_dimensions[0..],
788 };
789
790 var graph = switch (instance.dtype) {
791 .f32 => try DeviceScanAddBaseRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
792 .f16 => try DeviceScanAddBaseRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
793 .u32 => try DeviceScanAddBaseRuntimeFamilyU32.buildNamed(allocator, options.limits, entry_name, instance),
794 else => return error.UnsupportedDType,
795 };
796 defer graph.deinit();
797 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
798 .target = target,
799 .version = device_scan_family_version,
800 .format = options.format,
801 .kernel_plan = options.kernel_plan,
802 .element_count_argument = options.element_count_argument,
803 .shape_family_fingerprint = family_fingerprint,
804 .shape_profile = shape_profile,
805 .launch = options.launch orelse try deviceScanLaunch(instance),
806 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
807 .static_arguments = options.static_arguments,
808 });
809 }
810
811 pub fn prefixSumInstanceTarget(allocator: std.mem.Allocator, instance: PrefixSum) ![]u8 {
812 return std.fmt.allocPrint(
813 allocator,
814 "accy.kernel.scan.prefix_sum{d}_{d}_{s}",
815 .{ instance.extent, instance.threads, instance.dtype.name() },
816 );
817 }
818
819 pub fn prefixSumInstanceEntryName(allocator: std.mem.Allocator, instance: PrefixSum) ![]u8 {
820 return std.fmt.allocPrint(
821 allocator,
822 "accy_kernel_scan_prefix_sum{d}_{d}_{s}",
823 .{ instance.extent, instance.threads, instance.dtype.name() },
824 );
825 }
826
827 fn prefixSumModePrefix(mode: PrefixSumMode) []const u8 {
828 return switch (mode) {
829 .inclusive => "prefix_sum",
830 .exclusive => "prefix_sum_exclusive",
831 };
832 }
833
834 pub fn prefixSumFamilyTarget(allocator: std.mem.Allocator, instance: PrefixSum) ![]u8 {
835 return std.fmt.allocPrint(
836 allocator,
837 "accy.kernel.scan.{s}_family_{d}_{s}",
838 .{ prefixSumModePrefix(instance.mode), instance.threads, instance.dtype.name() },
839 );
840 }
841
842 pub fn prefixSumFamilyEntryName(allocator: std.mem.Allocator, instance: PrefixSum) ![]u8 {
843 return std.fmt.allocPrint(
844 allocator,
845 "accy_kernel_scan_{s}_family_{d}_{s}",
846 .{ prefixSumModePrefix(instance.mode), instance.threads, instance.dtype.name() },
847 );
848 }
849
850 pub fn prefixSumTuningExtents(instance: PrefixSum) [1]u64 {
851 return .{instance.extent};
852 }
853
854 pub fn prefixSumTuningOperation(instance: PrefixSum) entry.Operation {
855 return .{ .scan = instance.mode.operation() };
856 }
857
858 pub fn prefixSumFamilyTuningKey(
859 backing_allocator: std.mem.Allocator,
860 device_fingerprint: u64,
861 instance: PrefixSum,
862 ) !tuning.FamilyTuningKey {
863 const family_fingerprint = try prefixSumFamilyFingerprint(backing_allocator, instance);
864 const extents = prefixSumTuningExtents(instance);
865 return tuning.FamilyTuningKey.init(
866 device_fingerprint,
867 family_fingerprint,
868 entry.operationFingerprint(prefixSumTuningOperation(instance)),
869 instance.dtype,
870 prefix_sum_family_version,
871 extents[0..],
872 ) orelse unreachable;
873 }
874
875 pub fn prefixSumRuntimeArguments(instance: PrefixSum) ![1]choir_abi.ScalarArgument {
876 return .{
877 .{ .u32 = try runtimeExtentArgument(instance.extent) },
878 };
879 }
880
881 pub fn prefixSumShapeProfileDimensions(instance: PrefixSum) [1]artifact_product.KernelCallShapeProfileDimension {
882 return .{
883 .{
884 .name = instance.element_axis,
885 .runtime_scalar_argument_index = 0,
886 .bounds = .{ .min = 1, .max = instance.threads },
887 },
888 };
889 }
890
891 fn prefixSumLaunch(instance: PrefixSum) !artifact_product.KernelCallLaunch {
892 if (!prefixSumInstanceValid(instance)) return error.UnsupportedPrefixSumInstance;
893 return .{ .derived = .{
894 .grid = .{
895 .{ .fixed = 1 },
896 .{ .fixed = 1 },
897 .{ .fixed = 1 },
898 },
899 .threadgroup = .{ instance.threads, 1, 1 },
900 } };
901 }
902
903 pub fn createPrefixSumFamilyArtifact(
904 allocator: std.mem.Allocator,
905 handle: kernel.BackendHandle,
906 instance: PrefixSum,
907 options: entry.ArtifactOptions,
908 ) !kernel.OwnedKernelCallArtifact {
909 if (!prefixSumInstanceValid(instance)) return error.UnsupportedPrefixSumInstance;
910 const target = try prefixSumFamilyTarget(allocator, instance);
911 defer allocator.free(target);
912 const entry_name = try prefixSumFamilyEntryName(allocator, instance);
913 defer allocator.free(entry_name);
914 const family_fingerprint = options.shape_family_fingerprint orelse try prefixSumFamilyFingerprint(allocator, instance);
915 const shape_profile_dimensions = prefixSumShapeProfileDimensions(instance);
916 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
917 .name = "prefix_sum",
918 .fingerprint = family_fingerprint,
919 .dimensions = shape_profile_dimensions[0..],
920 };
921
922 var graph = switch (instance.dtype) {
923 .f32 => try PrefixSumRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
924 .f16 => try PrefixSumRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
925 .u32 => try PrefixSumRuntimeFamilyU32.buildNamed(allocator, options.limits, entry_name, instance),
926 else => return error.UnsupportedDType,
927 };
928 defer graph.deinit();
929 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
930 .target = target,
931 .version = prefix_sum_family_version,
932 .format = options.format,
933 .kernel_plan = options.kernel_plan,
934 .element_count_argument = options.element_count_argument,
935 .shape_family_fingerprint = family_fingerprint,
936 .shape_profile = shape_profile,
937 .launch = options.launch orelse try prefixSumLaunch(instance),
938 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
939 .static_arguments = options.static_arguments,
940 });
941 }
942
943 pub fn prefixSumFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: PrefixSum) !u64 {
944 var family = try prefixSumShapeFamily(backing_allocator, instance);
945 defer family.deinit();
946 return shape.fingerprint(family);
947 }
948
949 pub fn prefixSumShapeFamily(backing_allocator: std.mem.Allocator, instance: PrefixSum) !shape.Family {
950 var builder = try shape.Builder.init(backing_allocator, "prefix_sum");
951 errdefer builder.deinit();
952
953 const elements = try builder.symbol(instance.element_axis);
954 const elements_expr = try builder.symbolExpression(elements);
955
956 _ = try builder.tensor("data", &.{elements_expr});
957 _ = try builder.tensor("out", &.{elements_expr});
958 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = instance.threads });
959
960 return builder.finish();
961 }
962
963 pub fn prefixSumFamilySpecialization(backing_allocator: std.mem.Allocator, instance: PrefixSum) !entry.OwnedSpecialization {
964 if (!prefixSumInstanceValid(instance)) return error.UnsupportedPrefixSumInstance;
965 var owned = entry.OwnedSpecialization.init(backing_allocator);
966 errdefer owned.deinit();
967 const lifetime_allocator = owned.allocator();
968
969 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
970 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
971
972 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
973 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
974
975 owned.value = .{
976 .dtype = instance.dtype,
977 .operation = .{ .scan = instance.mode.operation() },
978 .inputs = inputs,
979 .outputs = outputs,
980 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.element_axis, instance.threads, instance.threads),
981 };
982 owned.value.launch = owned.value.schedule.?.launch();
983 var family = try prefixSumShapeFamily(backing_allocator, instance);
984 errdefer family.deinit();
985 try owned.takeShapeFamily(&family);
986 return owned;
987 }
988
989 pub fn prefixSumInstanceFromSpecialization(specialization: entry.Specialization) ?PrefixSum {
990 if (!specialization.scheduleMatchesLaunch()) return null;
991 const mode: PrefixSumMode = if (specialization.operationIs(.{ .scan = .prefix_sum }))
992 .inclusive
993 else if (specialization.operationIs(.{ .scan = .prefix_sum_exclusive }))
994 .exclusive
995 else
996 return null;
997 const dtype = specialization.dtype orelse return null;
998 if (!prefixSumDTypeSupported(dtype)) return null;
999 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
1000 if (specialization.reductions.len != 0) return null;
1001 const data = specialization.inputs[0];
1002 const output = specialization.outputs[0];
1003 if (data.axes.len != 1 or output.axes.len != 1) return null;
1004 const extent = data.axes[0].extent;
1005 if (output.axes[0].extent != extent) return null;
1006 const launch = specialization.launch orelse return null;
1007 if (launch.grid[0] != 1) return null;
1008 const instance = PrefixSum{
1009 .extent = extent,
1010 .dtype = dtype,
1011 .mode = mode,
1012 .threads = launch.threadgroup[0],
1013 .element_axis = data.axes[0].name,
1014 };
1015 if (!prefixSumInstanceValid(instance)) return null;
1016 return instance;
1017 }
1018
1019 fn prefixSumSpecialization(comptime spec: PrefixSum) entry.Specialization {
1020 return .{
1021 .dtype = spec.dtype,
1022 .operation = .{ .scan = spec.mode.operation() },
1023 .inputs = &.{entry.shape1D(spec.element_axis, spec.extent)},
1024 .outputs = &.{entry.shape1D(spec.element_axis, spec.extent)},
1025 .launch = entry.launch1D(1, spec.threads),
1026 .schedule = entry.threadBlocks1D(spec.element_axis, spec.threads, spec.threads),
1027 };
1028 }
1029
1030 fn prefixSumProgram(comptime spec: PrefixSum) type {
1031 const Body = struct {
1032 fn run(k: anytype, args: anytype) !void {
1033 try prefixSumBody(k, spec, args);
1034 }
1035 };
1036
1037 return kernel.logical.Program(.{
1038 .name = std.fmt.comptimePrint(
1039 "accy_kernel_scan_prefix_sum{}_{}_{s}",
1040 .{ spec.extent, spec.threads, spec.dtype.name() },
1041 ),
1042 .parameters = .{
1043 .dst = kernel.dynamicBuffer(spec.dtype),
1044 .data = kernel.dynamicBuffer(spec.dtype),
1045 },
1046 .body = Body.run,
1047 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
1048 }
1049
1050 pub fn prefixSumF32(comptime spec: PrefixSum) type {
1051 return entry.Entry(prefixSumProgram(spec), .{
1052 .target = std.fmt.comptimePrint(
1053 "accy.kernel.scan.prefix_sum{}_{}_{s}",
1054 .{ spec.extent, spec.threads, spec.dtype.name() },
1055 ),
1056 .layer = .logical,
1057 .category = .scan,
1058 .specialization = prefixSumSpecialization(spec),
1059 });
1060 }
1061
1062 pub const PrefixSum8F32 = prefixSumF32(.{ .extent = 8, .threads = 32 });
1063
1064 const testing = std.testing;
1065
1066 test "scan prefix sum entry runs on CPU" {
1067 const allocator = std.testing.allocator;
1068 var data = [_]f32{ 1, 2, 3, 4, 5, 6, 7, 8 };
1069 var dst = @as([8]f32, @splat(0));
1070
1071 const ProgramType = prefixSumProgram(.{ .extent = 8, .threads = 32 });
1072 var graph = try ProgramType.build(allocator, ProgramType.Limits.testing);
1073 defer graph.deinit();
1074 try graph.runCpuWithLaunch(allocator, &.{
1075 kernel.argumentBuffer(f32, dst[0..]),
1076 kernel.argumentBuffer(f32, data[0..]),
1077 }, .{
1078 .grid = .{ 1, 1, 1 },
1079 .block = .{ 32, 1, 1 },
1080 });
1081 try std.testing.expectEqualSlices(f32, &.{ 1, 3, 6, 10, 15, 21, 28, 36 }, dst[0..]);
1082 }
1083
1084 fn expectPrefixSumMatchesOracle(extent: usize, threads: u32) !void {
1085 try expectPrefixSumModeMatchesOracle(extent, threads, .inclusive);
1086 }
1087
1088 fn expectPrefixSumModeMatchesOracle(extent: usize, threads: u32, mode: PrefixSumMode) !void {
1089 const allocator = testing.allocator;
1090 const compiled = PrefixSum{ .extent = 1, .threads = threads, .mode = mode };
1091 const runtime = PrefixSum{ .extent = extent, .threads = threads, .mode = mode };
1092
1093 var graph = try PrefixSumRuntimeFamilyF32.build(allocator, PrefixSumRuntimeFamilyF32.Limits.testing, compiled);
1094 defer graph.deinit();
1095
1096 const data = try allocator.alloc(f32, extent);
1097 defer allocator.free(data);
1098 for (data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
1099 const dst = try allocator.alloc(f32, extent);
1100 defer allocator.free(dst);
1101 @memset(dst, 0);
1102
1103 const expected = try allocator.alloc(f32, extent);
1104 defer allocator.free(expected);
1105 var running: f32 = 0;
1106 for (data, 0..) |value, index| {
1107 switch (mode) {
1108 .inclusive => {
1109 running += value;
1110 expected[index] = running;
1111 },
1112 .exclusive => {
1113 expected[index] = running;
1114 running += value;
1115 },
1116 }
1117 }
1118
1119 try graph.runCpuWithLaunch(allocator, &.{
1120 kernel.argumentBuffer(f32, dst),
1121 kernel.argumentBuffer(f32, data),
1122 kernel.argumentI32(@intCast(runtime.extent)),
1123 }, .{
1124 .grid = .{ 1, 1, 1 },
1125 .block = .{ runtime.threads, 1, 1 },
1126 });
1127 for (expected, dst) |want, got| {
1128 try testing.expectApproxEqAbs(want, got, 0.001);
1129 }
1130 }
1131
1132 test "scan prefix sum runtime family matches oracle across warp boundaries" {
1133 try expectPrefixSumMatchesOracle(40, 64);
1134 try expectPrefixSumMatchesOracle(5, 32);
1135 try expectPrefixSumMatchesOracle(128, 128);
1136 try expectPrefixSumMatchesOracle(100, 256);
1137 }
1138
1139 test "scan exclusive prefix sum runtime family matches oracle across warp boundaries" {
1140 try expectPrefixSumModeMatchesOracle(40, 64, .exclusive);
1141 try expectPrefixSumModeMatchesOracle(5, 32, .exclusive);
1142 try expectPrefixSumModeMatchesOracle(128, 128, .exclusive);
1143 try expectPrefixSumModeMatchesOracle(100, 256, .exclusive);
1144 }
1145
1146 test "scan prefix sum f16 runtime family matches oracle through f32 accumulation" {
1147 const allocator = testing.allocator;
1148 const compiled = PrefixSum{ .extent = 1, .dtype = .f16, .threads = 64 };
1149 const runtime = PrefixSum{ .extent = 40, .dtype = .f16, .threads = 64 };
1150
1151 var graph = try PrefixSumRuntimeFamilyF16.build(allocator, PrefixSumRuntimeFamilyF16.Limits.testing, compiled);
1152 defer graph.deinit();
1153
1154 var data: [40]f16 = undefined;
1155 for (&data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
1156 var dst = @as([40]f16, @splat(0));
1157
1158 var expected: [40]f32 = undefined;
1159 var running: f32 = 0;
1160 for (data, 0..) |value, index| {
1161 running += @floatCast(value);
1162 expected[index] = running;
1163 }
1164
1165 const launch_value = try entry.runtimeLaunch1D(runtime.threads, runtime.threads);
1166 try graph.runCpuWithLaunch(allocator, &.{
1167 kernel.argumentBuffer(f16, dst[0..]),
1168 kernel.argumentBuffer(f16, data[0..]),
1169 kernel.argumentI32(@intCast(runtime.extent)),
1170 }, .{
1171 .grid = launch_value.grid,
1172 .block = launch_value.threadgroup,
1173 });
1174 for (expected, dst) |want, got| {
1175 try testing.expectApproxEqAbs(want, @as(f32, @floatCast(got)), 0.5);
1176 }
1177 }
1178
1179 fn expectPrefixSumU32ModeMatchesOracle(extent: usize, threads: u32, mode: PrefixSumMode) !void {
1180 const allocator = testing.allocator;
1181 const compiled = PrefixSum{ .extent = 1, .dtype = .u32, .threads = threads, .mode = mode };
1182 const runtime = PrefixSum{ .extent = extent, .dtype = .u32, .threads = threads, .mode = mode };
1183
1184 var graph = try PrefixSumRuntimeFamilyU32.build(allocator, PrefixSumRuntimeFamilyU32.Limits.testing, compiled);
1185 defer graph.deinit();
1186
1187 const data = try allocator.alloc(u32, extent);
1188 defer allocator.free(data);
1189 for (data, 0..) |*value, index| value.* = @intCast((index % 7) + 1);
1190 const dst = try allocator.alloc(u32, extent);
1191 defer allocator.free(dst);
1192 @memset(dst, 0);
1193
1194 const expected = try allocator.alloc(u32, extent);
1195 defer allocator.free(expected);
1196 var running: u32 = 0;
1197 for (data, 0..) |value, index| {
1198 switch (mode) {
1199 .inclusive => {
1200 running += value;
1201 expected[index] = running;
1202 },
1203 .exclusive => {
1204 expected[index] = running;
1205 running += value;
1206 },
1207 }
1208 }
1209
1210 const launch_value = try entry.runtimeLaunch1D(runtime.threads, runtime.threads);
1211 try graph.runCpuWithLaunch(allocator, &.{
1212 kernel.argumentBuffer(u32, dst),
1213 kernel.argumentBuffer(u32, data),
1214 kernel.argumentI32(@intCast(runtime.extent)),
1215 }, .{
1216 .grid = launch_value.grid,
1217 .block = launch_value.threadgroup,
1218 });
1219 try testing.expectEqualSlices(u32, expected, dst);
1220 }
1221
1222 test "scan prefix sum u32 runtime family matches oracle across warp boundaries" {
1223 try expectPrefixSumU32ModeMatchesOracle(40, 64, .inclusive);
1224 try expectPrefixSumU32ModeMatchesOracle(5, 32, .inclusive);
1225 try expectPrefixSumU32ModeMatchesOracle(128, 128, .inclusive);
1226 try expectPrefixSumU32ModeMatchesOracle(100, 256, .inclusive);
1227 try expectPrefixSumU32ModeMatchesOracle(40, 64, .exclusive);
1228 try expectPrefixSumU32ModeMatchesOracle(100, 256, .exclusive);
1229 }
1230
1231 test "scan prefix sum f16 identity carries the dtype" {
1232 const instance = PrefixSum{ .extent = 100, .dtype = .f16, .threads = 128 };
1233 const family_target = try prefixSumFamilyTarget(testing.allocator, instance);
1234 defer testing.allocator.free(family_target);
1235 try testing.expectEqualStrings("accy.kernel.scan.prefix_sum_family_128_f16", family_target);
1236 }
1237
1238 test "scan prefix sum u32 identity carries the dtype" {
1239 const instance = PrefixSum{ .extent = 100, .dtype = .u32, .threads = 128 };
1240 const family_target = try prefixSumFamilyTarget(testing.allocator, instance);
1241 defer testing.allocator.free(family_target);
1242 try testing.expectEqualStrings("accy.kernel.scan.prefix_sum_family_128_u32", family_target);
1243 }
1244
1245 test "scan exclusive prefix sum identity carries the mode" {
1246 const instance = PrefixSum{ .extent = 100, .mode = .exclusive, .threads = 128 };
1247 const family_target = try prefixSumFamilyTarget(testing.allocator, instance);
1248 defer testing.allocator.free(family_target);
1249 try testing.expectEqualStrings("accy.kernel.scan.prefix_sum_exclusive_family_128_f32", family_target);
1250 }
1251
1252 test "scan prefix sum family tuning keys discriminate modes" {
1253 const allocator = testing.allocator;
1254 const device = tuning.deviceFingerprint(.{ .identity = .{
1255 .backend = .cuda,
1256 .family = .nvidia_cuda,
1257 .name = "scan-family-tuning-test-device",
1258 .vendor_id = 0x10de,
1259 .device_id = 0x2684,
1260 } });
1261
1262 const inclusive = PrefixSum{ .extent = 64 };
1263 const exclusive = PrefixSum{ .extent = 64, .mode = .exclusive };
1264
1265 const inclusive_family = try prefixSumFamilyFingerprint(allocator, inclusive);
1266 const exclusive_family = try prefixSumFamilyFingerprint(allocator, exclusive);
1267 try testing.expectEqual(inclusive_family, exclusive_family);
1268
1269 const inclusive_key = try prefixSumFamilyTuningKey(allocator, device, inclusive);
1270 const exclusive_key = try prefixSumFamilyTuningKey(allocator, device, exclusive);
1271 try testing.expect(!inclusive_key.eql(exclusive_key));
1272 try testing.expectEqual(inclusive_key.family_fingerprint, exclusive_key.family_fingerprint);
1273 try testing.expect(inclusive_key.operation_fingerprint != exclusive_key.operation_fingerprint);
1274
1275 const repeat_key = try prefixSumFamilyTuningKey(allocator, device, inclusive);
1276 try testing.expect(inclusive_key.eql(repeat_key));
1277 }
1278
1279 test "scan exclusive prefix sum instance round-trips through specialization" {
1280 const instance = PrefixSum{ .extent = 100, .mode = .exclusive, .threads = 128 };
1281 var owned = try prefixSumFamilySpecialization(testing.allocator, instance);
1282 defer owned.deinit();
1283
1284 const recovered = prefixSumInstanceFromSpecialization(owned.value) orelse return error.TestExpectedPrefixSumInstance;
1285 try testing.expectEqual(PrefixSumMode.exclusive, recovered.mode);
1286 try testing.expectEqual(instance.extent, recovered.extent);
1287 try testing.expectEqual(instance.threads, recovered.threads);
1288 }
1289
1290 test "scan prefix sum instance validity bounds extent by threads" {
1291 try testing.expect(prefixSumInstanceValid(.{ .extent = 256, .threads = 256 }));
1292 try testing.expect(!prefixSumInstanceValid(.{ .extent = 257, .threads = 256 }));
1293 try testing.expect(!prefixSumInstanceValid(.{ .extent = 8, .threads = 24 }));
1294 try testing.expect(!prefixSumInstanceValid(.{ .extent = 0, .threads = 32 }));
1295 try testing.expect(!prefixSumInstanceValid(.{ .extent = 8, .threads = 2048 }));
1296 }
1297
1298 test "scan prefix sum thread selection rounds to warps" {
1299 try testing.expectEqual(@as(?u32, 32), prefixSumThreadsForExtent(5));
1300 try testing.expectEqual(@as(?u32, 64), prefixSumThreadsForExtent(40));
1301 try testing.expectEqual(@as(?u32, 1024), prefixSumThreadsForExtent(1024));
1302 try testing.expectEqual(@as(?u32, null), prefixSumThreadsForExtent(1025));
1303
1304 const candidates = prefixSumThreadCandidatesForExtent(40);
1305 try testing.expect(candidates.count >= 2);
1306 try testing.expectEqual(@as(u32, 64), candidates.items[0]);
1307 for (candidates.slice()) |threads| {
1308 try testing.expect(@as(u64, threads) >= 40);
1309 try testing.expect(threads % prefix_sum_warp_size == 0);
1310 }
1311 }
1312
1313 test "scan prefix sum family identity and artifact contract" {
1314 const allocator = testing.allocator;
1315 var state = gpu.recording.BackendState{
1316 .allocator = allocator,
1317 .kind = .cuda,
1318 .format = .cuda_ptx,
1319 };
1320 const instance = PrefixSum{ .extent = 40, .threads = 64 };
1321
1322 const family_target = try prefixSumFamilyTarget(allocator, instance);
1323 defer allocator.free(family_target);
1324 try testing.expectEqualStrings("accy.kernel.scan.prefix_sum_family_64_f32", family_target);
1325
1326 var family_artifact = try createPrefixSumFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
1327 defer family_artifact.deinit();
1328
1329 const family_entry = family_artifact.entry();
1330 try testing.expectEqualStrings("accy_kernel_scan_prefix_sum_family_64_f32", family_entry.entry_name);
1331 try testing.expectEqual(@as(u32, 3), family_entry.argument_count);
1332 try testing.expectEqual(@as(u32, 1), family_entry.runtime_scalar_argument_count);
1333 switch (family_entry.launch) {
1334 .derived => |launch| {
1335 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
1336 switch (launch.grid[0]) {
1337 .fixed => |value| try testing.expectEqual(@as(u32, 1), value),
1338 else => return error.TestExpectedFixedGrid,
1339 }
1340 },
1341 else => return error.TestExpectedDerivedLaunch,
1342 }
1343 }
1344
1345 test "scan prefix sum u32 recording artifacts cover native targets" {
1346 const allocator = testing.allocator;
1347 const instance = PrefixSum{ .extent = 40, .dtype = .u32, .threads = 64 };
1348 inline for (.{ gpu.ArtifactFormat.cuda_ptx, .vulkan_spirv, .metal_msl }) |format| {
1349 var state = gpu.recording.BackendState{
1350 .allocator = allocator,
1351 .kind = switch (format) {
1352 .cuda_ptx => .cuda,
1353 .vulkan_spirv => .vulkan,
1354 .metal_msl => .metal,
1355 else => .external,
1356 },
1357 .format = format,
1358 };
1359 var artifact = try createPrefixSumFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing, .format = format });
1360 defer artifact.deinit();
1361 const entry_value = artifact.entry();
1362 try testing.expectEqual(format, entry_value.format);
1363 try testing.expectEqualStrings("accy_kernel_scan_prefix_sum_family_64_u32", entry_value.entry_name);
1364 }
1365 }
1366
1367 test "scan prefix sum instance round-trips through specialization" {
1368 const instance = PrefixSum{ .extent = 100, .threads = 128 };
1369 var owned = try prefixSumFamilySpecialization(testing.allocator, instance);
1370 defer owned.deinit();
1371
1372 const recovered = prefixSumInstanceFromSpecialization(owned.value) orelse return error.TestExpectedPrefixSumInstance;
1373 try testing.expectEqual(instance.extent, recovered.extent);
1374 try testing.expectEqual(instance.threads, recovered.threads);
1375 try testing.expectEqual(instance.dtype, recovered.dtype);
1376
1377 const u32_instance = PrefixSum{ .extent = 100, .dtype = .u32, .threads = 128, .mode = .exclusive };
1378 var u32_owned = try prefixSumFamilySpecialization(testing.allocator, u32_instance);
1379 defer u32_owned.deinit();
1380
1381 const u32_recovered = prefixSumInstanceFromSpecialization(u32_owned.value) orelse return error.TestExpectedPrefixSumInstance;
1382 try testing.expectEqual(u32_instance.mode, u32_recovered.mode);
1383 try testing.expectEqual(u32_instance.extent, u32_recovered.extent);
1384 try testing.expectEqual(u32_instance.threads, u32_recovered.threads);
1385 try testing.expectEqual(u32_instance.dtype, u32_recovered.dtype);
1386
1387 try testing.expectEqual(@as(?PrefixSum, null), prefixSumInstanceFromSpecialization(.{}));
1388 }
1389
1390 fn expectDeviceScanMatchesOracle(extent: usize, threads: u32, mode: PrefixSumMode) !void {
1391 const allocator = testing.allocator;
1392 const compiled = DeviceScan{ .extent = 1, .threads = threads, .mode = mode };
1393 const runtime = DeviceScan{ .extent = extent, .threads = threads, .mode = mode };
1394 const stages = try deviceScanStages(runtime);
1395
1396 const data = try allocator.alloc(f32, extent);
1397 defer allocator.free(data);
1398 for (data, 0..) |*value, index| value.* = @floatFromInt((index % 7) + 1);
1399 const dst = try allocator.alloc(f32, extent);
1400 defer allocator.free(dst);
1401 @memset(dst, 0);
1402 const sums = try allocator.alloc(f32, stages.block_count);
1403 defer allocator.free(sums);
1404 @memset(sums, 0);
1405 const bases = try allocator.alloc(f32, stages.block_count);
1406 defer allocator.free(bases);
1407 @memset(bases, 0);
1408
1409 var block_scan = try DeviceScanBlockScanRuntimeFamilyF32.build(allocator, DeviceScanBlockScanRuntimeFamilyF32.Limits.testing, compiled);
1410 defer block_scan.deinit();
1411 try block_scan.runCpuWithLaunch(allocator, &.{
1412 kernel.argumentBuffer(f32, dst),
1413 kernel.argumentBuffer(f32, data),
1414 kernel.argumentBuffer(f32, sums),
1415 kernel.argumentI32(@intCast(runtime.extent)),
1416 }, .{
1417 .grid = .{ stages.block_count, 1, 1 },
1418 .block = .{ runtime.threads, 1, 1 },
1419 });
1420
1421 var sums_scan = try PrefixSumRuntimeFamilyF32.build(allocator, PrefixSumRuntimeFamilyF32.Limits.testing, .{
1422 .extent = 1,
1423 .mode = .exclusive,
1424 .threads = stages.sums_scan.threads,
1425 });
1426 defer sums_scan.deinit();
1427 try sums_scan.runCpuWithLaunch(allocator, &.{
1428 kernel.argumentBuffer(f32, bases),
1429 kernel.argumentBuffer(f32, sums),
1430 kernel.argumentI32(@intCast(stages.block_count)),
1431 }, .{
1432 .grid = .{ 1, 1, 1 },
1433 .block = .{ stages.sums_scan.threads, 1, 1 },
1434 });
1435
1436 var add_base = try DeviceScanAddBaseRuntimeFamilyF32.build(allocator, DeviceScanAddBaseRuntimeFamilyF32.Limits.testing, compiled);
1437 defer add_base.deinit();
1438 try add_base.runCpuWithLaunch(allocator, &.{
1439 kernel.argumentBuffer(f32, dst),
1440 kernel.argumentBuffer(f32, bases),
1441 kernel.argumentI32(@intCast(runtime.extent)),
1442 }, .{
1443 .grid = .{ stages.block_count, 1, 1 },
1444 .block = .{ runtime.threads, 1, 1 },
1445 });
1446
1447 var running: f32 = 0;
1448 for (data, dst) |value, got| {
1449 switch (mode) {
1450 .inclusive => {
1451 running += value;
1452 try testing.expectEqual(running, got);
1453 },
1454 .exclusive => {
1455 try testing.expectEqual(running, got);
1456 running += value;
1457 },
1458 }
1459 }
1460 }
1461
1462 fn expectDeviceScanF16MatchesOracle(extent: usize, threads: u32, mode: PrefixSumMode) !void {
1463 const allocator = testing.allocator;
1464 const compiled = DeviceScan{ .extent = 1, .dtype = .f16, .threads = threads, .mode = mode };
1465 const runtime = DeviceScan{ .extent = extent, .dtype = .f16, .threads = threads, .mode = mode };
1466 const stages = try deviceScanStages(runtime);
1467
1468 const data = try allocator.alloc(f16, extent);
1469 defer allocator.free(data);
1470 for (data, 0..) |*value, index| value.* = @floatFromInt((index % 5) + 1);
1471 const local = try allocator.alloc(f32, extent);
1472 defer allocator.free(local);
1473 @memset(local, 0);
1474 const dst = try allocator.alloc(f16, extent);
1475 defer allocator.free(dst);
1476 @memset(dst, 0);
1477 const sums = try allocator.alloc(f32, stages.block_count);
1478 defer allocator.free(sums);
1479 @memset(sums, 0);
1480 const bases = try allocator.alloc(f32, stages.block_count);
1481 defer allocator.free(bases);
1482 @memset(bases, 0);
1483
1484 var block_scan = try DeviceScanBlockScanRuntimeFamilyF16.build(allocator, DeviceScanBlockScanRuntimeFamilyF16.Limits.testing, compiled);
1485 defer block_scan.deinit();
1486 try block_scan.runCpuWithLaunch(allocator, &.{
1487 kernel.argumentBuffer(f32, local),
1488 kernel.argumentBuffer(f16, data),
1489 kernel.argumentBuffer(f32, sums),
1490 kernel.argumentI32(@intCast(runtime.extent)),
1491 }, .{
1492 .grid = .{ stages.block_count, 1, 1 },
1493 .block = .{ runtime.threads, 1, 1 },
1494 });
1495
1496 var sums_scan = try PrefixSumRuntimeFamilyF32.build(allocator, PrefixSumRuntimeFamilyF32.Limits.testing, .{
1497 .extent = 1,
1498 .mode = .exclusive,
1499 .threads = stages.sums_scan.threads,
1500 });
1501 defer sums_scan.deinit();
1502 try sums_scan.runCpuWithLaunch(allocator, &.{
1503 kernel.argumentBuffer(f32, bases),
1504 kernel.argumentBuffer(f32, sums),
1505 kernel.argumentI32(@intCast(stages.block_count)),
1506 }, .{
1507 .grid = .{ 1, 1, 1 },
1508 .block = .{ stages.sums_scan.threads, 1, 1 },
1509 });
1510
1511 var add_base = try DeviceScanAddBaseRuntimeFamilyF16.build(allocator, DeviceScanAddBaseRuntimeFamilyF16.Limits.testing, compiled);
1512 defer add_base.deinit();
1513 try add_base.runCpuWithLaunch(allocator, &.{
1514 kernel.argumentBuffer(f16, dst),
1515 kernel.argumentBuffer(f32, local),
1516 kernel.argumentBuffer(f32, bases),
1517 kernel.argumentI32(@intCast(runtime.extent)),
1518 }, .{
1519 .grid = .{ stages.block_count, 1, 1 },
1520 .block = .{ runtime.threads, 1, 1 },
1521 });
1522
1523 var running: f32 = 0;
1524 for (data, dst) |value, got| {
1525 const value_f32: f32 = @floatCast(value);
1526 switch (mode) {
1527 .inclusive => {
1528 running += value_f32;
1529 const expected: f16 = @floatCast(running);
1530 try testing.expectEqual(expected, got);
1531 },
1532 .exclusive => {
1533 const expected: f16 = @floatCast(running);
1534 try testing.expectEqual(expected, got);
1535 running += value_f32;
1536 },
1537 }
1538 }
1539 }
1540
1541 fn expectDeviceScanU32MatchesOracle(extent: usize, threads: u32, mode: PrefixSumMode) !void {
1542 const allocator = testing.allocator;
1543 const compiled = DeviceScan{ .extent = 1, .dtype = .u32, .threads = threads, .mode = mode };
1544 const runtime = DeviceScan{ .extent = extent, .dtype = .u32, .threads = threads, .mode = mode };
1545 const stages = try deviceScanStages(runtime);
1546
1547 const data = try allocator.alloc(u32, extent);
1548 defer allocator.free(data);
1549 for (data, 0..) |*value, index| value.* = @intCast((index % 7) + 1);
1550 const dst = try allocator.alloc(u32, extent);
1551 defer allocator.free(dst);
1552 @memset(dst, 0);
1553 const sums = try allocator.alloc(u32, stages.block_count);
1554 defer allocator.free(sums);
1555 @memset(sums, 0);
1556 const bases = try allocator.alloc(u32, stages.block_count);
1557 defer allocator.free(bases);
1558 @memset(bases, 0);
1559
1560 var block_scan = try DeviceScanBlockScanRuntimeFamilyU32.build(allocator, DeviceScanBlockScanRuntimeFamilyU32.Limits.testing, compiled);
1561 defer block_scan.deinit();
1562 try block_scan.runCpuWithLaunch(allocator, &.{
1563 kernel.argumentBuffer(u32, dst),
1564 kernel.argumentBuffer(u32, data),
1565 kernel.argumentBuffer(u32, sums),
1566 kernel.argumentI32(@intCast(runtime.extent)),
1567 }, .{
1568 .grid = .{ stages.block_count, 1, 1 },
1569 .block = .{ runtime.threads, 1, 1 },
1570 });
1571
1572 var sums_scan = try PrefixSumRuntimeFamilyU32.build(allocator, PrefixSumRuntimeFamilyU32.Limits.testing, .{
1573 .extent = 1,
1574 .dtype = .u32,
1575 .mode = .exclusive,
1576 .threads = stages.sums_scan.threads,
1577 });
1578 defer sums_scan.deinit();
1579 try sums_scan.runCpuWithLaunch(allocator, &.{
1580 kernel.argumentBuffer(u32, bases),
1581 kernel.argumentBuffer(u32, sums),
1582 kernel.argumentI32(@intCast(stages.block_count)),
1583 }, .{
1584 .grid = .{ 1, 1, 1 },
1585 .block = .{ stages.sums_scan.threads, 1, 1 },
1586 });
1587
1588 var add_base = try DeviceScanAddBaseRuntimeFamilyU32.build(allocator, DeviceScanAddBaseRuntimeFamilyU32.Limits.testing, compiled);
1589 defer add_base.deinit();
1590 try add_base.runCpuWithLaunch(allocator, &.{
1591 kernel.argumentBuffer(u32, dst),
1592 kernel.argumentBuffer(u32, bases),
1593 kernel.argumentI32(@intCast(runtime.extent)),
1594 }, .{
1595 .grid = .{ stages.block_count, 1, 1 },
1596 .block = .{ runtime.threads, 1, 1 },
1597 });
1598
1599 var running: u32 = 0;
1600 for (data, dst) |value, got| {
1601 switch (mode) {
1602 .inclusive => {
1603 running += value;
1604 try testing.expectEqual(running, got);
1605 },
1606 .exclusive => {
1607 try testing.expectEqual(running, got);
1608 running += value;
1609 },
1610 }
1611 }
1612 }
1613
1614 test "scan device-wide composition matches running-sum oracle across blocks" {
1615 try expectDeviceScanMatchesOracle(100, 32, .inclusive);
1616 try expectDeviceScanMatchesOracle(33, 32, .inclusive);
1617 try expectDeviceScanMatchesOracle(2048, 256, .inclusive);
1618 try expectDeviceScanMatchesOracle(1000, 64, .inclusive);
1619 }
1620
1621 test "scan device-wide exclusive composition matches running-sum oracle" {
1622 try expectDeviceScanMatchesOracle(100, 32, .exclusive);
1623 try expectDeviceScanMatchesOracle(1000, 64, .exclusive);
1624 }
1625
1626 test "scan device-wide f16 composition uses f32 intermediates across blocks" {
1627 try expectDeviceScanF16MatchesOracle(100, 32, .inclusive);
1628 try expectDeviceScanF16MatchesOracle(1000, 64, .inclusive);
1629 try expectDeviceScanF16MatchesOracle(100, 32, .exclusive);
1630 }
1631
1632 test "scan device-wide u32 composition matches running-sum oracle across blocks" {
1633 try expectDeviceScanU32MatchesOracle(100, 32, .inclusive);
1634 try expectDeviceScanU32MatchesOracle(33, 32, .inclusive);
1635 try expectDeviceScanU32MatchesOracle(2048, 256, .inclusive);
1636 try expectDeviceScanU32MatchesOracle(1000, 64, .inclusive);
1637 try expectDeviceScanU32MatchesOracle(100, 32, .exclusive);
1638 try expectDeviceScanU32MatchesOracle(1000, 64, .exclusive);
1639 }
1640
1641 test "scan device-wide instance validity caps blocks and accepts supported dtypes" {
1642 try testing.expect(deviceScanInstanceValid(.{ .extent = 1024 * 1024, .threads = 1024 }));
1643 try testing.expect(!deviceScanInstanceValid(.{ .extent = 1024 * 1024 + 1, .threads = 1024 }));
1644 try testing.expect(deviceScanInstanceValid(.{ .extent = 100, .dtype = .f16, .threads = 32 }));
1645 try testing.expect(deviceScanInstanceValid(.{ .extent = 100, .dtype = .u32, .threads = 32 }));
1646 try testing.expect(!deviceScanInstanceValid(.{ .extent = 0, .threads = 32 }));
1647 try testing.expect(!deviceScanInstanceValid(.{ .extent = 100, .threads = 48 }));
1648 try testing.expect(!deviceScanInstanceValid(.{ .extent = std.math.maxInt(u64), .threads = 32 }));
1649 }
1650
1651 test "scan device-wide stages derive the sums scan from the block count" {
1652 const stages = try deviceScanStages(.{ .extent = 1000, .threads = 64 });
1653 try testing.expectEqual(@as(u32, 16), stages.block_count);
1654 try testing.expectEqual(@as(u64, 16), stages.sums_scan.extent);
1655 try testing.expectEqual(PrefixSumMode.exclusive, stages.sums_scan.mode);
1656 try testing.expectEqual(@as(u32, 32), stages.sums_scan.threads);
1657 try testing.expectEqual(@as(u64, 1000), stages.block_scan.extent);
1658 try testing.expectEqual(@as(u64, 1000), stages.add_base.extent);
1659
1660 const u32_stages = try deviceScanStages(.{ .extent = 1000, .dtype = .u32, .threads = 64 });
1661 try testing.expectEqual(DType.u32, u32_stages.sums_scan.dtype);
1662 }
1663
1664 test "scan device-wide identity carries mode threads and dtype" {
1665 const block_target = try deviceScanBlockScanFamilyTarget(testing.allocator, .{ .extent = 1000, .threads = 64 });
1666 defer testing.allocator.free(block_target);
1667 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_block_scan_family_64_f32", block_target);
1668 const exclusive_target = try deviceScanBlockScanFamilyTarget(testing.allocator, .{ .extent = 1000, .mode = .exclusive, .threads = 64 });
1669 defer testing.allocator.free(exclusive_target);
1670 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_exclusive_block_scan_family_64_f32", exclusive_target);
1671 const add_target = try deviceScanAddBaseFamilyTarget(testing.allocator, .{ .extent = 1000, .threads = 64 });
1672 defer testing.allocator.free(add_target);
1673 try testing.expectEqualStrings("accy.kernel.scan.device_add_base_family_64_f32", add_target);
1674 const f16_target = try deviceScanFamilyTarget(testing.allocator, .{ .extent = 1000, .dtype = .f16, .threads = 64 });
1675 defer testing.allocator.free(f16_target);
1676 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_family_64_f16", f16_target);
1677 const u32_target = try deviceScanFamilyTarget(testing.allocator, .{ .extent = 1000, .dtype = .u32, .threads = 64 });
1678 defer testing.allocator.free(u32_target);
1679 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_family_64_u32", u32_target);
1680 }
1681
1682 test "scan device-wide u32 instance round-trips through specialization" {
1683 const instance = DeviceScan{ .extent = 2000, .dtype = .u32, .threads = 64, .mode = .exclusive };
1684 var owned = try deviceScanFamilySpecialization(testing.allocator, instance);
1685 defer owned.deinit();
1686
1687 const recovered = deviceScanInstanceFromSpecialization(owned.value) orelse return error.TestExpectedDeviceScanInstance;
1688 try testing.expectEqual(instance.mode, recovered.mode);
1689 try testing.expectEqual(instance.extent, recovered.extent);
1690 try testing.expectEqual(instance.threads, recovered.threads);
1691 try testing.expectEqual(instance.dtype, recovered.dtype);
1692
1693 try testing.expectEqual(@as(?DeviceScan, null), deviceScanInstanceFromSpecialization(.{}));
1694 }
1695
1696 test "scan device-wide family artifact contract derives the grid from the extent" {
1697 const allocator = testing.allocator;
1698 var state = gpu.recording.BackendState{
1699 .allocator = allocator,
1700 .kind = .cuda,
1701 .format = .cuda_ptx,
1702 };
1703 const instance = DeviceScan{ .extent = 1000, .threads = 64 };
1704
1705 var block_artifact = try createDeviceScanBlockScanFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
1706 defer block_artifact.deinit();
1707 const block_entry = block_artifact.entry();
1708 try testing.expectEqualStrings("accy_kernel_scan_device_prefix_sum_block_scan_family_64_f32", block_entry.entry_name);
1709 try testing.expectEqual(@as(u32, 4), block_entry.argument_count);
1710 try testing.expectEqual(@as(u32, 1), block_entry.runtime_scalar_argument_count);
1711 switch (block_entry.launch) {
1712 .derived => |launch| {
1713 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
1714 switch (launch.grid[0]) {
1715 .runtime_u32_ceil_div => |axis| {
1716 try testing.expectEqual(@as(u32, 0), axis.argument_index);
1717 try testing.expectEqual(@as(u32, 64), axis.divisor);
1718 },
1719 else => return error.TestExpectedDerivedGrid,
1720 }
1721 },
1722 else => return error.TestExpectedDerivedLaunch,
1723 }
1724
1725 var add_artifact = try createDeviceScanAddBaseFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
1726 defer add_artifact.deinit();
1727 const add_entry = add_artifact.entry();
1728 try testing.expectEqualStrings("accy_kernel_scan_device_add_base_family_64_f32", add_entry.entry_name);
1729 try testing.expectEqual(@as(u32, 3), add_entry.argument_count);
1730 try testing.expectEqual(@as(u32, 1), add_entry.runtime_scalar_argument_count);
1731 switch (add_entry.launch) {
1732 .derived => |launch| {
1733 try testing.expectEqual(@as(u32, 64), launch.threadgroup[0]);
1734 switch (launch.grid[0]) {
1735 .runtime_u32_ceil_div => |axis| try testing.expectEqual(@as(u32, 64), axis.divisor),
1736 else => return error.TestExpectedDerivedGrid,
1737 }
1738 },
1739 else => return error.TestExpectedDerivedLaunch,
1740 }
1741
1742 const f16_instance = DeviceScan{ .extent = 1000, .dtype = .f16, .threads = 64 };
1743 var f16_block_artifact = try createDeviceScanBlockScanFamilyArtifact(allocator, state.handle(), f16_instance, .{ .limits = .testing });
1744 defer f16_block_artifact.deinit();
1745 const f16_block_entry = f16_block_artifact.entry();
1746 try testing.expectEqualStrings("accy_kernel_scan_device_prefix_sum_block_scan_family_64_f16", f16_block_entry.entry_name);
1747 try testing.expectEqual(@as(u32, 4), f16_block_entry.argument_count);
1748
1749 var f16_add_artifact = try createDeviceScanAddBaseFamilyArtifact(allocator, state.handle(), f16_instance, .{ .limits = .testing });
1750 defer f16_add_artifact.deinit();
1751 const f16_add_entry = f16_add_artifact.entry();
1752 try testing.expectEqualStrings("accy_kernel_scan_device_add_base_family_64_f16", f16_add_entry.entry_name);
1753 try testing.expectEqual(@as(u32, 4), f16_add_entry.argument_count);
1754
1755 const u32_instance = DeviceScan{ .extent = 1000, .dtype = .u32, .threads = 64 };
1756 var u32_block_artifact = try createDeviceScanBlockScanFamilyArtifact(allocator, state.handle(), u32_instance, .{ .limits = .testing });
1757 defer u32_block_artifact.deinit();
1758 const u32_block_entry = u32_block_artifact.entry();
1759 try testing.expectEqualStrings("accy_kernel_scan_device_prefix_sum_block_scan_family_64_u32", u32_block_entry.entry_name);
1760 try testing.expectEqual(@as(u32, 4), u32_block_entry.argument_count);
1761
1762 var u32_add_artifact = try createDeviceScanAddBaseFamilyArtifact(allocator, state.handle(), u32_instance, .{ .limits = .testing });
1763 defer u32_add_artifact.deinit();
1764 const u32_add_entry = u32_add_artifact.entry();
1765 try testing.expectEqualStrings("accy_kernel_scan_device_add_base_family_64_u32", u32_add_entry.entry_name);
1766 try testing.expectEqual(@as(u32, 3), u32_add_entry.argument_count);
1767 }
1768
1769 test "scan device-wide u32 recording artifacts cover native targets" {
1770 const allocator = testing.allocator;
1771 const instance = DeviceScan{ .extent = 1000, .dtype = .u32, .threads = 64 };
1772 inline for (.{ gpu.ArtifactFormat.cuda_ptx, .vulkan_spirv, .metal_msl }) |format| {
1773 var state = gpu.recording.BackendState{
1774 .allocator = allocator,
1775 .kind = switch (format) {
1776 .cuda_ptx => .cuda,
1777 .vulkan_spirv => .vulkan,
1778 .metal_msl => .metal,
1779 else => .external,
1780 },
1781 .format = format,
1782 };
1783 var block_artifact = try createDeviceScanBlockScanFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing, .format = format });
1784 defer block_artifact.deinit();
1785 try testing.expectEqual(format, block_artifact.entry().format);
1786 try testing.expectEqualStrings("accy_kernel_scan_device_prefix_sum_block_scan_family_64_u32", block_artifact.entry().entry_name);
1787
1788 var add_artifact = try createDeviceScanAddBaseFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing, .format = format });
1789 defer add_artifact.deinit();
1790 try testing.expectEqual(format, add_artifact.entry().format);
1791 try testing.expectEqualStrings("accy_kernel_scan_device_add_base_family_64_u32", add_artifact.entry().entry_name);
1792 }
1793 }
1794
1795 test "scan device-wide pipeline descriptor binds the family artifacts" {
1796 const allocator = testing.allocator;
1797 var state = gpu.recording.BackendState{
1798 .allocator = allocator,
1799 .kind = .cuda,
1800 .format = .cuda_ptx,
1801 };
1802 const instance = DeviceScan{ .extent = 5000, .threads = 64 };
1803
1804 var artifacts = try createDeviceScanPipelineArtifacts(allocator, state.handle(), instance, .{
1805 .limits = kernel.Limits.testing,
1806 });
1807 defer artifacts.deinit();
1808 const entries = artifacts.entries();
1809 const registry = artifact_product.KernelCallRegistry{ .entries = entries[0..] };
1810
1811 var owned = try deviceScanPipeline(allocator, instance);
1812 defer owned.deinit();
1813 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_family_64_f32", owned.value.target);
1814 try testing.expectEqual(@as(u32, 1), owned.value.operand_count);
1815 try testing.expectEqual(@as(usize, 1), owned.value.runtime_scalar_bounds.len);
1816 try testing.expectEqual(@as(u32, 0), owned.value.runtime_scalar_bounds[0].argument_index);
1817 try testing.expectEqual(@as(u32, 6144), owned.value.runtime_scalar_bounds[0].max_u32);
1818 try testing.expectEqual(@as(usize, 2), owned.value.intermediates.len);
1819 try testing.expectEqual(@as(usize, 3), owned.value.stages.len);
1820 try owned.value.validate(registry, .cuda_ptx);
1821 const capacity_args = [_]choir_abi.ScalarArgument{.{ .u32 = 6144 }};
1822 try owned.value.validateRuntimeScalarArguments(capacity_args[0..]);
1823 const too_large_args = [_]choir_abi.ScalarArgument{.{ .u32 = 10000 }};
1824 try testing.expectError(error.LaunchArgumentMismatch, owned.value.validateRuntimeScalarArguments(too_large_args[0..]));
1825
1826 const exclusive = DeviceScan{ .extent = 5000, .mode = .exclusive, .threads = 64 };
1827 var owned_exclusive = try deviceScanPipeline(allocator, exclusive);
1828 defer owned_exclusive.deinit();
1829 try testing.expectEqualStrings(
1830 "accy.kernel.scan.device_prefix_sum_exclusive_family_64_f32",
1831 owned_exclusive.value.target,
1832 );
1833 try testing.expectError(error.InvalidArtifact, owned_exclusive.value.validate(registry, .cuda_ptx));
1834
1835 try testing.expectError(
1836 error.UnsupportedDeviceScanInstance,
1837 deviceScanPipeline(allocator, .{ .extent = 0, .threads = 64 }),
1838 );
1839
1840 const f16_instance = DeviceScan{ .extent = 5000, .dtype = .f16, .threads = 64 };
1841 var f16_artifacts = try createDeviceScanPipelineArtifacts(allocator, state.handle(), f16_instance, .{
1842 .limits = kernel.Limits.testing,
1843 });
1844 defer f16_artifacts.deinit();
1845 const f16_entries = f16_artifacts.entries();
1846 const f16_registry = artifact_product.KernelCallRegistry{ .entries = f16_entries[0..] };
1847
1848 var f16_pipeline = try deviceScanPipeline(allocator, f16_instance);
1849 defer f16_pipeline.deinit();
1850 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_family_64_f16", f16_pipeline.value.target);
1851 try testing.expectEqual(@as(usize, 3), f16_pipeline.value.intermediates.len);
1852 try testing.expectEqual(DType.f32, f16_pipeline.value.intermediates[2].dtype);
1853 try testing.expectEqual(@as(usize, 3), f16_pipeline.value.stages[2].buffers.len);
1854 try f16_pipeline.value.validate(f16_registry, .cuda_ptx);
1855
1856 const u32_instance = DeviceScan{ .extent = 5000, .dtype = .u32, .threads = 64 };
1857 var u32_artifacts = try createDeviceScanPipelineArtifacts(allocator, state.handle(), u32_instance, .{
1858 .limits = kernel.Limits.testing,
1859 });
1860 defer u32_artifacts.deinit();
1861 const u32_entries = u32_artifacts.entries();
1862 const u32_registry = artifact_product.KernelCallRegistry{ .entries = u32_entries[0..] };
1863
1864 var u32_pipeline = try deviceScanPipeline(allocator, u32_instance);
1865 defer u32_pipeline.deinit();
1866 try testing.expectEqualStrings("accy.kernel.scan.device_prefix_sum_family_64_u32", u32_pipeline.value.target);
1867 try testing.expectEqual(@as(usize, 2), u32_pipeline.value.intermediates.len);
1868 try testing.expectEqual(DType.u32, u32_pipeline.value.intermediates[0].dtype);
1869 try testing.expectEqual(DType.u32, u32_pipeline.value.intermediates[1].dtype);
1870 try u32_pipeline.value.validate(u32_registry, .cuda_ptx);
1871 }