lib/accy/src/preparation/indexing.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const gpu = @import("gpu");
   3 const choir_abi = @import("choir_abi");
   4 const alloc_arena = @import("alloc_arena");
   5 const choir = @import("choir");
   6 const accy_root = @import("../root.zig");
   7 const accy_choir = @import("../choir/root.zig");
   8 const kernel_library = @import("../kernel/library/root.zig");
   9 const call_preparation = @import("call.zig");
  10 const library_preparation = @import("library.zig");
  11 const dialect_mod = accy_choir.dialect;
  12 
  13 const ir = choir.ir;
  14 const rewrite = ir.rewrite;
  15 const passes = choir.passes;
  16 const work = passes.pass.work;
  17 
  18 pub const KernelLibraryLowering = library_preparation.KernelLibraryLowering;
  19 
  20 pub const Options = struct {
  21     kernel_library: KernelLibraryLowering = .disabled,
  22     gather_schedule: ?kernel_library.GatherSchedule = null,
  23     scatter_schedule: ?kernel_library.ScatterSchedule = null,
  24     scatter_add_schedule: ?kernel_library.ScatterAddSchedule = null,
  25     family_tuning: ?*const kernel_library.tuning.FamilyTuningReader = null,
  26 
  27     pub fn eql(self: Options, other: Options) bool {
  28         return self.kernel_library == other.kernel_library and
  29             std.meta.eql(self.gather_schedule, other.gather_schedule) and
  30             std.meta.eql(self.scatter_schedule, other.scatter_schedule) and
  31             std.meta.eql(self.scatter_add_schedule, other.scatter_add_schedule) and
  32             self.family_tuning == other.family_tuning;
  33     }
  34 };
  35 
  36 pub const indexing_lowering_pass_name = "accy-choir-indexing-lower";
  37 pub const indexing_lowering_pass_description =
  38     "Lower semantic Accy indexing operations onto kernel library catalog calls";
  39 
  40 const kernel_library_option_choices = [_]passes.PassOptionChoice{
  41     .{ .name = "disabled" },
  42     .{ .name = "enabled" },
  43 };
  44 
  45 pub const indexing_lowering_pass_options = [_]passes.PassOptionSpec{
  46     .{
  47         .name = "kernel-library",
  48         .description = "Use kernel library calls for supported indexing operations",
  49         .kind = .choice,
  50         .choices = &kernel_library_option_choices,
  51         .default_value = "disabled",
  52     },
  53     .{
  54         .name = "gather-thread-blocks",
  55         .description = "Thread blocks for selected gather kernels",
  56         .kind = .unsigned,
  57     },
  58     .{
  59         .name = "scatter-thread-blocks",
  60         .description = "Thread blocks for selected scatter kernels",
  61         .kind = .unsigned,
  62     },
  63     .{
  64         .name = "scatter-add-thread-blocks",
  65         .description = "Thread blocks for selected scatter-add kernels",
  66         .kind = .unsigned,
  67     },
  68 };
  69 
  70 pub fn indexingLoweringPass() passes.Pass {
  71     return .{
  72         .name = indexing_lowering_pass_name,
  73         .description = indexing_lowering_pass_description,
  74         .run_fn = runIndexingLoweringPass,
  75         .work_contract = indexing_work_contract,
  76     };
  77 }
  78 
  79 pub fn indexingLoweringPassWithOptions(options: *const Options) passes.Pass {
  80     return .{
  81         .name = indexing_lowering_pass_name,
  82         .description = indexing_lowering_pass_description,
  83         .state = @constCast(options),
  84         .run_with_state_fn = runIndexingLoweringPassWithState,
  85         .work_contract = indexing_work_contract,
  86     };
  87 }
  88 
  89 pub fn indexingLoweringPassFromOptions(allocator: std.mem.Allocator, set: passes.PassOptionSet) anyerror!passes.Pass {
  90     const options = try allocator.create(Options);
  91     errdefer allocator.destroy(options);
  92     options.* = .{
  93         .kernel_library = try kernelLibraryLoweringFromText(set.choiceValue("kernel-library", "disabled")),
  94         .gather_schedule = if (set.get("gather-thread-blocks")) |value| .{
  95             .thread_blocks = try parseU32Option(value),
  96         } else null,
  97         .scatter_schedule = if (set.get("scatter-thread-blocks")) |value| .{
  98             .thread_blocks = try parseU32Option(value),
  99         } else null,
 100         .scatter_add_schedule = if (set.get("scatter-add-thread-blocks")) |value| .{
 101             .thread_blocks = try parseU32Option(value),
 102         } else null,
 103     };
 104     var pass = indexingLoweringPassWithOptions(options);
 105     pass.state_deinit_fn = destroyOptions;
 106     return pass;
 107 }
 108 
 109 const indexing_work_contract: work.Contract = .{
 110     .identity = .{ .name = indexing_lowering_pass_name, .version = 1 },
 111     .estimate = indexingWork,
 112 };
 113 
 114 const IndexingKind = enum {
 115     gather,
 116     scatter,
 117     scatter_add,
 118 
 119     fn fromOperation(op: *ir.Operation) ?IndexingKind {
 120         inline for (comptime std.meta.tags(IndexingKind)) |kind| {
 121             if (std.mem.eql(u8, op.name.name, "accy." ++ @tagName(kind))) return kind;
 122         }
 123         return null;
 124     }
 125 
 126     fn scheduled(self: IndexingKind, options: Options) bool {
 127         return switch (self) {
 128             .gather => options.gather_schedule != null,
 129             .scatter => options.scatter_schedule != null,
 130             .scatter_add => options.scatter_add_schedule != null,
 131         };
 132     }
 133 };
 134 
 135 const IndexingWork = struct {
 136     options: Options,
 137     candidates: u64 = 0,
 138     tuned: u64 = 0,
 139     type_bytes: u64 = 0,
 140 
 141     fn visit(self: *IndexingWork, op: *ir.Operation) !ir.WalkResult {
 142         const kind = IndexingKind.fromOperation(op) orelse return .advance;
 143         const scheduled = kind.scheduled(self.options);
 144         if (!scheduled and self.options.family_tuning == null) return .advance;
 145         if (op.getNumResults() != 1) return .advance;
 146         if (op.getNumOperands() != @as(usize, if (kind == .gather) 2 else 3)) return .advance;
 147         self.candidates = try work.add(self.candidates, 1);
 148         self.tuned = try work.add(self.tuned, @intFromBool(!scheduled));
 149         for (op.getOperandValues()) |operand| try self.typeBytes(operand.type);
 150         try self.typeBytes(op.getResult(0).?.type);
 151         return .advance;
 152     }
 153 
 154     fn typeBytes(self: *IndexingWork, typ: ir.Type) !void {
 155         if (typ.getDialectParamKey()) |key| {
 156             self.type_bytes = try work.add(self.type_bytes, key.len);
 157         }
 158     }
 159 };
 160 
 161 fn indexingDescriptorStorage() u64 {
 162     const entry = kernel_library.entry;
 163     const shape = accy_choir.shape;
 164     const arrays = 4 * std.ArrayList(shape.Symbol).growCapacity(4) * @sizeOf(shape.Symbol) +
 165         4 * std.ArrayList(shape.Tensor).growCapacity(4) * @sizeOf(shape.Tensor) +
 166         4 * std.ArrayList(shape.Fact).growCapacity(4) * @sizeOf(shape.Fact);
 167     const records = 4 * @sizeOf(entry.Shape) + 10 * @sizeOf(entry.Axis) +
 168         2 * @sizeOf(entry.ScheduleBinding) + 10 * @sizeOf(shape.Expression) +
 169         18 * @sizeOf(shape.Term);
 170     const arena_traffic = 8 * (arrays + records + 512 + 96 * 128);
 171     return arena_traffic + @sizeOf(alloc_arena.Arena) + @sizeOf(shape.Family) + 64;
 172 }
 173 
 174 const TuningWork = struct { input_bytes: u64 = 0, visits: u64 = 0 };
 175 
 176 fn indexingTuningWork(options: Options, queries: u64) !TuningWork {
 177     if (queries == 0) return .{};
 178     const reader = options.family_tuning.?;
 179     var bytes = try work.multiply(
 180         reader.table.records.len,
 181         @sizeOf(kernel_library.tuning.FamilyTuningRecord),
 182     );
 183     for (reader.table.records) |record| bytes = try work.add(bytes, record.target.len);
 184     return .{
 185         .input_bytes = bytes,
 186         .visits = try work.multiply(128, try work.multiply(queries, try work.add(bytes, 1))),
 187     };
 188 }
 189 
 190 fn indexingWork(input: work.Input) !work.Bounds {
 191     const options: *const Options = if (input.state) |state| @ptrCast(@alignCast(state)) else &.{};
 192     if (options.kernel_library != .enabled or (options.gather_schedule == null and
 193         options.scatter_schedule == null and options.scatter_add_schedule == null and
 194         options.family_tuning == null)) return .{ .work = .{ .structural_visits = 1 } };
 195     const counts = try work.Census.inspect(input.operation);
 196     var facts: IndexingWork = .{ .options = options.* };
 197     _ = try input.operation.walk(.{ .order = .pre_order }, &facts, IndexingWork.visit);
 198     const descriptors = try work.multiply(
 199         try work.add(facts.candidates, facts.tuned),
 200         indexingDescriptorStorage(),
 201     );
 202     const spellings = try work.multiply(
 203         facts.tuned,
 204         2 * kernel_library.geometry.max_thread_candidates * (128 + 8),
 205     );
 206     const decoding = try work.multiply(8, try work.add(
 207         try work.multiply(facts.type_bytes, @sizeOf(i64)),
 208         try work.multiply(facts.candidates, 4 * 128 + 64),
 209     ));
 210     const queues = try work.multiply(2, try work.arrayListGrowth(*ir.Operation, facts.candidates));
 211     const bytes = try work.add(
 212         try work.add(queues, decoding),
 213         try work.add(descriptors, spellings),
 214     );
 215     const tuning = try indexingTuningWork(options.*, facts.tuned);
 216     const units = try work.add(try work.add(counts.atoms, counts.input_bytes), 1);
 217     const uses = try work.add(try work.add(counts.values, counts.operands), 1);
 218     const traversal = try work.multiply(128, try work.multiply(units, uses));
 219     const processing = try work.add(tuning.visits, try work.multiply(2, descriptors));
 220     const nodes = try work.multiply(facts.candidates, @sizeOf(ir.Operation) + @sizeOf(ir.Value) +
 221         3 * @sizeOf(ir.OpOperand) + 8 * @sizeOf(ir.NamedAttribute) + 256 +
 222         5 * @sizeOf(dialect_mod.AccyDialect.KernelCallScalar));
 223     return .{
 224         .work = .{
 225             .input_bytes = try work.add(counts.input_bytes, tuning.input_bytes),
 226             .output_bytes = nodes,
 227             .structural_visits = try work.add(traversal, processing),
 228             .allocation_capacity = bytes,
 229         },
 230         .workspace = bytes,
 231     };
 232 }
 233 
 234 fn destroyOptions(raw: ?*anyopaque, allocator: std.mem.Allocator) void {
 235     const options: *Options = @ptrCast(@alignCast(raw orelse return));
 236     allocator.destroy(options);
 237 }
 238 
 239 fn kernelLibraryLoweringFromText(value: []const u8) !KernelLibraryLowering {
 240     if (std.mem.eql(u8, value, "disabled")) return .disabled;
 241     if (std.mem.eql(u8, value, "enabled")) return .enabled;
 242     return error.InvalidPassOptionValue;
 243 }
 244 
 245 fn parseU32Option(value: []const u8) !u32 {
 246     return std.fmt.parseUnsigned(u32, value, 10) catch return error.InvalidPassOptionValue;
 247 }
 248 
 249 fn runIndexingLoweringPass(pass_ctx: *passes.PassContext) passes.PassResult {
 250     return runIndexingLoweringWithOptions(pass_ctx, .{});
 251 }
 252 
 253 fn runIndexingLoweringPassWithState(raw: ?*anyopaque, pass_ctx: *passes.PassContext) passes.PassResult {
 254     const options: *const Options = @ptrCast(@alignCast(raw orelse return .failure));
 255     return runIndexingLoweringWithOptions(pass_ctx, options.*);
 256 }
 257 
 258 fn runIndexingLoweringWithOptions(pass_ctx: *passes.PassContext, options: Options) passes.PassResult {
 259     if (options.kernel_library != .enabled or
 260         (options.gather_schedule == null and
 261             options.scatter_schedule == null and
 262             options.scatter_add_schedule == null and
 263             options.family_tuning == null))
 264     {
 265         pass_ctx.preserveAllAnalyses();
 266         return .success;
 267     }
 268 
 269     var rewriter = rewrite.PatternRewriter.init(pass_ctx.allocator, pass_ctx.ir_ctx);
 270     defer rewriter.deinit();
 271 
 272     var lowered_count: usize = 0;
 273     lowerOnOp(pass_ctx.op, &rewriter, options, &lowered_count) catch return .failure;
 274     if (lowered_count == 0) {
 275         pass_ctx.preserveAllAnalyses();
 276     } else {
 277         rewriter.finalize(pass_ctx.op);
 278         pass_ctx.markModified();
 279     }
 280     return .success;
 281 }
 282 
 283 fn lowerOnOp(
 284     op: *ir.Operation,
 285     rewriter: *rewrite.PatternRewriter,
 286     options: Options,
 287     lowered_count: *usize,
 288 ) !void {
 289     for (op.regions.items) |*region| {
 290         var block_iter = region.getBlocks();
 291         while (block_iter.next()) |block| {
 292             var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
 293             while (current) |current_op| {
 294                 const next = current_op.next_op;
 295                 if (current_op.regions.items.len > 0) {
 296                     try lowerOnOp(current_op, rewriter, options, lowered_count);
 297                 }
 298                 if ((options.gather_schedule != null or options.family_tuning != null) and
 299                     std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name))
 300                 {
 301                     var guard = rewriter.insertionGuard();
 302                     defer guard.deinit();
 303                     rewriter.setInsertionPointBefore(current_op);
 304                     if (try lowerKnownKernelLibraryGather(current_op, rewriter, options)) {
 305                         lowered_count.* += 1;
 306                     }
 307                 } else if ((options.scatter_schedule != null or options.family_tuning != null) and
 308                     std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.ScatterOp.operation_name))
 309                 {
 310                     var guard = rewriter.insertionGuard();
 311                     defer guard.deinit();
 312                     rewriter.setInsertionPointBefore(current_op);
 313                     if (try lowerKnownKernelLibraryScatter(current_op, rewriter, options)) {
 314                         lowered_count.* += 1;
 315                     }
 316                 } else if ((options.scatter_add_schedule != null or options.family_tuning != null) and
 317                     std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.ScatterAddOp.operation_name))
 318                 {
 319                     var guard = rewriter.insertionGuard();
 320                     defer guard.deinit();
 321                     rewriter.setInsertionPointBefore(current_op);
 322                     if (try lowerKnownKernelLibraryScatterAdd(current_op, rewriter, options)) {
 323                         lowered_count.* += 1;
 324                     }
 325                 }
 326                 current = next;
 327             }
 328         }
 329     }
 330 }
 331 
 332 fn lowerKnownKernelLibraryGather(
 333     op: *ir.Operation,
 334     rewriter: *rewrite.PatternRewriter,
 335     options: Options,
 336 ) !bool {
 337     if (op.getNumResults() != 1) return false;
 338     const result = op.getResult(0) orelse return false;
 339     const operands = op.getOperandValues();
 340     if (operands.len != 2) return false;
 341 
 342     var arena_state = alloc_arena.Arena.init(rewriter.allocator);
 343     defer arena_state.deinit();
 344     const arena = arena_state.allocator();
 345 
 346     const data_type = try dialect_mod.decodeTensorType(arena, operands[0].type);
 347     const indices_type = try dialect_mod.decodeTensorType(arena, operands[1].type);
 348     const output_type = try dialect_mod.decodeTensorType(arena, result.type);
 349     if (indices_type.dtype != .i32) return false;
 350     if (output_type.dtype != data_type.dtype) return false;
 351     if (data_type.dims.len == 0 or indices_type.dims.len != 1) return false;
 352     if (output_type.dims.len != data_type.dims.len) return false;
 353 
 354     const gather_op = dialect_mod.AccyDialect.GatherOp{ .op = op };
 355     const axis_value = gather_op.getAxis() orelse return false;
 356     if (axis_value < 0) return false;
 357     const axis = std.math.cast(usize, axis_value) orelse return false;
 358     if (axis >= data_type.dims.len) return false;
 359 
 360     if (!std.mem.eql(i64, output_type.dims[0..axis], data_type.dims[0..axis])) return false;
 361     if (output_type.dims[axis] != indices_type.dims[0]) return false;
 362     if (!std.mem.eql(i64, output_type.dims[axis + 1 ..], data_type.dims[axis + 1 ..])) return false;
 363 
 364     const outer = dimsProduct(data_type.dims[0..axis]) orelse return false;
 365     const axis_size = dimExtent(data_type.dims[axis]) orelse return false;
 366     const gathered = dimExtent(indices_type.dims[0]) orelse return false;
 367     const inner = dimsProduct(data_type.dims[axis + 1 ..]) orelse return false;
 368 
 369     var gather_schedule = options.gather_schedule;
 370     if (gather_schedule == null) {
 371         const family_tuning = options.family_tuning orelse return false;
 372         const thread_blocks = (try kernel_library.indexing.resolveGatherSchedule(rewriter.allocator, family_tuning.*, .{
 373             .outer = outer,
 374             .axis_size = axis_size,
 375             .gathered = gathered,
 376             .inner = inner,
 377             .dtype = data_type.dtype,
 378         })) orelse return false;
 379         gather_schedule = .{ .thread_blocks = thread_blocks };
 380     }
 381 
 382     var selected = (try kernel_library.selectOwned(rewriter.allocator, .{ .gather = .{
 383         .dtype = data_type.dtype,
 384         .outer = outer,
 385         .axis_size = axis_size,
 386         .gathered = gathered,
 387         .inner = inner,
 388         .schedule = gather_schedule,
 389     } })) orelse return false;
 390     defer selected.deinit();
 391 
 392     const instance = kernel_library.indexing.gatherInstanceFromSpecialization(
 393         selected.descriptor.metadata.specialization,
 394     ) orelse return false;
 395     const runtime_scalars = call_preparation.catalogCallScalars(
 396         5,
 397         try kernel_library.indexing.gatherRuntimeArguments(instance),
 398     );
 399     const result_types = [_]ir.Type{result.type};
 400     const call = try call_preparation.insertCatalogCall(rewriter, .{
 401         .descriptor = selected.descriptor,
 402         .operands = operands,
 403         .result_types = &result_types,
 404         .options = .{ .runtime_scalars = runtime_scalars[0..] },
 405     });
 406     try rewriter.replaceOpWithValue(op, call.getFirstResult());
 407     return true;
 408 }
 409 
 410 fn lowerKnownKernelLibraryScatter(
 411     op: *ir.Operation,
 412     rewriter: *rewrite.PatternRewriter,
 413     options: Options,
 414 ) !bool {
 415     if (op.getNumResults() != 1) return false;
 416     const result = op.getResult(0) orelse return false;
 417     const operands = op.getOperandValues();
 418     if (operands.len != 3) return false;
 419 
 420     var arena_state = alloc_arena.Arena.init(rewriter.allocator);
 421     defer arena_state.deinit();
 422     const arena = arena_state.allocator();
 423 
 424     const data_type = try dialect_mod.decodeTensorType(arena, operands[0].type);
 425     const indices_type = try dialect_mod.decodeTensorType(arena, operands[1].type);
 426     const updates_type = try dialect_mod.decodeTensorType(arena, operands[2].type);
 427     const output_type = try dialect_mod.decodeTensorType(arena, result.type);
 428     if (indices_type.dtype != .i32) return false;
 429     if (updates_type.dtype != data_type.dtype) return false;
 430     if (output_type.dtype != data_type.dtype) return false;
 431     if (data_type.dims.len == 0 or indices_type.dims.len != 1) return false;
 432     if (updates_type.dims.len != data_type.dims.len) return false;
 433     if (output_type.dims.len != data_type.dims.len) return false;
 434 
 435     const scatter_op = dialect_mod.AccyDialect.ScatterOp{ .op = op };
 436     const axis_value = scatter_op.getAxis() orelse return false;
 437     if (axis_value < 0) return false;
 438     const axis = std.math.cast(usize, axis_value) orelse return false;
 439     if (axis >= data_type.dims.len) return false;
 440 
 441     if (!std.mem.eql(i64, output_type.dims, data_type.dims)) return false;
 442     if (!std.mem.eql(i64, updates_type.dims[0..axis], data_type.dims[0..axis])) return false;
 443     if (updates_type.dims[axis] != indices_type.dims[0]) return false;
 444     if (!std.mem.eql(i64, updates_type.dims[axis + 1 ..], data_type.dims[axis + 1 ..])) return false;
 445 
 446     const outer = dimsProduct(data_type.dims[0..axis]) orelse return false;
 447     const axis_size = dimExtent(data_type.dims[axis]) orelse return false;
 448     const updates = dimExtent(indices_type.dims[0]) orelse return false;
 449     const inner = dimsProduct(data_type.dims[axis + 1 ..]) orelse return false;
 450 
 451     var scatter_schedule = options.scatter_schedule;
 452     if (scatter_schedule == null) {
 453         const family_tuning = options.family_tuning orelse return false;
 454         const thread_blocks = (try kernel_library.indexing.resolveScatterSchedule(rewriter.allocator, family_tuning.*, .{
 455             .outer = outer,
 456             .axis_size = axis_size,
 457             .updates = updates,
 458             .inner = inner,
 459             .dtype = data_type.dtype,
 460         })) orelse return false;
 461         scatter_schedule = .{ .thread_blocks = thread_blocks };
 462     }
 463 
 464     var selected = (try kernel_library.selectOwned(rewriter.allocator, .{ .scatter = .{
 465         .dtype = data_type.dtype,
 466         .outer = outer,
 467         .axis_size = axis_size,
 468         .updates = updates,
 469         .inner = inner,
 470         .schedule = scatter_schedule,
 471     } })) orelse return false;
 472     defer selected.deinit();
 473 
 474     const instance = kernel_library.indexing.scatterInstanceFromSpecialization(
 475         selected.descriptor.metadata.specialization,
 476     ) orelse return false;
 477     const runtime_scalars = call_preparation.catalogCallScalars(
 478         5,
 479         try kernel_library.indexing.scatterRuntimeArguments(instance),
 480     );
 481     const result_types = [_]ir.Type{result.type};
 482     const call = try call_preparation.insertCatalogCall(rewriter, .{
 483         .descriptor = selected.descriptor,
 484         .operands = operands,
 485         .result_types = &result_types,
 486         .options = .{ .runtime_scalars = runtime_scalars[0..] },
 487     });
 488     try rewriter.replaceOpWithValue(op, call.getFirstResult());
 489     return true;
 490 }
 491 
 492 fn lowerKnownKernelLibraryScatterAdd(
 493     op: *ir.Operation,
 494     rewriter: *rewrite.PatternRewriter,
 495     options: Options,
 496 ) !bool {
 497     if (op.getNumResults() != 1) return false;
 498     const result = op.getResult(0) orelse return false;
 499     const operands = op.getOperandValues();
 500     if (operands.len != 3) return false;
 501 
 502     var arena_state = alloc_arena.Arena.init(rewriter.allocator);
 503     defer arena_state.deinit();
 504     const arena = arena_state.allocator();
 505 
 506     const data_type = try dialect_mod.decodeTensorType(arena, operands[0].type);
 507     const indices_type = try dialect_mod.decodeTensorType(arena, operands[1].type);
 508     const updates_type = try dialect_mod.decodeTensorType(arena, operands[2].type);
 509     const output_type = try dialect_mod.decodeTensorType(arena, result.type);
 510     if (indices_type.dtype != .i32) return false;
 511     if (updates_type.dtype != data_type.dtype) return false;
 512     if (output_type.dtype != data_type.dtype) return false;
 513     if (data_type.dims.len == 0 or indices_type.dims.len != 1) return false;
 514     if (updates_type.dims.len != data_type.dims.len) return false;
 515     if (output_type.dims.len != data_type.dims.len) return false;
 516 
 517     const scatter_add_op = dialect_mod.AccyDialect.ScatterAddOp{ .op = op };
 518     const axis_value = scatter_add_op.getAxis() orelse return false;
 519     if (axis_value < 0) return false;
 520     const axis = std.math.cast(usize, axis_value) orelse return false;
 521     if (axis >= data_type.dims.len) return false;
 522 
 523     if (!std.mem.eql(i64, output_type.dims, data_type.dims)) return false;
 524     if (!std.mem.eql(i64, updates_type.dims[0..axis], data_type.dims[0..axis])) return false;
 525     if (updates_type.dims[axis] != indices_type.dims[0]) return false;
 526     if (!std.mem.eql(i64, updates_type.dims[axis + 1 ..], data_type.dims[axis + 1 ..])) return false;
 527 
 528     const outer = dimsProduct(data_type.dims[0..axis]) orelse return false;
 529     const axis_size = dimExtent(data_type.dims[axis]) orelse return false;
 530     const updates = dimExtent(indices_type.dims[0]) orelse return false;
 531     const inner = dimsProduct(data_type.dims[axis + 1 ..]) orelse return false;
 532 
 533     var scatter_add_schedule = options.scatter_add_schedule;
 534     if (scatter_add_schedule == null) {
 535         const family_tuning = options.family_tuning orelse return false;
 536         const resolved = (try kernel_library.indexing.resolveScatterAddSchedule(rewriter.allocator, family_tuning.*, .{
 537             .outer = outer,
 538             .axis_size = axis_size,
 539             .updates = updates,
 540             .inner = inner,
 541             .dtype = data_type.dtype,
 542         })) orelse return false;
 543         scatter_add_schedule = switch (resolved.variant) {
 544             .direct => .{ .thread_blocks = resolved.threads },
 545             .shared_bins => .{ .shared_bins = resolved.threads },
 546         };
 547     }
 548 
 549     var selected = (try kernel_library.selectOwned(rewriter.allocator, .{ .scatter_add = .{
 550         .dtype = data_type.dtype,
 551         .outer = outer,
 552         .axis_size = axis_size,
 553         .updates = updates,
 554         .inner = inner,
 555         .schedule = scatter_add_schedule,
 556     } })) orelse return false;
 557     defer selected.deinit();
 558 
 559     const instance = kernel_library.indexing.scatterAddInstanceFromSpecialization(
 560         selected.descriptor.metadata.specialization,
 561     ) orelse return false;
 562     const runtime_scalars = call_preparation.catalogCallScalars(
 563         5,
 564         try kernel_library.indexing.scatterAddRuntimeArguments(instance),
 565     );
 566     const result_types = [_]ir.Type{result.type};
 567     const call = try call_preparation.insertCatalogCall(rewriter, .{
 568         .descriptor = selected.descriptor,
 569         .operands = operands,
 570         .result_types = &result_types,
 571         .options = .{
 572             .operand_effects = &.{ .read_write, .read, .read },
 573             .result_aliases = &.{0},
 574             .runtime_scalars = runtime_scalars[0..],
 575         },
 576     });
 577     try rewriter.replaceOpWithValue(op, call.getFirstResult());
 578     return true;
 579 }
 580 
 581 fn dimExtent(dim: i64) ?u64 {
 582     if (dim <= 0) return null;
 583     return @intCast(dim);
 584 }
 585 
 586 fn dimsProduct(dims: []const i64) ?u64 {
 587     var product: u64 = 1;
 588     for (dims) |dim| {
 589         const extent = dimExtent(dim) orelse return null;
 590         product = std.math.mul(u64, product, extent) catch return null;
 591     }
 592     return product;
 593 }
 594 
 595 const testing = std.testing;
 596 const semantic = accy_choir.semantic;
 597 
 598 const IndexingCase = struct {
 599     kind: IndexingKind = .gather,
 600     copies: u32 = 1,
 601     dtype: choir_abi.DType = .f32,
 602     threads: u32 = 4,
 603     shared: bool = false,
 604 
 605     fn outer(self: IndexingCase) u64 {
 606         return if (self.shared) 1 else 2;
 607     }
 608     fn inner(self: IndexingCase) u64 {
 609         return if (self.shared) 1 else 3;
 610     }
 611     fn entries(self: IndexingCase) u64 {
 612         return if (self.shared) 128 else 5;
 613     }
 614 
 615     fn module(self: IndexingCase) !*semantic.SemanticModule {
 616         var builder = try semantic.Builder.init(testing.allocator, .standard);
 617         defer builder.deinit();
 618         const data = try builder.tensor(
 619             self.dtype,
 620             &.{ @intCast(self.outer()), 8, @intCast(self.inner()) },
 621         );
 622         const indices = try builder.tensor(.i32, &.{@intCast(self.entries())});
 623         const updates = try builder.tensor(
 624             self.dtype,
 625             &.{ @intCast(self.outer()), @intCast(self.entries()), @intCast(self.inner()) },
 626         );
 627         const result = if (self.kind == .gather) updates else data;
 628         const inputs: []const ir.Type = if (self.kind == .gather)
 629             &.{ data, indices }
 630         else
 631             &.{ data, indices, updates };
 632         var function = try builder.beginFunction("indexing_accounting", inputs, &.{result});
 633         var value: *ir.Value = undefined;
 634         for (0..self.copies) |_| {
 635             value = switch (self.kind) {
 636                 .gather => try function.gather(
 637                     function.parameter(0),
 638                     function.parameter(1),
 639                     result,
 640                     1,
 641                 ),
 642                 .scatter => try function.scatter(
 643                     function.parameter(0),
 644                     function.parameter(1),
 645                     function.parameter(2),
 646                     result,
 647                     1,
 648                 ),
 649                 .scatter_add => try function.scatterAdd(
 650                     function.parameter(0),
 651                     function.parameter(1),
 652                     function.parameter(2),
 653                     result,
 654                     1,
 655                 ),
 656             };
 657         }
 658         try function.return_(&.{value});
 659         try function.finish();
 660         return builder.finish();
 661     }
 662 
 663     fn options(self: IndexingCase) Options {
 664         var result: Options = .{ .kernel_library = .enabled };
 665         switch (self.kind) {
 666             .gather => result.gather_schedule = .{ .thread_blocks = self.threads },
 667             .scatter => result.scatter_schedule = .{ .thread_blocks = self.threads },
 668             .scatter_add => result.scatter_add_schedule = if (self.shared)
 669                 .{ .shared_bins = self.threads }
 670             else
 671                 .{ .thread_blocks = self.threads },
 672         }
 673         return result;
 674     }
 675 
 676     fn gather(self: IndexingCase) kernel_library.indexing.Gather {
 677         return .{
 678             .dtype = self.dtype,
 679             .outer = self.outer(),
 680             .axis_size = 8,
 681             .gathered = self.entries(),
 682             .inner = self.inner(),
 683             .threads = self.threads,
 684         };
 685     }
 686 
 687     fn scatter(self: IndexingCase) kernel_library.indexing.Scatter {
 688         return .{
 689             .dtype = self.dtype,
 690             .outer = self.outer(),
 691             .axis_size = 8,
 692             .updates = self.entries(),
 693             .inner = self.inner(),
 694             .threads = self.threads,
 695         };
 696     }
 697 
 698     fn scatterAdd(self: IndexingCase) kernel_library.indexing.ScatterAdd {
 699         return .{
 700             .dtype = self.dtype,
 701             .outer = self.outer(),
 702             .axis_size = 8,
 703             .updates = self.entries(),
 704             .inner = self.inner(),
 705             .threads = self.threads,
 706             .variant = if (self.shared) .shared_bins else .direct,
 707         };
 708     }
 709 
 710     fn query(self: IndexingCase) kernel_library.CatalogQuery {
 711         const config = self.options();
 712         return switch (self.kind) {
 713             .gather => .{ .gather = .{
 714                 .dtype = self.dtype,
 715                 .outer = self.outer(),
 716                 .axis_size = 8,
 717                 .gathered = self.entries(),
 718                 .inner = self.inner(),
 719                 .schedule = config.gather_schedule,
 720             } },
 721             .scatter => .{ .scatter = .{
 722                 .dtype = self.dtype,
 723                 .outer = self.outer(),
 724                 .axis_size = 8,
 725                 .updates = self.entries(),
 726                 .inner = self.inner(),
 727                 .schedule = config.scatter_schedule,
 728             } },
 729             .scatter_add => .{ .scatter_add = .{
 730                 .dtype = self.dtype,
 731                 .outer = self.outer(),
 732                 .axis_size = 8,
 733                 .updates = self.entries(),
 734                 .inner = self.inner(),
 735                 .schedule = config.scatter_add_schedule,
 736             } },
 737         };
 738     }
 739 
 740     fn target(self: IndexingCase) ![]u8 {
 741         return switch (self.kind) {
 742             .gather => kernel_library.indexing.gatherFamilyTarget(testing.allocator, self.gather()),
 743             .scatter => kernel_library.indexing.scatterFamilyTarget(
 744                 testing.allocator,
 745                 self.scatter(),
 746             ),
 747             .scatter_add => kernel_library.indexing.scatterAddFamilyTarget(
 748                 testing.allocator,
 749                 self.scatterAdd(),
 750             ),
 751         };
 752     }
 753 
 754     fn tuningKey(self: IndexingCase, device: u64) !kernel_library.tuning.FamilyTuningKey {
 755         return switch (self.kind) {
 756             .gather => kernel_library.indexing.gatherFamilyTuningKey(
 757                 testing.allocator,
 758                 device,
 759                 self.gather(),
 760             ),
 761             .scatter => kernel_library.indexing.scatterFamilyTuningKey(
 762                 testing.allocator,
 763                 device,
 764                 self.scatter(),
 765             ),
 766             .scatter_add => kernel_library.indexing.scatterAddFamilyTuningKey(
 767                 testing.allocator,
 768                 device,
 769                 self.scatterAdd(),
 770             ),
 771         };
 772     }
 773 
 774     fn lastCandidate(self: IndexingCase) IndexingCase {
 775         const candidates = switch (self.kind) {
 776             .gather => kernel_library.indexing.gatherThreadCandidatesForTotal(
 777                 self.gather().total(),
 778             ),
 779             .scatter => kernel_library.indexing.scatterThreadCandidatesForTotal(
 780                 self.scatter().total(),
 781             ),
 782             .scatter_add => kernel_library.indexing.scatterAddThreadCandidatesForTotal(
 783                 self.scatterAdd().total(),
 784             ),
 785         };
 786         var result = self;
 787         result.threads = candidates.slice()[candidates.count - 1];
 788         return result;
 789     }
 790 
 791     fn check(self: IndexingCase, module_: *semantic.SemanticModule, lowered: bool) !void {
 792         try module_.verify();
 793         var witness: IndexingWitness = .{ .case = self };
 794         _ = try module_.choir_module.walk(
 795             .{ .order = .pre_order },
 796             &witness,
 797             IndexingWitness.visit,
 798         );
 799         try testing.expectEqual(@as(u64, if (lowered) 0 else self.copies), witness.originals);
 800         try testing.expectEqual(@as(u64, if (lowered) self.copies else 0), witness.calls);
 801     }
 802 };
 803 
 804 const IndexingWitness = struct {
 805     case: IndexingCase,
 806     originals: u64 = 0,
 807     calls: u64 = 0,
 808 
 809     fn visit(self: *IndexingWitness, op: *ir.Operation) !ir.WalkResult {
 810         if (IndexingKind.fromOperation(op) != null) self.originals += 1;
 811         if (!std.mem.eql(u8, op.name.name, "accy.kernel_call")) return .advance;
 812         self.calls += 1;
 813         const expected = try self.case.target();
 814         defer testing.allocator.free(expected);
 815         const target = op.getAttr("target").?.cast(ir.Attribute.DialectAttr).?.payload;
 816         try testing.expectEqualStrings(expected, target);
 817         const scalars = (try dialect_mod.AccyDialect.kernelCallRuntimeScalars(op)).?;
 818         const total = self.case.outer() * self.case.inner() *
 819             @as(u64, if (self.case.kind == .scatter) 8 else self.case.entries());
 820         const values = [_]u64{
 821             self.case.outer(), 8, self.case.entries(), self.case.inner(), total,
 822         };
 823         try testing.expectEqual(values.len, scalars.count);
 824         for (values, scalars.slice()) |value, scalar| {
 825             try testing.expectEqual(.u32, scalar.kind);
 826             try testing.expectEqual(value, scalar.bits);
 827         }
 828         const effects = op.getAttr("operand_effects").?.cast(ir.Attribute.DialectAttr).?.payload;
 829         const aliases = op.getAttr("result_aliases").?.cast(ir.Attribute.DialectAttr).?.payload;
 830         const write = self.case.kind == .scatter_add;
 831         const effect: accy_choir.semantics.KernelOperandEffect = if (write) .read_write else .read;
 832         try testing.expectEqual(@backingInt(effect), effects[0]);
 833         const alias = std.mem.bytesToValue(i64, aliases[0..8]);
 834         try testing.expectEqual(@as(i64, if (write) 0 else -1), alias);
 835         return .advance;
 836     }
 837 };
 838 
 839 fn checkIndexingAccounting(admitted: bool, constructor: u32) !void {
 840     const allocator = testing.allocator;
 841     const revision = choir.product.revision;
 842     const fixture: IndexingCase = .{};
 843     const module = try fixture.module();
 844     defer module.deinit();
 845     const root = module.choir_module;
 846     const before = try choir.bytecode.encodeModule(allocator, root);
 847     defer allocator.free(before);
 848     const options: Options = switch (constructor) {
 849         0 => .{},
 850         3 => .{ .kernel_library = .enabled },
 851         else => fixture.options(),
 852     };
 853     const bounds = try indexingWork(.{ .operation = root, .state = &options });
 854     var allowance = revision.WorkVector.uniform(1 << 40);
 855     if (!admitted) allowance.structural_visits = bounds.work.structural_visits - 1;
 856     const ledger = try revision.AccountingV1.create(allocator, .{
 857         .allowance = allowance,
 858         .workspace = 1 << 24,
 859         .events = 4,
 860     }, &.{.{ .name = indexing_lowering_pass_name, .version = 1 }});
 861     defer ledger.destroy();
 862     var cache = try passes.AnalysisCache.initAccounted(
 863         allocator,
 864         null,
 865         ledger,
 866         .{ .context = module.context() },
 867         0,
 868     );
 869     defer cache.deinit();
 870     var manager = passes.PassManager.init(allocator);
 871     defer manager.deinit();
 872     try manager.addPass(switch (constructor) {
 873         0 => indexingLoweringPass(),
 874         1, 3 => indexingLoweringPassWithOptions(&options),
 875         2 => try indexingLoweringPassFromOptions(allocator, .{ .assignments = &.{
 876             .{ .name = "kernel-library", .value = "enabled" },
 877             .{ .name = "gather-thread-blocks", .value = "4" },
 878         } }),
 879         else => unreachable,
 880     });
 881     const result = manager.runWithAnalysisCache(root, module.context(), &cache, .{});
 882     try testing.expectEqual(if (admitted) passes.PassResult.success else .failure, result);
 883     if (admitted) {
 884         try ledger.producersComplete();
 885         try testing.expect(!ledger.view().missing_work_contract);
 886         try fixture.check(module, constructor == 1 or constructor == 2);
 887     } else {
 888         try testing.expectEqual(.exhausted, ledger.view().outcome);
 889         try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
 890     }
 891     if (!admitted or constructor == 0 or constructor == 3) {
 892         const after = try choir.bytecode.encodeModule(allocator, root);
 893         defer allocator.free(after);
 894         try testing.expectEqualSlices(u8, before, after);
 895     }
 896 }
 897 
 898 test "indexing lowering accounts constructors and refuses before mutation" {
 899     for (0..4) |constructor| {
 900         try checkIndexingAccounting(false, @intCast(constructor));
 901         try checkIndexingAccounting(true, @intCast(constructor));
 902     }
 903 }
 904 
 905 fn checkIndexingStorage(fixture: IndexingCase, options: Options, lowered: bool) !void {
 906     const allocator = testing.allocator;
 907     const module = try fixture.module();
 908     defer module.deinit();
 909     const bounds = try indexingWork(.{ .operation = module.choir_module, .state = &options });
 910     const bytes = try allocator.alloc(u8, @intCast(bounds.workspace));
 911     defer allocator.free(bytes);
 912     var storage = @import("alloc_fixed").Tracked.init(bytes);
 913     var cache = passes.AnalysisCache.init(allocator, null);
 914     defer cache.deinit();
 915     var context = passes.PassContext.init(module.choir_module, module.context(), allocator, &cache);
 916     defer context.deinit();
 917     context.allocator = storage.allocator();
 918     defer context.allocator = allocator;
 919     const result = runIndexingLoweringWithOptions(&context, options);
 920     try testing.expect(!storage.exhausted);
 921     try testing.expectEqual(null, module.context().exhaustedSegment());
 922     try testing.expectEqual(.success, result);
 923     try testing.expect(storage.status().high_water_bytes <= bounds.workspace);
 924     try fixture.check(module, lowered);
 925 }
 926 
 927 test "indexing lowering scratch covers scheduled producers and scatter add aliases" {
 928     for (comptime std.meta.tags(IndexingKind)) |kind| {
 929         for ([_]u32{ 1, 17 }) |copies| {
 930             const fixture: IndexingCase = .{ .kind = kind, .copies = copies };
 931             try checkIndexingStorage(fixture, fixture.options(), true);
 932         }
 933     }
 934     const shared: IndexingCase = .{ .kind = .scatter_add, .shared = true, .copies = 17 };
 935     try checkIndexingStorage(shared, shared.options(), true);
 936 }
 937 
 938 const indexing_descriptor_cases = [_]IndexingCase{
 939     .{},
 940     .{ .kind = .scatter },
 941     .{ .kind = .scatter_add },
 942     .{ .kind = .scatter_add, .shared = true },
 943     .{ .dtype = .f16 },
 944     .{ .kind = .scatter, .dtype = .f16 },
 945     .{ .kind = .scatter_add, .dtype = .i32 },
 946     .{ .kind = .scatter_add, .dtype = .i32, .shared = true },
 947 };
 948 
 949 test "indexing lowering descriptor storage covers selected families and element types" {
 950     for (indexing_descriptor_cases) |case| {
 951         const bytes = try testing.allocator.alloc(u8, @intCast(indexingDescriptorStorage()));
 952         defer testing.allocator.free(bytes);
 953         var storage = @import("alloc_fixed").Tracked.init(bytes);
 954         var selected = (try kernel_library.selectOwned(storage.allocator(), case.query())) orelse
 955             return error.TestExpectedDescriptor;
 956         defer selected.deinit();
 957         const target = try case.target();
 958         defer testing.allocator.free(target);
 959         try testing.expectEqualStrings(target, selected.descriptor.metadata.target);
 960         try testing.expect(!storage.exhausted);
 961         try testing.expect(storage.status().high_water_bytes <= bytes.len);
 962         try testing.expect(storage.status().high_water_bytes > 0);
 963         try checkIndexingStorage(case, case.options(), true);
 964     }
 965 }
 966 
 967 fn checkIndexingTuning(fixture_: IndexingCase, mode: enum { hit, miss, stale, precedence }) !void {
 968     const fixture = fixture_.lastCandidate();
 969     const caps = familyTuningTestCapabilities();
 970     const device = kernel_library.tuning.deviceFingerprint(caps);
 971     const key = try fixture.tuningKey(device);
 972     const target = try fixture.target();
 973     defer testing.allocator.free(target);
 974     var records: [129]kernel_library.tuning.FamilyTuningRecord = undefined;
 975     for (&records, 0..) |*record, index| {
 976         record.* = .{
 977             .key = key,
 978             .target = target,
 979             .winner_median_ns = 1,
 980             .runner_up_median_ns = 2,
 981             .sample_count = 3,
 982         };
 983         record.key.device_fingerprint = device +% index +% 1;
 984     }
 985     records[128].key = key;
 986     if (mode == .stale) records[128].target = "unavailable.target";
 987     const reader = kernel_library.tuning.FamilyTuningReader.init(
 988         caps,
 989         .{ .records = if (mode == .miss) records[0..128] else &records },
 990     );
 991     var options: Options = .{ .kernel_library = .enabled, .family_tuning = &reader };
 992     if (mode == .precedence) {
 993         options = fixture_.options();
 994         options.family_tuning = &reader;
 995     }
 996     if (mode == .hit and fixture.kind == .gather) {
 997         try checkTuningCharge(fixture, options, false);
 998         try checkTuningCharge(fixture, options, true);
 999     }
1000     const expected = if (mode == .precedence) fixture_ else fixture;
1001     try checkIndexingStorage(expected, options, mode == .hit or mode == .precedence);
1002 }
1003 
1004 fn checkTuningCharge(fixture: IndexingCase, options: Options, admitted: bool) !void {
1005     const allocator = testing.allocator;
1006     const module = try fixture.module();
1007     defer module.deinit();
1008     const root = module.choir_module;
1009     const before = try choir.bytecode.encodeModule(allocator, root);
1010     defer allocator.free(before);
1011     var small_reader = options.family_tuning.?.*;
1012     small_reader.table.records = small_reader.table.records[0..1];
1013     var small_options = options;
1014     small_options.family_tuning = &small_reader;
1015     const small = try indexingWork(.{ .operation = root, .state = &small_options });
1016     const revision = choir.product.revision;
1017     var allowance = revision.WorkVector.uniform(1 << 40);
1018     if (!admitted) allowance.structural_visits = small.work.structural_visits;
1019     const ledger = try revision.AccountingV1.create(allocator, .{
1020         .allowance = allowance,
1021         .workspace = 1 << 24,
1022         .events = 4,
1023     }, &.{.{ .name = indexing_lowering_pass_name, .version = 1 }});
1024     defer ledger.destroy();
1025     var cache = try passes.AnalysisCache.initAccounted(
1026         allocator,
1027         null,
1028         ledger,
1029         .{ .context = module.context() },
1030         0,
1031     );
1032     defer cache.deinit();
1033     var manager = passes.PassManager.init(allocator);
1034     defer manager.deinit();
1035     try manager.addPass(indexingLoweringPassWithOptions(&options));
1036     const result = manager.runWithAnalysisCache(root, module.context(), &cache, .{});
1037     try testing.expectEqual(if (admitted) passes.PassResult.success else .failure, result);
1038     if (admitted) {
1039         try ledger.producersComplete();
1040         try testing.expect(!ledger.view().missing_work_contract);
1041         try testing.expectEqual(@as(u64, 1), ledger.view().executed.counters.pass_runs);
1042         try fixture.check(module, true);
1043     } else {
1044         try testing.expectEqual(.exhausted, ledger.view().outcome);
1045         try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
1046         const after = try choir.bytecode.encodeModule(allocator, root);
1047         defer allocator.free(after);
1048         try testing.expectEqualSlices(u8, before, after);
1049     }
1050 }
1051 
1052 test "indexing lowering scratch covers tuning table hits misses and schedule precedence" {
1053     for (comptime std.meta.tags(IndexingKind)) |kind| {
1054         const fixture: IndexingCase = .{ .kind = kind, .copies = 17 };
1055         try checkIndexingTuning(fixture, .hit);
1056         try checkIndexingTuning(fixture, .miss);
1057         try checkIndexingTuning(fixture, .stale);
1058         try checkIndexingTuning(fixture, .precedence);
1059     }
1060     try checkIndexingTuning(.{ .kind = .scatter_add, .shared = true, .copies = 17 }, .hit);
1061 }
1062 
1063 fn gatherTestModule(allocator: std.mem.Allocator, dtype: choir_abi.DType) !*semantic.SemanticModule {
1064     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1065     defer builder.deinit();
1066     const data_ty = try builder.tensor(dtype, &.{ 2, 8, 3 });
1067     const indices_ty = try builder.tensor(.i32, &.{5});
1068     const out_ty = try builder.tensor(dtype, &.{ 2, 5, 3 });
1069     var fb = try builder.beginFunction("indexing_lowering_gather", &.{ data_ty, indices_ty }, &.{out_ty});
1070     const out = try fb.gather(fb.parameter(0), fb.parameter(1), out_ty, 1);
1071     try fb.return_(&.{out});
1072     try fb.finish();
1073     return try builder.finish();
1074 }
1075 
1076 test "indexing lowering pass keeps gather generic by default" {
1077     const allocator = testing.allocator;
1078 
1079     const module = try gatherTestModule(allocator, .f32);
1080     defer module.deinit();
1081 
1082     var pm = passes.PassManager.init(allocator);
1083     defer pm.deinit();
1084     try pm.addPass(indexingLoweringPass());
1085 
1086     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1087     try module.verify();
1088     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1089     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1090     try testing.expectEqual(@as(u64, 0), pm.stats.passes_modified);
1091 }
1092 
1093 test "indexing lowering pass keeps gather generic without schedule" {
1094     const allocator = testing.allocator;
1095 
1096     const module = try gatherTestModule(allocator, .f32);
1097     defer module.deinit();
1098 
1099     var options = Options{ .kernel_library = .enabled };
1100     var pm = passes.PassManager.init(allocator);
1101     defer pm.deinit();
1102     try pm.addPass(indexingLoweringPassWithOptions(&options));
1103 
1104     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1105     try module.verify();
1106     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1107     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1108 }
1109 
1110 test "indexing lowering pass selects scheduled kernel library gather family" {
1111     const allocator = testing.allocator;
1112 
1113     const module = try gatherTestModule(allocator, .f32);
1114     defer module.deinit();
1115 
1116     var options = Options{
1117         .kernel_library = .enabled,
1118         .gather_schedule = .{ .thread_blocks = 16 },
1119     };
1120     var pm = passes.PassManager.init(allocator);
1121     defer pm.deinit();
1122     try pm.addPass(indexingLoweringPassWithOptions(&options));
1123 
1124     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1125     try module.verify();
1126     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1127     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1128         return error.TestExpectedKernelCall;
1129     };
1130     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1131     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1132     try testing.expectEqualStrings("accy.kernel.indexing.gather_family_16_f32", target.payload);
1133 }
1134 
1135 test "indexing lowering pass keeps integer gather generic with schedule" {
1136     const allocator = testing.allocator;
1137 
1138     const module = try gatherTestModule(allocator, .i32);
1139     defer module.deinit();
1140 
1141     var options = Options{
1142         .kernel_library = .enabled,
1143         .gather_schedule = .{ .thread_blocks = 16 },
1144     };
1145     var pm = passes.PassManager.init(allocator);
1146     defer pm.deinit();
1147     try pm.addPass(indexingLoweringPassWithOptions(&options));
1148 
1149     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1150     try module.verify();
1151     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1152     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1153 }
1154 
1155 fn scatterTestModule(allocator: std.mem.Allocator, dtype: choir_abi.DType) !*semantic.SemanticModule {
1156     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1157     defer builder.deinit();
1158     const data_ty = try builder.tensor(dtype, &.{ 2, 8, 3 });
1159     const indices_ty = try builder.tensor(.i32, &.{5});
1160     const updates_ty = try builder.tensor(dtype, &.{ 2, 5, 3 });
1161     var fb = try builder.beginFunction("indexing_lowering_scatter", &.{ data_ty, indices_ty, updates_ty }, &.{data_ty});
1162     const out = try fb.scatter(fb.parameter(0), fb.parameter(1), fb.parameter(2), data_ty, 1);
1163     try fb.return_(&.{out});
1164     try fb.finish();
1165     return try builder.finish();
1166 }
1167 
1168 fn scatterAddTestModule(allocator: std.mem.Allocator, dtype: choir_abi.DType) !*semantic.SemanticModule {
1169     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
1170     defer builder.deinit();
1171     const data_ty = try builder.tensor(dtype, &.{ 2, 8, 3 });
1172     const indices_ty = try builder.tensor(.i32, &.{5});
1173     const updates_ty = try builder.tensor(dtype, &.{ 2, 5, 3 });
1174     var fb = try builder.beginFunction("indexing_lowering_scatter_add", &.{ data_ty, indices_ty, updates_ty }, &.{data_ty});
1175     const out = try fb.scatterAdd(fb.parameter(0), fb.parameter(1), fb.parameter(2), data_ty, 1);
1176     try fb.return_(&.{out});
1177     try fb.finish();
1178     return try builder.finish();
1179 }
1180 
1181 test "indexing lowering pass keeps scatter generic without schedule" {
1182     const allocator = testing.allocator;
1183 
1184     const module = try scatterTestModule(allocator, .f32);
1185     defer module.deinit();
1186 
1187     var options = Options{ .kernel_library = .enabled, .gather_schedule = .{ .thread_blocks = 16 } };
1188     var pm = passes.PassManager.init(allocator);
1189     defer pm.deinit();
1190     try pm.addPass(indexingLoweringPassWithOptions(&options));
1191 
1192     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1193     try module.verify();
1194     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterOp.operation_name));
1195     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1196 }
1197 
1198 test "indexing lowering pass selects scheduled kernel library scatter family" {
1199     const allocator = testing.allocator;
1200 
1201     const module = try scatterTestModule(allocator, .f32);
1202     defer module.deinit();
1203 
1204     var options = Options{
1205         .kernel_library = .enabled,
1206         .scatter_schedule = .{ .thread_blocks = 16 },
1207     };
1208     var pm = passes.PassManager.init(allocator);
1209     defer pm.deinit();
1210     try pm.addPass(indexingLoweringPassWithOptions(&options));
1211 
1212     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1213     try module.verify();
1214     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterOp.operation_name));
1215     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1216         return error.TestExpectedKernelCall;
1217     };
1218     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1219     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1220     try testing.expectEqualStrings("accy.kernel.indexing.scatter_family_16_f32", target.payload);
1221 }
1222 
1223 test "indexing lowering pass keeps integer scatter generic with schedule" {
1224     const allocator = testing.allocator;
1225 
1226     const module = try scatterTestModule(allocator, .i32);
1227     defer module.deinit();
1228 
1229     var options = Options{
1230         .kernel_library = .enabled,
1231         .scatter_schedule = .{ .thread_blocks = 16 },
1232     };
1233     var pm = passes.PassManager.init(allocator);
1234     defer pm.deinit();
1235     try pm.addPass(indexingLoweringPassWithOptions(&options));
1236 
1237     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1238     try module.verify();
1239     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterOp.operation_name));
1240     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1241 }
1242 
1243 test "indexing lowering pass keeps scatter add generic without schedule" {
1244     const allocator = testing.allocator;
1245 
1246     const module = try scatterAddTestModule(allocator, .f32);
1247     defer module.deinit();
1248 
1249     var options = Options{ .kernel_library = .enabled, .gather_schedule = .{ .thread_blocks = 16 } };
1250     var pm = passes.PassManager.init(allocator);
1251     defer pm.deinit();
1252     try pm.addPass(indexingLoweringPassWithOptions(&options));
1253 
1254     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1255     try module.verify();
1256     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterAddOp.operation_name));
1257     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1258 }
1259 
1260 test "indexing lowering pass selects scheduled kernel library scatter add family" {
1261     const allocator = testing.allocator;
1262 
1263     const module = try scatterAddTestModule(allocator, .f32);
1264     defer module.deinit();
1265 
1266     var options = Options{
1267         .kernel_library = .enabled,
1268         .scatter_add_schedule = .{ .thread_blocks = 16 },
1269     };
1270     var pm = passes.PassManager.init(allocator);
1271     defer pm.deinit();
1272     try pm.addPass(indexingLoweringPassWithOptions(&options));
1273 
1274     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1275     try module.verify();
1276     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterAddOp.operation_name));
1277     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1278         return error.TestExpectedKernelCall;
1279     };
1280     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1281     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1282     try testing.expectEqualStrings("accy.kernel.indexing.scatter_add_family_16_f32", target.payload);
1283 }
1284 
1285 test "indexing lowering pass keeps half scatter add generic with schedule" {
1286     const allocator = testing.allocator;
1287 
1288     const module = try scatterAddTestModule(allocator, .f16);
1289     defer module.deinit();
1290 
1291     var options = Options{
1292         .kernel_library = .enabled,
1293         .scatter_add_schedule = .{ .thread_blocks = 16 },
1294     };
1295     var pm = passes.PassManager.init(allocator);
1296     defer pm.deinit();
1297     try pm.addPass(indexingLoweringPassWithOptions(&options));
1298 
1299     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1300     try module.verify();
1301     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterAddOp.operation_name));
1302     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1303 }
1304 
1305 test "indexing lowering pass rejects oversized gather schedule" {
1306     const allocator = testing.allocator;
1307 
1308     const module = try gatherTestModule(allocator, .f32);
1309     defer module.deinit();
1310 
1311     var options = Options{
1312         .kernel_library = .enabled,
1313         .gather_schedule = .{ .thread_blocks = 64 },
1314     };
1315     var pm = passes.PassManager.init(allocator);
1316     defer pm.deinit();
1317     try pm.addPass(indexingLoweringPassWithOptions(&options));
1318 
1319     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1320     try module.verify();
1321     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1322     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1323 }
1324 
1325 fn familyTuningTestCapabilities() gpu.BackendCapabilities {
1326     return .{ .identity = .{
1327         .backend = .cuda,
1328         .family = .nvidia_cuda,
1329         .name = "pass-test-device",
1330         .vendor_id = 0x10de,
1331         .device_id = 0x2684,
1332     } };
1333 }
1334 
1335 test "indexing lowering pass consults family tuning for gather" {
1336     const allocator = testing.allocator;
1337 
1338     const module = try gatherTestModule(allocator, .f32);
1339     defer module.deinit();
1340 
1341     const caps = familyTuningTestCapabilities();
1342     const device = kernel_library.tuning.deviceFingerprint(caps);
1343     const probe = kernel_library.indexing.Gather{ .outer = 2, .axis_size = 8, .gathered = 5, .inner = 3 };
1344     const thread_candidates = kernel_library.indexing.gatherThreadCandidatesForTotal(30);
1345     try testing.expect(thread_candidates.slice().len >= 1);
1346     var winner = probe;
1347     winner.threads = thread_candidates.slice()[0];
1348     const winner_target = try kernel_library.indexing.gatherFamilyTarget(allocator, winner);
1349     defer allocator.free(winner_target);
1350 
1351     const records = [_]kernel_library.tuning.FamilyTuningRecord{.{
1352         .key = try kernel_library.indexing.gatherFamilyTuningKey(allocator, device, probe),
1353         .target = winner_target,
1354         .winner_median_ns = 500,
1355         .runner_up_median_ns = 900,
1356         .sample_count = 30,
1357     }};
1358     const reader = kernel_library.tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
1359 
1360     var options = Options{
1361         .kernel_library = .enabled,
1362         .family_tuning = &reader,
1363     };
1364     var pm = passes.PassManager.init(allocator);
1365     defer pm.deinit();
1366     try pm.addPass(indexingLoweringPassWithOptions(&options));
1367 
1368     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1369     try module.verify();
1370     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1371     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1372         return error.TestExpectedKernelCall;
1373     };
1374     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1375     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1376     try testing.expectEqualStrings(winner_target, target.payload);
1377 }
1378 
1379 test "indexing lowering pass consults family tuning for scatter" {
1380     const allocator = testing.allocator;
1381 
1382     const module = try scatterTestModule(allocator, .f32);
1383     defer module.deinit();
1384 
1385     const caps = familyTuningTestCapabilities();
1386     const device = kernel_library.tuning.deviceFingerprint(caps);
1387     const probe = kernel_library.indexing.Scatter{ .outer = 2, .axis_size = 8, .updates = 5, .inner = 3 };
1388     const thread_candidates = kernel_library.indexing.scatterThreadCandidatesForTotal(48);
1389     try testing.expect(thread_candidates.slice().len >= 1);
1390     var winner = probe;
1391     winner.threads = thread_candidates.slice()[0];
1392     const winner_target = try kernel_library.indexing.scatterFamilyTarget(allocator, winner);
1393     defer allocator.free(winner_target);
1394 
1395     const records = [_]kernel_library.tuning.FamilyTuningRecord{.{
1396         .key = try kernel_library.indexing.scatterFamilyTuningKey(allocator, device, probe),
1397         .target = winner_target,
1398         .winner_median_ns = 700,
1399         .runner_up_median_ns = 1300,
1400         .sample_count = 30,
1401     }};
1402     const reader = kernel_library.tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
1403 
1404     var options = Options{
1405         .kernel_library = .enabled,
1406         .family_tuning = &reader,
1407     };
1408     var pm = passes.PassManager.init(allocator);
1409     defer pm.deinit();
1410     try pm.addPass(indexingLoweringPassWithOptions(&options));
1411 
1412     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1413     try module.verify();
1414     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterOp.operation_name));
1415     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1416         return error.TestExpectedKernelCall;
1417     };
1418     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1419     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1420     try testing.expectEqualStrings(winner_target, target.payload);
1421 }
1422 
1423 test "indexing lowering pass consults family tuning for scatter add" {
1424     const allocator = testing.allocator;
1425 
1426     const module = try scatterAddTestModule(allocator, .f32);
1427     defer module.deinit();
1428 
1429     const caps = familyTuningTestCapabilities();
1430     const device = kernel_library.tuning.deviceFingerprint(caps);
1431     const probe = kernel_library.indexing.ScatterAdd{ .outer = 2, .axis_size = 8, .updates = 5, .inner = 3, .dtype = .f32 };
1432     const thread_candidates = kernel_library.indexing.scatterAddThreadCandidatesForTotal(30);
1433     try testing.expect(thread_candidates.slice().len >= 1);
1434     var winner = probe;
1435     winner.threads = thread_candidates.slice()[0];
1436     const winner_target = try kernel_library.indexing.scatterAddFamilyTarget(allocator, winner);
1437     defer allocator.free(winner_target);
1438 
1439     const records = [_]kernel_library.tuning.FamilyTuningRecord{.{
1440         .key = try kernel_library.indexing.scatterAddFamilyTuningKey(allocator, device, probe),
1441         .target = winner_target,
1442         .winner_median_ns = 800,
1443         .runner_up_median_ns = 1500,
1444         .sample_count = 30,
1445     }};
1446     const reader = kernel_library.tuning.FamilyTuningReader.init(caps, .{ .records = records[0..] });
1447 
1448     var options = Options{
1449         .kernel_library = .enabled,
1450         .family_tuning = &reader,
1451     };
1452     var pm = passes.PassManager.init(allocator);
1453     defer pm.deinit();
1454     try pm.addPass(indexingLoweringPassWithOptions(&options));
1455 
1456     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1457     try module.verify();
1458     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.ScatterAddOp.operation_name));
1459     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
1460         return error.TestExpectedKernelCall;
1461     };
1462     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
1463     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
1464     try testing.expectEqualStrings(winner_target, target.payload);
1465 }
1466 
1467 test "indexing lowering pass keeps gather generic on family tuning miss" {
1468     const allocator = testing.allocator;
1469 
1470     const module = try gatherTestModule(allocator, .f32);
1471     defer module.deinit();
1472 
1473     const reader = kernel_library.tuning.FamilyTuningReader.init(familyTuningTestCapabilities(), .{});
1474 
1475     var options = Options{
1476         .kernel_library = .enabled,
1477         .family_tuning = &reader,
1478     };
1479     var pm = passes.PassManager.init(allocator);
1480     defer pm.deinit();
1481     try pm.addPass(indexingLoweringPassWithOptions(&options));
1482 
1483     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
1484     try module.verify();
1485     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.GatherOp.operation_name));
1486     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
1487 }