lib/accy/src/kernel/library/compaction.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4
5 const artifact_product = @import("../../artifact/model/root.zig");
6 const shape = @import("../../choir/shape/root.zig");
7 const entry = @import("entry.zig");
8 const extent_mod = @import("extent.zig");
9 const kernel = @import("../root.zig");
10 const tuning = @import("tuning.zig");
11
12 const DType = choir_abi.DType;
13 const indexExtent = extent_mod.indexExtent;
14 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
15
16 pub const Filter = struct {
17 extent: u64,
18 dtype: DType = .f32,
19 predicate: entry.CompactionPredicate = .nonzero,
20 threads: u32 = 256,
21 element_axis: []const u8 = "n",
22 segment_axis: []const u8 = "b",
23
24 pub fn segments(self: Filter) u64 {
25 return ceilDiv(self.extent, self.threads);
26 }
27
28 pub fn padded(self: Filter) u64 {
29 return self.extent + self.segments();
30 }
31 };
32
33 pub fn filterPredicateName(predicate: entry.CompactionPredicate) []const u8 {
34 return switch (predicate) {
35 .nonzero => "nonzero",
36 .greater_than => "greater",
37 };
38 }
39
40 pub const filter_family_version: u32 = 1;
41 pub const filter_warp_size: u32 = 32;
42 pub const filter_max_threads: u32 = 1024;
43
44 pub fn filterDTypeSupported(dtype: DType) bool {
45 return switch (dtype) {
46 .f32, .i32 => true,
47 else => false,
48 };
49 }
50
51 pub fn filterInstanceValid(instance: Filter) bool {
52 if (!filterDTypeSupported(instance.dtype)) return false;
53 if (instance.extent == 0) return false;
54 if (instance.threads == 0 or instance.threads > filter_max_threads) return false;
55 return instance.threads % filter_warp_size == 0;
56 }
57
58 fn ceilDiv(numerator: u64, denominator: u64) u64 {
59 return numerator / denominator + @intFromBool(numerator % denominator != 0);
60 }
61
62 pub fn filterBlocksExpectedF32(data: []const f32, threads: u32, dst: []f32) void {
63 const segment_count = ceilDiv(data.len, threads);
64 for (0..segment_count) |segment| {
65 const begin = segment * threads;
66 const end = @min(begin + threads, data.len);
67 var survivors: usize = 0;
68 for (data[begin..end]) |value| {
69 if (value != 0.0) {
70 dst[begin + survivors] = value;
71 survivors += 1;
72 }
73 }
74 dst[data.len + segment] = @floatFromInt(survivors);
75 }
76 }
77
78 pub fn filterBlocksExpectedI32(data: []const i32, threads: u32, dst: []i32) void {
79 const segment_count = ceilDiv(data.len, threads);
80 for (0..segment_count) |segment| {
81 const begin = segment * threads;
82 const end = @min(begin + threads, data.len);
83 var survivors: usize = 0;
84 for (data[begin..end]) |value| {
85 if (value != 0) {
86 dst[begin + survivors] = value;
87 survivors += 1;
88 }
89 }
90 dst[data.len + segment] = @intCast(survivors);
91 }
92 }
93
94 pub fn filterBlocksGreaterExpectedF32(data: []const f32, threads: u32, threshold: f32, dst: []f32) void {
95 const segment_count = ceilDiv(data.len, threads);
96 for (0..segment_count) |segment| {
97 const begin = segment * threads;
98 const end = @min(begin + threads, data.len);
99 var survivors: usize = 0;
100 for (data[begin..end]) |value| {
101 if (value > threshold) {
102 dst[begin + survivors] = value;
103 survivors += 1;
104 }
105 }
106 dst[data.len + segment] = @floatFromInt(survivors);
107 }
108 }
109
110 pub fn filterBlocksGreaterExpectedI32(data: []const i32, threads: u32, threshold: i32, dst: []i32) void {
111 const segment_count = ceilDiv(data.len, threads);
112 for (0..segment_count) |segment| {
113 const begin = segment * threads;
114 const end = @min(begin + threads, data.len);
115 var survivors: usize = 0;
116 for (data[begin..end]) |value| {
117 if (value > threshold) {
118 dst[begin + survivors] = value;
119 survivors += 1;
120 }
121 }
122 dst[data.len + segment] = @intCast(survivors);
123 }
124 }
125
126 fn predicateFlag(
127 k: anytype,
128 comptime dtype: DType,
129 comptime predicate: entry.CompactionPredicate,
130 args: anytype,
131 element: kernel.Value,
132 ) !kernel.Value {
133 const survives = switch (predicate) {
134 .nonzero => switch (dtype) {
135 .f32 => try k.compare(.ne, element, try k.constantFloat(.f32, 0.0)),
136 .i32 => try k.compare(.ne, element, try k.constantInt(.i32, 0)),
137 else => @compileError("filter kernels support dtype .f32 or .i32"),
138 },
139 .greater_than => try k.compare(.gt, element, args.param(.threshold).raw()),
140 };
141 const one = try k.constantInt(.i32, 1);
142 const zero = try k.constantInt(.i32, 0);
143 return k.select(survives, one, zero);
144 }
145
146 fn zeroElement(k: anytype, comptime dtype: DType) !kernel.Value {
147 return switch (dtype) {
148 .f32 => k.constantFloat(.f32, 0.0),
149 .i32 => k.constantInt(.i32, 0),
150 else => @compileError("filter kernels support dtype .f32 or .i32"),
151 };
152 }
153
154 fn filter_scan_core_seeds_zero(inner: anytype, ctx: anytype) !void {
155 try inner.storeIndex(ctx.zero_flag, ctx.warp_sums, ctx.local);
156 }
157
158 fn filter_scan_core_is_last_lane(inner: anytype, ctx: anytype) !void {
159 try inner.storeIndex(ctx.scanned, ctx.warp_sums, ctx.warp);
160 }
161
162 fn filter_scan_core_is_first_warp(inner: anytype, ctx: anytype) !void {
163 const warp_sum = try inner.loadIndex(ctx.warp_sums, ctx.lane);
164 const warp_scan = try inner.warpScan(.add, .inclusive, warp_sum);
165 try inner.storeIndex(warp_scan, ctx.warp_sums, ctx.lane);
166 }
167
168 fn filter_scan_core_survives(inner: anytype, ctx: anytype) !void {
169 try ctx.args.param(.dst).store(inner, ctx.element, ctx.destination);
170 }
171
172 fn filter_scan_core_is_last_thread(inner: anytype, ctx: anytype) !void {
173 try ctx.args.param(.dst).store(inner, ctx.count_value, ctx.count_slot);
174 }
175
176 fn filterScanCore(
177 k: anytype,
178 comptime dtype: DType,
179 comptime predicate: entry.CompactionPredicate,
180 args: anytype,
181 extent: kernel.Value,
182 ) !void {
183 const tid = try k.globalId(.x);
184 const local = try k.threadId(.x);
185 const block = try k.blockId(.x);
186 const block_threads = try k.blockDim(.x);
187 const lane = try k.laneId();
188 const warp = try k.warpId();
189 const zero = try k.constantIndex(0);
190 const one = try k.constantIndex(1);
191
192 const in_range = try k.compare(.lt, tid, extent);
193 const extent_minus_one = try k.sub(extent, one);
194 const clamped_tid = try k.min(tid, extent_minus_one);
195 const loaded = try args.param(.data).load(k, clamped_tid);
196 const element = try k.select(in_range, loaded.raw(), try zeroElement(k, dtype));
197
198 const live = try predicateFlag(k, dtype, predicate, args, element);
199 const flag = try k.select(in_range, live, try k.constantInt(.i32, 0));
200
201 const scanned = try k.warpScan(.add, .inclusive, flag);
202
203 const warp_sums = try k.sharedBuffer(.i32, filter_warp_size);
204 const lane_limit = try k.constantIndex(filter_warp_size - 1);
205 const warp_count_value = try k.constantIndex(filter_warp_size);
206
207 const seeds_zero = try k.compare(.lt, local, warp_count_value);
208 try k.guardDo(seeds_zero, .{ .warp_sums = warp_sums, .local = local, .zero_flag = try k.constantInt(.i32, 0) }, filter_scan_core_seeds_zero);
209 try k.barrier(.block);
210
211 const is_last_lane = try k.compare(.eq, lane, lane_limit);
212 try k.guardDo(is_last_lane, .{ .warp_sums = warp_sums, .warp = warp, .scanned = scanned }, filter_scan_core_is_last_lane);
213 try k.barrier(.block);
214
215 const is_first_warp = try k.compare(.eq, warp, zero);
216 try k.guardDo(is_first_warp, .{ .warp_sums = warp_sums, .lane = lane }, filter_scan_core_is_first_warp);
217 try k.barrier(.block);
218
219 const has_base = try k.compare(.gt, warp, zero);
220 const warp_minus_one = try k.sub(warp, one);
221 const base_index = try k.select(has_base, warp_minus_one, zero);
222 const base_loaded = try k.loadIndex(warp_sums, base_index);
223 const base = try k.select(has_base, base_loaded, try k.constantInt(.i32, 0));
224 const inclusive = try k.add(scanned, base);
225
226 const segment_base = try k.mul(block, block_threads);
227 const survives = try k.compare(.gt, flag, try k.constantInt(.i32, 0));
228 const offset = try k.castIndex(try k.sub(inclusive, try k.constantInt(.i32, 1)));
229 const destination = try k.add(segment_base, offset);
230 try k.guardDo(survives, .{ .args = args, .element = element, .destination = destination }, filter_scan_core_survives);
231
232 const block_threads_minus_one = try k.sub(block_threads, one);
233 const is_last_thread = try k.compare(.eq, local, block_threads_minus_one);
234 const count_value = switch (dtype) {
235 .f32 => try k.cast(inclusive, .f32),
236 .i32 => inclusive,
237 else => @compileError("filter kernels support dtype .f32 or .i32"),
238 };
239 const count_slot = try k.add(extent, block);
240 try k.guardDo(is_last_thread, .{ .args = args, .count_value = count_value, .count_slot = count_slot }, filter_scan_core_is_last_thread);
241 }
242
243 fn filterBody(
244 k: anytype,
245 comptime dtype: DType,
246 comptime predicate: entry.CompactionPredicate,
247 spec: Filter,
248 args: anytype,
249 ) !void {
250 if (!filterInstanceValid(spec)) return error.UnsupportedFilterInstance;
251 const extent = try k.constantIndex(try indexExtent(spec.extent));
252 try filterScanCore(k, dtype, predicate, args, extent);
253 }
254
255 fn filterRuntimeBody(
256 k: anytype,
257 comptime dtype: DType,
258 comptime predicate: entry.CompactionPredicate,
259 spec: Filter,
260 args: anytype,
261 ) !void {
262 if (!filterInstanceValid(spec)) return error.UnsupportedFilterInstance;
263 const extent = try k.castIndex(args.param(.extent).raw());
264 try filterScanCore(k, dtype, predicate, args, extent);
265 }
266
267 fn filterFamilySchedule(instance: Filter) kernel.logical.schedule.ThreadBlocks {
268 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
269 }
270
271 fn filter_runtime_family_body_f32(k: anytype, spec: Filter, args: anytype) !void {
272 try filterRuntimeBody(k, .f32, .nonzero, spec, args);
273 }
274
275 fn filter_runtime_family_body_i32(k: anytype, spec: Filter, args: anytype) !void {
276 try filterRuntimeBody(k, .i32, .nonzero, spec, args);
277 }
278
279 fn filterRuntimeFamily(comptime dtype: DType) type {
280 return kernel.logical.Family(.{
281 .name = std.fmt.comptimePrint("accy_kernel_compaction_filter_runtime_{s}", .{dtype.name()}),
282 .parameters = .{
283 .dst = kernel.dynamicBuffer(dtype),
284 .data = kernel.dynamicBuffer(dtype),
285 .extent = kernel.scalar(.i32),
286 },
287 .Instance = Filter,
288 .schedule = filterFamilySchedule,
289 .body = switch (dtype) {
290 .f32 => filter_runtime_family_body_f32,
291 .i32 => filter_runtime_family_body_i32,
292 else => @compileError("runtime filter supports dtype .f32 or .i32"),
293 },
294 });
295 }
296
297 fn filter_greater_runtime_family_body_f32(k: anytype, spec: Filter, args: anytype) !void {
298 try filterRuntimeBody(k, .f32, .greater_than, spec, args);
299 }
300
301 fn filter_greater_runtime_family_body_i32(k: anytype, spec: Filter, args: anytype) !void {
302 try filterRuntimeBody(k, .i32, .greater_than, spec, args);
303 }
304
305 fn filterGreaterRuntimeFamily(comptime dtype: DType) type {
306 return kernel.logical.Family(.{
307 .name = std.fmt.comptimePrint("accy_kernel_compaction_filter_greater_runtime_{s}", .{dtype.name()}),
308 .parameters = .{
309 .dst = kernel.dynamicBuffer(dtype),
310 .data = kernel.dynamicBuffer(dtype),
311 .extent = kernel.scalar(.i32),
312 .threshold = kernel.scalar(dtype),
313 },
314 .Instance = Filter,
315 .schedule = filterFamilySchedule,
316 .body = switch (dtype) {
317 .f32 => filter_greater_runtime_family_body_f32,
318 .i32 => filter_greater_runtime_family_body_i32,
319 else => @compileError("runtime greater-than filter supports dtype .f32 or .i32"),
320 },
321 });
322 }
323
324 pub const FilterRuntimeFamilyF32 = filterRuntimeFamily(.f32);
325 pub const FilterRuntimeFamilyI32 = filterRuntimeFamily(.i32);
326 pub const FilterGreaterRuntimeFamilyF32 = filterGreaterRuntimeFamily(.f32);
327 pub const FilterGreaterRuntimeFamilyI32 = filterGreaterRuntimeFamily(.i32);
328
329 pub fn filterThreadsForExtent(extent: u64) u32 {
330 if (extent >= 256) return 256;
331 const wide: u64 = extent + filter_warp_size - 1;
332 const rounded: u32 = @intCast((wide / filter_warp_size) * filter_warp_size);
333 return @max(rounded, filter_warp_size);
334 }
335
336 pub const FilterThreadCandidates = struct {
337 count: usize = 0,
338 items: [6]u32 = @as([6]u32, @splat(0)),
339
340 pub fn slice(self: *const FilterThreadCandidates) []const u32 {
341 return self.items[0..self.count];
342 }
343 };
344
345 pub fn filterThreadCandidatesForExtent(extent: u64) FilterThreadCandidates {
346 var result = FilterThreadCandidates{};
347 if (extent == 0) return result;
348 const base = filterThreadsForExtent(extent);
349 result.items[result.count] = base;
350 result.count += 1;
351 var threads: u32 = filter_warp_size;
352 while (threads <= filter_max_threads) : (threads *= 2) {
353 if (threads == base) continue;
354 if (result.count >= result.items.len) break;
355 result.items[result.count] = threads;
356 result.count += 1;
357 }
358 return result;
359 }
360
361 pub fn filterInstanceTarget(allocator: std.mem.Allocator, instance: Filter) ![]u8 {
362 return std.fmt.allocPrint(
363 allocator,
364 "accy.kernel.compaction.filter{d}_{d}_{s}",
365 .{ instance.extent, instance.threads, instance.dtype.name() },
366 );
367 }
368
369 pub fn filterInstanceEntryName(allocator: std.mem.Allocator, instance: Filter) ![]u8 {
370 return std.fmt.allocPrint(
371 allocator,
372 "accy_kernel_compaction_filter{d}_{d}_{s}",
373 .{ instance.extent, instance.threads, instance.dtype.name() },
374 );
375 }
376
377 pub fn filterFamilyTarget(allocator: std.mem.Allocator, instance: Filter) ![]u8 {
378 return std.fmt.allocPrint(
379 allocator,
380 "accy.kernel.compaction.filter_family_{s}_{d}_{s}",
381 .{ filterPredicateName(instance.predicate), instance.threads, instance.dtype.name() },
382 );
383 }
384
385 pub fn filterFamilyEntryName(allocator: std.mem.Allocator, instance: Filter) ![]u8 {
386 return std.fmt.allocPrint(
387 allocator,
388 "accy_kernel_compaction_filter_family_{s}_{d}_{s}",
389 .{ filterPredicateName(instance.predicate), instance.threads, instance.dtype.name() },
390 );
391 }
392
393 pub fn filterTuningExtents(instance: Filter) [1]u64 {
394 return .{instance.extent};
395 }
396
397 pub fn filterTuningOperation(instance: Filter) entry.Operation {
398 return .{ .compaction = .{ .blocks = instance.predicate } };
399 }
400
401 pub fn filterFamilyTuningKey(
402 backing_allocator: std.mem.Allocator,
403 device_fingerprint: u64,
404 instance: Filter,
405 ) !tuning.FamilyTuningKey {
406 const family_fingerprint = try filterFamilyFingerprint(backing_allocator, instance);
407 const extents = filterTuningExtents(instance);
408 return tuning.FamilyTuningKey.init(
409 device_fingerprint,
410 family_fingerprint,
411 entry.operationFingerprint(filterTuningOperation(instance)),
412 instance.dtype,
413 filter_family_version,
414 extents[0..],
415 ) orelse unreachable;
416 }
417
418 pub fn filterRuntimeArguments(instance: Filter) ![1]choir_abi.ScalarArgument {
419 return .{
420 .{ .u32 = try runtimeExtentArgument(instance.extent) },
421 };
422 }
423
424 pub fn filterGreaterRuntimeArguments(
425 instance: Filter,
426 threshold: choir_abi.ScalarArgument,
427 ) ![2]choir_abi.ScalarArgument {
428 return .{
429 .{ .u32 = try runtimeExtentArgument(instance.extent) },
430 threshold,
431 };
432 }
433
434 pub fn filterRuntimeScalarArgumentCount(instance: Filter) u32 {
435 return switch (instance.predicate) {
436 .nonzero => 1,
437 .greater_than => 2,
438 };
439 }
440
441 fn filterRuntimeExtentBounds() shape.Bounds {
442 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
443 }
444
445 pub fn filterShapeProfileDimensions(instance: Filter) [1]artifact_product.KernelCallShapeProfileDimension {
446 return .{
447 .{ .name = instance.element_axis, .runtime_scalar_argument_index = 0, .bounds = filterRuntimeExtentBounds() },
448 };
449 }
450
451 fn filterDerivedLaunch(instance: Filter) !artifact_product.KernelCallLaunch {
452 if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
453 return .{ .derived = .{
454 .grid = .{
455 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = instance.threads } },
456 .{ .fixed = 1 },
457 .{ .fixed = 1 },
458 },
459 .threadgroup = .{ instance.threads, 1, 1 },
460 } };
461 }
462
463 pub fn createFilterFamilyArtifact(
464 allocator: std.mem.Allocator,
465 handle: kernel.BackendHandle,
466 instance: Filter,
467 options: entry.ArtifactOptions,
468 ) !kernel.OwnedKernelCallArtifact {
469 const target = try filterFamilyTarget(allocator, instance);
470 defer allocator.free(target);
471 const entry_name = try filterFamilyEntryName(allocator, instance);
472 defer allocator.free(entry_name);
473 const family_fingerprint = options.shape_family_fingerprint orelse try filterFamilyFingerprint(allocator, instance);
474 const shape_profile_dimensions = filterShapeProfileDimensions(instance);
475 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
476 .name = "filter",
477 .fingerprint = family_fingerprint,
478 .dimensions = shape_profile_dimensions[0..],
479 };
480
481 var graph = switch (instance.predicate) {
482 .nonzero => switch (instance.dtype) {
483 .f32 => try FilterRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
484 .i32 => try FilterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),
485 else => return error.UnsupportedDType,
486 },
487 .greater_than => switch (instance.dtype) {
488 .f32 => try FilterGreaterRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
489 .i32 => try FilterGreaterRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),
490 else => return error.UnsupportedDType,
491 },
492 };
493 defer graph.deinit();
494 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
495 .target = target,
496 .version = filter_family_version,
497 .format = options.format,
498 .kernel_plan = options.kernel_plan,
499 .element_count_argument = options.element_count_argument,
500 .shape_family_fingerprint = family_fingerprint,
501 .shape_profile = shape_profile,
502 .launch = options.launch orelse try filterDerivedLaunch(instance),
503 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0)
504 filterRuntimeScalarArgumentCount(instance)
505 else
506 options.runtime_scalar_argument_count,
507 .static_arguments = options.static_arguments,
508 });
509 }
510
511 pub fn filterFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Filter) !u64 {
512 var family = try filterShapeFamily(backing_allocator, instance);
513 defer family.deinit();
514 return shape.fingerprint(family);
515 }
516
517 pub fn filterShapeFamily(backing_allocator: std.mem.Allocator, instance: Filter) !shape.Family {
518 var builder = try shape.Builder.init(backing_allocator, "filter");
519 errdefer builder.deinit();
520
521 const element = try builder.symbol(instance.element_axis);
522 const segment = try builder.symbol(instance.segment_axis);
523 const element_expr = try builder.symbolExpression(element);
524 const segment_expr = try builder.symbolExpression(segment);
525
526 const padded_expr = try builder.addExpression(element_expr, segment_expr);
527 _ = try builder.tensor("data", &.{element_expr});
528 _ = try builder.tensor("out", &.{padded_expr});
529 try builder.assumeBounds(element_expr, filterRuntimeExtentBounds());
530 try builder.assumeBounds(segment_expr, filterRuntimeExtentBounds());
531
532 return builder.finish();
533 }
534
535 pub fn filterFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Filter) !entry.OwnedSpecialization {
536 var owned = entry.OwnedSpecialization.init(backing_allocator);
537 errdefer owned.deinit();
538 const lifetime_allocator = owned.allocator();
539
540 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
541 inputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.element_axis, instance.extent);
542
543 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
544 outputs[0] = try entry.runtimeShape1D(lifetime_allocator, instance.segment_axis, instance.padded());
545
546 owned.value = .{
547 .dtype = instance.dtype,
548 .operation = .{ .compaction = .{ .blocks = instance.predicate } },
549 .inputs = inputs,
550 .outputs = outputs,
551 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.extent, instance.threads),
552 };
553 owned.value.launch = owned.value.schedule.?.launch();
554 var family = try filterShapeFamily(backing_allocator, instance);
555 errdefer family.deinit();
556 try owned.takeShapeFamily(&family);
557 return owned;
558 }
559
560 pub fn filterInstanceFromSpecialization(specialization: entry.Specialization) ?Filter {
561 if (!specialization.scheduleMatchesLaunch()) return null;
562 const operation = specialization.operation orelse return null;
563 const predicate = switch (operation) {
564 .compaction => |compaction_operation| switch (compaction_operation) {
565 .blocks => |predicate| predicate,
566 },
567 else => return null,
568 };
569 const dtype = specialization.dtype orelse return null;
570 if (!filterDTypeSupported(dtype)) return null;
571 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
572 if (specialization.reductions.len != 0) return null;
573 const data = specialization.inputs[0];
574 const packed_output = specialization.outputs[0];
575 if (data.axes.len != 1 or packed_output.axes.len != 1) return null;
576 const launch = specialization.launch orelse return null;
577 if (launch.threadgroup[0] == 0) return null;
578 const instance = Filter{
579 .extent = data.axes[0].extent,
580 .dtype = dtype,
581 .predicate = predicate,
582 .threads = launch.threadgroup[0],
583 .element_axis = data.axes[0].name,
584 .segment_axis = packed_output.axes[0].name,
585 };
586 if (!filterInstanceValid(instance)) return null;
587 if (packed_output.axes[0].extent != instance.padded()) return null;
588 return instance;
589 }
590
591 fn ceilDivComptime(comptime numerator: u64, comptime denominator: u32) u32 {
592 return @intCast(numerator / denominator + @as(u64, @intFromBool(numerator % denominator != 0)));
593 }
594
595 fn filterSpecialization(comptime spec: Filter) entry.Specialization {
596 return .{
597 .dtype = spec.dtype,
598 .operation = .{ .compaction = .{ .blocks = spec.predicate } },
599 .inputs = &.{entry.shape1D(spec.element_axis, spec.extent)},
600 .outputs = &.{entry.shape1D(spec.segment_axis, spec.padded())},
601 .launch = entry.launch1D(ceilDivComptime(spec.extent, spec.threads), spec.threads),
602 .schedule = entry.threadBlocks1D("e", spec.extent, spec.threads),
603 };
604 }
605
606 fn filterProgram(comptime spec: Filter) type {
607 const Body = struct {
608 fn run(k: anytype, args: anytype) !void {
609 try filterBody(k, spec.dtype, spec.predicate, spec, args);
610 }
611 };
612
613 return kernel.logical.Program(.{
614 .name = std.fmt.comptimePrint(
615 "accy_kernel_compaction_filter{}_{}_{s}",
616 .{ spec.extent, spec.threads, spec.dtype.name() },
617 ),
618 .parameters = .{
619 .dst = kernel.dynamicBuffer(spec.dtype),
620 .data = kernel.dynamicBuffer(spec.dtype),
621 },
622 .body = Body.run,
623 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
624 }
625
626 pub fn filterEntry(comptime spec: Filter) type {
627 return entry.Entry(filterProgram(spec), .{
628 .target = std.fmt.comptimePrint(
629 "accy.kernel.compaction.filter{}_{}_{s}",
630 .{ spec.extent, spec.threads, spec.dtype.name() },
631 ),
632 .layer = .logical,
633 .category = .compaction,
634 .specialization = filterSpecialization(spec),
635 });
636 }
637
638 pub const Filter8F32 = filterEntry(.{ .extent = 8, .threads = 32 });
639
640 test "compaction filter entry compacts one block on CPU" {
641 const allocator = std.testing.allocator;
642 var data = [_]f32{ 0.0, 3.5, 0.0, -1.25, 2.0, 0.0, 0.0, 7.0 };
643 var dst = @as([9]f32, @splat(-99.0));
644
645 const ProgramType = filterProgram(.{ .extent = 8, .threads = 32 });
646 var graph = try ProgramType.build(allocator, ProgramType.Limits.testing);
647 defer graph.deinit();
648 try graph.runCpuWithLaunch(allocator, &.{
649 kernel.argumentBuffer(f32, dst[0..]),
650 kernel.argumentBuffer(f32, data[0..]),
651 }, .{
652 .grid = .{ 1, 1, 1 },
653 .block = .{ 32, 1, 1 },
654 });
655
656 try std.testing.expectEqual(@as(f32, 4.0), dst[8]);
657 try std.testing.expectEqualSlices(f32, &.{ 3.5, -1.25, 2.0, 7.0 }, dst[0..4]);
658 for (dst[4..8]) |value| try std.testing.expectEqual(@as(f32, -99.0), value);
659 }
660
661 test "compaction filter runtime family compacts segments with tail" {
662 const allocator = std.testing.allocator;
663 const compiled = Filter{ .extent = 1, .threads = 32 };
664 const runtime = Filter{ .extent = 70, .threads = 32 };
665
666 var graph = try FilterRuntimeFamilyF32.build(allocator, FilterRuntimeFamilyF32.Limits.testing, compiled);
667 defer graph.deinit();
668
669 var data: [70]f32 = undefined;
670 var seed: u32 = 0x2545f491;
671 for (&data, 0..) |*value, index| {
672 seed ^= seed << 13;
673 seed ^= seed >> 17;
674 seed ^= seed << 5;
675 value.* = if (seed % 3 == 0) 0.0 else @floatFromInt(index + 1);
676 }
677
678 var dst = @as([73]f32, @splat(-1.0));
679
680 const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);
681 try graph.runCpuWithLaunch(allocator, &.{
682 kernel.argumentBuffer(f32, dst[0..]),
683 kernel.argumentBuffer(f32, data[0..]),
684 kernel.argumentI32(@intCast(runtime.extent)),
685 }, .{
686 .grid = launch_value.grid,
687 .block = launch_value.threadgroup,
688 });
689
690 var expected_dst = @as([73]f32, @splat(-1.0));
691 filterBlocksExpectedF32(data[0..], runtime.threads, expected_dst[0..]);
692
693 try std.testing.expectEqualSlices(f32, expected_dst[70..], dst[70..]);
694 for (0..3) |segment| {
695 const begin = segment * 32;
696 const survivors: usize = @intFromFloat(expected_dst[70 + segment]);
697 try std.testing.expectEqualSlices(f32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);
698 }
699 }
700
701 test "compaction filter runtime family compacts i32 data" {
702 const allocator = std.testing.allocator;
703 const compiled = Filter{ .extent = 1, .dtype = .i32, .threads = 32 };
704 const runtime = Filter{ .extent = 40, .dtype = .i32, .threads = 32 };
705
706 var graph = try FilterRuntimeFamilyI32.build(allocator, FilterRuntimeFamilyI32.Limits.testing, compiled);
707 defer graph.deinit();
708
709 var data: [40]i32 = undefined;
710 for (&data, 0..) |*value, index| value.* = if (index % 2 == 0) 0 else @intCast(index);
711
712 var dst = @as([42]i32, @splat(-1));
713
714 const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);
715 try graph.runCpuWithLaunch(allocator, &.{
716 kernel.argumentBuffer(i32, dst[0..]),
717 kernel.argumentBuffer(i32, data[0..]),
718 kernel.argumentI32(@intCast(runtime.extent)),
719 }, .{
720 .grid = launch_value.grid,
721 .block = launch_value.threadgroup,
722 });
723
724 var expected_dst = @as([42]i32, @splat(-1));
725 filterBlocksExpectedI32(data[0..], runtime.threads, expected_dst[0..]);
726
727 try std.testing.expectEqualSlices(i32, expected_dst[40..], dst[40..]);
728 for (0..2) |segment| {
729 const begin = segment * 32;
730 const survivors: usize = @intCast(expected_dst[40 + segment]);
731 try std.testing.expectEqualSlices(i32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);
732 }
733 }
734
735 test "compaction filter greater runtime family keeps survivors above the threshold" {
736 const allocator = std.testing.allocator;
737 const compiled = Filter{ .extent = 1, .predicate = .greater_than, .threads = 32 };
738 const runtime = Filter{ .extent = 70, .predicate = .greater_than, .threads = 32 };
739
740 var graph = try FilterGreaterRuntimeFamilyF32.build(allocator, FilterGreaterRuntimeFamilyF32.Limits.testing, compiled);
741 defer graph.deinit();
742
743 var data: [70]f32 = undefined;
744 var seed: u32 = 0x2545f491;
745 for (&data, 0..) |*value, index| {
746 seed ^= seed << 13;
747 seed ^= seed >> 17;
748 seed ^= seed << 5;
749 const magnitude: f32 = @floatFromInt(index + 1);
750 value.* = if (seed % 2 == 0) -magnitude else magnitude;
751 }
752
753 const threshold: f32 = 20.0;
754 var dst = @as([73]f32, @splat(-999.0));
755
756 const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);
757 try graph.runCpuWithLaunch(allocator, &.{
758 kernel.argumentBuffer(f32, dst[0..]),
759 kernel.argumentBuffer(f32, data[0..]),
760 kernel.argumentI32(@intCast(runtime.extent)),
761 kernel.argumentF32(threshold),
762 }, .{
763 .grid = launch_value.grid,
764 .block = launch_value.threadgroup,
765 });
766
767 var expected_dst = @as([73]f32, @splat(-999.0));
768 filterBlocksGreaterExpectedF32(data[0..], runtime.threads, threshold, expected_dst[0..]);
769
770 try std.testing.expectEqualSlices(f32, expected_dst[70..], dst[70..]);
771 for (0..3) |segment| {
772 const begin = segment * 32;
773 const survivors: usize = @intFromFloat(expected_dst[70 + segment]);
774 for (dst[begin .. begin + survivors]) |value| try std.testing.expect(value > threshold);
775 try std.testing.expectEqualSlices(f32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);
776 }
777 }
778
779 test "compaction filter greater runtime family filters i32 data" {
780 const allocator = std.testing.allocator;
781 const compiled = Filter{ .extent = 1, .dtype = .i32, .predicate = .greater_than, .threads = 32 };
782 const runtime = Filter{ .extent = 40, .dtype = .i32, .predicate = .greater_than, .threads = 32 };
783
784 var graph = try FilterGreaterRuntimeFamilyI32.build(allocator, FilterGreaterRuntimeFamilyI32.Limits.testing, compiled);
785 defer graph.deinit();
786
787 var data: [40]i32 = undefined;
788 for (&data, 0..) |*value, index| {
789 const magnitude: i32 = @intCast(index);
790 value.* = if (index % 3 == 0) -magnitude else magnitude;
791 }
792
793 const threshold: i32 = 11;
794 var dst = @as([42]i32, @splat(-999));
795
796 const launch_value = try entry.runtimeLaunch1D(runtime.extent, runtime.threads);
797 try graph.runCpuWithLaunch(allocator, &.{
798 kernel.argumentBuffer(i32, dst[0..]),
799 kernel.argumentBuffer(i32, data[0..]),
800 kernel.argumentI32(@intCast(runtime.extent)),
801 kernel.argumentI32(threshold),
802 }, .{
803 .grid = launch_value.grid,
804 .block = launch_value.threadgroup,
805 });
806
807 var expected_dst = @as([42]i32, @splat(-999));
808 filterBlocksGreaterExpectedI32(data[0..], runtime.threads, threshold, expected_dst[0..]);
809
810 try std.testing.expectEqualSlices(i32, expected_dst[40..], dst[40..]);
811 for (0..2) |segment| {
812 const begin = segment * 32;
813 const survivors: usize = @intCast(expected_dst[40 + segment]);
814 try std.testing.expectEqualSlices(i32, expected_dst[begin .. begin + survivors], dst[begin .. begin + survivors]);
815 }
816 }
817
818 test "compaction filter family instance identity matches fixed entry strings" {
819 const instance = Filter{ .extent = 8, .threads = 32 };
820
821 const target = try filterInstanceTarget(std.testing.allocator, instance);
822 defer std.testing.allocator.free(target);
823 try std.testing.expectEqualStrings(Filter8F32.target, target);
824
825 const entry_name = try filterInstanceEntryName(std.testing.allocator, instance);
826 defer std.testing.allocator.free(entry_name);
827 try std.testing.expectEqualStrings(Filter8F32.name, entry_name);
828
829 try std.testing.expectEqual(Filter8F32.version, filter_family_version);
830
831 const fresh = Filter{ .extent = 1 << 20, .threads = 128, .dtype = .i32 };
832 const family_target = try filterFamilyTarget(std.testing.allocator, fresh);
833 defer std.testing.allocator.free(family_target);
834 try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_nonzero_128_i32", family_target);
835 }
836
837 test "compaction filter family tuning keys discriminate predicates" {
838 const allocator = std.testing.allocator;
839 const device = tuning.deviceFingerprint(.{ .identity = .{
840 .backend = .cuda,
841 .family = .nvidia_cuda,
842 .name = "compaction-family-tuning-test-device",
843 .vendor_id = 0x10de,
844 .device_id = 0x2684,
845 } });
846
847 const nonzero = Filter{ .extent = 4096, .predicate = .nonzero };
848 const greater = Filter{ .extent = 4096, .predicate = .greater_than };
849 const nonzero_key = try filterFamilyTuningKey(allocator, device, nonzero);
850 const greater_key = try filterFamilyTuningKey(allocator, device, greater);
851 try std.testing.expect(!nonzero_key.eql(greater_key));
852 try std.testing.expectEqual(nonzero_key.family_fingerprint, greater_key.family_fingerprint);
853 try std.testing.expect(nonzero_key.operation_fingerprint != greater_key.operation_fingerprint);
854
855 const other_dtype = try filterFamilyTuningKey(allocator, device, .{ .extent = 4096, .dtype = .i32 });
856 try std.testing.expect(!nonzero_key.eql(other_dtype));
857
858 const repeat_key = try filterFamilyTuningKey(allocator, device, nonzero);
859 try std.testing.expect(nonzero_key.eql(repeat_key));
860 }
861
862 test "compaction filter family artifact carries runtime launch contract" {
863 const allocator = std.testing.allocator;
864 var state = gpu.recording.BackendState{
865 .allocator = allocator,
866 .kind = .cuda,
867 .format = .cuda_ptx,
868 };
869 const instance = Filter{ .extent = 4096, .threads = 128 };
870
871 var family_artifact = try createFilterFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
872 defer family_artifact.deinit();
873
874 const family_entry = family_artifact.entry();
875 try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_nonzero_128_f32", family_entry.target);
876 try std.testing.expectEqualStrings("accy_kernel_compaction_filter_family_nonzero_128_f32", family_entry.entry_name);
877 try std.testing.expectEqual(@as(u32, 3), family_entry.argument_count);
878 try std.testing.expectEqual(@as(u32, 1), family_entry.runtime_scalar_argument_count);
879 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
880 try std.testing.expect(family_entry.shape_family_fingerprint != null);
881 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
882 try std.testing.expectEqualStrings("filter", profile.name);
883 switch (family_entry.launch) {
884 .derived => |launch| {
885 try std.testing.expectEqual(@as(u32, 128), launch.threadgroup[0]);
886 switch (launch.grid[0]) {
887 .runtime_u32_ceil_div => |term| {
888 try std.testing.expectEqual(@as(usize, 0), term.argument_index);
889 try std.testing.expectEqual(@as(u32, 128), term.divisor);
890 },
891 else => return error.TestExpectedDerivedLaunch,
892 }
893 },
894 else => return error.TestExpectedDerivedLaunch,
895 }
896 }
897
898 test "compaction filter instance round-trips through specialization" {
899 const instance = Filter{ .extent = 1000, .threads = 64, .dtype = .i32 };
900 var owned = try filterFamilySpecialization(std.testing.allocator, instance);
901 defer owned.deinit();
902
903 const recovered = filterInstanceFromSpecialization(owned.value) orelse return error.TestExpectedFilterInstance;
904 try std.testing.expectEqual(instance.extent, recovered.extent);
905 try std.testing.expectEqual(instance.dtype, recovered.dtype);
906 try std.testing.expectEqual(instance.predicate, recovered.predicate);
907 try std.testing.expectEqual(instance.threads, recovered.threads);
908 try std.testing.expectEqual(@as(u64, 16), recovered.segments());
909
910 try std.testing.expectEqual(@as(?Filter, null), filterInstanceFromSpecialization(.{}));
911 }
912
913 test "compaction filter greater instance keeps its predicate through identity and specialization" {
914 const instance = Filter{ .extent = 4096, .predicate = .greater_than, .threads = 128 };
915
916 const family_target = try filterFamilyTarget(std.testing.allocator, instance);
917 defer std.testing.allocator.free(family_target);
918 try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_greater_128_f32", family_target);
919
920 const family_entry_name = try filterFamilyEntryName(std.testing.allocator, instance);
921 defer std.testing.allocator.free(family_entry_name);
922 try std.testing.expectEqualStrings("accy_kernel_compaction_filter_family_greater_128_f32", family_entry_name);
923
924 try std.testing.expectEqual(@as(u32, 2), filterRuntimeScalarArgumentCount(instance));
925 try std.testing.expectEqual(@as(u32, 1), filterRuntimeScalarArgumentCount(.{ .extent = 4096 }));
926
927 var owned = try filterFamilySpecialization(std.testing.allocator, instance);
928 defer owned.deinit();
929 const recovered = filterInstanceFromSpecialization(owned.value) orelse return error.TestExpectedFilterInstance;
930 try std.testing.expectEqual(entry.CompactionPredicate.greater_than, recovered.predicate);
931
932 const arguments = try filterGreaterRuntimeArguments(instance, .{ .f32 = 0.5 });
933 try std.testing.expectEqual(@as(u32, 4096), arguments[0].u32);
934 try std.testing.expectEqual(@as(f32, 0.5), arguments[1].f32);
935 }
936
937 test "compaction filter greater family artifact requires two runtime scalars" {
938 const allocator = std.testing.allocator;
939 var state = gpu.recording.BackendState{
940 .allocator = allocator,
941 .kind = .cuda,
942 .format = .cuda_ptx,
943 };
944 const instance = Filter{ .extent = 4096, .predicate = .greater_than, .threads = 128 };
945
946 var family_artifact = try createFilterFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
947 defer family_artifact.deinit();
948
949 const family_entry = family_artifact.entry();
950 try std.testing.expectEqualStrings("accy.kernel.compaction.filter_family_greater_128_f32", family_entry.target);
951 try std.testing.expectEqual(@as(u32, 4), family_entry.argument_count);
952 try std.testing.expectEqual(@as(u32, 2), family_entry.runtime_scalar_argument_count);
953 }
954
955 test "compaction filter thread candidates stay bounded and lead with the default" {
956 const candidates = filterThreadCandidatesForExtent(100_000);
957 try std.testing.expect(candidates.count > 2);
958 try std.testing.expectEqual(filterThreadsForExtent(100_000), candidates.items[0]);
959 for (candidates.slice(), 0..) |candidate, index| {
960 try std.testing.expect(candidate != 0);
961 try std.testing.expect(candidate % filter_warp_size == 0);
962 for (candidates.slice()[0..index]) |previous| try std.testing.expect(previous != candidate);
963 }
964 }