lib/accy/src/kernel/library/indexing.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 geometry_mod = @import("geometry.zig");
10 const kernel = @import("../root.zig");
11 const tuning = @import("tuning.zig");
12
13 const DType = choir_abi.DType;
14 const indexExtent = extent_mod.indexExtent;
15 const runtimeExtentArgument = extent_mod.runtimeExtentArgument;
16
17 pub const Gather = struct {
18 outer: u64 = 1,
19 axis_size: u64,
20 gathered: u64,
21 inner: u64 = 1,
22 dtype: DType = .f32,
23 index_dtype: DType = .i32,
24 threads: u32 = 256,
25 outer_axis: []const u8 = "o",
26 source_axis: []const u8 = "s",
27 gathered_axis: []const u8 = "g",
28 inner_axis: []const u8 = "i",
29
30 pub fn total(self: Gather) u64 {
31 return self.outer * self.gathered * self.inner;
32 }
33 };
34
35 pub const gather_family_version: u32 = 1;
36 const gather_thread_caps = geometry_mod.ThreadCaps1D{};
37
38 pub fn gatherDTypeSupported(dtype: DType) bool {
39 return switch (dtype) {
40 .f32, .f16 => true,
41 else => false,
42 };
43 }
44
45 fn gatherSourceIndexValue(
46 inner_builder: anytype,
47 indices: anytype,
48 element: kernel.Value,
49 axis_size: kernel.Value,
50 gathered: kernel.Value,
51 inner_extent: kernel.Value,
52 ) !kernel.Value {
53 const gathered_inner = try inner_builder.mul(gathered, inner_extent);
54 const outer_coord = try inner_builder.div(element, gathered_inner);
55 const outer_consumed = try inner_builder.mul(outer_coord, gathered_inner);
56 const rem = try inner_builder.sub(element, outer_consumed);
57 const position = try inner_builder.div(rem, inner_extent);
58 const position_consumed = try inner_builder.mul(position, inner_extent);
59 const within = try inner_builder.sub(rem, position_consumed);
60
61 const loaded = try indices.load(inner_builder, position);
62 const zero_i32 = try inner_builder.constantInt(.i32, 0);
63 const one_i32 = try inner_builder.constantInt(.i32, 1);
64 const axis_size_i32 = try inner_builder.cast(axis_size, .i32);
65 const limit_i32 = try inner_builder.sub(axis_size_i32, one_i32);
66 const lower_clamped_i32 = try inner_builder.max(loaded.raw(), zero_i32);
67 const clamped_i32 = try inner_builder.min(lower_clamped_i32, limit_i32);
68 const clamped = try inner_builder.castIndex(clamped_i32);
69
70 const axis_block = try inner_builder.mul(axis_size, inner_extent);
71 const outer_offset = try inner_builder.mul(outer_coord, axis_block);
72 const gathered_offset = try inner_builder.mul(clamped, inner_extent);
73 const partial = try inner_builder.add(outer_offset, gathered_offset);
74 return inner_builder.add(partial, within);
75 }
76
77 fn gather_body_each(inner_builder: anytype, index: kernel.Index1D, ctx: anytype) !void {
78 const axis_size = try inner_builder.constantIndex(try indexExtent(ctx.spec.axis_size));
79 const gathered = try inner_builder.constantIndex(try indexExtent(ctx.spec.gathered));
80 const inner_extent = try inner_builder.constantIndex(try indexExtent(ctx.spec.inner));
81 const src = try gatherSourceIndexValue(
82 inner_builder,
83 ctx.args.param(.indices),
84 index.index,
85 axis_size,
86 gathered,
87 inner_extent,
88 );
89 const value = try ctx.args.param(.data).load(inner_builder, src);
90 try ctx.args.param(.dst).store(inner_builder, value.raw(), index);
91 }
92
93 fn gatherBody(k: anytype, spec: Gather, args: anytype) !void {
94 _ = try k.forEach1D(spec.gathered_axis, spec.total(), .{ .spec = spec, .args = args }, gather_body_each);
95 }
96
97 fn gather_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
98 const src = try gatherSourceIndexValue(
99 inner_builder,
100 ctx.args.param(.indices),
101 ctx.element,
102 ctx.axis_size,
103 ctx.gathered,
104 ctx.inner_extent,
105 );
106 const value = try ctx.args.param(.data).load(inner_builder, src);
107 try ctx.args.param(.dst).store(inner_builder, value.raw(), ctx.element);
108 }
109
110 fn gatherRuntimeBody(k: anytype, spec: Gather, args: anytype) !void {
111 _ = spec;
112 const element = try k.globalId(.x);
113 const axis_size = try k.castIndex(args.param(.axis_size).raw());
114 const gathered = try k.castIndex(args.param(.gathered).raw());
115 const inner_extent = try k.castIndex(args.param(.inner).raw());
116 const total = try k.castIndex(args.param(.total).raw());
117 const active = try k.compare(.lt, element, total);
118 try k.guardDo(active, .{
119 .args = args,
120 .element = element,
121 .axis_size = axis_size,
122 .gathered = gathered,
123 .inner_extent = inner_extent,
124 }, gather_runtime_body_active);
125 }
126
127 fn gatherFamilySchedule(instance: Gather) kernel.logical.schedule.ThreadBlocks {
128 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
129 }
130
131 fn gatherFamily(comptime dtype: DType) type {
132 return kernel.logical.Family(.{
133 .name = std.fmt.comptimePrint("accy_kernel_indexing_gather_{s}", .{dtype.name()}),
134 .parameters = .{
135 .dst = kernel.dynamicBuffer(dtype),
136 .data = kernel.dynamicBuffer(dtype),
137 .indices = kernel.dynamicBuffer(.i32),
138 },
139 .Instance = Gather,
140 .schedule = gatherFamilySchedule,
141 .body = gatherBody,
142 });
143 }
144
145 fn gatherRuntimeFamily(comptime dtype: DType) type {
146 return kernel.logical.Family(.{
147 .name = std.fmt.comptimePrint("accy_kernel_indexing_gather_runtime_{s}", .{dtype.name()}),
148 .parameters = .{
149 .dst = kernel.dynamicBuffer(dtype),
150 .data = kernel.dynamicBuffer(dtype),
151 .indices = kernel.dynamicBuffer(.i32),
152 .outer = kernel.scalar(.i32),
153 .axis_size = kernel.scalar(.i32),
154 .gathered = kernel.scalar(.i32),
155 .inner = kernel.scalar(.i32),
156 .total = kernel.scalar(.i32),
157 },
158 .Instance = Gather,
159 .schedule = gatherFamilySchedule,
160 .body = gatherRuntimeBody,
161 });
162 }
163
164 pub const GatherFamilyF32 = gatherFamily(.f32);
165 pub const GatherFamilyF16 = gatherFamily(.f16);
166 pub const GatherRuntimeFamilyF32 = gatherRuntimeFamily(.f32);
167 pub const GatherRuntimeFamilyF16 = gatherRuntimeFamily(.f16);
168
169 pub fn gatherThreadsForTotal(total: u64) u32 {
170 return geometry_mod.threadsForExtent(total, gather_thread_caps);
171 }
172
173 pub fn gatherThreadCandidatesForTotal(total: u64) geometry_mod.Thread1DCandidates {
174 return geometry_mod.threadCandidatesForExtent(total, gather_thread_caps);
175 }
176
177 pub fn gatherInstanceTarget(allocator: std.mem.Allocator, instance: Gather) ![]u8 {
178 return std.fmt.allocPrint(
179 allocator,
180 "accy.kernel.indexing.gather{d}x{d}x{d}x{d}_{d}_{s}",
181 .{ instance.outer, instance.axis_size, instance.gathered, instance.inner, instance.threads, instance.dtype.name() },
182 );
183 }
184
185 pub fn gatherInstanceEntryName(allocator: std.mem.Allocator, instance: Gather) ![]u8 {
186 return std.fmt.allocPrint(
187 allocator,
188 "accy_kernel_indexing_gather{d}x{d}x{d}x{d}_{d}_{s}",
189 .{ instance.outer, instance.axis_size, instance.gathered, instance.inner, instance.threads, instance.dtype.name() },
190 );
191 }
192
193 pub fn gatherFamilyTarget(allocator: std.mem.Allocator, instance: Gather) ![]u8 {
194 return std.fmt.allocPrint(
195 allocator,
196 "accy.kernel.indexing.gather_family_{d}_{s}",
197 .{ instance.threads, instance.dtype.name() },
198 );
199 }
200
201 pub fn gatherFamilyEntryName(allocator: std.mem.Allocator, instance: Gather) ![]u8 {
202 return std.fmt.allocPrint(
203 allocator,
204 "accy_kernel_indexing_gather_family_{d}_{s}",
205 .{ instance.threads, instance.dtype.name() },
206 );
207 }
208
209 pub fn gatherTuningExtents(instance: Gather) [4]u64 {
210 return .{ instance.outer, instance.axis_size, instance.gathered, instance.inner };
211 }
212
213 pub fn gatherTuningOperation(instance: Gather) entry.Operation {
214 _ = instance;
215 return .{ .indexing = .gather };
216 }
217
218 pub fn gatherFamilyTuningKey(
219 backing_allocator: std.mem.Allocator,
220 device_fingerprint: u64,
221 instance: Gather,
222 ) !tuning.FamilyTuningKey {
223 const family_fingerprint = try gatherFamilyFingerprint(backing_allocator, instance);
224 const extents = gatherTuningExtents(instance);
225 return tuning.FamilyTuningKey.init(
226 device_fingerprint,
227 family_fingerprint,
228 entry.operationFingerprint(gatherTuningOperation(instance)),
229 instance.dtype,
230 gather_family_version,
231 extents[0..],
232 ) orelse unreachable;
233 }
234
235 pub fn resolveGatherSchedule(
236 backing_allocator: std.mem.Allocator,
237 reader: tuning.FamilyTuningReader,
238 instance: Gather,
239 ) !?u32 {
240 const key = try gatherFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
241 const record = reader.table.find(key) orelse return null;
242 const thread_candidates = gatherThreadCandidatesForTotal(instance.total());
243 for (thread_candidates.slice()) |threads| {
244 var candidate = instance;
245 candidate.threads = threads;
246 const target = try gatherFamilyTarget(backing_allocator, candidate);
247 defer backing_allocator.free(target);
248 if (std.mem.eql(u8, target, record.target)) return threads;
249 }
250 return null;
251 }
252
253 pub fn gatherRuntimeArguments(instance: Gather) ![5]choir_abi.ScalarArgument {
254 return .{
255 .{ .u32 = try runtimeExtentArgument(instance.outer) },
256 .{ .u32 = try runtimeExtentArgument(instance.axis_size) },
257 .{ .u32 = try runtimeExtentArgument(instance.gathered) },
258 .{ .u32 = try runtimeExtentArgument(instance.inner) },
259 .{ .u32 = try runtimeExtentArgument(instance.total()) },
260 };
261 }
262
263 pub fn gatherShapeProfileDimensions(instance: Gather) [5]artifact_product.KernelCallShapeProfileDimension {
264 const bounds = gatherRuntimeExtentBounds();
265 return .{
266 .{ .name = instance.outer_axis, .runtime_scalar_argument_index = 0, .bounds = bounds },
267 .{ .name = instance.source_axis, .runtime_scalar_argument_index = 1, .bounds = bounds },
268 .{ .name = instance.gathered_axis, .runtime_scalar_argument_index = 2, .bounds = bounds },
269 .{ .name = instance.inner_axis, .runtime_scalar_argument_index = 3, .bounds = bounds },
270 .{ .name = "e", .runtime_scalar_argument_index = 4, .bounds = bounds },
271 };
272 }
273
274 fn gatherRuntimeExtentBounds() shape.Bounds {
275 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
276 }
277
278 fn gatherDerivedLaunch(instance: Gather) !artifact_product.KernelCallLaunch {
279 if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
280 return .{ .derived = .{
281 .grid = .{
282 .{ .runtime_u32_ceil_div = .{ .argument_index = 4, .divisor = instance.threads } },
283 .{ .fixed = 1 },
284 .{ .fixed = 1 },
285 },
286 .threadgroup = .{ instance.threads, 1, 1 },
287 } };
288 }
289
290 pub fn createGatherFamilyArtifact(
291 allocator: std.mem.Allocator,
292 handle: kernel.BackendHandle,
293 instance: Gather,
294 options: entry.ArtifactOptions,
295 ) !kernel.OwnedKernelCallArtifact {
296 const target = try gatherFamilyTarget(allocator, instance);
297 defer allocator.free(target);
298 const entry_name = try gatherFamilyEntryName(allocator, instance);
299 defer allocator.free(entry_name);
300 const family_fingerprint = options.shape_family_fingerprint orelse try gatherFamilyFingerprint(allocator, instance);
301 const shape_profile_dimensions = gatherShapeProfileDimensions(instance);
302 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
303 .name = "gather",
304 .fingerprint = family_fingerprint,
305 .dimensions = shape_profile_dimensions[0..],
306 };
307
308 var graph = switch (instance.dtype) {
309 .f32 => try GatherRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
310 .f16 => try GatherRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
311 else => return error.UnsupportedDType,
312 };
313 defer graph.deinit();
314 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
315 .target = target,
316 .version = gather_family_version,
317 .format = options.format,
318 .kernel_plan = options.kernel_plan,
319 .element_count_argument = options.element_count_argument,
320 .shape_family_fingerprint = family_fingerprint,
321 .shape_profile = shape_profile,
322 .launch = options.launch orelse try gatherDerivedLaunch(instance),
323 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 5 else options.runtime_scalar_argument_count,
324 .static_arguments = options.static_arguments,
325 });
326 }
327
328 pub fn gatherFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Gather) !u64 {
329 var family = try gatherShapeFamily(backing_allocator, instance);
330 defer family.deinit();
331 return shape.fingerprint(family);
332 }
333
334 pub fn gatherShapeFamily(backing_allocator: std.mem.Allocator, instance: Gather) !shape.Family {
335 var builder = try shape.Builder.init(backing_allocator, "gather");
336 errdefer builder.deinit();
337
338 const outer = try builder.symbol(instance.outer_axis);
339 const source = try builder.symbol(instance.source_axis);
340 const gathered = try builder.symbol(instance.gathered_axis);
341 const inner = try builder.symbol(instance.inner_axis);
342
343 const outer_expr = try builder.symbolExpression(outer);
344 const source_expr = try builder.symbolExpression(source);
345 const gathered_expr = try builder.symbolExpression(gathered);
346 const inner_expr = try builder.symbolExpression(inner);
347
348 _ = try builder.tensor("data", &.{ outer_expr, source_expr, inner_expr });
349 _ = try builder.tensor("indices", &.{gathered_expr});
350 _ = try builder.tensor("out", &.{ outer_expr, gathered_expr, inner_expr });
351 try builder.assumeBounds(outer_expr, gatherRuntimeExtentBounds());
352 try builder.assumeBounds(source_expr, gatherRuntimeExtentBounds());
353 try builder.assumeBounds(gathered_expr, gatherRuntimeExtentBounds());
354 try builder.assumeBounds(inner_expr, gatherRuntimeExtentBounds());
355
356 return builder.finish();
357 }
358
359 pub fn gatherFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Gather) !entry.OwnedSpecialization {
360 var owned = entry.OwnedSpecialization.init(backing_allocator);
361 errdefer owned.deinit();
362 const lifetime_allocator = owned.allocator();
363
364 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
365 inputs[0] = try entry.runtimeShape3D(
366 lifetime_allocator,
367 instance.outer_axis,
368 instance.outer,
369 instance.source_axis,
370 instance.axis_size,
371 instance.inner_axis,
372 instance.inner,
373 );
374 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.gathered_axis, instance.gathered);
375
376 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
377 outputs[0] = try entry.runtimeShape3D(
378 lifetime_allocator,
379 instance.outer_axis,
380 instance.outer,
381 instance.gathered_axis,
382 instance.gathered,
383 instance.inner_axis,
384 instance.inner,
385 );
386
387 owned.value = .{
388 .dtype = instance.dtype,
389 .operation = .{ .indexing = .gather },
390 .inputs = inputs,
391 .outputs = outputs,
392 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.total(), instance.threads),
393 };
394 owned.value.launch = owned.value.schedule.?.launch();
395 var family = try gatherShapeFamily(backing_allocator, instance);
396 errdefer family.deinit();
397 try owned.takeShapeFamily(&family);
398 return owned;
399 }
400
401 pub fn gatherInstanceFromSpecialization(specialization: entry.Specialization) ?Gather {
402 if (!specialization.scheduleMatchesLaunch()) return null;
403 if (!specialization.operationIs(.{ .indexing = .gather })) return null;
404 const dtype = specialization.dtype orelse return null;
405 if (!gatherDTypeSupported(dtype)) return null;
406 if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return null;
407 if (specialization.reductions.len != 0) return null;
408 const data = specialization.inputs[0];
409 const indices = specialization.inputs[1];
410 const output = specialization.outputs[0];
411 if (data.axes.len != 3 or indices.axes.len != 1 or output.axes.len != 3) return null;
412 const outer = data.axes[0].extent;
413 const axis_size = data.axes[1].extent;
414 const inner = data.axes[2].extent;
415 const gathered = indices.axes[0].extent;
416 if (output.axes[0].extent != outer or output.axes[1].extent != gathered or output.axes[2].extent != inner) return null;
417 if (!std.mem.eql(u8, data.axes[0].name, output.axes[0].name)) return null;
418 if (!std.mem.eql(u8, indices.axes[0].name, output.axes[1].name)) return null;
419 if (!std.mem.eql(u8, data.axes[2].name, output.axes[2].name)) return null;
420 const launch = specialization.launch orelse return null;
421 if (launch.threadgroup[0] == 0) return null;
422 return .{
423 .outer = outer,
424 .axis_size = axis_size,
425 .gathered = gathered,
426 .inner = inner,
427 .dtype = dtype,
428 .threads = launch.threadgroup[0],
429 .outer_axis = data.axes[0].name,
430 .source_axis = data.axes[1].name,
431 .gathered_axis = indices.axes[0].name,
432 .inner_axis = data.axes[2].name,
433 };
434 }
435
436 fn gatherSpecialization(comptime spec: Gather) entry.Specialization {
437 return .{
438 .dtype = spec.dtype,
439 .operation = .{ .indexing = .gather },
440 .inputs = &.{
441 entry.shape3D(spec.outer_axis, spec.outer, spec.source_axis, spec.axis_size, spec.inner_axis, spec.inner),
442 entry.shape1D(spec.gathered_axis, spec.gathered),
443 },
444 .outputs = &.{entry.shape3D(spec.outer_axis, spec.outer, spec.gathered_axis, spec.gathered, spec.inner_axis, spec.inner)},
445 .launch = entry.launch1D(ceilDivComptime(spec.total(), spec.threads), spec.threads),
446 .schedule = entry.threadBlocks1D("e", spec.total(), spec.threads),
447 };
448 }
449
450 fn ceilDivComptime(comptime numerator: u64, comptime denominator: u32) u32 {
451 return @intCast(numerator / denominator + @as(u64, @intFromBool(numerator % denominator != 0)));
452 }
453
454 fn gatherProgram(comptime spec: Gather) type {
455 const Body = struct {
456 fn run(k: anytype, args: anytype) !void {
457 try gatherBody(k, spec, args);
458 }
459 };
460
461 return kernel.logical.Program(.{
462 .name = std.fmt.comptimePrint(
463 "accy_kernel_indexing_gather{}x{}x{}x{}_{}_{s}",
464 .{ spec.outer, spec.axis_size, spec.gathered, spec.inner, spec.threads, spec.dtype.name() },
465 ),
466 .parameters = .{
467 .dst = kernel.dynamicBuffer(spec.dtype),
468 .data = kernel.dynamicBuffer(spec.dtype),
469 .indices = kernel.dynamicBuffer(.i32),
470 },
471 .body = Body.run,
472 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
473 }
474
475 pub fn gatherF32(comptime spec: Gather) type {
476 return entry.Entry(gatherProgram(spec), .{
477 .target = std.fmt.comptimePrint(
478 "accy.kernel.indexing.gather{}x{}x{}x{}_{}_{s}",
479 .{ spec.outer, spec.axis_size, spec.gathered, spec.inner, spec.threads, spec.dtype.name() },
480 ),
481 .layer = .logical,
482 .category = .indexing,
483 .specialization = gatherSpecialization(spec),
484 });
485 }
486
487 pub const Gather8F32 = gatherF32(.{ .axis_size = 8, .gathered = 8, .threads = 8 });
488
489 test "indexing gather entry runs on CPU with clamped indices" {
490 var data = [_]f32{ 10, 11, 12, 13, 14, 15, 16, 17 };
491 var indices = [_]i32{ 3, 0, 7, 2, -1, 9, 5, 1 };
492 var dst = @as([8]f32, @splat(0));
493
494 try Gather8F32.runCpu(std.testing.allocator, Gather8F32.Limits.testing, &.{
495 kernel.argumentBuffer(f32, dst[0..]),
496 kernel.argumentBuffer(f32, data[0..]),
497 kernel.argumentBuffer(i32, indices[0..]),
498 });
499 try std.testing.expectEqualSlices(f32, &.{ 13, 10, 17, 12, 10, 17, 15, 11 }, dst[0..]);
500 }
501
502 test "indexing gather runtime family executes explicit runtime extents" {
503 const allocator = std.testing.allocator;
504 const compiled = Gather{ .axis_size = 1, .gathered = 1, .threads = 4 };
505 const runtime = Gather{ .outer = 2, .axis_size = 4, .gathered = 3, .inner = 2, .threads = 4 };
506
507 var graph = try GatherRuntimeFamilyF32.build(allocator, GatherRuntimeFamilyF32.Limits.testing, compiled);
508 defer graph.deinit();
509
510 var data: [16]f32 = undefined;
511 for (&data, 0..) |*value, index| value.* = @floatFromInt(index);
512 var indices = [_]i32{ 2, 0, 3 };
513 var dst = @as([12]f32, @splat(0));
514
515 var expected: [12]f32 = undefined;
516 for (0..2) |outer| {
517 for (0..3) |position| {
518 const clamped: usize = @intCast(@max(@min(indices[position], 3), 0));
519 for (0..2) |within| {
520 expected[outer * 6 + position * 2 + within] = data[outer * 8 + clamped * 2 + within];
521 }
522 }
523 }
524
525 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
526 try graph.runCpuWithLaunch(allocator, &.{
527 kernel.argumentBuffer(f32, dst[0..]),
528 kernel.argumentBuffer(f32, data[0..]),
529 kernel.argumentBuffer(i32, indices[0..]),
530 kernel.argumentI32(@intCast(runtime.outer)),
531 kernel.argumentI32(@intCast(runtime.axis_size)),
532 kernel.argumentI32(@intCast(runtime.gathered)),
533 kernel.argumentI32(@intCast(runtime.inner)),
534 kernel.argumentI32(@intCast(runtime.total())),
535 }, .{
536 .grid = launch_value.grid,
537 .block = launch_value.threadgroup,
538 });
539 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
540 }
541
542 test "indexing gather family instance identity matches fixed entry strings" {
543 const instance = Gather{ .axis_size = 8, .gathered = 8, .threads = 8 };
544
545 const target = try gatherInstanceTarget(std.testing.allocator, instance);
546 defer std.testing.allocator.free(target);
547 try std.testing.expectEqualStrings(Gather8F32.target, target);
548
549 const entry_name = try gatherInstanceEntryName(std.testing.allocator, instance);
550 defer std.testing.allocator.free(entry_name);
551 try std.testing.expectEqualStrings(Gather8F32.name, entry_name);
552
553 try std.testing.expectEqual(Gather8F32.version, gather_family_version);
554
555 const fresh = Gather{ .outer = 4, .axis_size = 1024, .gathered = 256, .inner = 8, .threads = 128 };
556 const family_target = try gatherFamilyTarget(std.testing.allocator, fresh);
557 defer std.testing.allocator.free(family_target);
558 try std.testing.expectEqualStrings("accy.kernel.indexing.gather_family_128_f32", family_target);
559
560 const family_entry = try gatherFamilyEntryName(std.testing.allocator, fresh);
561 defer std.testing.allocator.free(family_entry);
562 try std.testing.expectEqualStrings("accy_kernel_indexing_gather_family_128_f32", family_entry);
563 }
564
565 test "indexing gather family artifact carries runtime launch contract" {
566 const allocator = std.testing.allocator;
567 var state = gpu.recording.BackendState{
568 .allocator = allocator,
569 .kind = .cuda,
570 .format = .cuda_ptx,
571 };
572 const instance = Gather{ .axis_size = 8, .gathered = 8, .threads = 8 };
573
574 var family_artifact = try createGatherFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
575 defer family_artifact.deinit();
576
577 const family_entry = family_artifact.entry();
578 try std.testing.expectEqualStrings("accy.kernel.indexing.gather_family_8_f32", family_entry.target);
579 try std.testing.expectEqualStrings("accy_kernel_indexing_gather_family_8_f32", family_entry.entry_name);
580 try std.testing.expectEqual(@as(u32, 8), family_entry.argument_count);
581 try std.testing.expectEqual(@as(u32, 5), family_entry.runtime_scalar_argument_count);
582 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
583 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
584 try std.testing.expect(family_entry.shape_family_fingerprint != null);
585 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
586 try std.testing.expectEqualStrings("gather", profile.name);
587 try std.testing.expectEqual(@as(usize, 5), profile.dimensions.len);
588 switch (family_entry.launch) {
589 .derived => |launch| {
590 try std.testing.expectEqual(@as(u32, 8), launch.threadgroup[0]);
591 switch (launch.grid[0]) {
592 .runtime_u32_ceil_div => |term| {
593 try std.testing.expectEqual(@as(usize, 4), term.argument_index);
594 try std.testing.expectEqual(@as(u32, 8), term.divisor);
595 },
596 else => return error.TestExpectedDerivedLaunch,
597 }
598 },
599 else => return error.TestExpectedDerivedLaunch,
600 }
601 }
602
603 test "indexing gather instance round-trips through specialization" {
604 const instance = Gather{ .outer = 2, .axis_size = 16, .gathered = 5, .inner = 3, .threads = 16 };
605 var owned = try gatherFamilySpecialization(std.testing.allocator, instance);
606 defer owned.deinit();
607
608 const recovered = gatherInstanceFromSpecialization(owned.value) orelse return error.TestExpectedGatherInstance;
609 try std.testing.expectEqual(instance.outer, recovered.outer);
610 try std.testing.expectEqual(instance.axis_size, recovered.axis_size);
611 try std.testing.expectEqual(instance.gathered, recovered.gathered);
612 try std.testing.expectEqual(instance.inner, recovered.inner);
613 try std.testing.expectEqual(instance.dtype, recovered.dtype);
614 try std.testing.expectEqual(instance.threads, recovered.threads);
615
616 try std.testing.expectEqual(@as(?Gather, null), gatherInstanceFromSpecialization(.{}));
617 }
618
619 test "indexing gather thread candidates stay bounded and lead with the default" {
620 const candidates = gatherThreadCandidatesForTotal(100_000);
621 try std.testing.expect(candidates.count > 2);
622 try std.testing.expectEqual(gatherThreadsForTotal(100_000), candidates.items[0]);
623 for (candidates.slice(), 0..) |candidate, index| {
624 try std.testing.expect(candidate != 0);
625 for (candidates.slice()[0..index]) |previous| try std.testing.expect(previous != candidate);
626 }
627 }
628
629 pub const Scatter = struct {
630 outer: u64 = 1,
631 axis_size: u64,
632 updates: u64,
633 inner: u64 = 1,
634 dtype: DType = .f32,
635 index_dtype: DType = .i32,
636 threads: u32 = 256,
637 outer_axis: []const u8 = "o",
638 source_axis: []const u8 = "s",
639 update_axis: []const u8 = "u",
640 inner_axis: []const u8 = "i",
641
642 pub fn total(self: Scatter) u64 {
643 return self.outer * self.axis_size * self.inner;
644 }
645 };
646
647 pub const scatter_family_version: u32 = 1;
648 const scatter_thread_caps = geometry_mod.ThreadCaps1D{};
649
650 pub fn scatterDTypeSupported(dtype: DType) bool {
651 return switch (dtype) {
652 .f32, .f16 => true,
653 else => false,
654 };
655 }
656
657 fn scatter_output_value_apply(fold_builder: anytype, update_position: kernel.Value, current: kernel.Value, ctx: anytype) !kernel.Value {
658 const loaded = try ctx.args.param(.indices).load(fold_builder, update_position);
659 const index_value = try fold_builder.castIndex(loaded.raw());
660 const matches = try fold_builder.compare(.eq, index_value, ctx.axis_coord);
661 const update_offset = try fold_builder.mul(update_position, ctx.inner_extent);
662 const update_partial = try fold_builder.add(ctx.update_base, update_offset);
663 const update_index = try fold_builder.add(update_partial, ctx.within);
664 const candidate = try ctx.args.param(.updates).load(fold_builder, update_index);
665 return fold_builder.select(matches, candidate.raw(), current);
666 }
667
668 fn scatterOutputValue(
669 inner_builder: anytype,
670 args: anytype,
671 element: kernel.Value,
672 axis_size: kernel.Value,
673 update_count: kernel.Value,
674 inner_extent: kernel.Value,
675 ) !kernel.Value {
676 const axis_block = try inner_builder.mul(axis_size, inner_extent);
677 const outer_coord = try inner_builder.div(element, axis_block);
678 const outer_consumed = try inner_builder.mul(outer_coord, axis_block);
679 const axis_rem = try inner_builder.sub(element, outer_consumed);
680 const axis_coord = try inner_builder.div(axis_rem, inner_extent);
681 const axis_consumed = try inner_builder.mul(axis_coord, inner_extent);
682 const within = try inner_builder.sub(axis_rem, axis_consumed);
683
684 const update_block = try inner_builder.mul(update_count, inner_extent);
685 const update_base = try inner_builder.mul(outer_coord, update_block);
686
687 const zero = try inner_builder.constantIndex(0);
688 const one = try inner_builder.constantIndex(1);
689 const initial = try args.param(.data).load(inner_builder, element);
690
691 return inner_builder.fold(zero, update_count, one, initial.raw(), .{
692 .args = args,
693 .axis_coord = axis_coord,
694 .within = within,
695 .update_base = update_base,
696 .inner_extent = inner_extent,
697 }, scatter_output_value_apply);
698 }
699
700 fn scatter_body_each(inner_builder: anytype, index: kernel.Index1D, ctx: anytype) !void {
701 const axis_size = try inner_builder.constantIndex(try indexExtent(ctx.spec.axis_size));
702 const update_count = try inner_builder.constantIndex(try indexExtent(ctx.spec.updates));
703 const inner_extent = try inner_builder.constantIndex(try indexExtent(ctx.spec.inner));
704 const value = try scatterOutputValue(
705 inner_builder,
706 ctx.args,
707 index.index,
708 axis_size,
709 update_count,
710 inner_extent,
711 );
712 try ctx.args.param(.dst).store(inner_builder, value, index);
713 }
714
715 fn scatterBody(k: anytype, spec: Scatter, args: anytype) !void {
716 _ = try k.forEach1D(spec.source_axis, spec.total(), .{ .spec = spec, .args = args }, scatter_body_each);
717 }
718
719 fn scatter_runtime_body_active(inner_builder: anytype, ctx: anytype) !void {
720 const value = try scatterOutputValue(
721 inner_builder,
722 ctx.args,
723 ctx.element,
724 ctx.axis_size,
725 ctx.update_count,
726 ctx.inner_extent,
727 );
728 try ctx.args.param(.dst).store(inner_builder, value, ctx.element);
729 }
730
731 fn scatterRuntimeBody(k: anytype, spec: Scatter, args: anytype) !void {
732 _ = spec;
733 const element = try k.globalId(.x);
734 const axis_size = try k.castIndex(args.param(.axis_size).raw());
735 const update_count = try k.castIndex(args.param(.update_count).raw());
736 const inner_extent = try k.castIndex(args.param(.inner).raw());
737 const total = try k.castIndex(args.param(.total).raw());
738 const active = try k.compare(.lt, element, total);
739 try k.guardDo(active, .{
740 .args = args,
741 .element = element,
742 .axis_size = axis_size,
743 .update_count = update_count,
744 .inner_extent = inner_extent,
745 }, scatter_runtime_body_active);
746 }
747
748 fn scatterFamilySchedule(instance: Scatter) kernel.logical.schedule.ThreadBlocks {
749 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
750 }
751
752 fn scatterFamily(comptime dtype: DType) type {
753 return kernel.logical.Family(.{
754 .name = std.fmt.comptimePrint("accy_kernel_indexing_scatter_{s}", .{dtype.name()}),
755 .parameters = .{
756 .dst = kernel.dynamicBuffer(dtype),
757 .data = kernel.dynamicBuffer(dtype),
758 .indices = kernel.dynamicBuffer(.i32),
759 .updates = kernel.dynamicBuffer(dtype),
760 },
761 .Instance = Scatter,
762 .schedule = scatterFamilySchedule,
763 .body = scatterBody,
764 });
765 }
766
767 fn scatterRuntimeFamily(comptime dtype: DType) type {
768 return kernel.logical.Family(.{
769 .name = std.fmt.comptimePrint("accy_kernel_indexing_scatter_runtime_{s}", .{dtype.name()}),
770 .parameters = .{
771 .dst = kernel.dynamicBuffer(dtype),
772 .data = kernel.dynamicBuffer(dtype),
773 .indices = kernel.dynamicBuffer(.i32),
774 .updates = kernel.dynamicBuffer(dtype),
775 .outer = kernel.scalar(.i32),
776 .axis_size = kernel.scalar(.i32),
777 .update_count = kernel.scalar(.i32),
778 .inner = kernel.scalar(.i32),
779 .total = kernel.scalar(.i32),
780 },
781 .Instance = Scatter,
782 .schedule = scatterFamilySchedule,
783 .body = scatterRuntimeBody,
784 });
785 }
786
787 pub const ScatterFamilyF32 = scatterFamily(.f32);
788 pub const ScatterFamilyF16 = scatterFamily(.f16);
789 pub const ScatterRuntimeFamilyF32 = scatterRuntimeFamily(.f32);
790 pub const ScatterRuntimeFamilyF16 = scatterRuntimeFamily(.f16);
791
792 pub fn scatterThreadsForTotal(total: u64) u32 {
793 return geometry_mod.threadsForExtent(total, scatter_thread_caps);
794 }
795
796 pub fn scatterThreadCandidatesForTotal(total: u64) geometry_mod.Thread1DCandidates {
797 return geometry_mod.threadCandidatesForExtent(total, scatter_thread_caps);
798 }
799
800 pub fn scatterAddThreadsForTotal(total: u64) u32 {
801 return geometry_mod.threadsForExtent(total, scatter_thread_caps);
802 }
803
804 pub fn scatterAddThreadCandidatesForTotal(total: u64) geometry_mod.Thread1DCandidates {
805 return geometry_mod.threadCandidatesForExtent(total, scatter_thread_caps);
806 }
807
808 pub fn scatterInstanceTarget(allocator: std.mem.Allocator, instance: Scatter) ![]u8 {
809 return std.fmt.allocPrint(
810 allocator,
811 "accy.kernel.indexing.scatter{d}x{d}x{d}x{d}_{d}_{s}",
812 .{ instance.outer, instance.axis_size, instance.updates, instance.inner, instance.threads, instance.dtype.name() },
813 );
814 }
815
816 pub fn scatterInstanceEntryName(allocator: std.mem.Allocator, instance: Scatter) ![]u8 {
817 return std.fmt.allocPrint(
818 allocator,
819 "accy_kernel_indexing_scatter{d}x{d}x{d}x{d}_{d}_{s}",
820 .{ instance.outer, instance.axis_size, instance.updates, instance.inner, instance.threads, instance.dtype.name() },
821 );
822 }
823
824 pub fn scatterFamilyTarget(allocator: std.mem.Allocator, instance: Scatter) ![]u8 {
825 return std.fmt.allocPrint(
826 allocator,
827 "accy.kernel.indexing.scatter_family_{d}_{s}",
828 .{ instance.threads, instance.dtype.name() },
829 );
830 }
831
832 pub fn scatterFamilyEntryName(allocator: std.mem.Allocator, instance: Scatter) ![]u8 {
833 return std.fmt.allocPrint(
834 allocator,
835 "accy_kernel_indexing_scatter_family_{d}_{s}",
836 .{ instance.threads, instance.dtype.name() },
837 );
838 }
839
840 pub fn scatterTuningExtents(instance: Scatter) [4]u64 {
841 return .{ instance.outer, instance.axis_size, instance.updates, instance.inner };
842 }
843
844 pub fn scatterTuningOperation(instance: Scatter) entry.Operation {
845 _ = instance;
846 return .{ .indexing = .scatter };
847 }
848
849 pub fn scatterFamilyTuningKey(
850 backing_allocator: std.mem.Allocator,
851 device_fingerprint: u64,
852 instance: Scatter,
853 ) !tuning.FamilyTuningKey {
854 const family_fingerprint = try scatterFamilyFingerprint(backing_allocator, instance);
855 const extents = scatterTuningExtents(instance);
856 return tuning.FamilyTuningKey.init(
857 device_fingerprint,
858 family_fingerprint,
859 entry.operationFingerprint(scatterTuningOperation(instance)),
860 instance.dtype,
861 scatter_family_version,
862 extents[0..],
863 ) orelse unreachable;
864 }
865
866 pub fn resolveScatterSchedule(
867 backing_allocator: std.mem.Allocator,
868 reader: tuning.FamilyTuningReader,
869 instance: Scatter,
870 ) !?u32 {
871 const key = try scatterFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
872 const record = reader.table.find(key) orelse return null;
873 const thread_candidates = scatterThreadCandidatesForTotal(instance.total());
874 for (thread_candidates.slice()) |threads| {
875 var candidate = instance;
876 candidate.threads = threads;
877 const target = try scatterFamilyTarget(backing_allocator, candidate);
878 defer backing_allocator.free(target);
879 if (std.mem.eql(u8, target, record.target)) return threads;
880 }
881 return null;
882 }
883
884 pub fn scatterRuntimeArguments(instance: Scatter) ![5]choir_abi.ScalarArgument {
885 return .{
886 .{ .u32 = try runtimeExtentArgument(instance.outer) },
887 .{ .u32 = try runtimeExtentArgument(instance.axis_size) },
888 .{ .u32 = try runtimeExtentArgument(instance.updates) },
889 .{ .u32 = try runtimeExtentArgument(instance.inner) },
890 .{ .u32 = try runtimeExtentArgument(instance.total()) },
891 };
892 }
893
894 pub fn scatterShapeProfileDimensions(instance: Scatter) [5]artifact_product.KernelCallShapeProfileDimension {
895 const bounds = scatterRuntimeExtentBounds();
896 return .{
897 .{ .name = instance.outer_axis, .runtime_scalar_argument_index = 0, .bounds = bounds },
898 .{ .name = instance.source_axis, .runtime_scalar_argument_index = 1, .bounds = bounds },
899 .{ .name = instance.update_axis, .runtime_scalar_argument_index = 2, .bounds = bounds },
900 .{ .name = instance.inner_axis, .runtime_scalar_argument_index = 3, .bounds = bounds },
901 .{ .name = "e", .runtime_scalar_argument_index = 4, .bounds = bounds },
902 };
903 }
904
905 fn scatterRuntimeExtentBounds() shape.Bounds {
906 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
907 }
908
909 fn scatterDerivedLaunch(instance: Scatter) !artifact_product.KernelCallLaunch {
910 if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
911 return .{ .derived = .{
912 .grid = .{
913 .{ .runtime_u32_ceil_div = .{ .argument_index = 4, .divisor = instance.threads } },
914 .{ .fixed = 1 },
915 .{ .fixed = 1 },
916 },
917 .threadgroup = .{ instance.threads, 1, 1 },
918 } };
919 }
920
921 pub fn createScatterFamilyArtifact(
922 allocator: std.mem.Allocator,
923 handle: kernel.BackendHandle,
924 instance: Scatter,
925 options: entry.ArtifactOptions,
926 ) !kernel.OwnedKernelCallArtifact {
927 const target = try scatterFamilyTarget(allocator, instance);
928 defer allocator.free(target);
929 const entry_name = try scatterFamilyEntryName(allocator, instance);
930 defer allocator.free(entry_name);
931 const family_fingerprint = options.shape_family_fingerprint orelse try scatterFamilyFingerprint(allocator, instance);
932 const shape_profile_dimensions = scatterShapeProfileDimensions(instance);
933 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
934 .name = "scatter",
935 .fingerprint = family_fingerprint,
936 .dimensions = shape_profile_dimensions[0..],
937 };
938
939 var graph = switch (instance.dtype) {
940 .f32 => try ScatterRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
941 .f16 => try ScatterRuntimeFamilyF16.buildNamed(allocator, options.limits, entry_name, instance),
942 else => return error.UnsupportedDType,
943 };
944 defer graph.deinit();
945 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
946 .target = target,
947 .version = scatter_family_version,
948 .format = options.format,
949 .kernel_plan = options.kernel_plan,
950 .element_count_argument = options.element_count_argument,
951 .shape_family_fingerprint = family_fingerprint,
952 .shape_profile = shape_profile,
953 .launch = options.launch orelse try scatterDerivedLaunch(instance),
954 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 5 else options.runtime_scalar_argument_count,
955 .static_arguments = options.static_arguments,
956 });
957 }
958
959 pub fn scatterFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Scatter) !u64 {
960 var family = try scatterShapeFamily(backing_allocator, instance);
961 defer family.deinit();
962 return shape.fingerprint(family);
963 }
964
965 pub fn scatterShapeFamily(backing_allocator: std.mem.Allocator, instance: Scatter) !shape.Family {
966 var builder = try shape.Builder.init(backing_allocator, "scatter");
967 errdefer builder.deinit();
968
969 const outer = try builder.symbol(instance.outer_axis);
970 const source = try builder.symbol(instance.source_axis);
971 const update = try builder.symbol(instance.update_axis);
972 const inner = try builder.symbol(instance.inner_axis);
973
974 const outer_expr = try builder.symbolExpression(outer);
975 const source_expr = try builder.symbolExpression(source);
976 const update_expr = try builder.symbolExpression(update);
977 const inner_expr = try builder.symbolExpression(inner);
978
979 _ = try builder.tensor("data", &.{ outer_expr, source_expr, inner_expr });
980 _ = try builder.tensor("indices", &.{update_expr});
981 _ = try builder.tensor("updates", &.{ outer_expr, update_expr, inner_expr });
982 _ = try builder.tensor("out", &.{ outer_expr, source_expr, inner_expr });
983 try builder.assumeBounds(outer_expr, scatterRuntimeExtentBounds());
984 try builder.assumeBounds(source_expr, scatterRuntimeExtentBounds());
985 try builder.assumeBounds(update_expr, scatterRuntimeExtentBounds());
986 try builder.assumeBounds(inner_expr, scatterRuntimeExtentBounds());
987
988 return builder.finish();
989 }
990
991 pub fn scatterFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Scatter) !entry.OwnedSpecialization {
992 var owned = entry.OwnedSpecialization.init(backing_allocator);
993 errdefer owned.deinit();
994 const lifetime_allocator = owned.allocator();
995
996 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
997 inputs[0] = try entry.runtimeShape3D(
998 lifetime_allocator,
999 instance.outer_axis,
1000 instance.outer,
1001 instance.source_axis,
1002 instance.axis_size,
1003 instance.inner_axis,
1004 instance.inner,
1005 );
1006 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.update_axis, instance.updates);
1007 inputs[2] = try entry.runtimeShape3D(
1008 lifetime_allocator,
1009 instance.outer_axis,
1010 instance.outer,
1011 instance.update_axis,
1012 instance.updates,
1013 instance.inner_axis,
1014 instance.inner,
1015 );
1016
1017 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1018 outputs[0] = try entry.runtimeShape3D(
1019 lifetime_allocator,
1020 instance.outer_axis,
1021 instance.outer,
1022 instance.source_axis,
1023 instance.axis_size,
1024 instance.inner_axis,
1025 instance.inner,
1026 );
1027
1028 owned.value = .{
1029 .dtype = instance.dtype,
1030 .operation = .{ .indexing = .scatter },
1031 .inputs = inputs,
1032 .outputs = outputs,
1033 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.total(), instance.threads),
1034 };
1035 owned.value.launch = owned.value.schedule.?.launch();
1036 var family = try scatterShapeFamily(backing_allocator, instance);
1037 errdefer family.deinit();
1038 try owned.takeShapeFamily(&family);
1039 return owned;
1040 }
1041
1042 pub fn scatterInstanceFromSpecialization(specialization: entry.Specialization) ?Scatter {
1043 if (!specialization.scheduleMatchesLaunch()) return null;
1044 if (!specialization.operationIs(.{ .indexing = .scatter })) return null;
1045 const dtype = specialization.dtype orelse return null;
1046 if (!scatterDTypeSupported(dtype)) return null;
1047 if (specialization.inputs.len != 3 or specialization.outputs.len != 1) return null;
1048 if (specialization.reductions.len != 0) return null;
1049 const data = specialization.inputs[0];
1050 const indices = specialization.inputs[1];
1051 const update_values = specialization.inputs[2];
1052 const output = specialization.outputs[0];
1053 if (data.axes.len != 3 or indices.axes.len != 1 or update_values.axes.len != 3 or output.axes.len != 3) return null;
1054 const outer = data.axes[0].extent;
1055 const axis_size = data.axes[1].extent;
1056 const inner = data.axes[2].extent;
1057 const updates = indices.axes[0].extent;
1058 if (update_values.axes[0].extent != outer or update_values.axes[1].extent != updates or update_values.axes[2].extent != inner) return null;
1059 if (output.axes[0].extent != outer or output.axes[1].extent != axis_size or output.axes[2].extent != inner) return null;
1060 if (!std.mem.eql(u8, data.axes[1].name, output.axes[1].name)) return null;
1061 if (!std.mem.eql(u8, indices.axes[0].name, update_values.axes[1].name)) return null;
1062 const launch = specialization.launch orelse return null;
1063 if (launch.threadgroup[0] == 0) return null;
1064 return .{
1065 .outer = outer,
1066 .axis_size = axis_size,
1067 .updates = updates,
1068 .inner = inner,
1069 .dtype = dtype,
1070 .threads = launch.threadgroup[0],
1071 .outer_axis = data.axes[0].name,
1072 .source_axis = data.axes[1].name,
1073 .update_axis = indices.axes[0].name,
1074 .inner_axis = data.axes[2].name,
1075 };
1076 }
1077
1078 fn scatterSpecialization(comptime spec: Scatter) entry.Specialization {
1079 return .{
1080 .dtype = spec.dtype,
1081 .operation = .{ .indexing = .scatter },
1082 .inputs = &.{
1083 entry.shape3D(spec.outer_axis, spec.outer, spec.source_axis, spec.axis_size, spec.inner_axis, spec.inner),
1084 entry.shape1D(spec.update_axis, spec.updates),
1085 entry.shape3D(spec.outer_axis, spec.outer, spec.update_axis, spec.updates, spec.inner_axis, spec.inner),
1086 },
1087 .outputs = &.{entry.shape3D(spec.outer_axis, spec.outer, spec.source_axis, spec.axis_size, spec.inner_axis, spec.inner)},
1088 .launch = entry.launch1D(ceilDivComptime(spec.total(), spec.threads), spec.threads),
1089 .schedule = entry.threadBlocks1D("e", spec.total(), spec.threads),
1090 };
1091 }
1092
1093 fn scatterProgram(comptime spec: Scatter) type {
1094 const Body = struct {
1095 fn run(k: anytype, args: anytype) !void {
1096 try scatterBody(k, spec, args);
1097 }
1098 };
1099
1100 return kernel.logical.Program(.{
1101 .name = std.fmt.comptimePrint(
1102 "accy_kernel_indexing_scatter{}x{}x{}x{}_{}_{s}",
1103 .{ spec.outer, spec.axis_size, spec.updates, spec.inner, spec.threads, spec.dtype.name() },
1104 ),
1105 .parameters = .{
1106 .dst = kernel.dynamicBuffer(spec.dtype),
1107 .data = kernel.dynamicBuffer(spec.dtype),
1108 .indices = kernel.dynamicBuffer(.i32),
1109 .updates = kernel.dynamicBuffer(spec.dtype),
1110 },
1111 .body = Body.run,
1112 }).withSchedule(kernel.logical.schedule.threadBlocks(.{ .x = spec.threads }));
1113 }
1114
1115 pub fn scatterF32(comptime spec: Scatter) type {
1116 return entry.Entry(scatterProgram(spec), .{
1117 .target = std.fmt.comptimePrint(
1118 "accy.kernel.indexing.scatter{}x{}x{}x{}_{}_{s}",
1119 .{ spec.outer, spec.axis_size, spec.updates, spec.inner, spec.threads, spec.dtype.name() },
1120 ),
1121 .layer = .logical,
1122 .category = .indexing,
1123 .specialization = scatterSpecialization(spec),
1124 });
1125 }
1126
1127 pub const Scatter8F32 = scatterF32(.{ .axis_size = 8, .updates = 4, .threads = 8 });
1128
1129 test "indexing scatter entry runs on CPU with last match wins" {
1130 var data = [_]f32{ 10, 11, 12, 13, 14, 15, 16, 17 };
1131 var indices = [_]i32{ 3, 0, 3, 9 };
1132 var updates = [_]f32{ 100, 200, 300, 400 };
1133 var dst = @as([8]f32, @splat(0));
1134
1135 try Scatter8F32.runCpu(std.testing.allocator, Scatter8F32.Limits.testing, &.{
1136 kernel.argumentBuffer(f32, dst[0..]),
1137 kernel.argumentBuffer(f32, data[0..]),
1138 kernel.argumentBuffer(i32, indices[0..]),
1139 kernel.argumentBuffer(f32, updates[0..]),
1140 });
1141 try std.testing.expectEqualSlices(f32, &.{ 200, 11, 12, 300, 14, 15, 16, 17 }, dst[0..]);
1142 }
1143
1144 test "indexing scatter runtime family executes explicit runtime extents" {
1145 const allocator = std.testing.allocator;
1146 const compiled = Scatter{ .axis_size = 1, .updates = 1, .threads = 4 };
1147 const runtime = Scatter{ .outer = 2, .axis_size = 4, .updates = 3, .inner = 2, .threads = 4 };
1148
1149 var graph = try ScatterRuntimeFamilyF32.build(allocator, ScatterRuntimeFamilyF32.Limits.testing, compiled);
1150 defer graph.deinit();
1151
1152 var data: [16]f32 = undefined;
1153 for (&data, 0..) |*value, index| value.* = @floatFromInt(index);
1154 var indices = [_]i32{ 2, 0, 2 };
1155 var updates: [12]f32 = undefined;
1156 for (&updates, 0..) |*value, index| value.* = @floatFromInt(100 + index);
1157 var dst = @as([16]f32, @splat(0));
1158
1159 var expected: [16]f32 = undefined;
1160 @memcpy(expected[0..], data[0..]);
1161 for (0..2) |outer| {
1162 for (0..3) |update_position| {
1163 const target_axis: usize = @intCast(indices[update_position]);
1164 for (0..2) |within| {
1165 expected[outer * 8 + target_axis * 2 + within] = updates[outer * 6 + update_position * 2 + within];
1166 }
1167 }
1168 }
1169
1170 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
1171 try graph.runCpuWithLaunch(allocator, &.{
1172 kernel.argumentBuffer(f32, dst[0..]),
1173 kernel.argumentBuffer(f32, data[0..]),
1174 kernel.argumentBuffer(i32, indices[0..]),
1175 kernel.argumentBuffer(f32, updates[0..]),
1176 kernel.argumentI32(@intCast(runtime.outer)),
1177 kernel.argumentI32(@intCast(runtime.axis_size)),
1178 kernel.argumentI32(@intCast(runtime.updates)),
1179 kernel.argumentI32(@intCast(runtime.inner)),
1180 kernel.argumentI32(@intCast(runtime.total())),
1181 }, .{
1182 .grid = launch_value.grid,
1183 .block = launch_value.threadgroup,
1184 });
1185 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
1186 }
1187
1188 test "indexing scatter family instance identity matches fixed entry strings" {
1189 const instance = Scatter{ .axis_size = 8, .updates = 4, .threads = 8 };
1190
1191 const target = try scatterInstanceTarget(std.testing.allocator, instance);
1192 defer std.testing.allocator.free(target);
1193 try std.testing.expectEqualStrings(Scatter8F32.target, target);
1194
1195 const entry_name = try scatterInstanceEntryName(std.testing.allocator, instance);
1196 defer std.testing.allocator.free(entry_name);
1197 try std.testing.expectEqualStrings(Scatter8F32.name, entry_name);
1198
1199 try std.testing.expectEqual(Scatter8F32.version, scatter_family_version);
1200
1201 const fresh = Scatter{ .outer = 4, .axis_size = 1024, .updates = 256, .inner = 8, .threads = 128 };
1202 const family_target = try scatterFamilyTarget(std.testing.allocator, fresh);
1203 defer std.testing.allocator.free(family_target);
1204 try std.testing.expectEqualStrings("accy.kernel.indexing.scatter_family_128_f32", family_target);
1205 }
1206
1207 test "indexing scatter family artifact carries runtime launch contract" {
1208 const allocator = std.testing.allocator;
1209 var state = gpu.recording.BackendState{
1210 .allocator = allocator,
1211 .kind = .cuda,
1212 .format = .cuda_ptx,
1213 };
1214 const instance = Scatter{ .axis_size = 8, .updates = 4, .threads = 8 };
1215
1216 var family_artifact = try createScatterFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
1217 defer family_artifact.deinit();
1218
1219 const family_entry = family_artifact.entry();
1220 try std.testing.expectEqualStrings("accy.kernel.indexing.scatter_family_8_f32", family_entry.target);
1221 try std.testing.expectEqual(@as(u32, 9), family_entry.argument_count);
1222 try std.testing.expectEqual(@as(u32, 5), family_entry.runtime_scalar_argument_count);
1223 try std.testing.expect(family_entry.required_dtypes.contains(.f32));
1224 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
1225 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
1226 try std.testing.expectEqualStrings("scatter", profile.name);
1227 switch (family_entry.launch) {
1228 .derived => |launch| {
1229 try std.testing.expectEqual(@as(u32, 8), launch.threadgroup[0]);
1230 switch (launch.grid[0]) {
1231 .runtime_u32_ceil_div => |term| {
1232 try std.testing.expectEqual(@as(usize, 4), term.argument_index);
1233 try std.testing.expectEqual(@as(u32, 8), term.divisor);
1234 },
1235 else => return error.TestExpectedDerivedLaunch,
1236 }
1237 },
1238 else => return error.TestExpectedDerivedLaunch,
1239 }
1240 }
1241
1242 test "indexing scatter instance round-trips through specialization" {
1243 const instance = Scatter{ .outer = 2, .axis_size = 16, .updates = 5, .inner = 3, .threads = 16 };
1244 var owned = try scatterFamilySpecialization(std.testing.allocator, instance);
1245 defer owned.deinit();
1246
1247 const recovered = scatterInstanceFromSpecialization(owned.value) orelse return error.TestExpectedScatterInstance;
1248 try std.testing.expectEqual(instance.outer, recovered.outer);
1249 try std.testing.expectEqual(instance.axis_size, recovered.axis_size);
1250 try std.testing.expectEqual(instance.updates, recovered.updates);
1251 try std.testing.expectEqual(instance.inner, recovered.inner);
1252 try std.testing.expectEqual(instance.dtype, recovered.dtype);
1253 try std.testing.expectEqual(instance.threads, recovered.threads);
1254
1255 try std.testing.expectEqual(@as(?Scatter, null), scatterInstanceFromSpecialization(.{}));
1256 }
1257
1258 pub const ScatterAddVariant = enum {
1259 direct,
1260 shared_bins,
1261 };
1262
1263 pub const ScatterAdd = struct {
1264 outer: u64 = 1,
1265 axis_size: u64,
1266 updates: u64,
1267 inner: u64 = 1,
1268 dtype: DType = .i32,
1269 variant: ScatterAddVariant = .direct,
1270 threads: u32 = 256,
1271 outer_axis: []const u8 = "o",
1272 source_axis: []const u8 = "s",
1273 update_axis: []const u8 = "u",
1274 inner_axis: []const u8 = "i",
1275
1276 pub fn total(self: ScatterAdd) u64 {
1277 return self.outer * self.updates * self.inner;
1278 }
1279 };
1280
1281 pub const scatter_add_shared_bins_cap: u64 = 4096;
1282
1283 pub const scatter_add_family_version: u32 = 3;
1284
1285 pub const ScatterAddResolvedSchedule = struct {
1286 variant: ScatterAddVariant,
1287 threads: u32,
1288 };
1289
1290 pub fn scatterAddDTypeSupported(dtype: DType) bool {
1291 return switch (dtype) {
1292 .i32, .f32 => true,
1293 else => false,
1294 };
1295 }
1296
1297 pub fn scatterAddInstanceValid(instance: ScatterAdd) bool {
1298 if (!scatterAddDTypeSupported(instance.dtype)) return false;
1299 if (instance.outer == 0 or instance.axis_size == 0 or instance.updates == 0 or instance.inner == 0) return false;
1300 if (instance.variant == .shared_bins and
1301 (instance.axis_size > scatter_add_shared_bins_cap or instance.outer != 1 or instance.inner != 1))
1302 {
1303 return false;
1304 }
1305 return instance.threads != 0;
1306 }
1307
1308 pub fn scatterAddFamilyTarget(allocator: std.mem.Allocator, instance: ScatterAdd) ![]u8 {
1309 return switch (instance.variant) {
1310 .direct => std.fmt.allocPrint(
1311 allocator,
1312 "accy.kernel.indexing.scatter_add_family_{d}_{s}",
1313 .{ instance.threads, instance.dtype.name() },
1314 ),
1315 .shared_bins => std.fmt.allocPrint(
1316 allocator,
1317 "accy.kernel.indexing.scatter_add_family_shared{d}_{d}_{s}",
1318 .{ instance.axis_size, instance.threads, instance.dtype.name() },
1319 ),
1320 };
1321 }
1322
1323 pub fn scatterAddFamilyEntryName(allocator: std.mem.Allocator, instance: ScatterAdd) ![]u8 {
1324 return switch (instance.variant) {
1325 .direct => std.fmt.allocPrint(
1326 allocator,
1327 "accy_kernel_indexing_scatter_add_family_{d}_{s}",
1328 .{ instance.threads, instance.dtype.name() },
1329 ),
1330 .shared_bins => std.fmt.allocPrint(
1331 allocator,
1332 "accy_kernel_indexing_scatter_add_family_shared{d}_{d}_{s}",
1333 .{ instance.axis_size, instance.threads, instance.dtype.name() },
1334 ),
1335 };
1336 }
1337
1338 pub fn scatterAddTuningExtents(instance: ScatterAdd) [4]u64 {
1339 return .{ instance.outer, instance.axis_size, instance.updates, instance.inner };
1340 }
1341
1342 pub fn scatterAddTuningOperation(instance: ScatterAdd) entry.Operation {
1343 _ = instance;
1344 return .{ .indexing = .scatter_add };
1345 }
1346
1347 pub fn scatterAddFamilyTuningKey(
1348 backing_allocator: std.mem.Allocator,
1349 device_fingerprint: u64,
1350 instance: ScatterAdd,
1351 ) !tuning.FamilyTuningKey {
1352 const family_fingerprint = try scatterAddFamilyFingerprint(backing_allocator, instance);
1353 const extents = scatterAddTuningExtents(instance);
1354 return tuning.FamilyTuningKey.init(
1355 device_fingerprint,
1356 family_fingerprint,
1357 entry.operationFingerprint(scatterAddTuningOperation(instance)),
1358 instance.dtype,
1359 scatter_add_family_version,
1360 extents[0..],
1361 ) orelse unreachable;
1362 }
1363
1364 pub fn resolveScatterAddSchedule(
1365 backing_allocator: std.mem.Allocator,
1366 reader: tuning.FamilyTuningReader,
1367 instance: ScatterAdd,
1368 ) !?ScatterAddResolvedSchedule {
1369 const key = try scatterAddFamilyTuningKey(backing_allocator, reader.device_fingerprint, instance);
1370 const record = reader.table.find(key) orelse return null;
1371 const thread_candidates = scatterAddThreadCandidatesForTotal(instance.total());
1372 const variants = [_]ScatterAddVariant{ .direct, .shared_bins };
1373 for (variants) |variant| {
1374 for (thread_candidates.slice()) |threads| {
1375 var candidate = instance;
1376 candidate.variant = variant;
1377 candidate.threads = threads;
1378 if (!scatterAddInstanceValid(candidate)) continue;
1379 const target = try scatterAddFamilyTarget(backing_allocator, candidate);
1380 defer backing_allocator.free(target);
1381 if (std.mem.eql(u8, target, record.target)) {
1382 return .{ .variant = variant, .threads = threads };
1383 }
1384 }
1385 }
1386 return null;
1387 }
1388
1389 pub fn scatterAddRuntimeArguments(instance: ScatterAdd) ![5]choir_abi.ScalarArgument {
1390 return .{
1391 .{ .u32 = try runtimeExtentArgument(instance.outer) },
1392 .{ .u32 = try runtimeExtentArgument(instance.axis_size) },
1393 .{ .u32 = try runtimeExtentArgument(instance.updates) },
1394 .{ .u32 = try runtimeExtentArgument(instance.inner) },
1395 .{ .u32 = try runtimeExtentArgument(instance.total()) },
1396 };
1397 }
1398
1399 pub fn scatterAddShapeProfileDimensions(instance: ScatterAdd) [5]artifact_product.KernelCallShapeProfileDimension {
1400 const bounds = scatterRuntimeExtentBounds();
1401 return .{
1402 .{ .name = instance.outer_axis, .runtime_scalar_argument_index = 0, .bounds = bounds },
1403 .{ .name = instance.source_axis, .runtime_scalar_argument_index = 1, .bounds = bounds },
1404 .{ .name = instance.update_axis, .runtime_scalar_argument_index = 2, .bounds = bounds },
1405 .{ .name = instance.inner_axis, .runtime_scalar_argument_index = 3, .bounds = bounds },
1406 .{ .name = "e", .runtime_scalar_argument_index = 4, .bounds = bounds },
1407 };
1408 }
1409
1410 fn scatterAddDerivedLaunch(instance: ScatterAdd) !artifact_product.KernelCallLaunch {
1411 if (instance.threads == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
1412 return .{ .derived = .{
1413 .grid = .{
1414 .{ .runtime_u32_ceil_div = .{ .argument_index = 4, .divisor = instance.threads } },
1415 .{ .fixed = 1 },
1416 .{ .fixed = 1 },
1417 },
1418 .threadgroup = .{ instance.threads, 1, 1 },
1419 } };
1420 }
1421
1422 pub fn scatterAddShapeFamily(backing_allocator: std.mem.Allocator, instance: ScatterAdd) !shape.Family {
1423 var builder = try shape.Builder.init(backing_allocator, "scatter_add");
1424 errdefer builder.deinit();
1425
1426 const outer = try builder.symbol(instance.outer_axis);
1427 const source = try builder.symbol(instance.source_axis);
1428 const update = try builder.symbol(instance.update_axis);
1429 const inner = try builder.symbol(instance.inner_axis);
1430 const outer_expr = try builder.symbolExpression(outer);
1431 const source_expr = try builder.symbolExpression(source);
1432 const update_expr = try builder.symbolExpression(update);
1433 const inner_expr = try builder.symbolExpression(inner);
1434
1435 _ = try builder.tensor("dst", &.{ outer_expr, source_expr, inner_expr });
1436 _ = try builder.tensor("indices", &.{update_expr});
1437 _ = try builder.tensor("updates", &.{ outer_expr, update_expr, inner_expr });
1438 _ = try builder.tensor("out", &.{ outer_expr, source_expr, inner_expr });
1439 try builder.assumeBounds(outer_expr, scatterRuntimeExtentBounds());
1440 try builder.assumeBounds(source_expr, scatterRuntimeExtentBounds());
1441 try builder.assumeBounds(update_expr, scatterRuntimeExtentBounds());
1442 try builder.assumeBounds(inner_expr, scatterRuntimeExtentBounds());
1443
1444 return builder.finish();
1445 }
1446
1447 pub fn scatterAddFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: ScatterAdd) !u64 {
1448 var family = try scatterAddShapeFamily(backing_allocator, instance);
1449 defer family.deinit();
1450 return shape.fingerprint(family);
1451 }
1452
1453 pub fn scatterAddFamilySpecialization(backing_allocator: std.mem.Allocator, instance: ScatterAdd) !entry.OwnedSpecialization {
1454 var owned = entry.OwnedSpecialization.init(backing_allocator);
1455 errdefer owned.deinit();
1456 const lifetime_allocator = owned.allocator();
1457
1458 const inputs = try lifetime_allocator.alloc(entry.Shape, 3);
1459 inputs[0] = try entry.runtimeShape3D(
1460 lifetime_allocator,
1461 instance.outer_axis,
1462 instance.outer,
1463 instance.source_axis,
1464 instance.axis_size,
1465 instance.inner_axis,
1466 instance.inner,
1467 );
1468 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, instance.update_axis, instance.updates);
1469 inputs[2] = try entry.runtimeShape3D(
1470 lifetime_allocator,
1471 instance.outer_axis,
1472 instance.outer,
1473 instance.update_axis,
1474 instance.updates,
1475 instance.inner_axis,
1476 instance.inner,
1477 );
1478
1479 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
1480 outputs[0] = try entry.runtimeShape3D(
1481 lifetime_allocator,
1482 instance.outer_axis,
1483 instance.outer,
1484 instance.source_axis,
1485 instance.axis_size,
1486 instance.inner_axis,
1487 instance.inner,
1488 );
1489
1490 owned.value = .{
1491 .dtype = instance.dtype,
1492 .operation = .{ .indexing = .scatter_add },
1493 .inputs = inputs,
1494 .outputs = outputs,
1495 .schedule = try entry.runtimeThreadBlocks1D(lifetime_allocator, "e", instance.total(), instance.threads),
1496 .structure = @tagName(instance.variant),
1497 };
1498 owned.value.launch = owned.value.schedule.?.launch();
1499 var family = try scatterAddShapeFamily(backing_allocator, instance);
1500 errdefer family.deinit();
1501 try owned.takeShapeFamily(&family);
1502 return owned;
1503 }
1504
1505 pub fn scatterAddInstanceFromSpecialization(specialization: entry.Specialization) ?ScatterAdd {
1506 if (!specialization.scheduleMatchesLaunch()) return null;
1507 if (!specialization.operationIs(.{ .indexing = .scatter_add })) return null;
1508 const dtype = specialization.dtype orelse return null;
1509 if (!scatterAddDTypeSupported(dtype)) return null;
1510 if (specialization.inputs.len != 3 or specialization.outputs.len != 1) return null;
1511 if (specialization.reductions.len != 0) return null;
1512 const seed = specialization.inputs[0];
1513 const indices = specialization.inputs[1];
1514 const update_values = specialization.inputs[2];
1515 const output = specialization.outputs[0];
1516 if (seed.axes.len != 3 or indices.axes.len != 1 or update_values.axes.len != 3 or output.axes.len != 3) return null;
1517 const outer = seed.axes[0].extent;
1518 const axis_size = seed.axes[1].extent;
1519 const inner = seed.axes[2].extent;
1520 const updates = indices.axes[0].extent;
1521 if (update_values.axes[0].extent != outer or update_values.axes[1].extent != updates or update_values.axes[2].extent != inner) return null;
1522 if (output.axes[0].extent != outer or output.axes[1].extent != axis_size or output.axes[2].extent != inner) return null;
1523 if (!std.mem.eql(u8, seed.axes[0].name, output.axes[0].name)) return null;
1524 if (!std.mem.eql(u8, seed.axes[1].name, output.axes[1].name)) return null;
1525 if (!std.mem.eql(u8, seed.axes[2].name, output.axes[2].name)) return null;
1526 if (!std.mem.eql(u8, indices.axes[0].name, update_values.axes[1].name)) return null;
1527 const launch = specialization.launch orelse return null;
1528 if (launch.threadgroup[0] == 0) return null;
1529 const variant: ScatterAddVariant = if (specialization.structure) |structure| blk: {
1530 if (std.mem.eql(u8, structure, "direct")) break :blk .direct;
1531 if (std.mem.eql(u8, structure, "shared_bins")) break :blk .shared_bins;
1532 return null;
1533 } else .direct;
1534 return .{
1535 .outer = outer,
1536 .axis_size = axis_size,
1537 .updates = updates,
1538 .inner = inner,
1539 .dtype = dtype,
1540 .variant = variant,
1541 .threads = launch.threadgroup[0],
1542 .outer_axis = seed.axes[0].name,
1543 .source_axis = seed.axes[1].name,
1544 .update_axis = indices.axes[0].name,
1545 .inner_axis = seed.axes[2].name,
1546 };
1547 }
1548
1549 fn scatterAddFamilySchedule(instance: ScatterAdd) kernel.logical.schedule.ThreadBlocks {
1550 return kernel.logical.schedule.threadBlocks(.{ .x = instance.threads });
1551 }
1552
1553 fn scatter_add_direct_runtime_body_active(guard_builder: anytype, ctx: anytype) !void {
1554 const axis_size = try guard_builder.castIndex(ctx.args.param(.axis_size).raw());
1555 const update_block = try guard_builder.mul(ctx.update_count, ctx.inner_extent);
1556 const outer = try guard_builder.div(ctx.element, update_block);
1557 const outer_consumed = try guard_builder.mul(outer, update_block);
1558 const rem = try guard_builder.sub(ctx.element, outer_consumed);
1559 const update_position = try guard_builder.div(rem, ctx.inner_extent);
1560 const update_consumed = try guard_builder.mul(update_position, ctx.inner_extent);
1561 const within = try guard_builder.sub(rem, update_consumed);
1562 const loaded = try ctx.args.param(.indices).load(guard_builder, update_position);
1563 const target = try guard_builder.castIndex(loaded.raw());
1564 const zero = try guard_builder.constantIndex(0);
1565 const non_negative = try guard_builder.compare(.ge, target, zero);
1566 try guard_builder.guardDo(non_negative, .{
1567 .args = ctx.args,
1568 .element = ctx.element,
1569 .outer = outer,
1570 .target = target,
1571 .axis_size = axis_size,
1572 .inner_extent = ctx.inner_extent,
1573 .within = within,
1574 }, scatter_add_direct_runtime_body_non_negative);
1575 }
1576
1577 fn scatter_add_direct_runtime_body_non_negative(range_builder: anytype, range_ctx: anytype) !void {
1578 const in_range = try range_builder.compare(.lt, range_ctx.target, range_ctx.axis_size);
1579 try range_builder.guardDo(in_range, .{
1580 .args = range_ctx.args,
1581 .element = range_ctx.element,
1582 .outer = range_ctx.outer,
1583 .target = range_ctx.target,
1584 .axis_size = range_ctx.axis_size,
1585 .inner_extent = range_ctx.inner_extent,
1586 .within = range_ctx.within,
1587 }, scatter_add_direct_runtime_body_in_range);
1588 }
1589
1590 fn scatter_add_direct_runtime_body_in_range(atomic_builder: anytype, atomic_ctx: anytype) !void {
1591 const axis_block = try atomic_builder.mul(atomic_ctx.axis_size, atomic_ctx.inner_extent);
1592 const outer_offset = try atomic_builder.mul(atomic_ctx.outer, axis_block);
1593 const target_offset = try atomic_builder.mul(atomic_ctx.target, atomic_ctx.inner_extent);
1594 const partial = try atomic_builder.add(outer_offset, target_offset);
1595 const dst_index = try atomic_builder.add(partial, atomic_ctx.within);
1596 const value = try atomic_ctx.args.param(.updates).load(atomic_builder, atomic_ctx.element);
1597 _ = try atomic_ctx.args.param(.dst).atomicRmw(atomic_builder, .add, value, dst_index);
1598 }
1599
1600 fn scatterAddDirectRuntimeBody(k: anytype, spec: ScatterAdd, args: anytype) !void {
1601 _ = spec;
1602 const element = try k.globalId(.x);
1603 const total = try k.castIndex(args.param(.total).raw());
1604 const update_count = try k.castIndex(args.param(.update_count).raw());
1605 const inner_extent = try k.castIndex(args.param(.inner).raw());
1606 const active = try k.compare(.lt, element, total);
1607 try k.guardDo(active, .{
1608 .args = args,
1609 .element = element,
1610 .update_count = update_count,
1611 .inner_extent = inner_extent,
1612 }, scatter_add_direct_runtime_body_active);
1613 }
1614
1615 fn scatterAddRuntimeBody(k: anytype, spec: ScatterAdd, args: anytype) !void {
1616 switch (spec.variant) {
1617 .direct => try scatterAddDirectRuntimeBody(k, spec, args),
1618 .shared_bins => try scatterAddSharedRuntimeBody(k, spec, args),
1619 }
1620 }
1621
1622 fn scatter_add_shared_runtime_body_zero_bin(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
1623 try loop_builder.storeIndex(ctx.zero_value, ctx.bins, bin);
1624 return acc;
1625 }
1626
1627 fn scatter_add_shared_runtime_body_active(guard_builder: anytype, ctx: anytype) !void {
1628 const loaded = try ctx.args.param(.indices).load(guard_builder, ctx.update);
1629 const target = try guard_builder.castIndex(loaded.raw());
1630 const zero = try guard_builder.constantIndex(0);
1631 const non_negative = try guard_builder.compare(.ge, target, zero);
1632 try guard_builder.guardDo(non_negative, .{
1633 .args = ctx.args,
1634 .update = ctx.update,
1635 .target = target,
1636 .axis_size = ctx.axis_size,
1637 .bins = ctx.bins,
1638 }, scatter_add_shared_runtime_body_non_negative);
1639 }
1640
1641 fn scatter_add_shared_runtime_body_non_negative(range_builder: anytype, range_ctx: anytype) !void {
1642 const in_range = try range_builder.compare(.lt, range_ctx.target, range_ctx.axis_size);
1643 try range_builder.guardDo(in_range, .{
1644 .args = range_ctx.args,
1645 .update = range_ctx.update,
1646 .target = range_ctx.target,
1647 .bins = range_ctx.bins,
1648 }, scatter_add_shared_runtime_body_in_range);
1649 }
1650
1651 fn scatter_add_shared_runtime_body_in_range(atomic_builder: anytype, atomic_ctx: anytype) !void {
1652 const value = try atomic_ctx.args.param(.updates).load(atomic_builder, atomic_ctx.update);
1653 _ = try atomic_builder.atomicRmwIndex(.add, value.raw(), atomic_ctx.bins, atomic_ctx.target);
1654 }
1655
1656 fn scatter_add_shared_runtime_body_value(loop_builder: anytype, bin: kernel.Value, acc: kernel.Value, ctx: anytype) !kernel.Value {
1657 const partial = try loop_builder.loadIndex(ctx.bins, bin);
1658 _ = try loop_builder.atomicRmwIndex(.add, partial, ctx.args.param(.dst).raw(), bin);
1659 return acc;
1660 }
1661
1662 fn scatterAddSharedRuntimeBody(k: anytype, spec: ScatterAdd, args: anytype) !void {
1663 const bins = try k.sharedBuffer(spec.dtype, spec.axis_size);
1664 const zero_value = switch (spec.dtype) {
1665 .i32 => try k.constantInt(.i32, 0),
1666 .f32 => try k.constantFloat(.f32, 0),
1667 else => return error.UnsupportedDType,
1668 };
1669 const thread = try k.castIndex(try k.threadId(.x));
1670 const stride = try k.castIndex(try k.blockDim(.x));
1671 const axis_size = try k.castIndex(args.param(.axis_size).raw());
1672
1673 _ = try k.fold(thread, axis_size, stride, zero_value, .{
1674 .bins = bins,
1675 .zero_value = zero_value,
1676 }, scatter_add_shared_runtime_body_zero_bin);
1677 try k.barrier(.block);
1678
1679 const update = try k.globalId(.x);
1680 const update_count = try k.castIndex(args.param(.update_count).raw());
1681 const active = try k.compare(.lt, update, update_count);
1682 try k.guardDo(active, .{
1683 .args = args,
1684 .update = update,
1685 .bins = bins,
1686 .axis_size = axis_size,
1687 }, scatter_add_shared_runtime_body_active);
1688 try k.barrier(.block);
1689
1690 _ = try k.fold(thread, axis_size, stride, zero_value, .{
1691 .args = args,
1692 .bins = bins,
1693 }, scatter_add_shared_runtime_body_value);
1694 }
1695
1696 fn scatterAddRuntimeFamily(comptime dtype: DType) type {
1697 return kernel.logical.Family(.{
1698 .name = std.fmt.comptimePrint("accy_kernel_indexing_scatter_add_runtime_{s}", .{dtype.name()}),
1699 .parameters = .{
1700 .dst = kernel.dynamicBuffer(dtype),
1701 .src = kernel.dynamicBuffer(dtype),
1702 .indices = kernel.dynamicBuffer(.i32),
1703 .updates = kernel.dynamicBuffer(dtype),
1704 .outer = kernel.scalar(.i32),
1705 .axis_size = kernel.scalar(.i32),
1706 .update_count = kernel.scalar(.i32),
1707 .inner = kernel.scalar(.i32),
1708 .total = kernel.scalar(.i32),
1709 },
1710 .Instance = ScatterAdd,
1711 .schedule = scatterAddFamilySchedule,
1712 .body = scatterAddRuntimeBody,
1713 });
1714 }
1715
1716 pub const ScatterAddRuntimeFamilyI32 = scatterAddRuntimeFamily(.i32);
1717 pub const ScatterAddRuntimeFamilyF32 = scatterAddRuntimeFamily(.f32);
1718
1719 pub fn createScatterAddFamilyArtifact(
1720 allocator: std.mem.Allocator,
1721 handle: kernel.BackendHandle,
1722 instance: ScatterAdd,
1723 options: entry.ArtifactOptions,
1724 ) !kernel.OwnedKernelCallArtifact {
1725 if (!scatterAddInstanceValid(instance)) return error.InvalidKernelLibraryEntry;
1726 const target = try scatterAddFamilyTarget(allocator, instance);
1727 defer allocator.free(target);
1728 const entry_name = try scatterAddFamilyEntryName(allocator, instance);
1729 defer allocator.free(entry_name);
1730 const family_fingerprint = options.shape_family_fingerprint orelse try scatterAddFamilyFingerprint(allocator, instance);
1731 const shape_profile_dimensions = scatterAddShapeProfileDimensions(instance);
1732 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
1733 .name = "scatter_add",
1734 .fingerprint = family_fingerprint,
1735 .dimensions = shape_profile_dimensions[0..],
1736 };
1737
1738 var graph = switch (instance.dtype) {
1739 .i32 => try ScatterAddRuntimeFamilyI32.buildNamed(allocator, options.limits, entry_name, instance),
1740 .f32 => try ScatterAddRuntimeFamilyF32.buildNamed(allocator, options.limits, entry_name, instance),
1741 else => return error.UnsupportedDType,
1742 };
1743 defer graph.deinit();
1744 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
1745 .target = target,
1746 .version = scatter_add_family_version,
1747 .format = options.format,
1748 .kernel_plan = options.kernel_plan,
1749 .element_count_argument = options.element_count_argument,
1750 .shape_family_fingerprint = family_fingerprint,
1751 .shape_profile = shape_profile,
1752 .launch = options.launch orelse try scatterAddDerivedLaunch(instance),
1753 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 5 else options.runtime_scalar_argument_count,
1754 .static_arguments = options.static_arguments,
1755 });
1756 }
1757
1758 test "indexing scatter add runtime family accumulates on the oracle" {
1759 const allocator = std.testing.allocator;
1760 const compiled = ScatterAdd{ .axis_size = 1, .updates = 1, .threads = 32 };
1761 const runtime = ScatterAdd{ .axis_size = 8, .updates = 5, .threads = 32 };
1762
1763 var graph = try ScatterAddRuntimeFamilyI32.build(allocator, ScatterAddRuntimeFamilyI32.Limits.testing, compiled);
1764 defer graph.deinit();
1765
1766 var dst = [_]i32{ 5, 5, 5, 5, 5, 5, 5, 5 };
1767 var indices = [_]i32{ 3, 0, 3, 9, 1 };
1768 var updates = [_]i32{ 100, 200, 300, 400, 500 };
1769 const expected = [_]i32{ 205, 505, 5, 405, 5, 5, 5, 5 };
1770
1771 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
1772 try graph.runCpuWithLaunch(allocator, &.{
1773 kernel.argumentBuffer(i32, dst[0..]),
1774 kernel.argumentBuffer(i32, dst[0..]),
1775 kernel.argumentBuffer(i32, indices[0..]),
1776 kernel.argumentBuffer(i32, updates[0..]),
1777 kernel.argumentI32(@intCast(runtime.outer)),
1778 kernel.argumentI32(@intCast(runtime.axis_size)),
1779 kernel.argumentI32(@intCast(runtime.updates)),
1780 kernel.argumentI32(@intCast(runtime.inner)),
1781 kernel.argumentI32(@intCast(runtime.total())),
1782 }, .{
1783 .grid = launch_value.grid,
1784 .block = launch_value.threadgroup,
1785 });
1786 try std.testing.expectEqualSlices(i32, expected[0..], dst[0..]);
1787 }
1788
1789 test "indexing scatter add runtime family accumulates shaped updates on the oracle" {
1790 const allocator = std.testing.allocator;
1791 const compiled = ScatterAdd{ .axis_size = 1, .updates = 1, .threads = 32 };
1792 const runtime = ScatterAdd{ .outer = 2, .axis_size = 4, .updates = 3, .inner = 2, .threads = 32 };
1793
1794 var graph = try ScatterAddRuntimeFamilyI32.build(allocator, ScatterAddRuntimeFamilyI32.Limits.testing, compiled);
1795 defer graph.deinit();
1796
1797 var dst = [_]i32{ 1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 13, 14, 15, 16, 17, 18 };
1798 var indices = [_]i32{ 2, 0, 2 };
1799 var updates = [_]i32{ 10, 20, 30, 40, 50, 60, 100, 200, 300, 400, 500, 600 };
1800 var expected = dst;
1801 for (0..2) |outer| {
1802 for (0..3) |update_position| {
1803 const target: usize = @intCast(indices[update_position]);
1804 for (0..2) |within| {
1805 const out_index = outer * 8 + target * 2 + within;
1806 const update_index = outer * 6 + update_position * 2 + within;
1807 expected[out_index] += updates[update_index];
1808 }
1809 }
1810 }
1811
1812 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
1813 try graph.runCpuWithLaunch(allocator, &.{
1814 kernel.argumentBuffer(i32, dst[0..]),
1815 kernel.argumentBuffer(i32, dst[0..]),
1816 kernel.argumentBuffer(i32, indices[0..]),
1817 kernel.argumentBuffer(i32, updates[0..]),
1818 kernel.argumentI32(@intCast(runtime.outer)),
1819 kernel.argumentI32(@intCast(runtime.axis_size)),
1820 kernel.argumentI32(@intCast(runtime.updates)),
1821 kernel.argumentI32(@intCast(runtime.inner)),
1822 kernel.argumentI32(@intCast(runtime.total())),
1823 }, .{
1824 .grid = launch_value.grid,
1825 .block = launch_value.threadgroup,
1826 });
1827 try std.testing.expectEqualSlices(i32, expected[0..], dst[0..]);
1828 }
1829
1830 test "indexing scatter add instance round-trips through specialization" {
1831 const instance = ScatterAdd{ .outer = 2, .axis_size = 16, .updates = 100, .inner = 3, .threads = 64 };
1832 var owned = try scatterAddFamilySpecialization(std.testing.allocator, instance);
1833 defer owned.deinit();
1834
1835 const recovered = scatterAddInstanceFromSpecialization(owned.value) orelse return error.TestExpectedScatterAddInstance;
1836 try std.testing.expectEqual(instance.outer, recovered.outer);
1837 try std.testing.expectEqual(instance.axis_size, recovered.axis_size);
1838 try std.testing.expectEqual(instance.updates, recovered.updates);
1839 try std.testing.expectEqual(instance.inner, recovered.inner);
1840 try std.testing.expectEqual(instance.dtype, recovered.dtype);
1841 try std.testing.expectEqual(instance.threads, recovered.threads);
1842
1843 try std.testing.expectEqual(@as(?ScatterAdd, null), scatterAddInstanceFromSpecialization(.{}));
1844 }
1845
1846 test "indexing scatter add family identity carries the operation" {
1847 const instance = ScatterAdd{ .axis_size = 1024, .updates = 4096, .threads = 128 };
1848
1849 const family_target = try scatterAddFamilyTarget(std.testing.allocator, instance);
1850 defer std.testing.allocator.free(family_target);
1851 try std.testing.expectEqualStrings("accy.kernel.indexing.scatter_add_family_128_i32", family_target);
1852
1853 const replace_instance = Scatter{ .axis_size = 1024, .updates = 4096, .threads = 128 };
1854 try std.testing.expect(entry.operationFingerprint(scatterAddTuningOperation(instance)) !=
1855 entry.operationFingerprint(scatterTuningOperation(replace_instance)));
1856
1857 try std.testing.expect(!scatterAddInstanceValid(.{ .axis_size = 8, .updates = 4, .dtype = .f16 }));
1858 try std.testing.expect(scatterAddInstanceValid(.{ .axis_size = 8, .updates = 4, .dtype = .f32 }));
1859 try std.testing.expect(scatterAddInstanceValid(.{ .axis_size = 8, .updates = 4 }));
1860 }
1861
1862 test "indexing scatter add family artifact carries runtime launch contract" {
1863 const allocator = std.testing.allocator;
1864 var state = gpu.recording.BackendState{
1865 .allocator = allocator,
1866 .kind = .cuda,
1867 .format = .cuda_ptx,
1868 };
1869 const instance = ScatterAdd{ .axis_size = 8, .updates = 4, .threads = 8 };
1870
1871 var family_artifact = try createScatterAddFamilyArtifact(allocator, state.handle(), instance, .{ .limits = .testing });
1872 defer family_artifact.deinit();
1873
1874 const family_entry = family_artifact.entry();
1875 try std.testing.expectEqualStrings("accy.kernel.indexing.scatter_add_family_8_i32", family_entry.target);
1876 try std.testing.expectEqual(@as(u32, 5), family_entry.runtime_scalar_argument_count);
1877 try std.testing.expect(family_entry.required_dtypes.contains(.i32));
1878 const profile = family_entry.shape_profile orelse return error.TestExpectedShapeProfile;
1879 try std.testing.expectEqualStrings("scatter_add", profile.name);
1880 try std.testing.expectEqual(@as(usize, 5), profile.dimensions.len);
1881 switch (family_entry.launch) {
1882 .derived => |launch| {
1883 try std.testing.expectEqual(@as(u32, 8), launch.threadgroup[0]);
1884 switch (launch.grid[0]) {
1885 .runtime_u32_ceil_div => |term| {
1886 try std.testing.expectEqual(@as(usize, 4), term.argument_index);
1887 try std.testing.expectEqual(@as(u32, 8), term.divisor);
1888 },
1889 else => return error.TestExpectedDerivedLaunch,
1890 }
1891 },
1892 else => return error.TestExpectedDerivedLaunch,
1893 }
1894 }
1895
1896 test "indexing scatter add f32 runtime family accumulates exactly on the sequential oracle" {
1897 const allocator = std.testing.allocator;
1898 const compiled = ScatterAdd{ .axis_size = 1, .updates = 1, .dtype = .f32, .threads = 32 };
1899 const runtime = ScatterAdd{ .axis_size = 8, .updates = 5, .dtype = .f32, .threads = 32 };
1900
1901 var graph = try ScatterAddRuntimeFamilyF32.build(allocator, ScatterAddRuntimeFamilyF32.Limits.testing, compiled);
1902 defer graph.deinit();
1903
1904 var dst = [_]f32{ 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5 };
1905 var indices = [_]i32{ 3, 0, 3, 9, 1 };
1906 var updates = [_]f32{ 0.125, 2.5, 0.25, 99.0, 7.75 };
1907 const expected = [_]f32{ 3.0, 8.25, 0.5, 0.875, 0.5, 0.5, 0.5, 0.5 };
1908
1909 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
1910 try graph.runCpuWithLaunch(allocator, &.{
1911 kernel.argumentBuffer(f32, dst[0..]),
1912 kernel.argumentBuffer(f32, dst[0..]),
1913 kernel.argumentBuffer(i32, indices[0..]),
1914 kernel.argumentBuffer(f32, updates[0..]),
1915 kernel.argumentI32(@intCast(runtime.outer)),
1916 kernel.argumentI32(@intCast(runtime.axis_size)),
1917 kernel.argumentI32(@intCast(runtime.updates)),
1918 kernel.argumentI32(@intCast(runtime.inner)),
1919 kernel.argumentI32(@intCast(runtime.total())),
1920 }, .{
1921 .grid = launch_value.grid,
1922 .block = launch_value.threadgroup,
1923 });
1924 try std.testing.expectEqualSlices(f32, expected[0..], dst[0..]);
1925 }
1926
1927 test "indexing scatter add shared bins variant matches the direct arm on the oracle" {
1928 const allocator = std.testing.allocator;
1929 const compiled = ScatterAdd{ .axis_size = 8, .updates = 1, .variant = .shared_bins, .threads = 4 };
1930 const runtime = ScatterAdd{ .axis_size = 8, .updates = 10, .variant = .shared_bins, .threads = 4 };
1931
1932 var graph = try ScatterAddRuntimeFamilyI32.build(allocator, ScatterAddRuntimeFamilyI32.Limits.testing, compiled);
1933 defer graph.deinit();
1934
1935 var dst = [_]i32{ 1, 1, 1, 1, 1, 1, 1, 1 };
1936 var indices = [_]i32{ 3, 0, 3, 9, 1, 0, 7, 3, -2, 7 };
1937 var updates = [_]i32{ 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000 };
1938
1939 var expected = [_]i32{ 1, 1, 1, 1, 1, 1, 1, 1 };
1940 for (indices, updates) |index, update| {
1941 if (index < 0 or index >= 8) continue;
1942 expected[@intCast(index)] += update;
1943 }
1944
1945 const launch_value = try entry.runtimeLaunch1D(runtime.total(), runtime.threads);
1946 try graph.runCpuWithLaunch(allocator, &.{
1947 kernel.argumentBuffer(i32, dst[0..]),
1948 kernel.argumentBuffer(i32, dst[0..]),
1949 kernel.argumentBuffer(i32, indices[0..]),
1950 kernel.argumentBuffer(i32, updates[0..]),
1951 kernel.argumentI32(@intCast(runtime.outer)),
1952 kernel.argumentI32(@intCast(runtime.axis_size)),
1953 kernel.argumentI32(@intCast(runtime.updates)),
1954 kernel.argumentI32(@intCast(runtime.inner)),
1955 kernel.argumentI32(@intCast(runtime.total())),
1956 }, .{
1957 .grid = launch_value.grid,
1958 .block = launch_value.threadgroup,
1959 });
1960 try std.testing.expectEqualSlices(i32, expected[0..], dst[0..]);
1961 }
1962
1963 test "indexing scatter add shared bins identity and validity" {
1964 const shared_instance = ScatterAdd{ .axis_size = 256, .updates = 4096, .variant = .shared_bins, .threads = 128 };
1965 const shared_target = try scatterAddFamilyTarget(std.testing.allocator, shared_instance);
1966 defer std.testing.allocator.free(shared_target);
1967 try std.testing.expectEqualStrings("accy.kernel.indexing.scatter_add_family_shared256_128_i32", shared_target);
1968
1969 try std.testing.expect(scatterAddInstanceValid(shared_instance));
1970 try std.testing.expect(!scatterAddInstanceValid(.{
1971 .axis_size = scatter_add_shared_bins_cap + 1,
1972 .updates = 16,
1973 .variant = .shared_bins,
1974 }));
1975 try std.testing.expect(!scatterAddInstanceValid(.{
1976 .outer = 2,
1977 .axis_size = 16,
1978 .updates = 16,
1979 .variant = .shared_bins,
1980 }));
1981 try std.testing.expect(scatterAddInstanceValid(.{
1982 .axis_size = scatter_add_shared_bins_cap + 1,
1983 .updates = 16,
1984 }));
1985 }
1986
1987 test "indexing scatter add shared variant round-trips through specialization structure" {
1988 const instance = ScatterAdd{ .axis_size = 64, .updates = 1024, .variant = .shared_bins, .threads = 128 };
1989 var owned = try scatterAddFamilySpecialization(std.testing.allocator, instance);
1990 defer owned.deinit();
1991
1992 try std.testing.expect(owned.value.structureIs("shared_bins"));
1993 const recovered = scatterAddInstanceFromSpecialization(owned.value) orelse
1994 return error.TestExpectedScatterAddInstance;
1995 try std.testing.expectEqual(ScatterAddVariant.shared_bins, recovered.variant);
1996 try std.testing.expectEqual(instance.threads, recovered.threads);
1997 }
1998
1999 fn indexingFamilyTuningTestCapabilities() gpu.BackendCapabilities {
2000 return .{ .identity = .{
2001 .backend = .cuda,
2002 .family = .nvidia_cuda,
2003 .name = "indexing-family-tuning-test-device",
2004 .vendor_id = 0x10de,
2005 .device_id = 0x2684,
2006 } };
2007 }
2008
2009 test "indexing family tuning keys discriminate gather scatter and scatter add" {
2010 const allocator = std.testing.allocator;
2011 const device = tuning.deviceFingerprint(indexingFamilyTuningTestCapabilities());
2012
2013 const gather_key = try gatherFamilyTuningKey(allocator, device, .{ .axis_size = 16, .gathered = 8 });
2014 const scatter_key = try scatterFamilyTuningKey(allocator, device, .{ .axis_size = 16, .updates = 8 });
2015 const scatter_add_key = try scatterAddFamilyTuningKey(allocator, device, .{ .axis_size = 16, .updates = 4096 });
2016
2017 try std.testing.expect(!gather_key.eql(scatter_key));
2018 try std.testing.expect(!gather_key.eql(scatter_add_key));
2019 try std.testing.expect(!scatter_key.eql(scatter_add_key));
2020
2021 const replacement = Scatter{ .outer = 16, .axis_size = 4096, .updates = 16, .inner = 1 };
2022 const replacement_key = try scatterFamilyTuningKey(allocator, device, replacement);
2023 try std.testing.expect(!scatter_add_key.eql(replacement_key));
2024 try std.testing.expect(scatter_add_key.operation_fingerprint != replacement_key.operation_fingerprint);
2025 }
2026
2027 test "indexing family tuning resolves gather and stale scatter targets" {
2028 const allocator = std.testing.allocator;
2029 const caps = indexingFamilyTuningTestCapabilities();
2030 const device = tuning.deviceFingerprint(caps);
2031
2032 const gather_probe = Gather{ .axis_size = 128, .gathered = 64 };
2033 const gather_candidates = gatherThreadCandidatesForTotal(gather_probe.total());
2034 try std.testing.expect(gather_candidates.slice().len >= 1);
2035 var gather_winner = gather_probe;
2036 gather_winner.threads = gather_candidates.slice()[0];
2037 const gather_target = try gatherFamilyTarget(allocator, gather_winner);
2038 defer allocator.free(gather_target);
2039
2040 const gather_records = [_]tuning.FamilyTuningRecord{.{
2041 .key = try gatherFamilyTuningKey(allocator, device, gather_probe),
2042 .target = gather_target,
2043 .winner_median_ns = 500,
2044 .runner_up_median_ns = 700,
2045 .sample_count = 30,
2046 }};
2047 const gather_reader = tuning.FamilyTuningReader.init(caps, .{ .records = gather_records[0..] });
2048 const gather_resolved = (try resolveGatherSchedule(allocator, gather_reader, gather_probe)) orelse
2049 return error.TestExpectedSchedule;
2050 try std.testing.expectEqual(gather_winner.threads, gather_resolved);
2051
2052 const stale = [_]tuning.FamilyTuningRecord{.{
2053 .key = try scatterFamilyTuningKey(allocator, device, .{ .axis_size = 16, .updates = 8 }),
2054 .target = "accy.kernel.indexing.scatter_family_9999_f32",
2055 .winner_median_ns = 1,
2056 .runner_up_median_ns = 2,
2057 .sample_count = 1,
2058 }};
2059 const stale_reader = tuning.FamilyTuningReader.init(caps, .{ .records = stale[0..] });
2060 const unresolvable = try resolveScatterSchedule(allocator, stale_reader, .{ .axis_size = 16, .updates = 8 });
2061 try std.testing.expectEqual(@as(?u32, null), unresolvable);
2062 }
2063
2064 test "indexing family tuning resolves scatter add variants" {
2065 const allocator = std.testing.allocator;
2066 const caps = indexingFamilyTuningTestCapabilities();
2067 const device = tuning.deviceFingerprint(caps);
2068
2069 const direct_probe = ScatterAdd{ .axis_size = 16, .updates = 4096 };
2070 const direct_key = try scatterAddFamilyTuningKey(allocator, device, direct_probe);
2071 const direct_candidates = scatterAddThreadCandidatesForTotal(direct_probe.total());
2072 try std.testing.expect(direct_candidates.slice().len >= 2);
2073 var direct_winner = direct_probe;
2074 direct_winner.threads = direct_candidates.slice()[0];
2075 const direct_target = try scatterAddFamilyTarget(allocator, direct_winner);
2076 defer allocator.free(direct_target);
2077
2078 const direct_records = [_]tuning.FamilyTuningRecord{.{
2079 .key = direct_key,
2080 .target = direct_target,
2081 .winner_median_ns = 600,
2082 .runner_up_median_ns = 900,
2083 .sample_count = 30,
2084 }};
2085 const direct_reader = tuning.FamilyTuningReader.init(caps, .{ .records = direct_records[0..] });
2086 const direct_found = direct_reader.table.find(direct_key) orelse return error.TestExpectedTuningRecord;
2087 try std.testing.expectEqualStrings(direct_target, direct_found.target);
2088 const direct_resolved = (try resolveScatterAddSchedule(allocator, direct_reader, direct_probe)) orelse
2089 return error.TestExpectedSchedule;
2090 try std.testing.expectEqual(direct_winner.threads, direct_resolved.threads);
2091 try std.testing.expectEqual(ScatterAddVariant.direct, direct_resolved.variant);
2092
2093 const miss = try resolveScatterAddSchedule(allocator, direct_reader, .{ .axis_size = 16, .updates = 2048 });
2094 try std.testing.expectEqual(@as(?ScatterAddResolvedSchedule, null), miss);
2095
2096 const shared_probe = ScatterAdd{ .axis_size = 64, .updates = 4096 };
2097 var shared_winner = shared_probe;
2098 shared_winner.variant = .shared_bins;
2099 shared_winner.threads = direct_candidates.slice()[0];
2100 const shared_target = try scatterAddFamilyTarget(allocator, shared_winner);
2101 defer allocator.free(shared_target);
2102
2103 const shared_records = [_]tuning.FamilyTuningRecord{.{
2104 .key = try scatterAddFamilyTuningKey(allocator, device, shared_probe),
2105 .target = shared_target,
2106 .winner_median_ns = 400,
2107 .runner_up_median_ns = 900,
2108 .sample_count = 30,
2109 }};
2110 const shared_reader = tuning.FamilyTuningReader.init(caps, .{ .records = shared_records[0..] });
2111 const shared_resolved = (try resolveScatterAddSchedule(allocator, shared_reader, shared_probe)) orelse
2112 return error.TestExpectedSchedule;
2113 try std.testing.expectEqual(ScatterAddVariant.shared_bins, shared_resolved.variant);
2114 try std.testing.expectEqual(shared_winner.threads, shared_resolved.threads);
2115 }