lib/accy/src/kernel/library/sort.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 scan = @import("scan.zig");
11 const tuning = @import("tuning.zig");
12
13 const DType = choir_abi.DType;
14 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
15
16 pub const RadixSplit = struct {
17 extent: u64,
18 threads: u32 = 256,
19 element_axis: []const u8 = "e",
20 };
21
22 pub const BitonicBlock = struct {
23 extent: u64,
24 threads: u32 = 256,
25 element_axis: []const u8 = "e",
26 };
27
28 pub const TopKBlock = struct {
29 extent: u64,
30 k: u64,
31 threads: u32 = 256,
32 element_axis: []const u8 = "e",
33 };
34
35 pub const TopKBlockPairs = struct {
36 extent: u64,
37 k: u64,
38 threads: u32 = 256,
39 element_axis: []const u8 = "e",
40 };
41
42 pub const radix_split_family_version: u32 = 1;
43 pub const radix_split_warp_size: u32 = 32;
44 pub const radix_split_max_threads: u32 = 1024;
45 pub const radix_split_max_blocks: u32 = 1024;
46 pub const radix_split_key_bits: u32 = 32;
47 pub const radix_digit_bins: u32 = 16;
48 pub const radix_digit_bits: u32 = 4;
49 pub const bitonic_block_family_version: u32 = 1;
50 pub const bitonic_block_structure_name = "bitonic_block";
51 pub const top_k_block_structure_name = "top_k_block";
52 pub const top_k_block_pairs_structure_name = "top_k_block_pairs";
53 pub const bitonic_block_min_threads: u32 = 32;
54 pub const bitonic_block_max_threads: u32 = 1024;
55 pub const top_k_block_family_version: u32 = 1;
56 pub const top_k_block_pairs_family_version: u32 = 1;
57
58 pub const RadixSplitResolvedStructure = enum {
59 radix_split,
60 radix_digit,
61 };
62
63 pub fn radixSplitBlockCount(extent: u64, threads: u32) u64 {
64 return (extent + threads - 1) / threads;
65 }
66
67 pub fn radixSplitInstanceValid(instance: RadixSplit) bool {
68 if (instance.extent == 0) return false;
69 if (instance.threads == 0 or instance.threads > radix_split_max_threads) return false;
70 if (instance.threads % radix_split_warp_size != 0) return false;
71 return extent_mod.blockCountWithinLimit(instance.extent, instance.threads, radix_split_max_blocks);
72 }
73
74 fn powerOfTwo(value: u32) bool {
75 return value != 0 and (value & (value - 1)) == 0;
76 }
77
78 pub fn bitonicBlockInstanceValid(instance: BitonicBlock) bool {
79 if (instance.extent == 0) return false;
80 if (instance.threads < bitonic_block_min_threads or instance.threads > bitonic_block_max_threads) return false;
81 if (!powerOfTwo(instance.threads)) return false;
82 return instance.extent <= instance.threads;
83 }
84
85 pub fn bitonicBlockThreadsForExtent(extent: u64) ?u32 {
86 if (extent == 0 or extent > bitonic_block_max_threads) return null;
87 var threads = bitonic_block_min_threads;
88 while (@as(u64, threads) < extent) : (threads *= 2) {}
89 return threads;
90 }
91
92 pub fn topKBlockInstanceValid(instance: TopKBlock) bool {
93 if (instance.k == 0 or instance.k > instance.extent) return false;
94 return bitonicBlockInstanceValid(.{
95 .extent = instance.extent,
96 .threads = instance.threads,
97 .element_axis = instance.element_axis,
98 });
99 }
100
101 pub fn topKBlockPairsInstanceValid(instance: TopKBlockPairs) bool {
102 return topKBlockInstanceValid(.{
103 .extent = instance.extent,
104 .k = instance.k,
105 .threads = instance.threads,
106 .element_axis = instance.element_axis,
107 });
108 }
109
110 pub fn radixSplitTuningOperation(_: RadixSplit) entry.Operation {
111 return .{ .sort = .radix_ascending };
112 }
113
114 pub fn radixSplitTuningExtents(instance: RadixSplit) [1]u64 {
115 return .{instance.extent};
116 }
117
118 pub fn radixSplitFamilyTuningKey(
119 backing_allocator: std.mem.Allocator,
120 device_fingerprint: u64,
121 instance: RadixSplit,
122 ) !tuning.FamilyTuningKey {
123 const family_fingerprint = try radixSplitScatterFamilyFingerprint(backing_allocator, instance);
124 const extents = radixSplitTuningExtents(instance);
125 return tuning.FamilyTuningKey.init(
126 device_fingerprint,
127 family_fingerprint,
128 entry.operationFingerprint(radixSplitTuningOperation(instance)),
129 .i32,
130 radix_split_family_version,
131 extents[0..],
132 ) orelse unreachable;
133 }
134
135 pub fn resolveRadixSplitStructure(
136 backing_allocator: std.mem.Allocator,
137 reader: tuning.FamilyTuningReader,
138 instance: RadixSplit,
139 ) !?RadixSplitResolvedStructure {
140 const key = try radixSplitFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
141 const record = reader.table.find(key) orelse return null;
142 const split_target = try radixSplitPipelineTarget(backing_allocator, instance);
143 defer backing_allocator.free(split_target);
144 if (std.mem.eql(u8, split_target, record.target)) return .radix_split;
145 const digit_target = try radixDigitPipelineTarget(backing_allocator, instance);
146 defer backing_allocator.free(digit_target);
147 if (std.mem.eql(u8, digit_target, record.target)) return .radix_digit;
148 return null;
149 }
150
151 fn radix_split_flags_body_in_range(inner: anytype, ctx: anytype) !void {
152 const key = try ctx.args.param(.keys).load(inner, ctx.tid);
153 const shifted = try inner.shr(key.raw(), ctx.bit);
154 const one_i32 = try inner.constantInt(.i32, 1);
155 const masked = try inner.and_(shifted, one_i32);
156 const first_bucket = try inner.compare(.eq, masked, ctx.polarity);
157 const one_value = try inner.constantFloat(.f32, 1.0);
158 const zero_value = try inner.constantFloat(.f32, 0.0);
159 const flag = try inner.select(first_bucket, one_value, zero_value);
160 try ctx.args.param(.dst).store(inner, flag, ctx.tid);
161 }
162
163 fn radixSplitFlagsBody(k: anytype, spec: RadixSplit, args: anytype) !void {
164 if (!radixSplitInstanceValid(spec)) return error.UnsupportedRadixSplitInstance;
165 const extent = try k.castIndex(args.param(.extent).raw());
166 const tid = try k.globalId(.x);
167 const in_range = try k.compare(.lt, tid, extent);
168 const bit = args.param(.bit).raw();
169 const polarity = args.param(.polarity).raw();
170 try k.guardDo(in_range, .{ .args = args, .tid = tid, .bit = bit, .polarity = polarity }, radix_split_flags_body_in_range);
171 }
172
173 fn radix_split_scatter_body_in_range(inner: anytype, ctx: anytype) !void {
174 const key = try ctx.args.param(.keys).load(inner, ctx.tid);
175 const flag = try ctx.args.param(.flags).load(inner, ctx.tid);
176 const scanned = try ctx.args.param(.scanned).load(inner, ctx.tid);
177 const last_scanned = try ctx.args.param(.scanned).load(inner, ctx.last);
178 const last_flag = try ctx.args.param(.flags).load(inner, ctx.last);
179 const total_zeros_value = try inner.add(last_scanned.raw(), last_flag.raw());
180 const zeros_before = try inner.castIndex(try inner.cast(scanned.raw(), .i32));
181 const total_zeros = try inner.castIndex(try inner.cast(total_zeros_value, .i32));
182 const ones_before = try inner.sub(ctx.tid, zeros_before);
183 const ones_position = try inner.add(total_zeros, ones_before);
184 const one_value = try inner.constantFloat(.f32, 1.0);
185 const zeros_bucket = try inner.compare(.eq, flag.raw(), one_value);
186 const position = try inner.select(zeros_bucket, zeros_before, ones_position);
187 try ctx.args.param(.dst).store(inner, key.raw(), position);
188 }
189
190 fn radixSplitScatterBody(k: anytype, spec: RadixSplit, args: anytype) !void {
191 if (!radixSplitInstanceValid(spec)) return error.UnsupportedRadixSplitInstance;
192 const extent = try k.castIndex(args.param(.extent).raw());
193 const tid = try k.globalId(.x);
194 const one = try k.constantIndex(1);
195 const last = try k.sub(extent, one);
196 const in_range = try k.compare(.lt, tid, extent);
197 try k.guardDo(in_range, .{ .args = args, .tid = tid, .last = last }, radix_split_scatter_body_in_range);
198 }
199
200 fn radixSplitFamilySchedule(instance: RadixSplit) kernel.logical.schedule.ThreadBlocks {
201 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
202 }
203
204 fn radixSplitFlagsRuntimeFamily() type {
205 return kernel.logical.Family(.{
206 .name = "accy_kernel_sort_radix_split_flags_runtime_i32",
207 .parameters = .{
208 .dst = kernel.dynamicBuffer(.f32),
209 .keys = kernel.dynamicBuffer(.i32),
210 .extent = kernel.scalar(.i32),
211 .bit = kernel.scalar(.i32),
212 .polarity = kernel.scalar(.i32),
213 },
214 .Instance = RadixSplit,
215 .schedule = radixSplitFamilySchedule,
216 .body = radixSplitFlagsBody,
217 });
218 }
219
220 fn radixSplitScatterRuntimeFamily() type {
221 return kernel.logical.Family(.{
222 .name = "accy_kernel_sort_radix_split_scatter_runtime_i32",
223 .parameters = .{
224 .dst = kernel.dynamicBuffer(.i32),
225 .keys = kernel.dynamicBuffer(.i32),
226 .flags = kernel.dynamicBuffer(.f32),
227 .scanned = kernel.dynamicBuffer(.f32),
228 .extent = kernel.scalar(.i32),
229 },
230 .Instance = RadixSplit,
231 .schedule = radixSplitFamilySchedule,
232 .body = radixSplitScatterBody,
233 });
234 }
235
236 fn sort_shared_pair(inner: anytype, ctx: anytype) !void {
237 const lhs = try inner.loadIndex(ctx.shared, ctx.local);
238 const rhs = try inner.loadIndex(ctx.shared, ctx.partner);
239 const lower = try inner.min(lhs, rhs);
240 const upper = try inner.max(lhs, rhs);
241 const segment = try inner.and_(ctx.local_i32, ctx.direction_bit);
242 const ascending = try inner.compare(.eq, segment, ctx.zero_i32);
243 const first = try inner.select(ascending, lower, upper);
244 const second = try inner.select(ascending, upper, lower);
245 try inner.storeIndex(first, ctx.shared, ctx.local);
246 try inner.storeIndex(second, ctx.shared, ctx.partner);
247 }
248
249 fn bitonic_block_body_in_range(inner: anytype, ctx: anytype) !void {
250 try ctx.args.param(.dst).store(inner, ctx.sorted, ctx.local);
251 }
252
253 fn bitonicBlockBody(k: anytype, spec: BitonicBlock, args: anytype) !void {
254 if (!bitonicBlockInstanceValid(spec)) return error.UnsupportedBitonicBlockInstance;
255 const shared = try k.sharedBuffer(.i32, spec.threads);
256 const extent = try k.castIndex(args.param(.extent).raw());
257 const local = try k.castIndex(try k.threadId(.x));
258 const local_i32 = try k.cast(local, .i32);
259 const one = try k.constantIndex(1);
260 const last = try k.sub(extent, one);
261 const clamped = try k.min(local, last);
262 const in_range = try k.compare(.lt, local, extent);
263 const loaded = try args.param(.keys).load(k, clamped);
264 const padding = try k.constantInt(.i32, std.math.maxInt(i32));
265 const value = try k.select(in_range, loaded.raw(), padding);
266 try k.storeIndex(value, shared, local);
267 try k.barrier(.block);
268
269 var size: u32 = 2;
270 while (size <= spec.threads) : (size *= 2) {
271 var stride_value: u32 = size / 2;
272 while (stride_value > 0) : (stride_value /= 2) {
273 const stride_i32 = try k.constantInt(.i32, @as(i32, @intCast(stride_value)));
274 const partner_i32 = try k.xor(local_i32, stride_i32);
275 const partner = try k.castIndex(partner_i32);
276 const writes_pair = try k.compare(.lt, local_i32, partner_i32);
277 const direction_bit = try k.constantInt(.i32, @as(i32, @intCast(size)));
278 const zero_i32 = try k.constantInt(.i32, 0);
279 try k.guardDo(writes_pair, .{
280 .shared = shared,
281 .local = local,
282 .partner = partner,
283 .local_i32 = local_i32,
284 .direction_bit = direction_bit,
285 .zero_i32 = zero_i32,
286 }, sort_shared_pair);
287 try k.barrier(.block);
288 }
289 }
290
291 const sorted = try k.loadIndex(shared, local);
292 try k.guardDo(in_range, .{ .args = args, .sorted = sorted, .local = local }, bitonic_block_body_in_range);
293 }
294
295 fn top_k_block_body_writes_output(inner: anytype, ctx: anytype) !void {
296 try ctx.args.param(.dst).store(inner, ctx.selected, ctx.local);
297 }
298
299 fn topKBlockBody(k: anytype, spec: TopKBlock, args: anytype) !void {
300 if (!topKBlockInstanceValid(spec)) return error.UnsupportedTopKBlockInstance;
301 const shared = try k.sharedBuffer(.i32, spec.threads);
302 const extent = try k.castIndex(args.param(.extent).raw());
303 const local = try k.castIndex(try k.threadId(.x));
304 const local_i32 = try k.cast(local, .i32);
305 const one = try k.constantIndex(1);
306 const last = try k.sub(extent, one);
307 const clamped = try k.min(local, last);
308 const in_range = try k.compare(.lt, local, extent);
309 const loaded = try args.param(.keys).load(k, clamped);
310 const padding = try k.constantInt(.i32, std.math.maxInt(i32));
311 const value = try k.select(in_range, loaded.raw(), padding);
312 try k.storeIndex(value, shared, local);
313 try k.barrier(.block);
314
315 var size: u32 = 2;
316 while (size <= spec.threads) : (size *= 2) {
317 var stride_value: u32 = size / 2;
318 while (stride_value > 0) : (stride_value /= 2) {
319 const stride_i32 = try k.constantInt(.i32, @as(i32, @intCast(stride_value)));
320 const partner_i32 = try k.xor(local_i32, stride_i32);
321 const partner = try k.castIndex(partner_i32);
322 const writes_pair = try k.compare(.lt, local_i32, partner_i32);
323 const direction_bit = try k.constantInt(.i32, @as(i32, @intCast(size)));
324 const zero_i32 = try k.constantInt(.i32, 0);
325 try k.guardDo(writes_pair, .{
326 .shared = shared,
327 .local = local,
328 .partner = partner,
329 .local_i32 = local_i32,
330 .direction_bit = direction_bit,
331 .zero_i32 = zero_i32,
332 }, sort_shared_pair);
333 try k.barrier(.block);
334 }
335 }
336
337 const top_count = try k.constantIndex(@intCast(spec.k));
338 const writes_output = try k.compare(.lt, local, top_count);
339 const selected = try k.loadIndex(shared, local);
340 try k.guardDo(writes_output, .{ .args = args, .selected = selected, .local = local }, top_k_block_body_writes_output);
341 }
342
343 fn top_k_block_pairs_body_writes_pair(inner: anytype, ctx: anytype) !void {
344 const lhs_key = try inner.loadIndex(ctx.shared_keys, ctx.local);
345 const rhs_key = try inner.loadIndex(ctx.shared_keys, ctx.partner);
346 const lhs_value = try inner.loadIndex(ctx.shared_values, ctx.local);
347 const rhs_value = try inner.loadIndex(ctx.shared_values, ctx.partner);
348 const key_lt = try inner.compare(.lt, lhs_key, rhs_key);
349 const key_eq = try inner.compare(.eq, lhs_key, rhs_key);
350 const value_le = try inner.compare(.le, lhs_value, rhs_value);
351 const tie_before = try inner.and_(key_eq, value_le);
352 const lhs_first = try inner.or_(key_lt, tie_before);
353 const lower_key = try inner.select(lhs_first, lhs_key, rhs_key);
354 const upper_key = try inner.select(lhs_first, rhs_key, lhs_key);
355 const lower_value = try inner.select(lhs_first, lhs_value, rhs_value);
356 const upper_value = try inner.select(lhs_first, rhs_value, lhs_value);
357 const segment = try inner.and_(ctx.local_i32, ctx.direction_bit);
358 const ascending = try inner.compare(.eq, segment, ctx.zero_i32);
359 const first_key = try inner.select(ascending, lower_key, upper_key);
360 const second_key = try inner.select(ascending, upper_key, lower_key);
361 const first_value = try inner.select(ascending, lower_value, upper_value);
362 const second_value = try inner.select(ascending, upper_value, lower_value);
363 try inner.storeIndex(first_key, ctx.shared_keys, ctx.local);
364 try inner.storeIndex(second_key, ctx.shared_keys, ctx.partner);
365 try inner.storeIndex(first_value, ctx.shared_values, ctx.local);
366 try inner.storeIndex(second_value, ctx.shared_values, ctx.partner);
367 }
368
369 fn top_k_block_pairs_body_writes_output(inner: anytype, ctx: anytype) !void {
370 try ctx.args.param(.dst).store(inner, ctx.key, ctx.local);
371 try ctx.args.param(.dst_values).store(inner, ctx.value, ctx.local);
372 }
373
374 fn topKBlockPairsBody(k: anytype, spec: TopKBlockPairs, args: anytype) !void {
375 if (!topKBlockPairsInstanceValid(spec)) return error.UnsupportedTopKBlockPairsInstance;
376 const shared_keys = try k.sharedBuffer(.i32, spec.threads);
377 const shared_values = try k.sharedBuffer(.i32, spec.threads);
378 const extent = try k.castIndex(args.param(.extent).raw());
379 const local = try k.castIndex(try k.threadId(.x));
380 const local_i32 = try k.cast(local, .i32);
381 const one = try k.constantIndex(1);
382 const last = try k.sub(extent, one);
383 const clamped = try k.min(local, last);
384 const in_range = try k.compare(.lt, local, extent);
385 const loaded_key = try args.param(.keys).load(k, clamped);
386 const loaded_value = try args.param(.values).load(k, clamped);
387 const padding = try k.constantInt(.i32, std.math.maxInt(i32));
388 const key = try k.select(in_range, loaded_key.raw(), padding);
389 const payload = try k.select(in_range, loaded_value.raw(), padding);
390 try k.storeIndex(key, shared_keys, local);
391 try k.storeIndex(payload, shared_values, local);
392 try k.barrier(.block);
393
394 var size: u32 = 2;
395 while (size <= spec.threads) : (size *= 2) {
396 var stride_value: u32 = size / 2;
397 while (stride_value > 0) : (stride_value /= 2) {
398 const stride_i32 = try k.constantInt(.i32, @as(i32, @intCast(stride_value)));
399 const partner_i32 = try k.xor(local_i32, stride_i32);
400 const partner = try k.castIndex(partner_i32);
401 const writes_pair = try k.compare(.lt, local_i32, partner_i32);
402 const direction_bit = try k.constantInt(.i32, @as(i32, @intCast(size)));
403 const zero_i32 = try k.constantInt(.i32, 0);
404 try k.guardDo(writes_pair, .{
405 .shared_keys = shared_keys,
406 .shared_values = shared_values,
407 .local = local,
408 .partner = partner,
409 .local_i32 = local_i32,
410 .direction_bit = direction_bit,
411 .zero_i32 = zero_i32,
412 }, top_k_block_pairs_body_writes_pair);
413 try k.barrier(.block);
414 }
415 }
416
417 const top_count = try k.constantIndex(@intCast(spec.k));
418 const writes_output = try k.compare(.lt, local, top_count);
419 const selected_key = try k.loadIndex(shared_keys, local);
420 const selected_value = try k.loadIndex(shared_values, local);
421 try k.guardDo(writes_output, .{ .args = args, .key = selected_key, .value = selected_value, .local = local }, top_k_block_pairs_body_writes_output);
422 }
423
424 fn bitonicBlockFamilySchedule(instance: BitonicBlock) kernel.logical.schedule.ThreadBlocks {
425 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
426 }
427
428 fn topKBlockFamilySchedule(instance: TopKBlock) kernel.logical.schedule.ThreadBlocks {
429 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
430 }
431
432 fn topKBlockPairsFamilySchedule(instance: TopKBlockPairs) kernel.logical.schedule.ThreadBlocks {
433 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
434 }
435
436 fn bitonicBlockRuntimeFamily() type {
437 return kernel.logical.Family(.{
438 .name = "accy_kernel_sort_bitonic_block_runtime_i32",
439 .parameters = .{
440 .dst = kernel.dynamicBuffer(.i32),
441 .keys = kernel.dynamicBuffer(.i32),
442 .extent = kernel.scalar(.i32),
443 },
444 .Instance = BitonicBlock,
445 .schedule = bitonicBlockFamilySchedule,
446 .body = bitonicBlockBody,
447 });
448 }
449
450 fn topKBlockRuntimeFamily() type {
451 return kernel.logical.Family(.{
452 .name = "accy_kernel_sort_top_k_block_runtime_i32",
453 .parameters = .{
454 .dst = kernel.dynamicBuffer(.i32),
455 .keys = kernel.dynamicBuffer(.i32),
456 .extent = kernel.scalar(.i32),
457 },
458 .Instance = TopKBlock,
459 .schedule = topKBlockFamilySchedule,
460 .body = topKBlockBody,
461 });
462 }
463
464 fn topKBlockPairsRuntimeFamily() type {
465 return kernel.logical.Family(.{
466 .name = "accy_kernel_sort_top_k_block_pairs_runtime_i32",
467 .parameters = .{
468 .dst = kernel.dynamicBuffer(.i32),
469 .dst_values = kernel.dynamicBuffer(.i32),
470 .keys = kernel.dynamicBuffer(.i32),
471 .values = kernel.dynamicBuffer(.i32),
472 .extent = kernel.scalar(.i32),
473 },
474 .Instance = TopKBlockPairs,
475 .schedule = topKBlockPairsFamilySchedule,
476 .body = topKBlockPairsBody,
477 });
478 }
479
480 fn radix_digit_histogram_body_zero_bin(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
481 try loop_builder.storeIndex(ctx.zero_count, ctx.shared_bins, bin);
482 return acc;
483 }
484
485 fn radix_digit_histogram_body_active(inner: anytype, ctx: anytype) !void {
486 const key = try ctx.args.param(.keys).load(inner, ctx.element);
487 const shifted = try inner.shr(key.raw(), ctx.shift);
488 const mask = try inner.constantInt(.i32, radix_digit_bins - 1);
489 const masked = try inner.and_(shifted, mask);
490 const digit_i32 = try inner.xor(masked, ctx.bias);
491 const digit = try inner.castIndex(digit_i32);
492 const one = try inner.constantInt(.i32, 1);
493 _ = try inner.atomicRmwIndex(.add, one, ctx.shared_bins, digit);
494 }
495
496 fn radix_digit_histogram_body_grid(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
497 const partial = try loop_builder.loadIndex(ctx.shared_bins, bin);
498 const partial_value = try loop_builder.cast(partial, .f32);
499 const column = try loop_builder.mul(bin, ctx.grid);
500 const cell = try loop_builder.add(column, ctx.block);
501 try loop_builder.storeIndex(partial_value, ctx.args.param(.counts).raw(), cell);
502 return acc;
503 }
504
505 fn radixDigitHistogramBody(k: anytype, spec: RadixSplit, args: anytype) !void {
506 if (!radixSplitInstanceValid(spec)) return error.UnsupportedRadixSplitInstance;
507 const shared_bins = try k.sharedBuffer(.i32, radix_digit_bins);
508 const zero_count = try k.constantInt(.i32, 0);
509 const thread = try k.castIndex(try k.threadId(.x));
510 const stride = try k.castIndex(try k.blockDim(.x));
511 const bins = try k.constantIndex(radix_digit_bins);
512
513 _ = try k.fold(thread, bins, stride, zero_count, .{
514 .shared_bins = shared_bins,
515 .zero_count = zero_count,
516 }, radix_digit_histogram_body_zero_bin);
517 try k.barrier(.block);
518
519 const element = try k.globalId(.x);
520 const extent = try k.castIndex(args.param(.extent).raw());
521 const active = try k.compare(.lt, element, extent);
522 const shift = args.param(.shift).raw();
523 const bias = args.param(.bias).raw();
524 try k.guardDo(active, .{
525 .args = args,
526 .element = element,
527 .shift = shift,
528 .bias = bias,
529 .shared_bins = shared_bins,
530 }, radix_digit_histogram_body_active);
531 try k.barrier(.block);
532
533 const block = try k.blockId(.x);
534 const grid = try k.gridDim(.x);
535 _ = try k.fold(thread, bins, stride, zero_count, .{
536 .args = args,
537 .shared_bins = shared_bins,
538 .block = block,
539 .grid = grid,
540 }, radix_digit_histogram_body_grid);
541 }
542
543 fn radixDigitHistogramRuntimeFamily() type {
544 return kernel.logical.Family(.{
545 .name = "accy_kernel_sort_radix_digit_histogram_runtime_i32",
546 .parameters = .{
547 .counts = kernel.dynamicBuffer(.f32),
548 .keys = kernel.dynamicBuffer(.i32),
549 .extent = kernel.scalar(.i32),
550 .shift = kernel.scalar(.i32),
551 .bias = kernel.scalar(.i32),
552 },
553 .Instance = RadixSplit,
554 .schedule = radixSplitFamilySchedule,
555 .body = radixDigitHistogramBody,
556 });
557 }
558
559 fn radix_digit_rank_scatter_zero_cell(loop_builder: anytype, cell: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
560 try loop_builder.storeIndex(ctx.zero_i32, ctx.shared_counts, cell);
561 return acc;
562 }
563
564 fn radix_digit_rank_scatter_record_group(inner: anytype, ctx: anytype) !void {
565 try inner.guardDo(ctx.is_rank_zero, .{
566 .shared_counts = ctx.shared_counts,
567 .warp_cell = ctx.warp_cell,
568 .group_count = ctx.group_count,
569 }, radix_digit_rank_scatter_write_group_count);
570 }
571
572 fn radix_digit_rank_scatter_write_group_count(write_builder: anytype, write_ctx: anytype) !void {
573 try write_builder.storeIndex(write_ctx.group_count, write_ctx.shared_counts, write_ctx.warp_cell);
574 }
575
576 fn radix_digit_rank_scatter_key(inner: anytype, ctx: anytype) !void {
577 const cross_warp = try inner.fold(ctx.zero_index, ctx.warp, ctx.one_index, ctx.zero_i32, .{
578 .shared_counts = ctx.shared_counts,
579 .digit_index = ctx.digit_index,
580 .warps_value = ctx.warps_value,
581 }, radix_digit_rank_scatter_count_lower_warps);
582
583 const column = try inner.mul(ctx.digit_index, ctx.grid);
584 const base_cell = try inner.add(column, ctx.block);
585 const base_value = try inner.loadIndex(ctx.args.param(.scanned_counts).raw(), base_cell);
586 const base = try inner.cast(base_value, .i32);
587 const local = try inner.add(cross_warp, ctx.within_rank);
588 const position_i32 = try inner.add(base, local);
589 const position = try inner.castIndex(position_i32);
590 try ctx.args.param(.dst).store(inner, ctx.key, position);
591 }
592
593 fn radix_digit_rank_scatter_count_lower_warps(loop_builder: anytype, lower_warp: kernel.Value, acc: kernel.Value, fold_ctx: anytype) !kernel.Value {
594 const cell = try loop_builder.add(
595 try loop_builder.mul(fold_ctx.digit_index, fold_ctx.warps_value),
596 lower_warp,
597 );
598 const count = try loop_builder.loadIndex(fold_ctx.shared_counts, cell);
599 return try loop_builder.add(acc, count);
600 }
601
602 fn radixDigitRankScatterBody(k: anytype, spec: RadixSplit, args: anytype) !void {
603 if (!radixSplitInstanceValid(spec)) return error.UnsupportedRadixSplitInstance;
604 const warps_per_block = spec.threads / radix_split_warp_size;
605 const shared_counts = try k.sharedBuffer(.i32, radix_digit_bins * warps_per_block);
606 const zero_i32 = try k.constantInt(.i32, 0);
607 const zero_index = try k.constantIndex(0);
608 const one_index = try k.constantIndex(1);
609 const thread = try k.castIndex(try k.threadId(.x));
610 const stride = try k.castIndex(try k.blockDim(.x));
611 const cells = try k.constantIndex(radix_digit_bins * warps_per_block);
612
613 _ = try k.fold(thread, cells, stride, zero_i32, .{
614 .shared_counts = shared_counts,
615 .zero_i32 = zero_i32,
616 }, radix_digit_rank_scatter_zero_cell);
617 try k.barrier(.block);
618
619 const tid = try k.globalId(.x);
620 const extent = try k.castIndex(args.param(.extent).raw());
621 const in_range = try k.compare(.lt, tid, extent);
622 const last = try k.sub(extent, one_index);
623 const clamped = try k.min(tid, last);
624 const key = try args.param(.keys).load(k, clamped);
625 const shift = args.param(.shift).raw();
626 const bias = args.param(.bias).raw();
627 const shifted = try k.shr(key.raw(), shift);
628 const digit_mask = try k.constantInt(.i32, radix_digit_bins - 1);
629 const masked_digit = try k.and_(shifted, digit_mask);
630 const digit = try k.xor(masked_digit, bias);
631
632 var same_mask = try k.ballotSync(in_range);
633 inline for (0..radix_digit_bits) |bit_index| {
634 const bit_constant = try k.constantInt(.i32, @as(i32, 1) << bit_index);
635 const bit_value = try k.and_(digit, bit_constant);
636 const bit_set = try k.compare(.eq, bit_value, bit_constant);
637 const ballot = try k.ballotSync(bit_set);
638 const inverted = try k.not(ballot);
639 const matching = try k.select(bit_set, ballot, inverted);
640 same_mask = try k.and_(same_mask, matching);
641 }
642
643 const lane = try k.laneId();
644 const lane_i32 = try k.cast(lane, .i32);
645 const one_i32 = try k.constantInt(.i32, 1);
646 const lane_bit = try k.shl(one_i32, lane_i32);
647 const lower_mask = try k.sub(lane_bit, one_i32);
648 const below = try k.and_(same_mask, lower_mask);
649 const within_rank = try k.popcount(below);
650 const group_count = try k.popcount(same_mask);
651
652 const warp = try k.warpId();
653 const digit_index = try k.castIndex(digit);
654 const warps_value = try k.constantIndex(warps_per_block);
655 const warp_cell = try k.add(try k.mul(digit_index, warps_value), warp);
656 const is_rank_zero = try k.compare(.eq, within_rank, zero_i32);
657
658 try k.guardDo(in_range, .{
659 .shared_counts = shared_counts,
660 .warp_cell = warp_cell,
661 .group_count = group_count,
662 .is_rank_zero = is_rank_zero,
663 }, radix_digit_rank_scatter_record_group);
664 try k.barrier(.block);
665
666 const block = try k.blockId(.x);
667 const grid = try k.gridDim(.x);
668 try k.guardDo(in_range, .{
669 .args = args,
670 .key = key.raw(),
671 .digit_index = digit_index,
672 .warps_value = warps_value,
673 .warp = warp,
674 .block = block,
675 .grid = grid,
676 .within_rank = within_rank,
677 .shared_counts = shared_counts,
678 .zero_index = zero_index,
679 .one_index = one_index,
680 .zero_i32 = zero_i32,
681 }, radix_digit_rank_scatter_key);
682 }
683
684 fn radix_digit_rank_scatter_pair(inner: anytype, ctx: anytype) !void {
685 const cross_warp = try inner.fold(ctx.zero_index, ctx.warp, ctx.one_index, ctx.zero_i32, .{
686 .shared_counts = ctx.shared_counts,
687 .digit_index = ctx.digit_index,
688 .warps_value = ctx.warps_value,
689 }, radix_digit_rank_scatter_count_lower_warps);
690
691 const column = try inner.mul(ctx.digit_index, ctx.grid);
692 const base_cell = try inner.add(column, ctx.block);
693 const base_value = try inner.loadIndex(ctx.args.param(.scanned_counts).raw(), base_cell);
694 const base = try inner.cast(base_value, .i32);
695 const local = try inner.add(cross_warp, ctx.within_rank);
696 const position_i32 = try inner.add(base, local);
697 const position = try inner.castIndex(position_i32);
698 try ctx.args.param(.dst).store(inner, ctx.key, position);
699 try ctx.args.param(.dst_values).store(inner, ctx.payload, position);
700 }
701
702 fn radixDigitRankScatterPairsBody(k: anytype, spec: RadixSplit, args: anytype) !void {
703 if (!radixSplitInstanceValid(spec)) return error.UnsupportedRadixSplitInstance;
704 const warps_per_block = spec.threads / radix_split_warp_size;
705 const shared_counts = try k.sharedBuffer(.i32, radix_digit_bins * warps_per_block);
706 const zero_i32 = try k.constantInt(.i32, 0);
707 const zero_index = try k.constantIndex(0);
708 const one_index = try k.constantIndex(1);
709 const thread = try k.castIndex(try k.threadId(.x));
710 const stride = try k.castIndex(try k.blockDim(.x));
711 const cells = try k.constantIndex(radix_digit_bins * warps_per_block);
712
713 _ = try k.fold(thread, cells, stride, zero_i32, .{
714 .shared_counts = shared_counts,
715 .zero_i32 = zero_i32,
716 }, radix_digit_rank_scatter_zero_cell);
717 try k.barrier(.block);
718
719 const tid = try k.globalId(.x);
720 const extent = try k.castIndex(args.param(.extent).raw());
721 const in_range = try k.compare(.lt, tid, extent);
722 const last = try k.sub(extent, one_index);
723 const clamped = try k.min(tid, last);
724 const key = try args.param(.keys).load(k, clamped);
725 const payload = try args.param(.values).load(k, clamped);
726 const shift = args.param(.shift).raw();
727 const bias = args.param(.bias).raw();
728 const shifted = try k.shr(key.raw(), shift);
729 const digit_mask = try k.constantInt(.i32, radix_digit_bins - 1);
730 const masked_digit = try k.and_(shifted, digit_mask);
731 const digit = try k.xor(masked_digit, bias);
732
733 var same_mask = try k.ballotSync(in_range);
734 inline for (0..radix_digit_bits) |bit_index| {
735 const bit_constant = try k.constantInt(.i32, @as(i32, 1) << bit_index);
736 const bit_value = try k.and_(digit, bit_constant);
737 const bit_set = try k.compare(.eq, bit_value, bit_constant);
738 const ballot = try k.ballotSync(bit_set);
739 const inverted = try k.not(ballot);
740 const matching = try k.select(bit_set, ballot, inverted);
741 same_mask = try k.and_(same_mask, matching);
742 }
743
744 const lane = try k.laneId();
745 const lane_i32 = try k.cast(lane, .i32);
746 const one_i32 = try k.constantInt(.i32, 1);
747 const lane_bit = try k.shl(one_i32, lane_i32);
748 const lower_mask = try k.sub(lane_bit, one_i32);
749 const below = try k.and_(same_mask, lower_mask);
750 const within_rank = try k.popcount(below);
751 const group_count = try k.popcount(same_mask);
752
753 const warp = try k.warpId();
754 const digit_index = try k.castIndex(digit);
755 const warps_value = try k.constantIndex(warps_per_block);
756 const warp_cell = try k.add(try k.mul(digit_index, warps_value), warp);
757 const is_rank_zero = try k.compare(.eq, within_rank, zero_i32);
758
759 try k.guardDo(in_range, .{
760 .shared_counts = shared_counts,
761 .warp_cell = warp_cell,
762 .group_count = group_count,
763 .is_rank_zero = is_rank_zero,
764 }, radix_digit_rank_scatter_record_group);
765 try k.barrier(.block);
766
767 const block = try k.blockId(.x);
768 const grid = try k.gridDim(.x);
769 try k.guardDo(in_range, .{
770 .args = args,
771 .key = key.raw(),
772 .payload = payload.raw(),
773 .digit_index = digit_index,
774 .warps_value = warps_value,
775 .warp = warp,
776 .block = block,
777 .grid = grid,
778 .within_rank = within_rank,
779 .shared_counts = shared_counts,
780 .zero_index = zero_index,
781 .one_index = one_index,
782 .zero_i32 = zero_i32,
783 }, radix_digit_rank_scatter_pair);
784 }
785
786 fn radixDigitRankScatterPairsRuntimeFamily() type {
787 return kernel.logical.Family(.{
788 .name = "accy_kernel_sort_radix_digit_rank_scatter_pairs_runtime_i32",
789 .parameters = .{
790 .dst = kernel.dynamicBuffer(.i32),
791 .dst_values = kernel.dynamicBuffer(.i32),
792 .keys = kernel.dynamicBuffer(.i32),
793 .values = kernel.dynamicBuffer(.i32),
794 .scanned_counts = kernel.dynamicBuffer(.f32),
795 .extent = kernel.scalar(.i32),
796 .shift = kernel.scalar(.i32),
797 .bias = kernel.scalar(.i32),
798 },
799 .Instance = RadixSplit,
800 .schedule = radixSplitFamilySchedule,
801 .body = radixDigitRankScatterPairsBody,
802 });
803 }
804
805 pub const RadixDigitRankScatterPairsRuntimeFamilyI32 = radixDigitRankScatterPairsRuntimeFamily();
806
807 pub fn radixDigitRankScatterPairsFamilyTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
808 return std.fmt.allocPrint(
809 allocator,
810 "accy.kernel.sort.radix_digit_rank_scatter_pairs_family_{d}_i32",
811 .{instance.threads},
812 );
813 }
814
815 pub fn radixDigitRankScatterPairsFamilyEntryName(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
816 return std.fmt.allocPrint(
817 allocator,
818 "accy_kernel_sort_radix_digit_rank_scatter_pairs_family_{d}_i32",
819 .{instance.threads},
820 );
821 }
822
823 fn radixDigitRankScatterRuntimeFamily() type {
824 return kernel.logical.Family(.{
825 .name = "accy_kernel_sort_radix_digit_rank_scatter_runtime_i32",
826 .parameters = .{
827 .dst = kernel.dynamicBuffer(.i32),
828 .keys = kernel.dynamicBuffer(.i32),
829 .scanned_counts = kernel.dynamicBuffer(.f32),
830 .extent = kernel.scalar(.i32),
831 .shift = kernel.scalar(.i32),
832 .bias = kernel.scalar(.i32),
833 },
834 .Instance = RadixSplit,
835 .schedule = radixSplitFamilySchedule,
836 .body = radixDigitRankScatterBody,
837 });
838 }
839
840 pub const RadixDigitRankScatterRuntimeFamilyI32 = radixDigitRankScatterRuntimeFamily();
841
842 pub fn radixDigitRankScatterFamilyTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
843 return std.fmt.allocPrint(
844 allocator,
845 "accy.kernel.sort.radix_digit_rank_scatter_family_{d}_i32",
846 .{instance.threads},
847 );
848 }
849
850 pub fn radixDigitRankScatterFamilyEntryName(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
851 return std.fmt.allocPrint(
852 allocator,
853 "accy_kernel_sort_radix_digit_rank_scatter_family_{d}_i32",
854 .{instance.threads},
855 );
856 }
857
858 pub const RadixDigitHistogramRuntimeFamilyI32 = radixDigitHistogramRuntimeFamily();
859
860 pub fn radixDigitHistogramFamilyTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
861 return std.fmt.allocPrint(
862 allocator,
863 "accy.kernel.sort.radix_digit_histogram_family_{d}_i32",
864 .{instance.threads},
865 );
866 }
867
868 pub fn radixDigitHistogramFamilyEntryName(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
869 return std.fmt.allocPrint(
870 allocator,
871 "accy_kernel_sort_radix_digit_histogram_family_{d}_i32",
872 .{instance.threads},
873 );
874 }
875
876 pub fn radixDigitSignedPassBias(shift: u32) u32 {
877 return if (shift == radix_split_key_bits - radix_digit_bits) radix_digit_bins / 2 else 0;
878 }
879
880 pub fn radixDigitHistogramRuntimeArguments(instance: RadixSplit, shift: u32) ![3]choir_abi.ScalarArgument {
881 if (shift >= radix_split_key_bits) return error.UnsupportedRadixSplitInstance;
882 if (shift % radix_digit_bits != 0) return error.UnsupportedRadixSplitInstance;
883 return .{
884 .{ .u32 = try runtimeExtentArgument(instance.extent) },
885 .{ .u32 = shift },
886 .{ .u32 = radixDigitSignedPassBias(shift) },
887 };
888 }
889
890 pub const RadixSplitFlagsRuntimeFamilyI32 = radixSplitFlagsRuntimeFamily();
891 pub const RadixSplitScatterRuntimeFamilyI32 = radixSplitScatterRuntimeFamily();
892 pub const BitonicBlockRuntimeFamilyI32 = bitonicBlockRuntimeFamily();
893 pub const TopKBlockRuntimeFamilyI32 = topKBlockRuntimeFamily();
894 pub const TopKBlockPairsRuntimeFamilyI32 = topKBlockPairsRuntimeFamily();
895
896 pub fn radixSplitFlagsFamilyTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
897 return std.fmt.allocPrint(
898 allocator,
899 "accy.kernel.sort.radix_split_flags_family_{d}_i32",
900 .{instance.threads},
901 );
902 }
903
904 pub fn radixSplitFlagsFamilyEntryName(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
905 return std.fmt.allocPrint(
906 allocator,
907 "accy_kernel_sort_radix_split_flags_family_{d}_i32",
908 .{instance.threads},
909 );
910 }
911
912 pub fn radixSplitScatterFamilyTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
913 return std.fmt.allocPrint(
914 allocator,
915 "accy.kernel.sort.radix_split_scatter_family_{d}_i32",
916 .{instance.threads},
917 );
918 }
919
920 pub fn radixSplitScatterFamilyEntryName(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
921 return std.fmt.allocPrint(
922 allocator,
923 "accy_kernel_sort_radix_split_scatter_family_{d}_i32",
924 .{instance.threads},
925 );
926 }
927
928 pub fn radixSplitSignedPassPolarity(bit: u32) u32 {
929 return if (bit == radix_split_key_bits - 1) 1 else 0;
930 }
931
932 pub fn radixSplitFlagsRuntimeArguments(instance: RadixSplit, bit: u32) ![3]choir_abi.ScalarArgument {
933 if (bit >= radix_split_key_bits) return error.UnsupportedRadixSplitInstance;
934 return .{
935 .{ .u32 = try runtimeExtentArgument(instance.extent) },
936 .{ .u32 = bit },
937 .{ .u32 = radixSplitSignedPassPolarity(bit) },
938 };
939 }
940
941 pub fn radixSplitScatterRuntimeArguments(instance: RadixSplit) ![1]choir_abi.ScalarArgument {
942 return .{
943 .{ .u32 = try runtimeExtentArgument(instance.extent) },
944 };
945 }
946
947 pub fn bitonicBlockFamilyTarget(allocator: std.mem.Allocator, instance: BitonicBlock) ![]u8 {
948 return std.fmt.allocPrint(
949 allocator,
950 "accy.kernel.sort.bitonic_block_family_{d}_i32",
951 .{instance.threads},
952 );
953 }
954
955 pub fn bitonicBlockFamilyEntryName(allocator: std.mem.Allocator, instance: BitonicBlock) ![]u8 {
956 return std.fmt.allocPrint(
957 allocator,
958 "accy_kernel_sort_bitonic_block_family_{d}_i32",
959 .{instance.threads},
960 );
961 }
962
963 pub fn bitonicBlockRuntimeArguments(instance: BitonicBlock) ![1]choir_abi.ScalarArgument {
964 return .{
965 .{ .u32 = try runtimeExtentArgument(instance.extent) },
966 };
967 }
968
969 pub fn topKBlockFamilyTarget(allocator: std.mem.Allocator, instance: TopKBlock) ![]u8 {
970 return std.fmt.allocPrint(
971 allocator,
972 "accy.kernel.sort.top_k_block_family_{d}x{d}_i32",
973 .{ instance.threads, instance.k },
974 );
975 }
976
977 pub fn topKBlockFamilyEntryName(allocator: std.mem.Allocator, instance: TopKBlock) ![]u8 {
978 return std.fmt.allocPrint(
979 allocator,
980 "accy_kernel_sort_top_k_block_family_{d}x{d}_i32",
981 .{ instance.threads, instance.k },
982 );
983 }
984
985 pub fn topKBlockRuntimeArguments(instance: TopKBlock) ![1]choir_abi.ScalarArgument {
986 return .{
987 .{ .u32 = try runtimeExtentArgument(instance.extent) },
988 };
989 }
990
991 pub fn topKBlockPairsFamilyTarget(allocator: std.mem.Allocator, instance: TopKBlockPairs) ![]u8 {
992 return std.fmt.allocPrint(
993 allocator,
994 "accy.kernel.sort.top_k_block_pairs_family_{d}x{d}_i32",
995 .{ instance.threads, instance.k },
996 );
997 }
998
999 pub fn topKBlockPairsFamilyEntryName(allocator: std.mem.Allocator, instance: TopKBlockPairs) ![]u8 {
1000 return std.fmt.allocPrint(
1001 allocator,
1002 "accy_kernel_sort_top_k_block_pairs_family_{d}x{d}_i32",
1003 .{ instance.threads, instance.k },
1004 );
1005 }
1006
1007 pub fn topKBlockPairsRuntimeArguments(instance: TopKBlockPairs) ![1]choir_abi.ScalarArgument {
1008 return .{
1009 .{ .u32 = try runtimeExtentArgument(instance.extent) },
1010 };
1011 }
1012
1013 pub fn radixSplitThreadsForExtent(extent: u64) ?u32 {
1014 if (extent == 0) return null;
1015 const max_extent = @as(u64, radix_split_max_threads) * radix_split_max_blocks;
1016 if (extent > max_extent) return null;
1017 const needed = (extent + radix_split_max_blocks - 1) / radix_split_max_blocks;
1018 const wide = needed + radix_split_warp_size - 1;
1019 const rounded: u32 = @intCast((wide / radix_split_warp_size) * radix_split_warp_size);
1020 return @max(rounded, radix_split_warp_size);
1021 }
1022
1023 fn bitonicBlockScheduleMetadata(lifetime_allocator: std.mem.Allocator, instance: BitonicBlock) !entry.Schedule {
1024 const bindings = try lifetime_allocator.alloc(entry.ScheduleBinding, 1);
1025 bindings[0] = .{
1026 .axis = try std.fmt.allocPrint(lifetime_allocator, "{s}_lane", .{instance.element_axis}),
1027 .target = .thread_x,
1028 .extent = instance.threads,
1029 };
1030 return .{ .bindings = bindings };
1031 }
1032
1033 pub fn bitonicBlockShapeProfileDimensions(instance: BitonicBlock) [1]artifact_product.KernelCallShapeProfileDimension {
1034 return .{
1035 .{
1036 .name = instance.element_axis,
1037 .runtime_scalar_argument_index = 0,
1038 .bounds = .{ .min = 1, .max = instance.threads },
1039 },
1040 };
1041 }
1042
1043 pub fn topKBlockShapeProfileDimensions(instance: TopKBlock) [1]artifact_product.KernelCallShapeProfileDimension {
1044 return .{
1045 .{
1046 .name = instance.element_axis,
1047 .runtime_scalar_argument_index = 0,
1048 .bounds = .{ .min = 1, .max = instance.threads },
1049 },
1050 };
1051 }
1052
1053 pub fn topKBlockPairsShapeProfileDimensions(instance: TopKBlockPairs) [1]artifact_product.KernelCallShapeProfileDimension {
1054 return .{
1055 .{
1056 .name = instance.element_axis,
1057 .runtime_scalar_argument_index = 0,
1058 .bounds = .{ .min = 1, .max = instance.threads },
1059 },
1060 };
1061 }
1062
1063 fn bitonicBlockLaunch(instance: BitonicBlock) !artifact_product.KernelCallLaunch {
1064 if (!bitonicBlockInstanceValid(instance)) return error.UnsupportedBitonicBlockInstance;
1065 return .{ .derived = .{
1066 .grid = .{
1067 .{ .fixed = 1 },
1068 .{ .fixed = 1 },
1069 .{ .fixed = 1 },
1070 },
1071 .threadgroup = .{ instance.threads, 1, 1 },
1072 } };
1073 }
1074
1075 fn topKBlockLaunch(instance: TopKBlock) !artifact_product.KernelCallLaunch {
1076 if (!topKBlockInstanceValid(instance)) return error.UnsupportedTopKBlockInstance;
1077 return .{ .derived = .{
1078 .grid = .{
1079 .{ .fixed = 1 },
1080 .{ .fixed = 1 },
1081 .{ .fixed = 1 },
1082 },
1083 .threadgroup = .{ instance.threads, 1, 1 },
1084 } };
1085 }
1086
1087 fn topKBlockPairsLaunch(instance: TopKBlockPairs) !artifact_product.KernelCallLaunch {
1088 if (!topKBlockPairsInstanceValid(instance)) return error.UnsupportedTopKBlockPairsInstance;
1089 return .{ .derived = .{
1090 .grid = .{
1091 .{ .fixed = 1 },
1092 .{ .fixed = 1 },
1093 .{ .fixed = 1 },
1094 },
1095 .threadgroup = .{ instance.threads, 1, 1 },
1096 } };
1097 }
1098
1099 pub fn createBitonicBlockFamilyArtifact(
1100 allocator: std.mem.Allocator,
1101 handle: kernel.BackendHandle,
1102 instance: BitonicBlock,
1103 options: entry.ArtifactOptions,
1104 ) !kernel.OwnedKernelCallArtifact {
1105 if (!bitonicBlockInstanceValid(instance)) return error.UnsupportedBitonicBlockInstance;
1106 const target = try bitonicBlockFamilyTarget(allocator, instance);
1107 defer allocator.free(target);
1108 const entry_name = try bitonicBlockFamilyEntryName(allocator, instance);
1109 defer allocator.free(entry_name);
1110 const family_fingerprint = options.shape_family_fingerprint orelse try bitonicBlockFamilyFingerprint(allocator, instance);
1111 const shape_profile_dimensions = bitonicBlockShapeProfileDimensions(instance);
1112 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1113 .name = "bitonic_block",
1114 .fingerprint = family_fingerprint,
1115 .dimensions = shape_profile_dimensions[0..],
1116 };
1117
1118 var graph = try BitonicBlockRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1119 defer graph.deinit();
1120 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1121 .target = target,
1122 .version = bitonic_block_family_version,
1123 .format = options.format,
1124 .kernel_plan = options.kernel_plan,
1125 .element_count_argument = options.element_count_argument,
1126 .shape_family_fingerprint = family_fingerprint,
1127 .shape_profile = shape_profile,
1128 .launch = options.launch orelse try bitonicBlockLaunch(instance),
1129 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
1130 .static_arguments = options.static_arguments,
1131 });
1132 }
1133
1134 pub fn createTopKBlockFamilyArtifact(
1135 allocator: std.mem.Allocator,
1136 handle: kernel.BackendHandle,
1137 instance: TopKBlock,
1138 options: entry.ArtifactOptions,
1139 ) !kernel.OwnedKernelCallArtifact {
1140 if (!topKBlockInstanceValid(instance)) return error.UnsupportedTopKBlockInstance;
1141 const target = try topKBlockFamilyTarget(allocator, instance);
1142 defer allocator.free(target);
1143 const entry_name = try topKBlockFamilyEntryName(allocator, instance);
1144 defer allocator.free(entry_name);
1145 const family_fingerprint = options.shape_family_fingerprint orelse try topKBlockFamilyFingerprint(allocator, instance);
1146 const shape_profile_dimensions = topKBlockShapeProfileDimensions(instance);
1147 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1148 .name = "top_k_block",
1149 .fingerprint = family_fingerprint,
1150 .dimensions = shape_profile_dimensions[0..],
1151 };
1152
1153 var graph = try TopKBlockRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1154 defer graph.deinit();
1155 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1156 .target = target,
1157 .version = top_k_block_family_version,
1158 .format = options.format,
1159 .kernel_plan = options.kernel_plan,
1160 .element_count_argument = options.element_count_argument,
1161 .shape_family_fingerprint = family_fingerprint,
1162 .shape_profile = shape_profile,
1163 .launch = options.launch orelse try topKBlockLaunch(instance),
1164 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
1165 .static_arguments = options.static_arguments,
1166 });
1167 }
1168
1169 pub fn createTopKBlockPairsFamilyArtifact(
1170 allocator: std.mem.Allocator,
1171 handle: kernel.BackendHandle,
1172 instance: TopKBlockPairs,
1173 options: entry.ArtifactOptions,
1174 ) !kernel.OwnedKernelCallArtifact {
1175 if (!topKBlockPairsInstanceValid(instance)) return error.UnsupportedTopKBlockPairsInstance;
1176 const target = try topKBlockPairsFamilyTarget(allocator, instance);
1177 defer allocator.free(target);
1178 const entry_name = try topKBlockPairsFamilyEntryName(allocator, instance);
1179 defer allocator.free(entry_name);
1180 const family_fingerprint = options.shape_family_fingerprint orelse try topKBlockPairsFamilyFingerprint(allocator, instance);
1181 const shape_profile_dimensions = topKBlockPairsShapeProfileDimensions(instance);
1182 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1183 .name = "top_k_block_pairs",
1184 .fingerprint = family_fingerprint,
1185 .dimensions = shape_profile_dimensions[0..],
1186 };
1187
1188 var graph = try TopKBlockPairsRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1189 defer graph.deinit();
1190 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1191 .target = target,
1192 .version = top_k_block_pairs_family_version,
1193 .format = options.format,
1194 .kernel_plan = options.kernel_plan,
1195 .element_count_argument = options.element_count_argument,
1196 .shape_family_fingerprint = family_fingerprint,
1197 .shape_profile = shape_profile,
1198 .launch = options.launch orelse try topKBlockPairsLaunch(instance),
1199 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
1200 .static_arguments = options.static_arguments,
1201 });
1202 }
1203
1204 pub fn bitonicBlockFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BitonicBlock) !u64 {
1205 var family = try bitonicBlockShapeFamily(backing_allocator, instance);
1206 defer family.deinit();
1207 return shape.fingerprint(family);
1208 }
1209
1210 pub fn topKBlockFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: TopKBlock) !u64 {
1211 var family = try topKBlockShapeFamily(backing_allocator, instance);
1212 defer family.deinit();
1213 return shape.fingerprint(family);
1214 }
1215
1216 pub fn topKBlockPairsFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: TopKBlockPairs) !u64 {
1217 var family = try topKBlockPairsShapeFamily(backing_allocator, instance);
1218 defer family.deinit();
1219 return shape.fingerprint(family);
1220 }
1221
1222 pub fn bitonicBlockShapeFamily(backing_allocator: std.mem.Allocator, instance: BitonicBlock) !shape.Family {
1223 var builder = try shape.Builder.init(backing_allocator, "bitonic_block");
1224 errdefer builder.deinit();
1225 const elements = try builder.symbol(instance.element_axis);
1226 const elements_expr = try builder.symbolExpression(elements);
1227 _ = try builder.tensor("keys", &.{elements_expr});
1228 _ = try builder.tensor("out", &.{elements_expr});
1229 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = instance.threads });
1230 return builder.finish();
1231 }
1232
1233 pub fn topKBlockShapeFamily(backing_allocator: std.mem.Allocator, instance: TopKBlock) !shape.Family {
1234 var builder = try shape.Builder.init(backing_allocator, "top_k_block");
1235 errdefer builder.deinit();
1236 const elements = try builder.symbol(instance.element_axis);
1237 const elements_expr = try builder.symbolExpression(elements);
1238 const top_expr = builder.constantExpression(@intCast(instance.k));
1239 _ = try builder.tensor("keys", &.{elements_expr});
1240 _ = try builder.tensor("out", &.{top_expr});
1241 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = instance.threads });
1242 return builder.finish();
1243 }
1244
1245 pub fn topKBlockPairsShapeFamily(backing_allocator: std.mem.Allocator, instance: TopKBlockPairs) !shape.Family {
1246 var builder = try shape.Builder.init(backing_allocator, "top_k_block_pairs");
1247 errdefer builder.deinit();
1248 const elements = try builder.symbol(instance.element_axis);
1249 const elements_expr = try builder.symbolExpression(elements);
1250 const top_expr = builder.constantExpression(@intCast(instance.k));
1251 _ = try builder.tensor("keys", &.{elements_expr});
1252 _ = try builder.tensor("values", &.{elements_expr});
1253 _ = try builder.tensor("out_keys", &.{top_expr});
1254 _ = try builder.tensor("out_values", &.{top_expr});
1255 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = instance.threads });
1256 return builder.finish();
1257 }
1258
1259 pub fn bitonicBlockFamilySpecialization(backing_allocator: std.mem.Allocator, instance: BitonicBlock) !entry.OwnedSpecialization {
1260 if (!bitonicBlockInstanceValid(instance)) return error.UnsupportedBitonicBlockInstance;
1261 var owned = entry.OwnedSpecialization.init(backing_allocator);
1262 errdefer owned.deinit();
1263 const lifetime_allocator = owned.allocator();
1264
1265 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
1266 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1267
1268 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1269 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1270
1271 owned.value = .{
1272 .dtype = .i32,
1273 .operation = .{ .sort = .radix_ascending },
1274 .inputs = inputs,
1275 .outputs = outputs,
1276 .schedule = try bitonicBlockScheduleMetadata(lifetime_allocator, instance),
1277 .structure = bitonic_block_structure_name,
1278 };
1279 owned.value.launch = owned.value.schedule.?.launch();
1280 var family = try bitonicBlockShapeFamily(backing_allocator, instance);
1281 errdefer family.deinit();
1282 try owned.takeShapeFamily(&family);
1283 return owned;
1284 }
1285
1286 pub fn topKBlockFamilySpecialization(backing_allocator: std.mem.Allocator, instance: TopKBlock) !entry.OwnedSpecialization {
1287 if (!topKBlockInstanceValid(instance)) return error.UnsupportedTopKBlockInstance;
1288 var owned = entry.OwnedSpecialization.init(backing_allocator);
1289 errdefer owned.deinit();
1290 const lifetime_allocator = owned.allocator();
1291
1292 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
1293 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1294
1295 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1296 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, "k", instance.k);
1297
1298 owned.value = .{
1299 .dtype = .i32,
1300 .operation = .{ .sort = .top_k_smallest },
1301 .inputs = inputs,
1302 .outputs = outputs,
1303 .schedule = try bitonicBlockScheduleMetadata(lifetime_allocator, .{
1304 .extent = instance.extent,
1305 .threads = instance.threads,
1306 .element_axis = instance.element_axis,
1307 }),
1308 .structure = top_k_block_structure_name,
1309 };
1310 owned.value.launch = owned.value.schedule.?.launch();
1311 var family = try topKBlockShapeFamily(backing_allocator, instance);
1312 errdefer family.deinit();
1313 try owned.takeShapeFamily(&family);
1314 return owned;
1315 }
1316
1317 pub fn topKBlockPairsFamilySpecialization(backing_allocator: std.mem.Allocator, instance: TopKBlockPairs) !entry.OwnedSpecialization {
1318 if (!topKBlockPairsInstanceValid(instance)) return error.UnsupportedTopKBlockPairsInstance;
1319 var owned = entry.OwnedSpecialization.init(backing_allocator);
1320 errdefer owned.deinit();
1321 const lifetime_allocator = owned.allocator();
1322
1323 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
1324 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1325 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1326
1327 const outputs = try lifetime_allocator.alloc(entry.Shape, 2);
1328 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, "k", instance.k);
1329 outputs[1] = try entry.runtimeShape1D(lifetime_allocator, "k", instance.k);
1330
1331 owned.value = .{
1332 .dtype = .i32,
1333 .operation = .{ .sort = .top_k_smallest },
1334 .inputs = inputs,
1335 .outputs = outputs,
1336 .schedule = try bitonicBlockScheduleMetadata(lifetime_allocator, .{
1337 .extent = instance.extent,
1338 .threads = instance.threads,
1339 .element_axis = instance.element_axis,
1340 }),
1341 .structure = top_k_block_pairs_structure_name,
1342 };
1343 owned.value.launch = owned.value.schedule.?.launch();
1344 var family = try topKBlockPairsShapeFamily(backing_allocator, instance);
1345 errdefer family.deinit();
1346 try owned.takeShapeFamily(&family);
1347 return owned;
1348 }
1349
1350 pub fn bitonicBlockInstanceFromSpecialization(specialization: entry.Specialization) ?BitonicBlock {
1351 if (!specialization.scheduleMatchesLaunch()) return null;
1352 if (!specialization.operationIs(.{ .sort = .radix_ascending })) return null;
1353 if (!specialization.structureIs(bitonic_block_structure_name)) return null;
1354 const dtype = specialization.dtype orelse return null;
1355 if (dtype != .i32) return null;
1356 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
1357 if (specialization.reductions.len != 0) return null;
1358 const data = specialization.inputs[0];
1359 const output = specialization.outputs[0];
1360 if (data.axes.len != 1 or output.axes.len != 1) return null;
1361 const extent = data.axes[0].extent;
1362 if (output.axes[0].extent != extent) return null;
1363 const launch = specialization.launch orelse return null;
1364 if (launch.grid[0] != 1 or launch.grid[1] != 1 or launch.grid[2] != 1) return null;
1365 if (launch.threadgroup[1] != 1 or launch.threadgroup[2] != 1) return null;
1366 const instance = BitonicBlock{
1367 .extent = extent,
1368 .threads = launch.threadgroup[0],
1369 .element_axis = data.axes[0].name,
1370 };
1371 if (!bitonicBlockInstanceValid(instance)) return null;
1372 return instance;
1373 }
1374
1375 pub fn topKBlockInstanceFromSpecialization(specialization: entry.Specialization) ?TopKBlock {
1376 if (!specialization.scheduleMatchesLaunch()) return null;
1377 if (!specialization.operationIs(.{ .sort = .top_k_smallest })) return null;
1378 if (!specialization.structureIs(top_k_block_structure_name)) return null;
1379 const dtype = specialization.dtype orelse return null;
1380 if (dtype != .i32) return null;
1381 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
1382 if (specialization.reductions.len != 0) return null;
1383 const data = specialization.inputs[0];
1384 const output = specialization.outputs[0];
1385 if (data.axes.len != 1 or output.axes.len != 1) return null;
1386 const extent = data.axes[0].extent;
1387 const top_count = output.axes[0].extent;
1388 const launch = specialization.launch orelse return null;
1389 if (launch.grid[0] != 1 or launch.grid[1] != 1 or launch.grid[2] != 1) return null;
1390 if (launch.threadgroup[1] != 1 or launch.threadgroup[2] != 1) return null;
1391 const instance = TopKBlock{
1392 .extent = extent,
1393 .k = top_count,
1394 .threads = launch.threadgroup[0],
1395 .element_axis = data.axes[0].name,
1396 };
1397 if (!topKBlockInstanceValid(instance)) return null;
1398 return instance;
1399 }
1400
1401 pub fn topKBlockPairsInstanceFromSpecialization(specialization: entry.Specialization) ?TopKBlockPairs {
1402 if (!specialization.scheduleMatchesLaunch()) return null;
1403 if (!specialization.operationIs(.{ .sort = .top_k_smallest })) return null;
1404 if (!specialization.structureIs(top_k_block_pairs_structure_name)) return null;
1405 const dtype = specialization.dtype orelse return null;
1406 if (dtype != .i32) return null;
1407 if (specialization.inputs.len != 2 or specialization.outputs.len != 2) return null;
1408 if (specialization.reductions.len != 0) return null;
1409 const keys = specialization.inputs[0];
1410 const values = specialization.inputs[1];
1411 const out_keys = specialization.outputs[0];
1412 const out_values = specialization.outputs[1];
1413 if (keys.axes.len != 1 or values.axes.len != 1 or out_keys.axes.len != 1 or out_values.axes.len != 1) return null;
1414 const extent = keys.axes[0].extent;
1415 const top_count = out_keys.axes[0].extent;
1416 if (values.axes[0].extent != extent) return null;
1417 if (out_values.axes[0].extent != top_count) return null;
1418 const launch = specialization.launch orelse return null;
1419 if (launch.grid[0] != 1 or launch.grid[1] != 1 or launch.grid[2] != 1) return null;
1420 if (launch.threadgroup[1] != 1 or launch.threadgroup[2] != 1) return null;
1421 const instance = TopKBlockPairs{
1422 .extent = extent,
1423 .k = top_count,
1424 .threads = launch.threadgroup[0],
1425 .element_axis = keys.axes[0].name,
1426 };
1427 if (!topKBlockPairsInstanceValid(instance)) return null;
1428 return instance;
1429 }
1430
1431 pub fn radixSplitFamilySpecialization(backing_allocator: std.mem.Allocator, instance: RadixSplit) !entry.OwnedSpecialization {
1432 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1433 var owned = entry.OwnedSpecialization.init(backing_allocator);
1434 errdefer owned.deinit();
1435 const lifetime_allocator = owned.allocator();
1436
1437 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
1438 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1439
1440 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1441 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
1442
1443 owned.value = .{
1444 .dtype = .i32,
1445 .operation = .{ .sort = .radix_ascending },
1446 .inputs = inputs,
1447 .outputs = outputs,
1448 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, instance.element_axis, instance.extent, instance.threads),
1449 };
1450 owned.value.launch = owned.value.schedule.?.launch();
1451 var family = try radixSplitShapeFamily(backing_allocator, "radix_split", &.{ "keys", "out" }, instance);
1452 errdefer family.deinit();
1453 try owned.takeShapeFamily(&family);
1454 return owned;
1455 }
1456
1457 pub fn radixSplitInstanceFromSpecialization(specialization: entry.Specialization) ?RadixSplit {
1458 if (!specialization.scheduleMatchesLaunch()) return null;
1459 if (!specialization.operationIs(.{ .sort = .radix_ascending })) return null;
1460 const dtype = specialization.dtype orelse return null;
1461 if (dtype != .i32) return null;
1462 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
1463 if (specialization.reductions.len != 0) return null;
1464 const data = specialization.inputs[0];
1465 const output = specialization.outputs[0];
1466 if (data.axes.len != 1 or output.axes.len != 1) return null;
1467 const extent = data.axes[0].extent;
1468 if (output.axes[0].extent != extent) return null;
1469 const launch = specialization.launch orelse return null;
1470 const instance = RadixSplit{
1471 .extent = extent,
1472 .threads = launch.threadgroup[0],
1473 .element_axis = data.axes[0].name,
1474 };
1475 if (!radixSplitInstanceValid(instance)) return null;
1476 if (launch.grid[0] != radixSplitBlockCount(extent, instance.threads)) return null;
1477 return instance;
1478 }
1479
1480 pub fn radixSplitMaxExtent(instance: RadixSplit) u64 {
1481 return @as(u64, instance.threads) * radix_split_max_blocks;
1482 }
1483
1484 pub fn radixSplitShapeProfileDimensions(instance: RadixSplit) [1]artifact_product.KernelCallShapeProfileDimension {
1485 return .{
1486 .{
1487 .name = instance.element_axis,
1488 .runtime_scalar_argument_index = 0,
1489 .bounds = .{ .min = 1, .max = radixSplitMaxExtent(instance) },
1490 },
1491 };
1492 }
1493
1494 fn radixSplitLaunch(instance: RadixSplit) !artifact_product.KernelCallLaunch {
1495 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1496 return .{ .derived = .{
1497 .grid = .{
1498 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
1499 .{ .fixed = 1 },
1500 .{ .fixed = 1 },
1501 },
1502 .threadgroup = .{ instance.threads, 1, 1 },
1503 } };
1504 }
1505
1506 fn radixSplitShapeFamily(
1507 backing_allocator: std.mem.Allocator,
1508 comptime family_name: []const u8,
1509 comptime tensor_names: []const []const u8,
1510 instance: RadixSplit,
1511 ) !shape.Family {
1512 var builder = try shape.Builder.init(backing_allocator, family_name);
1513 errdefer builder.deinit();
1514 const elements = try builder.symbol(instance.element_axis);
1515 const elements_expr = try builder.symbolExpression(elements);
1516 inline for (tensor_names) |tensor_name| {
1517 _ = try builder.tensor(tensor_name, &.{elements_expr});
1518 }
1519 try builder.assumeBounds(elements_expr, .{ .min = 1, .max = radixSplitMaxExtent(instance) });
1520 return builder.finish();
1521 }
1522
1523 pub fn radixSplitFlagsFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: RadixSplit) !u64 {
1524 var family = try radixSplitShapeFamily(backing_allocator, "radix_split_flags", &.{ "keys", "flags" }, instance);
1525 defer family.deinit();
1526 return shape.fingerprint(family);
1527 }
1528
1529 pub fn radixSplitScatterFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: RadixSplit) !u64 {
1530 var family = try radixSplitShapeFamily(
1531 backing_allocator,
1532 "radix_split_scatter",
1533 &.{ "keys", "flags", "scanned", "out" },
1534 instance,
1535 );
1536 defer family.deinit();
1537 return shape.fingerprint(family);
1538 }
1539
1540 pub fn createRadixSplitFlagsFamilyArtifact(
1541 allocator: std.mem.Allocator,
1542 handle: kernel.BackendHandle,
1543 instance: RadixSplit,
1544 options: entry.ArtifactOptions,
1545 ) !kernel.OwnedKernelCallArtifact {
1546 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1547 const target = try radixSplitFlagsFamilyTarget(allocator, instance);
1548 defer allocator.free(target);
1549 const entry_name = try radixSplitFlagsFamilyEntryName(allocator, instance);
1550 defer allocator.free(entry_name);
1551 const family_fingerprint = options.shape_family_fingerprint orelse try radixSplitFlagsFamilyFingerprint(allocator, instance);
1552 const shape_profile_dimensions = radixSplitShapeProfileDimensions(instance);
1553 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1554 .name = "radix_split_flags",
1555 .fingerprint = family_fingerprint,
1556 .dimensions = shape_profile_dimensions[0..],
1557 };
1558
1559 var graph = try RadixSplitFlagsRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1560 defer graph.deinit();
1561 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1562 .target = target,
1563 .version = radix_split_family_version,
1564 .format = options.format,
1565 .kernel_plan = options.kernel_plan,
1566 .element_count_argument = options.element_count_argument,
1567 .shape_family_fingerprint = family_fingerprint,
1568 .shape_profile = shape_profile,
1569 .launch = options.launch orelse try radixSplitLaunch(instance),
1570 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1571 .static_arguments = options.static_arguments,
1572 });
1573 }
1574
1575 pub fn createRadixSplitScatterFamilyArtifact(
1576 allocator: std.mem.Allocator,
1577 handle: kernel.BackendHandle,
1578 instance: RadixSplit,
1579 options: entry.ArtifactOptions,
1580 ) !kernel.OwnedKernelCallArtifact {
1581 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1582 const target = try radixSplitScatterFamilyTarget(allocator, instance);
1583 defer allocator.free(target);
1584 const entry_name = try radixSplitScatterFamilyEntryName(allocator, instance);
1585 defer allocator.free(entry_name);
1586 const family_fingerprint = options.shape_family_fingerprint orelse try radixSplitScatterFamilyFingerprint(allocator, instance);
1587 const shape_profile_dimensions = radixSplitShapeProfileDimensions(instance);
1588 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1589 .name = "radix_split_scatter",
1590 .fingerprint = family_fingerprint,
1591 .dimensions = shape_profile_dimensions[0..],
1592 };
1593
1594 var graph = try RadixSplitScatterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1595 defer graph.deinit();
1596 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1597 .target = target,
1598 .version = radix_split_family_version,
1599 .format = options.format,
1600 .kernel_plan = options.kernel_plan,
1601 .element_count_argument = options.element_count_argument,
1602 .shape_family_fingerprint = family_fingerprint,
1603 .shape_profile = shape_profile,
1604 .launch = options.launch orelse try radixSplitLaunch(instance),
1605 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 1 else options.runtime_scalar_argument_count,
1606 .static_arguments = options.static_arguments,
1607 });
1608 }
1609
1610 pub const radix_digit_scan_threads: u32 = 1024;
1611
1612 pub fn createRadixDigitHistogramFamilyArtifact(
1613 allocator: std.mem.Allocator,
1614 handle: kernel.BackendHandle,
1615 instance: RadixSplit,
1616 options: entry.ArtifactOptions,
1617 ) !kernel.OwnedKernelCallArtifact {
1618 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1619 const target = try radixDigitHistogramFamilyTarget(allocator, instance);
1620 defer allocator.free(target);
1621 const entry_name = try radixDigitHistogramFamilyEntryName(allocator, instance);
1622 defer allocator.free(entry_name);
1623 const family_fingerprint = options.shape_family_fingerprint orelse try radixSplitFlagsFamilyFingerprint(allocator, instance);
1624 const shape_profile_dimensions = radixSplitShapeProfileDimensions(instance);
1625 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1626 .name = "radix_digit_histogram",
1627 .fingerprint = family_fingerprint,
1628 .dimensions = shape_profile_dimensions[0..],
1629 };
1630
1631 var graph = try RadixDigitHistogramRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1632 defer graph.deinit();
1633 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1634 .target = target,
1635 .version = radix_split_family_version,
1636 .format = options.format,
1637 .kernel_plan = options.kernel_plan,
1638 .element_count_argument = options.element_count_argument,
1639 .shape_family_fingerprint = family_fingerprint,
1640 .shape_profile = shape_profile,
1641 .launch = options.launch orelse try radixSplitLaunch(instance),
1642 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1643 .static_arguments = options.static_arguments,
1644 });
1645 }
1646
1647 pub fn createRadixDigitRankScatterFamilyArtifact(
1648 allocator: std.mem.Allocator,
1649 handle: kernel.BackendHandle,
1650 instance: RadixSplit,
1651 options: entry.ArtifactOptions,
1652 ) !kernel.OwnedKernelCallArtifact {
1653 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1654 const target = try radixDigitRankScatterFamilyTarget(allocator, instance);
1655 defer allocator.free(target);
1656 const entry_name = try radixDigitRankScatterFamilyEntryName(allocator, instance);
1657 defer allocator.free(entry_name);
1658 const family_fingerprint = options.shape_family_fingerprint orelse try radixSplitScatterFamilyFingerprint(allocator, instance);
1659 const shape_profile_dimensions = radixSplitShapeProfileDimensions(instance);
1660 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1661 .name = "radix_digit_rank_scatter",
1662 .fingerprint = family_fingerprint,
1663 .dimensions = shape_profile_dimensions[0..],
1664 };
1665
1666 var graph = try RadixDigitRankScatterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1667 defer graph.deinit();
1668 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1669 .target = target,
1670 .version = radix_split_family_version,
1671 .format = options.format,
1672 .kernel_plan = options.kernel_plan,
1673 .element_count_argument = options.element_count_argument,
1674 .shape_family_fingerprint = family_fingerprint,
1675 .shape_profile = shape_profile,
1676 .launch = options.launch orelse try radixSplitLaunch(instance),
1677 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1678 .static_arguments = options.static_arguments,
1679 });
1680 }
1681
1682 fn radixDigitCountsScan(instance: RadixSplit) scan.DeviceScan {
1683 const blocks = radixSplitBlockCount(instance.extent, instance.threads);
1684 return .{
1685 .extent = radix_digit_bins * blocks,
1686 .dtype = .f32,
1687 .mode = .exclusive,
1688 .threads = radix_digit_scan_threads,
1689 .element_axis = instance.element_axis,
1690 };
1691 }
1692
1693 pub const RadixDigitPipelineArtifacts = struct {
1694 histogram: kernel.OwnedKernelCallArtifact,
1695 block_scan: kernel.OwnedKernelCallArtifact,
1696 sums_scan: kernel.OwnedKernelCallArtifact,
1697 add_base: kernel.OwnedKernelCallArtifact,
1698 rank_scatter: kernel.OwnedKernelCallArtifact,
1699
1700 pub fn entries(self: *const RadixDigitPipelineArtifacts) [5]artifact_product.KernelCallArtifact {
1701 return .{
1702 self.histogram.entry(),
1703 self.block_scan.entry(),
1704 self.sums_scan.entry(),
1705 self.add_base.entry(),
1706 self.rank_scatter.entry(),
1707 };
1708 }
1709
1710 pub fn deinit(self: *RadixDigitPipelineArtifacts) void {
1711 self.histogram.deinit();
1712 self.block_scan.deinit();
1713 self.sums_scan.deinit();
1714 self.add_base.deinit();
1715 self.rank_scatter.deinit();
1716 self.* = undefined;
1717 }
1718 };
1719
1720 pub fn createRadixDigitPipelineArtifacts(
1721 allocator: std.mem.Allocator,
1722 handle: kernel.BackendHandle,
1723 instance: RadixSplit,
1724 options: entry.ArtifactOptions,
1725 ) !RadixDigitPipelineArtifacts {
1726 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1727 const counts_scan = radixDigitCountsScan(instance);
1728 const stages = try scan.deviceScanStages(counts_scan);
1729
1730 var histogram = try createRadixDigitHistogramFamilyArtifact(allocator, handle, instance, options);
1731 errdefer histogram.deinit();
1732 var block_scan = try scan.createDeviceScanBlockScanFamilyArtifact(allocator, handle, counts_scan, options);
1733 errdefer block_scan.deinit();
1734 var sums_scan = try scan.createPrefixSumFamilyArtifact(allocator, handle, stages.sums_scan, options);
1735 errdefer sums_scan.deinit();
1736 var add_base = try scan.createDeviceScanAddBaseFamilyArtifact(allocator, handle, counts_scan, options);
1737 errdefer add_base.deinit();
1738 const rank_scatter = try createRadixDigitRankScatterFamilyArtifact(allocator, handle, instance, options);
1739 return .{
1740 .histogram = histogram,
1741 .block_scan = block_scan,
1742 .sums_scan = sums_scan,
1743 .add_base = add_base,
1744 .rank_scatter = rank_scatter,
1745 };
1746 }
1747
1748 pub fn radixDigitPipelineTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
1749 return std.fmt.allocPrint(
1750 allocator,
1751 "accy.kernel.sort.radix_digit_family_{d}_i32",
1752 .{instance.threads},
1753 );
1754 }
1755
1756 pub fn radixDigitPipeline(
1757 backing_allocator: std.mem.Allocator,
1758 instance: RadixSplit,
1759 ) !artifact_product.OwnedKernelCallPipeline {
1760 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1761 const counts_scan = radixDigitCountsScan(instance);
1762 const stages = try scan.deviceScanStages(counts_scan);
1763 var owned = artifact_product.OwnedKernelCallPipeline.init(backing_allocator);
1764 errdefer owned.deinit();
1765 const arena = owned.allocator();
1766
1767 const cells_extent = artifact_product.PipelineScalarDerivation{
1768 .ceil_div_scaled = .{ .argument_index = 0, .divisor = instance.threads, .scale = radix_digit_bins },
1769 };
1770 const scan_blocks_extent = artifact_product.PipelineScalarDerivation{
1771 .ceil_div = .{ .argument_index = 0, .divisor = 64 * instance.threads },
1772 };
1773
1774 const intermediates = try arena.alloc(artifact_product.PipelineIntermediate, 4);
1775 intermediates[0] = .{ .dtype = .f32, .extent = cells_extent };
1776 intermediates[1] = .{ .dtype = .f32, .extent = scan_blocks_extent };
1777 intermediates[2] = .{ .dtype = .f32, .extent = scan_blocks_extent };
1778 intermediates[3] = .{ .dtype = .f32, .extent = cells_extent };
1779
1780 const pipeline_stages = try arena.alloc(artifact_product.PipelineStage, 5);
1781 pipeline_stages[0] = .{
1782 .target = try radixDigitHistogramFamilyTarget(arena, instance),
1783 .version = radix_split_family_version,
1784 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1785 .{ .intermediate = 0 }, .{ .operand = 0 },
1786 }),
1787 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1788 .{ .forward = 0 }, .{ .forward = 1 }, .{ .forward = 2 },
1789 }),
1790 };
1791 pipeline_stages[1] = .{
1792 .target = try scan.deviceScanBlockScanFamilyTarget(arena, counts_scan),
1793 .version = scan.device_scan_family_version,
1794 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1795 .{ .intermediate = 3 }, .{ .intermediate = 0 }, .{ .intermediate = 1 },
1796 }),
1797 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1798 cells_extent,
1799 }),
1800 };
1801 pipeline_stages[2] = .{
1802 .target = try scan.prefixSumFamilyTarget(arena, stages.sums_scan),
1803 .version = scan.prefix_sum_family_version,
1804 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1805 .{ .intermediate = 2 }, .{ .intermediate = 1 },
1806 }),
1807 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1808 scan_blocks_extent,
1809 }),
1810 };
1811 pipeline_stages[3] = .{
1812 .target = try scan.deviceScanAddBaseFamilyTarget(arena, counts_scan),
1813 .version = scan.device_scan_family_version,
1814 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1815 .{ .intermediate = 3 }, .{ .intermediate = 2 },
1816 }),
1817 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1818 cells_extent,
1819 }),
1820 };
1821 pipeline_stages[4] = .{
1822 .target = try radixDigitRankScatterFamilyTarget(arena, instance),
1823 .version = radix_split_family_version,
1824 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1825 .{ .result = 0 }, .{ .operand = 0 }, .{ .intermediate = 3 },
1826 }),
1827 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1828 .{ .forward = 0 }, .{ .forward = 1 }, .{ .forward = 2 },
1829 }),
1830 };
1831
1832 owned.value = .{
1833 .target = try radixDigitPipelineTarget(arena, instance),
1834 .version = radix_split_family_version,
1835 .operand_count = 1,
1836 .result_count = 1,
1837 .runtime_scalar_argument_count = 3,
1838 .intermediates = intermediates,
1839 .stages = pipeline_stages,
1840 };
1841 return owned;
1842 }
1843
1844 pub fn createRadixDigitRankScatterPairsFamilyArtifact(
1845 allocator: std.mem.Allocator,
1846 handle: kernel.BackendHandle,
1847 instance: RadixSplit,
1848 options: entry.ArtifactOptions,
1849 ) !kernel.OwnedKernelCallArtifact {
1850 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1851 const target = try radixDigitRankScatterPairsFamilyTarget(allocator, instance);
1852 defer allocator.free(target);
1853 const entry_name = try radixDigitRankScatterPairsFamilyEntryName(allocator, instance);
1854 defer allocator.free(entry_name);
1855 const family_fingerprint = options.shape_family_fingerprint orelse try radixSplitScatterFamilyFingerprint(allocator, instance);
1856 const shape_profile_dimensions = radixSplitShapeProfileDimensions(instance);
1857 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1858 .name = "radix_digit_rank_scatter_pairs",
1859 .fingerprint = family_fingerprint,
1860 .dimensions = shape_profile_dimensions[0..],
1861 };
1862
1863 var graph = try RadixDigitRankScatterPairsRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance);
1864 defer graph.deinit();
1865 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1866 .target = target,
1867 .version = radix_split_family_version,
1868 .format = options.format,
1869 .kernel_plan = options.kernel_plan,
1870 .element_count_argument = options.element_count_argument,
1871 .shape_family_fingerprint = family_fingerprint,
1872 .shape_profile = shape_profile,
1873 .launch = options.launch orelse try radixSplitLaunch(instance),
1874 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 3 else options.runtime_scalar_argument_count,
1875 .static_arguments = options.static_arguments,
1876 });
1877 }
1878
1879 pub const RadixDigitPairsPipelineArtifacts = struct {
1880 histogram: kernel.OwnedKernelCallArtifact,
1881 block_scan: kernel.OwnedKernelCallArtifact,
1882 sums_scan: kernel.OwnedKernelCallArtifact,
1883 add_base: kernel.OwnedKernelCallArtifact,
1884 rank_scatter_pairs: kernel.OwnedKernelCallArtifact,
1885
1886 pub fn entries(self: *const RadixDigitPairsPipelineArtifacts) [5]artifact_product.KernelCallArtifact {
1887 return .{
1888 self.histogram.entry(),
1889 self.block_scan.entry(),
1890 self.sums_scan.entry(),
1891 self.add_base.entry(),
1892 self.rank_scatter_pairs.entry(),
1893 };
1894 }
1895
1896 pub fn deinit(self: *RadixDigitPairsPipelineArtifacts) void {
1897 self.histogram.deinit();
1898 self.block_scan.deinit();
1899 self.sums_scan.deinit();
1900 self.add_base.deinit();
1901 self.rank_scatter_pairs.deinit();
1902 self.* = undefined;
1903 }
1904 };
1905
1906 pub fn createRadixDigitPairsPipelineArtifacts(
1907 allocator: std.mem.Allocator,
1908 handle: kernel.BackendHandle,
1909 instance: RadixSplit,
1910 options: entry.ArtifactOptions,
1911 ) !RadixDigitPairsPipelineArtifacts {
1912 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1913 const counts_scan = radixDigitCountsScan(instance);
1914 const stages = try scan.deviceScanStages(counts_scan);
1915
1916 var histogram = try createRadixDigitHistogramFamilyArtifact(allocator, handle, instance, options);
1917 errdefer histogram.deinit();
1918 var block_scan = try scan.createDeviceScanBlockScanFamilyArtifact(allocator, handle, counts_scan, options);
1919 errdefer block_scan.deinit();
1920 var sums_scan = try scan.createPrefixSumFamilyArtifact(allocator, handle, stages.sums_scan, options);
1921 errdefer sums_scan.deinit();
1922 var add_base = try scan.createDeviceScanAddBaseFamilyArtifact(allocator, handle, counts_scan, options);
1923 errdefer add_base.deinit();
1924 const rank_scatter_pairs = try createRadixDigitRankScatterPairsFamilyArtifact(allocator, handle, instance, options);
1925 return .{
1926 .histogram = histogram,
1927 .block_scan = block_scan,
1928 .sums_scan = sums_scan,
1929 .add_base = add_base,
1930 .rank_scatter_pairs = rank_scatter_pairs,
1931 };
1932 }
1933
1934 pub fn radixDigitPairsPipelineTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
1935 return std.fmt.allocPrint(
1936 allocator,
1937 "accy.kernel.sort.radix_digit_pairs_family_{d}_i32",
1938 .{instance.threads},
1939 );
1940 }
1941
1942 pub fn radixDigitPairsPipeline(
1943 backing_allocator: std.mem.Allocator,
1944 instance: RadixSplit,
1945 ) !artifact_product.OwnedKernelCallPipeline {
1946 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
1947 const counts_scan = radixDigitCountsScan(instance);
1948 const stages = try scan.deviceScanStages(counts_scan);
1949 var owned = artifact_product.OwnedKernelCallPipeline.init(backing_allocator);
1950 errdefer owned.deinit();
1951 const arena = owned.allocator();
1952
1953 const cells_extent = artifact_product.PipelineScalarDerivation{
1954 .ceil_div_scaled = .{ .argument_index = 0, .divisor = instance.threads, .scale = radix_digit_bins },
1955 };
1956 const scan_blocks_extent = artifact_product.PipelineScalarDerivation{
1957 .ceil_div = .{ .argument_index = 0, .divisor = 64 * instance.threads },
1958 };
1959
1960 const intermediates = try arena.alloc(artifact_product.PipelineIntermediate, 4);
1961 intermediates[0] = .{ .dtype = .f32, .extent = cells_extent };
1962 intermediates[1] = .{ .dtype = .f32, .extent = scan_blocks_extent };
1963 intermediates[2] = .{ .dtype = .f32, .extent = scan_blocks_extent };
1964 intermediates[3] = .{ .dtype = .f32, .extent = cells_extent };
1965
1966 const pipeline_stages = try arena.alloc(artifact_product.PipelineStage, 5);
1967 pipeline_stages[0] = .{
1968 .target = try radixDigitHistogramFamilyTarget(arena, instance),
1969 .version = radix_split_family_version,
1970 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1971 .{ .intermediate = 0 }, .{ .operand = 0 },
1972 }),
1973 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1974 .{ .forward = 0 }, .{ .forward = 1 }, .{ .forward = 2 },
1975 }),
1976 };
1977 pipeline_stages[1] = .{
1978 .target = try scan.deviceScanBlockScanFamilyTarget(arena, counts_scan),
1979 .version = scan.device_scan_family_version,
1980 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1981 .{ .intermediate = 3 }, .{ .intermediate = 0 }, .{ .intermediate = 1 },
1982 }),
1983 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1984 cells_extent,
1985 }),
1986 };
1987 pipeline_stages[2] = .{
1988 .target = try scan.prefixSumFamilyTarget(arena, stages.sums_scan),
1989 .version = scan.prefix_sum_family_version,
1990 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
1991 .{ .intermediate = 2 }, .{ .intermediate = 1 },
1992 }),
1993 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
1994 scan_blocks_extent,
1995 }),
1996 };
1997 pipeline_stages[3] = .{
1998 .target = try scan.deviceScanAddBaseFamilyTarget(arena, counts_scan),
1999 .version = scan.device_scan_family_version,
2000 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2001 .{ .intermediate = 3 }, .{ .intermediate = 2 },
2002 }),
2003 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2004 cells_extent,
2005 }),
2006 };
2007 pipeline_stages[4] = .{
2008 .target = try radixDigitRankScatterPairsFamilyTarget(arena, instance),
2009 .version = radix_split_family_version,
2010 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2011 .{ .result = 0 }, .{ .result = 1 }, .{ .operand = 0 }, .{ .operand = 1 }, .{ .intermediate = 3 },
2012 }),
2013 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2014 .{ .forward = 0 }, .{ .forward = 1 }, .{ .forward = 2 },
2015 }),
2016 };
2017
2018 owned.value = .{
2019 .target = try radixDigitPairsPipelineTarget(arena, instance),
2020 .version = radix_split_family_version,
2021 .operand_count = 2,
2022 .result_count = 2,
2023 .runtime_scalar_argument_count = 3,
2024 .intermediates = intermediates,
2025 .stages = pipeline_stages,
2026 };
2027 return owned;
2028 }
2029
2030 fn radixSplitDeviceScan(instance: RadixSplit) scan.DeviceScan {
2031 return .{
2032 .extent = instance.extent,
2033 .dtype = .f32,
2034 .mode = .exclusive,
2035 .threads = instance.threads,
2036 .element_axis = instance.element_axis,
2037 };
2038 }
2039
2040 pub const RadixSplitPipelineArtifacts = struct {
2041 flags: kernel.OwnedKernelCallArtifact,
2042 block_scan: kernel.OwnedKernelCallArtifact,
2043 sums_scan: kernel.OwnedKernelCallArtifact,
2044 add_base: kernel.OwnedKernelCallArtifact,
2045 scatter: kernel.OwnedKernelCallArtifact,
2046
2047 pub fn entries(self: *const RadixSplitPipelineArtifacts) [5]artifact_product.KernelCallArtifact {
2048 return .{
2049 self.flags.entry(),
2050 self.block_scan.entry(),
2051 self.sums_scan.entry(),
2052 self.add_base.entry(),
2053 self.scatter.entry(),
2054 };
2055 }
2056
2057 pub fn deinit(self: *RadixSplitPipelineArtifacts) void {
2058 self.flags.deinit();
2059 self.block_scan.deinit();
2060 self.sums_scan.deinit();
2061 self.add_base.deinit();
2062 self.scatter.deinit();
2063 self.* = undefined;
2064 }
2065 };
2066
2067 pub fn createRadixSplitPipelineArtifacts(
2068 allocator: std.mem.Allocator,
2069 handle: kernel.BackendHandle,
2070 instance: RadixSplit,
2071 options: entry.ArtifactOptions,
2072 ) !RadixSplitPipelineArtifacts {
2073 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
2074 const device_scan = radixSplitDeviceScan(instance);
2075 const stages = try scan.deviceScanStages(device_scan);
2076
2077 var flags = try createRadixSplitFlagsFamilyArtifact(allocator, handle, instance, options);
2078 errdefer flags.deinit();
2079 var block_scan = try scan.createDeviceScanBlockScanFamilyArtifact(allocator, handle, device_scan, options);
2080 errdefer block_scan.deinit();
2081 var sums_scan = try scan.createPrefixSumFamilyArtifact(allocator, handle, stages.sums_scan, options);
2082 errdefer sums_scan.deinit();
2083 var add_base = try scan.createDeviceScanAddBaseFamilyArtifact(allocator, handle, device_scan, options);
2084 errdefer add_base.deinit();
2085 const scatter = try createRadixSplitScatterFamilyArtifact(allocator, handle, instance, options);
2086 return .{
2087 .flags = flags,
2088 .block_scan = block_scan,
2089 .sums_scan = sums_scan,
2090 .add_base = add_base,
2091 .scatter = scatter,
2092 };
2093 }
2094
2095 pub fn radixSplitPipelineTarget(allocator: std.mem.Allocator, instance: RadixSplit) ![]u8 {
2096 return std.fmt.allocPrint(
2097 allocator,
2098 "accy.kernel.sort.radix_split_family_{d}_i32",
2099 .{instance.threads},
2100 );
2101 }
2102
2103 pub fn radixSplitPipeline(
2104 backing_allocator: std.mem.Allocator,
2105 instance: RadixSplit,
2106 ) !artifact_product.OwnedKernelCallPipeline {
2107 if (!radixSplitInstanceValid(instance)) return error.UnsupportedRadixSplitInstance;
2108 const device_scan = radixSplitDeviceScan(instance);
2109 const stages = try scan.deviceScanStages(device_scan);
2110 var owned = artifact_product.OwnedKernelCallPipeline.init(backing_allocator);
2111 errdefer owned.deinit();
2112 const arena = owned.allocator();
2113
2114 const block_count_extent = artifact_product.PipelineScalarDerivation{
2115 .ceil_div = .{ .argument_index = 0, .divisor = instance.threads },
2116 };
2117 const element_extent = artifact_product.PipelineScalarDerivation{ .forward = 0 };
2118
2119 const intermediates = try arena.alloc(artifact_product.PipelineIntermediate, 4);
2120 intermediates[0] = .{ .dtype = .f32, .extent = element_extent };
2121 intermediates[1] = .{ .dtype = .f32, .extent = block_count_extent };
2122 intermediates[2] = .{ .dtype = .f32, .extent = block_count_extent };
2123 intermediates[3] = .{ .dtype = .f32, .extent = element_extent };
2124
2125 const pipeline_stages = try arena.alloc(artifact_product.PipelineStage, 5);
2126 pipeline_stages[0] = .{
2127 .target = try radixSplitFlagsFamilyTarget(arena, instance),
2128 .version = radix_split_family_version,
2129 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2130 .{ .intermediate = 0 }, .{ .operand = 0 },
2131 }),
2132 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2133 .{ .forward = 0 }, .{ .forward = 1 }, .{ .forward = 2 },
2134 }),
2135 };
2136 pipeline_stages[1] = .{
2137 .target = try scan.deviceScanBlockScanFamilyTarget(arena, device_scan),
2138 .version = scan.device_scan_family_version,
2139 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2140 .{ .intermediate = 3 }, .{ .intermediate = 0 }, .{ .intermediate = 1 },
2141 }),
2142 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2143 .{ .forward = 0 },
2144 }),
2145 };
2146 pipeline_stages[2] = .{
2147 .target = try scan.prefixSumFamilyTarget(arena, stages.sums_scan),
2148 .version = scan.prefix_sum_family_version,
2149 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2150 .{ .intermediate = 2 }, .{ .intermediate = 1 },
2151 }),
2152 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2153 block_count_extent,
2154 }),
2155 };
2156 pipeline_stages[3] = .{
2157 .target = try scan.deviceScanAddBaseFamilyTarget(arena, device_scan),
2158 .version = scan.device_scan_family_version,
2159 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2160 .{ .intermediate = 3 }, .{ .intermediate = 2 },
2161 }),
2162 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2163 .{ .forward = 0 },
2164 }),
2165 };
2166 pipeline_stages[4] = .{
2167 .target = try radixSplitScatterFamilyTarget(arena, instance),
2168 .version = radix_split_family_version,
2169 .buffers = try arena.dupe(artifact_product.PipelineValueRef, &.{
2170 .{ .result = 0 }, .{ .operand = 0 }, .{ .intermediate = 0 }, .{ .intermediate = 3 },
2171 }),
2172 .scalars = try arena.dupe(artifact_product.PipelineScalarDerivation, &.{
2173 .{ .forward = 0 },
2174 }),
2175 };
2176
2177 owned.value = .{
2178 .target = try radixSplitPipelineTarget(arena, instance),
2179 .version = radix_split_family_version,
2180 .operand_count = 1,
2181 .result_count = 1,
2182 .runtime_scalar_argument_count = 3,
2183 .intermediates = intermediates,
2184 .stages = pipeline_stages,
2185 };
2186 return owned;
2187 }
2188
2189 pub fn radixSplitPipelineRuntimeArguments(instance: RadixSplit, bit: u32) ![3]choir_abi.ScalarArgument {
2190 return radixSplitFlagsRuntimeArguments(instance, bit);
2191 }
2192
2193 const testing = std.testing;
2194
2195 fn runRadixSplitPassOnOracle(
2196 allocator: std.mem.Allocator,
2197 instance: RadixSplit,
2198 bit: u32,
2199 keys: []i32,
2200 dst: []i32,
2201 ) !void {
2202 var flags_graph = try RadixSplitFlagsRuntimeFamilyI32.build(allocator, RadixSplitFlagsRuntimeFamilyI32.Limits.testing, instance);
2203 defer flags_graph.deinit();
2204 var scatter_graph = try RadixSplitScatterRuntimeFamilyI32.build(allocator, RadixSplitScatterRuntimeFamilyI32.Limits.testing, instance);
2205 defer scatter_graph.deinit();
2206 try runRadixSplitPassOnOracleWithGraphs(allocator, instance, bit, keys, dst, &flags_graph, &scatter_graph);
2207 }
2208
2209 fn runRadixSplitPassOnOracleWithGraphs(
2210 allocator: std.mem.Allocator,
2211 instance: RadixSplit,
2212 bit: u32,
2213 keys: []i32,
2214 dst: []i32,
2215 flags_graph: anytype,
2216 scatter_graph: anytype,
2217 ) !void {
2218 const extent = keys.len;
2219 const flags = try allocator.alloc(f32, extent);
2220 defer allocator.free(flags);
2221 @memset(flags, -1);
2222
2223 const blocks: u32 = @intCast(radixSplitBlockCount(instance.extent, instance.threads));
2224 try flags_graph.runCpuWithLaunch(allocator, &.{
2225 kernel.argumentBuffer(f32, flags),
2226 kernel.argumentBuffer(i32, keys),
2227 kernel.argumentI32(@intCast(extent)),
2228 kernel.argumentI32(@intCast(bit)),
2229 kernel.argumentI32(@intCast(radixSplitSignedPassPolarity(bit))),
2230 }, .{
2231 .grid = .{ blocks, 1, 1 },
2232 .block = .{ instance.threads, 1, 1 },
2233 });
2234
2235 const scanned = try allocator.alloc(f32, extent);
2236 defer allocator.free(scanned);
2237 var running: f32 = 0;
2238 for (flags, scanned) |flag, *value| {
2239 value.* = running;
2240 running += flag;
2241 }
2242
2243 try scatter_graph.runCpuWithLaunch(allocator, &.{
2244 kernel.argumentBuffer(i32, dst),
2245 kernel.argumentBuffer(i32, keys),
2246 kernel.argumentBuffer(f32, flags),
2247 kernel.argumentBuffer(f32, scanned),
2248 kernel.argumentI32(@intCast(extent)),
2249 }, .{
2250 .grid = .{ blocks, 1, 1 },
2251 .block = .{ instance.threads, 1, 1 },
2252 });
2253 }
2254
2255 fn expectStableSplit(keys: []const i32, dst: []const i32, bit: u32) !void {
2256 var expected = try testing.allocator.alloc(i32, keys.len);
2257 defer testing.allocator.free(expected);
2258 var count: usize = 0;
2259 const shift: u5 = @intCast(bit);
2260 for (keys) |key| {
2261 if ((key >> shift) & 1 == 0) {
2262 expected[count] = key;
2263 count += 1;
2264 }
2265 }
2266 for (keys) |key| {
2267 if ((key >> shift) & 1 == 1) {
2268 expected[count] = key;
2269 count += 1;
2270 }
2271 }
2272 try testing.expectEqualSlices(i32, expected, dst);
2273 }
2274
2275 test "sort radix split flags discriminate runtime bits on one compiled kernel" {
2276 const allocator = testing.allocator;
2277 const instance = RadixSplit{ .extent = 8, .threads = 32 };
2278 var keys = [_]i32{ 0, 1, 2, 3, 4, 5, 6, 7 };
2279 const flags = try allocator.alloc(f32, keys.len);
2280 defer allocator.free(flags);
2281
2282 var graph = try RadixSplitFlagsRuntimeFamilyI32.build(allocator, RadixSplitFlagsRuntimeFamilyI32.Limits.testing, instance);
2283 defer graph.deinit();
2284
2285 try graph.runCpuWithLaunch(allocator, &.{
2286 kernel.argumentBuffer(f32, flags),
2287 kernel.argumentBuffer(i32, keys[0..]),
2288 kernel.argumentI32(8),
2289 kernel.argumentI32(0),
2290 kernel.argumentI32(0),
2291 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 32, 1, 1 } });
2292 try testing.expectEqualSlices(f32, &.{ 1, 0, 1, 0, 1, 0, 1, 0 }, flags);
2293
2294 try graph.runCpuWithLaunch(allocator, &.{
2295 kernel.argumentBuffer(f32, flags),
2296 kernel.argumentBuffer(i32, keys[0..]),
2297 kernel.argumentI32(8),
2298 kernel.argumentI32(2),
2299 kernel.argumentI32(0),
2300 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 32, 1, 1 } });
2301 try testing.expectEqualSlices(f32, &.{ 1, 1, 1, 1, 0, 0, 0, 0 }, flags);
2302
2303 try graph.runCpuWithLaunch(allocator, &.{
2304 kernel.argumentBuffer(f32, flags),
2305 kernel.argumentBuffer(i32, keys[0..]),
2306 kernel.argumentI32(8),
2307 kernel.argumentI32(2),
2308 kernel.argumentI32(1),
2309 }, .{ .grid = .{ 1, 1, 1 }, .block = .{ 32, 1, 1 } });
2310 try testing.expectEqualSlices(f32, &.{ 0, 0, 0, 0, 1, 1, 1, 1 }, flags);
2311 }
2312
2313 test "sort radix split pass partitions stably across blocks" {
2314 const allocator = testing.allocator;
2315 const extent: usize = 90;
2316 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2317
2318 var keys: [extent]i32 = undefined;
2319 var seed: u32 = 0x2545f491;
2320 for (&keys, 0..) |*key, index| {
2321 seed ^= seed << 13;
2322 seed ^= seed >> 17;
2323 seed ^= seed << 5;
2324 key.* = @intCast((seed >> 8) % 1000 * 10 + index % 10);
2325 }
2326
2327 var dst = @as([extent]i32, @splat(-1));
2328 var flags_graph = try RadixSplitFlagsRuntimeFamilyI32.build(allocator, RadixSplitFlagsRuntimeFamilyI32.Limits.testing, instance);
2329 defer flags_graph.deinit();
2330 var scatter_graph = try RadixSplitScatterRuntimeFamilyI32.build(allocator, RadixSplitScatterRuntimeFamilyI32.Limits.testing, instance);
2331 defer scatter_graph.deinit();
2332
2333 try runRadixSplitPassOnOracleWithGraphs(allocator, instance, 0, keys[0..], dst[0..], &flags_graph, &scatter_graph);
2334 try expectStableSplit(keys[0..], dst[0..], 0);
2335
2336 var dst_bit5 = @as([extent]i32, @splat(-1));
2337 try runRadixSplitPassOnOracleWithGraphs(allocator, instance, 5, keys[0..], dst_bit5[0..], &flags_graph, &scatter_graph);
2338 try expectStableSplit(keys[0..], dst_bit5[0..], 5);
2339 }
2340
2341 test "sort radix split passes compose into a full sort on the oracle" {
2342 const allocator = testing.allocator;
2343 const extent: usize = 70;
2344 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2345
2346 var keys: [extent]i32 = undefined;
2347 var seed: u32 = 0x9e3779b9;
2348 for (&keys) |*key| {
2349 seed ^= seed << 13;
2350 seed ^= seed >> 17;
2351 seed ^= seed << 5;
2352 key.* = @intCast(seed % 100000);
2353 }
2354
2355 var current = keys;
2356 var scratch = @as([extent]i32, @splat(-1));
2357 var flags_graph = try RadixSplitFlagsRuntimeFamilyI32.build(allocator, RadixSplitFlagsRuntimeFamilyI32.Limits.testing, instance);
2358 defer flags_graph.deinit();
2359 var scatter_graph = try RadixSplitScatterRuntimeFamilyI32.build(allocator, RadixSplitScatterRuntimeFamilyI32.Limits.testing, instance);
2360 defer scatter_graph.deinit();
2361 var bit: u32 = 0;
2362 while (bit < 17) : (bit += 1) {
2363 try runRadixSplitPassOnOracleWithGraphs(allocator, instance, bit, current[0..], scratch[0..], &flags_graph, &scatter_graph);
2364 current = scratch;
2365 }
2366
2367 var expected = keys;
2368 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2369 try testing.expectEqualSlices(i32, expected[0..], current[0..]);
2370 }
2371
2372 test "sort radix split identity and validity" {
2373 const instance = RadixSplit{ .extent = 5000, .threads = 64 };
2374 const flags_target = try radixSplitFlagsFamilyTarget(testing.allocator, instance);
2375 defer testing.allocator.free(flags_target);
2376 try testing.expectEqualStrings("accy.kernel.sort.radix_split_flags_family_64_i32", flags_target);
2377 const scatter_target = try radixSplitScatterFamilyTarget(testing.allocator, instance);
2378 defer testing.allocator.free(scatter_target);
2379 try testing.expectEqualStrings("accy.kernel.sort.radix_split_scatter_family_64_i32", scatter_target);
2380
2381 try testing.expect(radixSplitInstanceValid(.{ .extent = 1024 * 1024, .threads = 1024 }));
2382 try testing.expect(!radixSplitInstanceValid(.{ .extent = 1024 * 1024 + 1, .threads = 1024 }));
2383 try testing.expect(!radixSplitInstanceValid(.{ .extent = 0, .threads = 32 }));
2384 try testing.expect(!radixSplitInstanceValid(.{ .extent = 100, .threads = 48 }));
2385 try testing.expect(!radixSplitInstanceValid(.{ .extent = std.math.maxInt(u64), .threads = 32 }));
2386
2387 const args = try radixSplitFlagsRuntimeArguments(instance, 31);
2388 try testing.expectEqual(@as(u32, 5000), args[0].u32);
2389 try testing.expectEqual(@as(u32, 31), args[1].u32);
2390 try testing.expectEqual(@as(u32, 1), args[2].u32);
2391 const low_bit_args = try radixSplitFlagsRuntimeArguments(instance, 7);
2392 try testing.expectEqual(@as(u32, 0), low_bit_args[2].u32);
2393 try testing.expectError(error.UnsupportedRadixSplitInstance, radixSplitFlagsRuntimeArguments(instance, 32));
2394 }
2395
2396 test "sort bitonic block identity and validity" {
2397 const instance = BitonicBlock{ .extent = 45, .threads = 64 };
2398 const target = try bitonicBlockFamilyTarget(testing.allocator, instance);
2399 defer testing.allocator.free(target);
2400 try testing.expectEqualStrings("accy.kernel.sort.bitonic_block_family_64_i32", target);
2401 const entry_name = try bitonicBlockFamilyEntryName(testing.allocator, instance);
2402 defer testing.allocator.free(entry_name);
2403 try testing.expectEqualStrings("accy_kernel_sort_bitonic_block_family_64_i32", entry_name);
2404
2405 try testing.expect(bitonicBlockInstanceValid(instance));
2406 try testing.expect(bitonicBlockInstanceValid(.{ .extent = 1024, .threads = 1024 }));
2407 try testing.expect(!bitonicBlockInstanceValid(.{ .extent = 0, .threads = 32 }));
2408 try testing.expect(!bitonicBlockInstanceValid(.{ .extent = 33, .threads = 48 }));
2409 try testing.expect(!bitonicBlockInstanceValid(.{ .extent = 65, .threads = 64 }));
2410 try testing.expect(!bitonicBlockInstanceValid(.{ .extent = 16, .threads = 16 }));
2411 try testing.expectEqual(@as(?u32, 32), bitonicBlockThreadsForExtent(1));
2412 try testing.expectEqual(@as(?u32, 32), bitonicBlockThreadsForExtent(32));
2413 try testing.expectEqual(@as(?u32, 64), bitonicBlockThreadsForExtent(33));
2414 try testing.expectEqual(@as(?u32, 1024), bitonicBlockThreadsForExtent(1024));
2415 try testing.expectEqual(@as(?u32, null), bitonicBlockThreadsForExtent(1025));
2416
2417 const args = try bitonicBlockRuntimeArguments(instance);
2418 try testing.expectEqual(@as(u32, 45), args[0].u32);
2419 }
2420
2421 test "sort bitonic block sorts a bounded tile on the oracle" {
2422 const allocator = testing.allocator;
2423 const extent: usize = 45;
2424 const instance = BitonicBlock{ .extent = extent, .threads = 64 };
2425
2426 var keys: [extent]i32 = undefined;
2427 for (&keys, 0..) |*key, index| {
2428 const raw: i32 = @intCast((index * 37 + 11) % 53);
2429 key.* = if (index % 3 == 0) -raw else raw - 19;
2430 }
2431 keys[7] = keys[4];
2432 keys[13] = std.math.maxInt(i32);
2433 keys[29] = std.math.minInt(i32);
2434
2435 var dst = @as([extent]i32, @splat(-7777));
2436 var graph = try BitonicBlockRuntimeFamilyI32.build(allocator, BitonicBlockRuntimeFamilyI32.Limits.testing, instance);
2437 defer graph.deinit();
2438 try graph.runCpuWithLaunch(allocator, &.{
2439 kernel.argumentBuffer(i32, dst[0..]),
2440 kernel.argumentBuffer(i32, keys[0..]),
2441 kernel.argumentI32(@intCast(extent)),
2442 }, .{
2443 .grid = .{ 1, 1, 1 },
2444 .block = .{ instance.threads, 1, 1 },
2445 });
2446
2447 var expected = keys;
2448 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2449 try testing.expectEqualSlices(i32, expected[0..], dst[0..]);
2450 }
2451
2452 test "sort bitonic block artifact records runtime family metadata" {
2453 const allocator = testing.allocator;
2454 var state = gpu.recording.BackendState{
2455 .allocator = allocator,
2456 .kind = .cuda,
2457 .format = .cuda_ptx,
2458 };
2459 const instance = BitonicBlock{ .extent = 45, .threads = 64 };
2460 var family_artifact = try createBitonicBlockFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
2461 defer family_artifact.deinit();
2462 const family_entry = family_artifact.entry();
2463 try testing.expectEqualStrings("accy.kernel.sort.bitonic_block_family_64_i32", family_entry.target);
2464 try testing.expectEqual(bitonic_block_family_version, family_entry.version);
2465 try testing.expectEqual(@as(u32, 1), family_entry.runtime_scalar_argument_count);
2466 try testing.expect(family_entry.shape_family_fingerprint != null);
2467
2468 var owned = try bitonicBlockFamilySpecialization(allocator, instance);
2469 defer owned.deinit();
2470 const recovered = bitonicBlockInstanceFromSpecialization(owned.value) orelse return error.TestExpectedBitonicBlockInstance;
2471 try testing.expectEqual(instance.extent, recovered.extent);
2472 try testing.expectEqual(instance.threads, recovered.threads);
2473 try testing.expect(owned.value.structureIs(bitonic_block_structure_name));
2474 }
2475
2476 test "sort top-k block identity and validity" {
2477 const instance = TopKBlock{ .extent = 45, .k = 8, .threads = 64 };
2478 const target = try topKBlockFamilyTarget(testing.allocator, instance);
2479 defer testing.allocator.free(target);
2480 try testing.expectEqualStrings("accy.kernel.sort.top_k_block_family_64x8_i32", target);
2481 const entry_name = try topKBlockFamilyEntryName(testing.allocator, instance);
2482 defer testing.allocator.free(entry_name);
2483 try testing.expectEqualStrings("accy_kernel_sort_top_k_block_family_64x8_i32", entry_name);
2484
2485 try testing.expect(topKBlockInstanceValid(instance));
2486 try testing.expect(topKBlockInstanceValid(.{ .extent = 1024, .k = 1024, .threads = 1024 }));
2487 try testing.expect(!topKBlockInstanceValid(.{ .extent = 0, .k = 1, .threads = 32 }));
2488 try testing.expect(!topKBlockInstanceValid(.{ .extent = 16, .k = 0, .threads = 32 }));
2489 try testing.expect(!topKBlockInstanceValid(.{ .extent = 16, .k = 17, .threads = 32 }));
2490 try testing.expect(!topKBlockInstanceValid(.{ .extent = 65, .k = 8, .threads = 64 }));
2491
2492 const args = try topKBlockRuntimeArguments(instance);
2493 try testing.expectEqual(@as(u32, 45), args[0].u32);
2494
2495 var owned = try topKBlockFamilySpecialization(testing.allocator, instance);
2496 defer owned.deinit();
2497 try testing.expect(owned.value.operationIs(.{ .sort = .top_k_smallest }));
2498 try testing.expect(owned.value.structureIs(top_k_block_structure_name));
2499 try testing.expect(owned.value.inputHasExtents(0, &.{45}));
2500 try testing.expect(owned.value.outputHasExtents(0, &.{8}));
2501 const recovered = topKBlockInstanceFromSpecialization(owned.value) orelse return error.TestExpectedTopKBlockInstance;
2502 try testing.expectEqual(@as(u64, 45), recovered.extent);
2503 try testing.expectEqual(@as(u64, 8), recovered.k);
2504 try testing.expectEqual(@as(u32, 64), recovered.threads);
2505
2506 var state = gpu.recording.BackendState{
2507 .allocator = testing.allocator,
2508 .kind = .cuda,
2509 .format = .cuda_ptx,
2510 };
2511 var artifact = try createTopKBlockFamilyArtifact(testing.allocator, state.handle(), instance, .{ .limits = .testing });
2512 defer artifact.deinit();
2513 const entry_value = artifact.entry();
2514 try testing.expectEqualStrings("accy.kernel.sort.top_k_block_family_64x8_i32", entry_value.target);
2515 try testing.expectEqual(top_k_block_family_version, entry_value.version);
2516 try testing.expectEqual(@as(u32, 1), entry_value.runtime_scalar_argument_count);
2517 }
2518
2519 test "sort top-k block selects the smallest sorted prefix on the oracle" {
2520 const allocator = testing.allocator;
2521 const extent: usize = 45;
2522 const top_count: usize = 8;
2523 const instance = TopKBlock{ .extent = extent, .k = top_count, .threads = 64 };
2524
2525 var keys: [extent]i32 = undefined;
2526 for (&keys, 0..) |*key, index| {
2527 const raw: i32 = @intCast((index * 41 + 5) % 67);
2528 key.* = if (index % 4 == 0) -raw else raw - 23;
2529 }
2530 keys[7] = keys[4];
2531 keys[13] = std.math.maxInt(i32);
2532 keys[29] = std.math.minInt(i32);
2533
2534 var dst = @as([top_count]i32, @splat(-7777));
2535 var graph = try TopKBlockRuntimeFamilyI32.build(allocator, TopKBlockRuntimeFamilyI32.Limits.testing, instance);
2536 defer graph.deinit();
2537 try graph.runCpuWithLaunch(allocator, &.{
2538 kernel.argumentBuffer(i32, dst[0..]),
2539 kernel.argumentBuffer(i32, keys[0..]),
2540 kernel.argumentI32(@intCast(extent)),
2541 }, .{
2542 .grid = .{ 1, 1, 1 },
2543 .block = .{ instance.threads, 1, 1 },
2544 });
2545
2546 var expected = keys;
2547 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2548 try testing.expectEqualSlices(i32, expected[0..top_count], dst[0..]);
2549 }
2550
2551 test "sort top-k block pairs identity and validity" {
2552 const instance = TopKBlockPairs{ .extent = 45, .k = 8, .threads = 64 };
2553 const target = try topKBlockPairsFamilyTarget(testing.allocator, instance);
2554 defer testing.allocator.free(target);
2555 try testing.expectEqualStrings("accy.kernel.sort.top_k_block_pairs_family_64x8_i32", target);
2556 const entry_name = try topKBlockPairsFamilyEntryName(testing.allocator, instance);
2557 defer testing.allocator.free(entry_name);
2558 try testing.expectEqualStrings("accy_kernel_sort_top_k_block_pairs_family_64x8_i32", entry_name);
2559
2560 try testing.expect(topKBlockPairsInstanceValid(instance));
2561 try testing.expect(topKBlockPairsInstanceValid(.{ .extent = 1024, .k = 1024, .threads = 1024 }));
2562 try testing.expect(!topKBlockPairsInstanceValid(.{ .extent = 0, .k = 1, .threads = 32 }));
2563 try testing.expect(!topKBlockPairsInstanceValid(.{ .extent = 16, .k = 0, .threads = 32 }));
2564 try testing.expect(!topKBlockPairsInstanceValid(.{ .extent = 16, .k = 17, .threads = 32 }));
2565 try testing.expect(!topKBlockPairsInstanceValid(.{ .extent = 65, .k = 8, .threads = 64 }));
2566
2567 const args = try topKBlockPairsRuntimeArguments(instance);
2568 try testing.expectEqual(@as(u32, 45), args[0].u32);
2569
2570 var owned = try topKBlockPairsFamilySpecialization(testing.allocator, instance);
2571 defer owned.deinit();
2572 try testing.expect(owned.value.operationIs(.{ .sort = .top_k_smallest }));
2573 try testing.expect(owned.value.structureIs(top_k_block_pairs_structure_name));
2574 try testing.expect(owned.value.inputHasExtents(0, &.{45}));
2575 try testing.expect(owned.value.inputHasExtents(1, &.{45}));
2576 try testing.expect(owned.value.outputHasExtents(0, &.{8}));
2577 try testing.expect(owned.value.outputHasExtents(1, &.{8}));
2578 const recovered = topKBlockPairsInstanceFromSpecialization(owned.value) orelse return error.TestExpectedTopKBlockPairsInstance;
2579 try testing.expectEqual(@as(u64, 45), recovered.extent);
2580 try testing.expectEqual(@as(u64, 8), recovered.k);
2581 try testing.expectEqual(@as(u32, 64), recovered.threads);
2582
2583 var state = gpu.recording.BackendState{
2584 .allocator = testing.allocator,
2585 .kind = .cuda,
2586 .format = .cuda_ptx,
2587 };
2588 var artifact = try createTopKBlockPairsFamilyArtifact(testing.allocator, state.handle(), instance, .{ .limits = .testing });
2589 defer artifact.deinit();
2590 const entry_value = artifact.entry();
2591 try testing.expectEqualStrings("accy.kernel.sort.top_k_block_pairs_family_64x8_i32", entry_value.target);
2592 try testing.expectEqual(top_k_block_pairs_family_version, entry_value.version);
2593 try testing.expectEqual(@as(u32, 1), entry_value.runtime_scalar_argument_count);
2594 }
2595
2596 const TopKOraclePair = struct {
2597 key: i32,
2598 value: i32,
2599
2600 fn lessThan(_: void, lhs: @This(), rhs: @This()) bool {
2601 return lhs.key < rhs.key or (lhs.key == rhs.key and lhs.value < rhs.value);
2602 }
2603 };
2604
2605 test "sort top-k block pairs selects key payload prefixes on the oracle" {
2606 const allocator = testing.allocator;
2607 const extent: usize = 45;
2608 const top_count: usize = 8;
2609 const instance = TopKBlockPairs{ .extent = extent, .k = top_count, .threads = 64 };
2610
2611 var keys: [extent]i32 = undefined;
2612 var values: [extent]i32 = undefined;
2613 for (&keys, &values, 0..) |*key, *value, index| {
2614 const raw: i32 = @intCast((index * 41 + 5) % 23);
2615 key.* = if (index % 4 == 0) -raw else raw - 11;
2616 value.* = @intCast(index);
2617 }
2618 keys[7] = keys[4];
2619 keys[13] = keys[4];
2620 keys[29] = std.math.minInt(i32);
2621
2622 var dst_keys = @as([top_count]i32, @splat(-7777));
2623 var dst_values = @as([top_count]i32, @splat(-7777));
2624 var graph = try TopKBlockPairsRuntimeFamilyI32.build(allocator, TopKBlockPairsRuntimeFamilyI32.Limits.testing, instance);
2625 defer graph.deinit();
2626 try graph.runCpuWithLaunch(allocator, &.{
2627 kernel.argumentBuffer(i32, dst_keys[0..]),
2628 kernel.argumentBuffer(i32, dst_values[0..]),
2629 kernel.argumentBuffer(i32, keys[0..]),
2630 kernel.argumentBuffer(i32, values[0..]),
2631 kernel.argumentI32(@intCast(extent)),
2632 }, .{
2633 .grid = .{ 1, 1, 1 },
2634 .block = .{ instance.threads, 1, 1 },
2635 });
2636
2637 var expected: [extent]TopKOraclePair = undefined;
2638 for (&expected, keys, values) |*pair, key, value| pair.* = .{ .key = key, .value = value };
2639 std.mem.sort(TopKOraclePair, expected[0..], {}, TopKOraclePair.lessThan);
2640 for (0..top_count) |index| {
2641 try testing.expectEqual(expected[index].key, dst_keys[index]);
2642 try testing.expectEqual(expected[index].value, dst_values[index]);
2643 }
2644 }
2645
2646 test "sort radix split pipeline descriptor binds the family artifacts" {
2647 const allocator = testing.allocator;
2648 var state = gpu.recording.BackendState{
2649 .allocator = allocator,
2650 .kind = .cuda,
2651 .format = .cuda_ptx,
2652 };
2653 const instance = RadixSplit{ .extent = 5000, .threads = 64 };
2654
2655 var artifacts = try createRadixSplitPipelineArtifacts(allocator, state.handle(), instance, .{ .limits = .testing });
2656 defer artifacts.deinit();
2657 const entries = artifacts.entries();
2658 const registry = artifact_product.KernelCallRegistry{ .entries = entries[0..] };
2659
2660 var owned = try radixSplitPipeline(allocator, instance);
2661 defer owned.deinit();
2662 try testing.expectEqualStrings("accy.kernel.sort.radix_split_family_64_i32", owned.value.target);
2663 try testing.expectEqual(@as(u32, 3), owned.value.runtime_scalar_argument_count);
2664 try testing.expectEqual(@as(usize, 4), owned.value.intermediates.len);
2665 try testing.expectEqual(@as(usize, 5), owned.value.stages.len);
2666 try owned.value.validate(registry, .cuda_ptx);
2667
2668 try testing.expectError(
2669 error.UnsupportedRadixSplitInstance,
2670 radixSplitPipeline(allocator, .{ .extent = 0, .threads = 64 }),
2671 );
2672 }
2673
2674 test "sort radix split passes sort signed keys on the oracle" {
2675 const allocator = testing.allocator;
2676 const extent: usize = 60;
2677 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2678
2679 var keys: [extent]i32 = undefined;
2680 var seed: u32 = 0xc0ffee11;
2681 for (&keys) |*key| {
2682 seed ^= seed << 13;
2683 seed ^= seed >> 17;
2684 seed ^= seed << 5;
2685 const magnitude: i32 = @intCast(seed % 50000);
2686 key.* = if (seed & 1 == 1) -magnitude else magnitude;
2687 }
2688
2689 var current = keys;
2690 var scratch = @as([extent]i32, @splat(-1));
2691 var flags_graph = try RadixSplitFlagsRuntimeFamilyI32.build(allocator, RadixSplitFlagsRuntimeFamilyI32.Limits.testing, instance);
2692 defer flags_graph.deinit();
2693 var scatter_graph = try RadixSplitScatterRuntimeFamilyI32.build(allocator, RadixSplitScatterRuntimeFamilyI32.Limits.testing, instance);
2694 defer scatter_graph.deinit();
2695 var bit: u32 = 0;
2696 while (bit < radix_split_key_bits) : (bit += 1) {
2697 try runRadixSplitPassOnOracleWithGraphs(allocator, instance, bit, current[0..], scratch[0..], &flags_graph, &scatter_graph);
2698 current = scratch;
2699 }
2700
2701 var expected = keys;
2702 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2703 try testing.expectEqualSlices(i32, expected[0..], current[0..]);
2704 }
2705
2706 test "sort radix digit histogram counts per block in column-major order" {
2707 const allocator = testing.allocator;
2708 const extent: usize = 90;
2709 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2710 const blocks: u32 = @intCast(radixSplitBlockCount(instance.extent, instance.threads));
2711 try testing.expectEqual(@as(u32, 3), blocks);
2712
2713 var keys: [extent]i32 = undefined;
2714 var seed: u32 = 0x2545f491;
2715 for (&keys) |*key| {
2716 seed ^= seed << 13;
2717 seed ^= seed >> 17;
2718 seed ^= seed << 5;
2719 key.* = @intCast(seed % 100000);
2720 }
2721
2722 var graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
2723 defer graph.deinit();
2724
2725 inline for (.{ 0, 4 }) |shift| {
2726 var counts = @as([(radix_digit_bins * 3)]f32, @splat(-1));
2727 try graph.runCpuWithLaunch(allocator, &.{
2728 kernel.argumentBuffer(f32, counts[0..]),
2729 kernel.argumentBuffer(i32, keys[0..]),
2730 kernel.argumentI32(@intCast(extent)),
2731 kernel.argumentI32(shift),
2732 kernel.argumentI32(0),
2733 }, .{
2734 .grid = .{ blocks, 1, 1 },
2735 .block = .{ instance.threads, 1, 1 },
2736 });
2737
2738 var expected = @as([(radix_digit_bins * 3)]f32, @splat(0));
2739 for (keys, 0..) |key, index| {
2740 const digit: usize = @intCast((key >> shift) & (radix_digit_bins - 1));
2741 const block = index / instance.threads;
2742 expected[digit * 3 + block] += 1;
2743 }
2744 try testing.expectEqualSlices(f32, expected[0..], counts[0..]);
2745 }
2746 }
2747
2748 fn runRadixDigitPassOnOracle(
2749 allocator: std.mem.Allocator,
2750 instance: RadixSplit,
2751 shift: u32,
2752 keys: []i32,
2753 dst: []i32,
2754 ) !void {
2755 var histogram_graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
2756 defer histogram_graph.deinit();
2757 var scatter_graph = try RadixDigitRankScatterRuntimeFamilyI32.build(allocator, RadixDigitRankScatterRuntimeFamilyI32.Limits.testing, instance);
2758 defer scatter_graph.deinit();
2759 try runRadixDigitPassOnOracleWithGraphs(allocator, instance, shift, keys, dst, &histogram_graph, &scatter_graph);
2760 }
2761
2762 fn runRadixDigitPassOnOracleWithGraphs(
2763 allocator: std.mem.Allocator,
2764 instance: RadixSplit,
2765 shift: u32,
2766 keys: []i32,
2767 dst: []i32,
2768 histogram_graph: anytype,
2769 scatter_graph: anytype,
2770 ) !void {
2771 const extent = keys.len;
2772 const blocks: u32 = @intCast(radixSplitBlockCount(instance.extent, instance.threads));
2773 const cell_count = radix_digit_bins * blocks;
2774
2775 const counts = try allocator.alloc(f32, cell_count);
2776 defer allocator.free(counts);
2777 @memset(counts, -1);
2778
2779 try histogram_graph.runCpuWithLaunch(allocator, &.{
2780 kernel.argumentBuffer(f32, counts),
2781 kernel.argumentBuffer(i32, keys),
2782 kernel.argumentI32(@intCast(extent)),
2783 kernel.argumentI32(@intCast(shift)),
2784 kernel.argumentI32(@intCast(radixDigitSignedPassBias(shift))),
2785 }, .{
2786 .grid = .{ blocks, 1, 1 },
2787 .block = .{ instance.threads, 1, 1 },
2788 });
2789
2790 const scanned = try allocator.alloc(f32, cell_count);
2791 defer allocator.free(scanned);
2792 var running: f32 = 0;
2793 for (counts, scanned) |count, *value| {
2794 value.* = running;
2795 running += count;
2796 }
2797
2798 try scatter_graph.runCpuWithLaunch(allocator, &.{
2799 kernel.argumentBuffer(i32, dst),
2800 kernel.argumentBuffer(i32, keys),
2801 kernel.argumentBuffer(f32, scanned),
2802 kernel.argumentI32(@intCast(extent)),
2803 kernel.argumentI32(@intCast(shift)),
2804 kernel.argumentI32(@intCast(radixDigitSignedPassBias(shift))),
2805 }, .{
2806 .grid = .{ blocks, 1, 1 },
2807 .block = .{ instance.threads, 1, 1 },
2808 });
2809 }
2810
2811 fn expectStableDigitPass(keys: []const i32, dst: []const i32, shift: u32) !void {
2812 const expected = try testing.allocator.alloc(i32, keys.len);
2813 defer testing.allocator.free(expected);
2814 var count: usize = 0;
2815 const shift_amount: u5 = @intCast(shift);
2816 var digit: i32 = 0;
2817 while (digit < radix_digit_bins) : (digit += 1) {
2818 for (keys) |key| {
2819 if ((key >> shift_amount) & (radix_digit_bins - 1) == digit) {
2820 expected[count] = key;
2821 count += 1;
2822 }
2823 }
2824 }
2825 try testing.expectEqualSlices(i32, expected, dst);
2826 }
2827
2828 test "sort radix digit pass partitions stably by digit across warps and blocks" {
2829 const allocator = testing.allocator;
2830 const extent: usize = 150;
2831 const instance = RadixSplit{ .extent = extent, .threads = 64 };
2832
2833 var keys: [extent]i32 = undefined;
2834 var seed: u32 = 0x2545f491;
2835 for (&keys, 0..) |*key, index| {
2836 seed ^= seed << 13;
2837 seed ^= seed >> 17;
2838 seed ^= seed << 5;
2839 key.* = @intCast((seed >> 8) % 10000 * 10 + index % 10);
2840 }
2841
2842 var dst = @as([extent]i32, @splat(-1));
2843 var histogram_graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
2844 defer histogram_graph.deinit();
2845 var scatter_graph = try RadixDigitRankScatterRuntimeFamilyI32.build(allocator, RadixDigitRankScatterRuntimeFamilyI32.Limits.testing, instance);
2846 defer scatter_graph.deinit();
2847
2848 try runRadixDigitPassOnOracleWithGraphs(allocator, instance, 0, keys[0..], dst[0..], &histogram_graph, &scatter_graph);
2849 try expectStableDigitPass(keys[0..], dst[0..], 0);
2850
2851 var dst_high = @as([extent]i32, @splat(-1));
2852 try runRadixDigitPassOnOracleWithGraphs(allocator, instance, 8, keys[0..], dst_high[0..], &histogram_graph, &scatter_graph);
2853 try expectStableDigitPass(keys[0..], dst_high[0..], 8);
2854 }
2855
2856 test "sort radix digit passes compose into a full sort on the oracle" {
2857 const allocator = testing.allocator;
2858 const extent: usize = 130;
2859 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2860
2861 var keys: [extent]i32 = undefined;
2862 var seed: u32 = 0x9e3779b9;
2863 for (&keys) |*key| {
2864 seed ^= seed << 13;
2865 seed ^= seed >> 17;
2866 seed ^= seed << 5;
2867 key.* = @intCast(seed % 1000000);
2868 }
2869
2870 var current = keys;
2871 var scratch = @as([extent]i32, @splat(-1));
2872 var histogram_graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
2873 defer histogram_graph.deinit();
2874 var scatter_graph = try RadixDigitRankScatterRuntimeFamilyI32.build(allocator, RadixDigitRankScatterRuntimeFamilyI32.Limits.testing, instance);
2875 defer scatter_graph.deinit();
2876 var shift: u32 = 0;
2877 while (shift < 20) : (shift += radix_digit_bits) {
2878 try runRadixDigitPassOnOracleWithGraphs(allocator, instance, shift, current[0..], scratch[0..], &histogram_graph, &scatter_graph);
2879 current = scratch;
2880 }
2881
2882 var expected = keys;
2883 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2884 try testing.expectEqualSlices(i32, expected[0..], current[0..]);
2885 }
2886
2887 test "sort radix digit pipeline descriptor binds the family artifacts" {
2888 const allocator = testing.allocator;
2889 var state = gpu.recording.BackendState{
2890 .allocator = allocator,
2891 .kind = .cuda,
2892 .format = .cuda_ptx,
2893 };
2894 const instance = RadixSplit{ .extent = 5000, .threads = 64 };
2895
2896 var artifacts = try createRadixDigitPipelineArtifacts(allocator, state.handle(), instance, .{ .limits = .testing });
2897 defer artifacts.deinit();
2898 const entries = artifacts.entries();
2899 const registry = artifact_product.KernelCallRegistry{ .entries = entries[0..] };
2900
2901 var owned = try radixDigitPipeline(allocator, instance);
2902 defer owned.deinit();
2903 try testing.expectEqualStrings("accy.kernel.sort.radix_digit_family_64_i32", owned.value.target);
2904 try testing.expectEqual(@as(u32, 3), owned.value.runtime_scalar_argument_count);
2905 try testing.expectEqual(@as(usize, 4), owned.value.intermediates.len);
2906 try testing.expectEqual(@as(usize, 5), owned.value.stages.len);
2907 try owned.value.validate(registry, .cuda_ptx);
2908
2909 const args = [_]choir_abi.ScalarArgument{ .{ .u32 = 5000 }, .{ .u32 = 0 }, .{ .u32 = 0 } };
2910 try testing.expectEqual(@as(u32, 79 * 16), try owned.value.intermediates[0].extent.resolveExtent(args[0..]));
2911 try testing.expectEqual(@as(u32, 2), try owned.value.intermediates[1].extent.resolveExtent(args[0..]));
2912 }
2913
2914 test "sort tuning resolves structure winners through pipeline targets" {
2915 const allocator = testing.allocator;
2916 const instance = RadixSplit{ .extent = 5000, .threads = 64 };
2917 const device: u64 = 0xfeed_dead_beef_0001;
2918
2919 var accumulator = tuning.FamilyMeasurementAccumulator.init(allocator);
2920 defer accumulator.deinit();
2921
2922 const key = try radixSplitFamilyTuningKey(allocator, device, instance);
2923 const digit_target = try radixDigitPipelineTarget(allocator, instance);
2924 defer allocator.free(digit_target);
2925 const split_target = try radixSplitPipelineTarget(allocator, instance);
2926 defer allocator.free(split_target);
2927
2928 try accumulator.append(key, split_target, 6_400_000, 50);
2929 try accumulator.append(key, digit_target, 1_550_000, 50);
2930
2931 var winners = try accumulator.selectWinners(allocator, tuning.family_tuning_default_margin_percent);
2932 defer winners.deinit();
2933 try testing.expectEqual(@as(usize, 1), winners.records.len);
2934
2935 const encoded_artifact = try tuning.encodeFamilyTuningArtifact(allocator, winners.records);
2936 defer allocator.free(encoded_artifact);
2937 var decoded = try tuning.decodeFamilyTuningArtifact(allocator, encoded_artifact);
2938 defer decoded.deinit();
2939
2940 const reader = tuning.FamilyTuningReader{
2941 .device_fingerprint = device,
2942 .table = decoded.table(),
2943 };
2944 const resolved = (try resolveRadixSplitStructure(allocator, reader, instance)) orelse {
2945 return error.TestExpectedSortStructure;
2946 };
2947 try testing.expectEqual(RadixSplitResolvedStructure.radix_digit, resolved);
2948
2949 const other_extent = RadixSplit{ .extent = 9000, .threads = 64 };
2950 try testing.expectEqual(
2951 @as(?RadixSplitResolvedStructure, null),
2952 try resolveRadixSplitStructure(allocator, reader, other_extent),
2953 );
2954 }
2955
2956 test "sort radix digit passes sort signed keys on the oracle" {
2957 const allocator = testing.allocator;
2958 const extent: usize = 96;
2959 const instance = RadixSplit{ .extent = extent, .threads = 32 };
2960
2961 var keys: [extent]i32 = undefined;
2962 var seed: u32 = 0xc0ffee11;
2963 for (&keys) |*key| {
2964 seed ^= seed << 13;
2965 seed ^= seed >> 17;
2966 seed ^= seed << 5;
2967 const magnitude: i32 = @intCast(seed % 500000);
2968 key.* = if (seed & 1 == 1) -magnitude else magnitude;
2969 }
2970
2971 var current = keys;
2972 var scratch = @as([extent]i32, @splat(-1));
2973 var histogram_graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
2974 defer histogram_graph.deinit();
2975 var scatter_graph = try RadixDigitRankScatterRuntimeFamilyI32.build(allocator, RadixDigitRankScatterRuntimeFamilyI32.Limits.testing, instance);
2976 defer scatter_graph.deinit();
2977 var shift: u32 = 0;
2978 while (shift < radix_split_key_bits) : (shift += radix_digit_bits) {
2979 try runRadixDigitPassOnOracleWithGraphs(allocator, instance, shift, current[0..], scratch[0..], &histogram_graph, &scatter_graph);
2980 current = scratch;
2981 }
2982
2983 var expected = keys;
2984 std.mem.sort(i32, expected[0..], {}, std.sort.asc(i32));
2985 try testing.expectEqualSlices(i32, expected[0..], current[0..]);
2986 }
2987
2988 test "sort radix digit pairs pipeline descriptor binds the family artifacts" {
2989 const allocator = testing.allocator;
2990 var state = gpu.recording.BackendState{
2991 .allocator = allocator,
2992 .kind = .cuda,
2993 .format = .cuda_ptx,
2994 };
2995 const instance = RadixSplit{ .extent = 5000, .threads = 64 };
2996
2997 var artifacts = try createRadixDigitPairsPipelineArtifacts(allocator, state.handle(), instance, .{ .limits = .testing });
2998 defer artifacts.deinit();
2999 const entries = artifacts.entries();
3000 const registry = artifact_product.KernelCallRegistry{ .entries = entries[0..] };
3001
3002 var owned = try radixDigitPairsPipeline(allocator, instance);
3003 defer owned.deinit();
3004 try testing.expectEqualStrings("accy.kernel.sort.radix_digit_pairs_family_64_i32", owned.value.target);
3005 try testing.expectEqual(@as(u32, 2), owned.value.operand_count);
3006 try testing.expectEqual(@as(u32, 2), owned.value.result_count);
3007 try testing.expectEqual(@as(u32, 3), owned.value.runtime_scalar_argument_count);
3008 try owned.value.validate(registry, .cuda_ptx);
3009 }
3010
3011 test "sort radix digit pairs pass carries payloads stably on the oracle" {
3012 const allocator = testing.allocator;
3013 const extent: usize = 96;
3014 const instance = RadixSplit{ .extent = extent, .threads = 32 };
3015 const blocks: u32 = @intCast(radixSplitBlockCount(instance.extent, instance.threads));
3016 const cell_count = radix_digit_bins * blocks;
3017
3018 var keys: [extent]i32 = undefined;
3019 var values: [extent]i32 = undefined;
3020 var seed: u32 = 0xc0ffee11;
3021 for (&keys, &values, 0..) |*key, *value, index| {
3022 seed ^= seed << 13;
3023 seed ^= seed >> 17;
3024 seed ^= seed << 5;
3025 const magnitude: i32 = @intCast(seed % 50000);
3026 key.* = if (seed & 1 == 1) -magnitude else magnitude;
3027 value.* = @intCast(index);
3028 }
3029
3030 const shift: u32 = 28;
3031 const counts = try allocator.alloc(f32, cell_count);
3032 defer allocator.free(counts);
3033 @memset(counts, -1);
3034 var histogram_graph = try RadixDigitHistogramRuntimeFamilyI32.build(allocator, RadixDigitHistogramRuntimeFamilyI32.Limits.testing, instance);
3035 defer histogram_graph.deinit();
3036 try histogram_graph.runCpuWithLaunch(allocator, &.{
3037 kernel.argumentBuffer(f32, counts),
3038 kernel.argumentBuffer(i32, keys[0..]),
3039 kernel.argumentI32(@intCast(extent)),
3040 kernel.argumentI32(@intCast(shift)),
3041 kernel.argumentI32(@intCast(radixDigitSignedPassBias(shift))),
3042 }, .{
3043 .grid = .{ blocks, 1, 1 },
3044 .block = .{ instance.threads, 1, 1 },
3045 });
3046
3047 const scanned = try allocator.alloc(f32, cell_count);
3048 defer allocator.free(scanned);
3049 var running: f32 = 0;
3050 for (counts, scanned) |count, *value| {
3051 value.* = running;
3052 running += count;
3053 }
3054
3055 var dst = @as([extent]i32, @splat(-1));
3056 var dst_values = @as([extent]i32, @splat(-1));
3057 var pairs_graph = try RadixDigitRankScatterPairsRuntimeFamilyI32.build(allocator, RadixDigitRankScatterPairsRuntimeFamilyI32.Limits.testing, instance);
3058 defer pairs_graph.deinit();
3059 try pairs_graph.runCpuWithLaunch(allocator, &.{
3060 kernel.argumentBuffer(i32, dst[0..]),
3061 kernel.argumentBuffer(i32, dst_values[0..]),
3062 kernel.argumentBuffer(i32, keys[0..]),
3063 kernel.argumentBuffer(i32, values[0..]),
3064 kernel.argumentBuffer(f32, scanned),
3065 kernel.argumentI32(@intCast(extent)),
3066 kernel.argumentI32(@intCast(shift)),
3067 kernel.argumentI32(@intCast(radixDigitSignedPassBias(shift))),
3068 }, .{
3069 .grid = .{ blocks, 1, 1 },
3070 .block = .{ instance.threads, 1, 1 },
3071 });
3072
3073 for (dst, dst_values) |key, original_index| {
3074 try testing.expectEqual(keys[@intCast(original_index)], key);
3075 }
3076 var previous_bucket: i32 = -1;
3077 for (dst) |key| {
3078 const shift_amount: u5 = @intCast(shift);
3079 const digit = ((key >> shift_amount) & (radix_digit_bins - 1)) ^ @as(i32, @intCast(radixDigitSignedPassBias(shift)));
3080 try testing.expect(digit >= previous_bucket);
3081 previous_bucket = digit;
3082 }
3083 }