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 }