lib/accy/src/preparation/loss.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const alloc_arena = @import("alloc_arena");
  3 const choir = @import("choir");
  4 const accy_choir = @import("../choir/root.zig");
  5 const kernel_library = @import("../kernel/library/root.zig");
  6 const call_preparation = @import("call.zig");
  7 const library_preparation = @import("library.zig");
  8 const dialect_mod = accy_choir.dialect;
  9 
 10 const ir = choir.ir;
 11 const rewrite = ir.rewrite;
 12 const passes = choir.passes;
 13 const work = passes.pass.work;
 14 
 15 pub const KernelLibraryLowering = library_preparation.KernelLibraryLowering;
 16 
 17 pub const Options = struct {
 18     kernel_library: KernelLibraryLowering = .disabled,
 19     row_sparse_cross_entropy_schedule: ?kernel_library.RowSparseCrossEntropySchedule = null,
 20 
 21     pub fn eql(self: Options, other: Options) bool {
 22         return self.kernel_library == other.kernel_library and
 23             std.meta.eql(self.row_sparse_cross_entropy_schedule, other.row_sparse_cross_entropy_schedule);
 24     }
 25 };
 26 
 27 pub const loss_lowering_pass_name = "accy-choir-loss-lower";
 28 pub const loss_lowering_pass_description =
 29     "Lower semantic Accy loss operations onto kernel library catalog calls";
 30 
 31 const kernel_library_option_choices = [_]passes.PassOptionChoice{
 32     .{ .name = "disabled" },
 33     .{ .name = "enabled" },
 34 };
 35 
 36 pub const loss_lowering_pass_options = [_]passes.PassOptionSpec{
 37     .{
 38         .name = "kernel-library",
 39         .description = "Use kernel library calls for supported loss operations",
 40         .kind = .choice,
 41         .choices = &kernel_library_option_choices,
 42         .default_value = "disabled",
 43     },
 44     .{
 45         .name = "row-sparse-cross-entropy-thread-blocks",
 46         .description = "Thread blocks for selected row sparse cross entropy kernels",
 47         .kind = .unsigned,
 48     },
 49 };
 50 
 51 pub fn lossLoweringPass() passes.Pass {
 52     return .{
 53         .name = loss_lowering_pass_name,
 54         .description = loss_lowering_pass_description,
 55         .run_fn = runLossLoweringPass,
 56         .work_contract = loss_work_contract,
 57     };
 58 }
 59 
 60 pub fn lossLoweringPassWithOptions(options: *const Options) passes.Pass {
 61     return .{
 62         .name = loss_lowering_pass_name,
 63         .description = loss_lowering_pass_description,
 64         .state = @constCast(options),
 65         .run_with_state_fn = runLossLoweringPassWithState,
 66         .work_contract = loss_work_contract,
 67     };
 68 }
 69 
 70 pub fn lossLoweringPassFromOptions(allocator: std.mem.Allocator, set: passes.PassOptionSet) anyerror!passes.Pass {
 71     const options = try allocator.create(Options);
 72     errdefer allocator.destroy(options);
 73     options.* = .{
 74         .kernel_library = try kernelLibraryLoweringFromText(set.choiceValue("kernel-library", "disabled")),
 75         .row_sparse_cross_entropy_schedule = if (set.get("row-sparse-cross-entropy-thread-blocks")) |value| .{
 76             .thread_blocks = try parseU32Option(value),
 77         } else null,
 78     };
 79     var pass = lossLoweringPassWithOptions(options);
 80     pass.state_deinit_fn = destroyOptions;
 81     return pass;
 82 }
 83 
 84 const loss_work_contract: work.Contract = .{
 85     .identity = .{ .name = loss_lowering_pass_name, .version = 1 },
 86     .estimate = lossWork,
 87 };
 88 
 89 const LossWork = struct {
 90     candidates: u64 = 0,
 91     type_bytes: u64 = 0,
 92 
 93     fn visit(self: *LossWork, op: *ir.Operation) !ir.WalkResult {
 94         if (!std.mem.eql(
 95             u8,
 96             op.name.name,
 97             dialect_mod.AccyDialect.SparseCrossEntropyOp.operation_name,
 98         )) return .advance;
 99         if (op.getNumResults() != 1 or op.getNumOperands() != 2) return .advance;
100         self.candidates = try work.add(self.candidates, 1);
101         for (op.getOperandValues()) |operand| try self.typeBytes(operand.type);
102         try self.typeBytes(op.getResult(0).?.type);
103         return .advance;
104     }
105 
106     fn typeBytes(self: *LossWork, typ: ir.Type) !void {
107         if (typ.getDialectParamKey()) |key| {
108             self.type_bytes = try work.add(self.type_bytes, key.len);
109         }
110     }
111 };
112 
113 fn lossDescriptorStorage() u64 {
114     const entry = kernel_library.entry;
115     const shape = accy_choir.shape;
116     const arrays = 2 * std.ArrayList(shape.Symbol).growCapacity(2) * @sizeOf(shape.Symbol) +
117         3 * std.ArrayList(shape.Tensor).growCapacity(3) * @sizeOf(shape.Tensor) +
118         2 * std.ArrayList(shape.Fact).growCapacity(2) * @sizeOf(shape.Fact);
119     const records = 3 * @sizeOf(entry.Shape) + 4 * @sizeOf(entry.Axis) +
120         2 * @sizeOf(entry.ScheduleBinding) + 4 * @sizeOf(shape.Expression) +
121         8 * @sizeOf(shape.Term);
122     const arena_traffic = 8 * (arrays + records + 256 + 64 * 128);
123     return arena_traffic + @sizeOf(alloc_arena.Arena) + @sizeOf(shape.Family) + 64;
124 }
125 
126 fn lossWork(input: work.Input) !work.Bounds {
127     const options: *const Options = if (input.state) |state| @ptrCast(@alignCast(state)) else &.{};
128     if (options.kernel_library != .enabled or options.row_sparse_cross_entropy_schedule == null) {
129         return .{ .work = .{ .structural_visits = 1 } };
130     }
131     const counts = try work.Census.inspect(input.operation);
132     var facts: LossWork = .{};
133     _ = try input.operation.walk(.{ .order = .pre_order }, &facts, LossWork.visit);
134     const descriptors = try work.multiply(facts.candidates, lossDescriptorStorage());
135     const decoding = try work.multiply(8, try work.add(
136         try work.multiply(facts.type_bytes, @sizeOf(i64)),
137         try work.multiply(facts.candidates, 3 * 128 + 64),
138     ));
139     const queues = try work.multiply(2, try work.arrayListGrowth(*ir.Operation, facts.candidates));
140     const bytes = try work.add(queues, try work.add(decoding, descriptors));
141     const units = try work.add(try work.add(counts.atoms, counts.input_bytes), 1);
142     const uses = try work.add(try work.add(counts.values, counts.operands), 1);
143     const traversal = try work.multiply(128, try work.multiply(units, uses));
144     const nodes = try work.multiply(facts.candidates, @sizeOf(ir.Operation) + @sizeOf(ir.Value) +
145         2 * @sizeOf(ir.OpOperand) + 8 * @sizeOf(ir.NamedAttribute) + 256 +
146         2 * @sizeOf(dialect_mod.AccyDialect.KernelCallScalar));
147     return .{
148         .work = .{
149             .input_bytes = counts.input_bytes,
150             .output_bytes = nodes,
151             .structural_visits = try work.add(traversal, try work.multiply(2, descriptors)),
152             .allocation_capacity = bytes,
153         },
154         .workspace = bytes,
155     };
156 }
157 
158 fn destroyOptions(raw: ?*anyopaque, allocator: std.mem.Allocator) void {
159     const options: *Options = @ptrCast(@alignCast(raw orelse return));
160     allocator.destroy(options);
161 }
162 
163 fn kernelLibraryLoweringFromText(value: []const u8) !KernelLibraryLowering {
164     if (std.mem.eql(u8, value, "disabled")) return .disabled;
165     if (std.mem.eql(u8, value, "enabled")) return .enabled;
166     return error.InvalidPassOptionValue;
167 }
168 
169 fn parseU32Option(value: []const u8) !u32 {
170     return std.fmt.parseUnsigned(u32, value, 10) catch return error.InvalidPassOptionValue;
171 }
172 
173 fn runLossLoweringPass(pass_ctx: *passes.PassContext) passes.PassResult {
174     return runLossLoweringWithOptions(pass_ctx, .{});
175 }
176 
177 fn runLossLoweringPassWithState(raw: ?*anyopaque, pass_ctx: *passes.PassContext) passes.PassResult {
178     const options: *const Options = @ptrCast(@alignCast(raw orelse return .failure));
179     return runLossLoweringWithOptions(pass_ctx, options.*);
180 }
181 
182 fn runLossLoweringWithOptions(pass_ctx: *passes.PassContext, options: Options) passes.PassResult {
183     if (options.kernel_library != .enabled or options.row_sparse_cross_entropy_schedule == null) {
184         pass_ctx.preserveAllAnalyses();
185         return .success;
186     }
187 
188     var rewriter = rewrite.PatternRewriter.init(pass_ctx.allocator, pass_ctx.ir_ctx);
189     defer rewriter.deinit();
190 
191     var lowered_count: usize = 0;
192     lowerOnOp(pass_ctx.op, &rewriter, options, &lowered_count) catch return .failure;
193     if (lowered_count == 0) {
194         pass_ctx.preserveAllAnalyses();
195     } else {
196         rewriter.finalize(pass_ctx.op);
197         pass_ctx.markModified();
198     }
199     return .success;
200 }
201 
202 fn lowerOnOp(
203     op: *ir.Operation,
204     rewriter: *rewrite.PatternRewriter,
205     options: Options,
206     lowered_count: *usize,
207 ) !void {
208     for (op.regions.items) |*region| {
209         var block_iter = region.getBlocks();
210         while (block_iter.next()) |block| {
211             var current: ?*ir.Operation = @ptrCast(@alignCast(block.operations.head));
212             while (current) |current_op| {
213                 const next = current_op.next_op;
214                 if (current_op.regions.items.len > 0) {
215                     try lowerOnOp(current_op, rewriter, options, lowered_count);
216                 }
217                 if (std.mem.eql(u8, current_op.name.name, dialect_mod.AccyDialect.SparseCrossEntropyOp.operation_name)) {
218                     var guard = rewriter.insertionGuard();
219                     defer guard.deinit();
220                     rewriter.setInsertionPointBefore(current_op);
221                     if (try lowerKnownKernelLibrarySparseCrossEntropy(current_op, rewriter, options)) {
222                         lowered_count.* += 1;
223                     }
224                 }
225                 current = next;
226             }
227         }
228     }
229 }
230 
231 fn lowerKnownKernelLibrarySparseCrossEntropy(
232     op: *ir.Operation,
233     rewriter: *rewrite.PatternRewriter,
234     options: Options,
235 ) !bool {
236     if (op.getNumResults() != 1) return false;
237     const result = op.getResult(0) orelse return false;
238     const operands = op.getOperandValues();
239     if (operands.len != 2) return false;
240 
241     var arena_state = alloc_arena.Arena.init(rewriter.allocator);
242     defer arena_state.deinit();
243     const arena = arena_state.allocator();
244 
245     const logits_type = try dialect_mod.decodeTensorType(arena, operands[0].type);
246     const targets_type = try dialect_mod.decodeTensorType(arena, operands[1].type);
247     const output_type = try dialect_mod.decodeTensorType(arena, result.type);
248     if (targets_type.dtype != .i32) return false;
249     if (output_type.dtype != logits_type.dtype) return false;
250     if (logits_type.dims.len != 2 or targets_type.dims.len != 1 or output_type.dims.len != 1) return false;
251     if (targets_type.dims[0] != logits_type.dims[0]) return false;
252     if (output_type.dims[0] != logits_type.dims[0]) return false;
253 
254     const rows = dimExtent(logits_type.dims[0]) orelse return false;
255     const classes = dimExtent(logits_type.dims[1]) orelse return false;
256 
257     var selected = (try kernel_library.selectOwned(rewriter.allocator, .{ .row_sparse_cross_entropy = .{
258         .dtype = logits_type.dtype,
259         .rows = rows,
260         .classes = classes,
261         .schedule = options.row_sparse_cross_entropy_schedule,
262     } })) orelse return false;
263     defer selected.deinit();
264 
265     const instance = kernel_library.loss.rowSparseCrossEntropyInstanceFromSpecialization(
266         selected.descriptor.metadata.specialization,
267     ) orelse return false;
268     const runtime_scalars = call_preparation.catalogCallScalars(
269         2,
270         try kernel_library.loss.rowSparseCrossEntropyRuntimeArguments(instance),
271     );
272     const result_types = [_]ir.Type{result.type};
273     const call = try call_preparation.insertCatalogCall(rewriter, .{
274         .descriptor = selected.descriptor,
275         .operands = operands,
276         .result_types = &result_types,
277         .options = .{ .runtime_scalars = runtime_scalars[0..] },
278     });
279     try rewriter.replaceOpWithValue(op, call.getFirstResult());
280     return true;
281 }
282 
283 fn dimExtent(dim: i64) ?u64 {
284     if (dim <= 0) return null;
285     return @intCast(dim);
286 }
287 
288 const testing = std.testing;
289 const semantic = accy_choir.semantic;
290 
291 const LossCase = struct {
292     rows: u32 = 6,
293     classes: u32 = 11,
294     threads: u32 = 4,
295     copies: u32 = 1,
296     dtype: accy_choir.semantics.DType = .f32,
297 
298     fn module(self: LossCase) !*semantic.SemanticModule {
299         var builder = try semantic.Builder.init(testing.allocator, .standard);
300         defer builder.deinit();
301         const logits = try builder.tensor(self.dtype, &.{ self.rows, self.classes });
302         const targets = try builder.tensor(.i32, &.{self.rows});
303         const losses = try builder.tensor(self.dtype, &.{self.rows});
304         var function = try builder.beginFunction(
305             "loss_accounting",
306             &.{ logits, targets },
307             &.{losses},
308         );
309         var value: *ir.Value = undefined;
310         for (0..self.copies) |_| {
311             value = try function.sparseCrossEntropy(
312                 function.parameter(0),
313                 function.parameter(1),
314                 losses,
315             );
316         }
317         try function.return_(&.{value});
318         try function.finish();
319         return builder.finish();
320     }
321 
322     fn options(self: LossCase) Options {
323         return .{
324             .kernel_library = .enabled,
325             .row_sparse_cross_entropy_schedule = .{ .thread_blocks = self.threads },
326         };
327     }
328 
329     fn check(self: LossCase, module_: *semantic.SemanticModule, lowered: bool) !void {
330         try module_.verify();
331         const root = module_.choir_module;
332         try testing.expectEqual(
333             @as(usize, if (lowered) 0 else self.copies),
334             ir.inspection.countOperationsNamed(root, "accy.sparse_cross_entropy"),
335         );
336         try testing.expectEqual(
337             @as(usize, if (lowered) self.copies else 0),
338             ir.inspection.countOperationsNamed(root, "accy.kernel_call"),
339         );
340         var witness: LossCallWitness = .{ .case = self };
341         _ = try root.walk(.{ .order = .pre_order }, &witness, LossCallWitness.visit);
342     }
343 };
344 
345 const LossCallWitness = struct {
346     case: LossCase,
347 
348     fn visit(self: *LossCallWitness, op: *ir.Operation) !ir.WalkResult {
349         if (!std.mem.eql(u8, op.name.name, "accy.kernel_call")) return .advance;
350         const target = op.getAttr("target").?.cast(ir.Attribute.DialectAttr).?.payload;
351         var name: [96]u8 = undefined;
352         const expected = try std.fmt.bufPrint(
353             &name,
354             "accy.kernel.loss.row_sparse_cross_entropy_family_{d}_f32",
355             .{self.case.threads},
356         );
357         try testing.expectEqualStrings(expected, target);
358         const scalars = (try dialect_mod.AccyDialect.kernelCallRuntimeScalars(op)).?;
359         try testing.expectEqual(@as(usize, 2), scalars.count);
360         try testing.expectEqual(.u32, scalars.items[0].kind);
361         try testing.expectEqual(.u32, scalars.items[1].kind);
362         try testing.expectEqual(@as(u64, self.case.rows), scalars.items[0].bits);
363         try testing.expectEqual(@as(u64, self.case.classes), scalars.items[1].bits);
364         return .advance;
365     }
366 };
367 
368 fn checkLossAccounting(admitted: bool, constructor: u32) !void {
369     const allocator = testing.allocator;
370     const revision = choir.product.revision;
371     const fixture: LossCase = .{};
372     const module = try fixture.module();
373     defer module.deinit();
374     const root = module.choir_module;
375     const before = try choir.bytecode.encodeModule(allocator, root);
376     defer allocator.free(before);
377     const options: Options = switch (constructor) {
378         0 => .{},
379         3 => .{ .kernel_library = .enabled },
380         else => fixture.options(),
381     };
382     const bounds = try lossWork(.{ .operation = root, .state = &options });
383     var allowance = revision.WorkVector.uniform(1 << 40);
384     if (!admitted) allowance.structural_visits = bounds.work.structural_visits - 1;
385     const ledger = try revision.AccountingV1.create(allocator, .{
386         .allowance = allowance,
387         .workspace = 1 << 24,
388         .events = 4,
389     }, &.{.{ .name = loss_lowering_pass_name, .version = 1 }});
390     defer ledger.destroy();
391     var cache = try passes.AnalysisCache.initAccounted(
392         allocator,
393         null,
394         ledger,
395         .{ .context = module.context() },
396         0,
397     );
398     defer cache.deinit();
399     var manager = passes.PassManager.init(allocator);
400     defer manager.deinit();
401     try manager.addPass(switch (constructor) {
402         0 => lossLoweringPass(),
403         1, 3 => lossLoweringPassWithOptions(&options),
404         2 => try lossLoweringPassFromOptions(allocator, .{ .assignments = &.{
405             .{ .name = "kernel-library", .value = "enabled" },
406             .{ .name = "row-sparse-cross-entropy-thread-blocks", .value = "4" },
407         } }),
408         else => unreachable,
409     });
410     const result = manager.runWithAnalysisCache(root, module.context(), &cache, .{});
411     try testing.expectEqual(if (admitted) passes.PassResult.success else .failure, result);
412     if (admitted) {
413         try ledger.producersComplete();
414         try testing.expect(!ledger.view().missing_work_contract);
415         try testing.expectEqual(@as(u64, 1), ledger.view().executed.counters.pass_runs);
416         try fixture.check(module, constructor == 1 or constructor == 2);
417     } else {
418         try testing.expectEqual(.exhausted, ledger.view().outcome);
419         try testing.expectEqual(@as(u64, 0), manager.stats.pass_runs);
420     }
421     if (!admitted or constructor == 0 or constructor == 3) {
422         const after = try choir.bytecode.encodeModule(allocator, root);
423         defer allocator.free(after);
424         try testing.expectEqualSlices(u8, before, after);
425     }
426 }
427 
428 test "loss lowering accounts all constructors and no-op paths before mutation" {
429     for (0..4) |constructor| {
430         try checkLossAccounting(false, @intCast(constructor));
431         try checkLossAccounting(true, @intCast(constructor));
432     }
433 }
434 
435 fn checkLossStorage(fixture: LossCase, lowered: bool) !void {
436     const allocator = testing.allocator;
437     const module = try fixture.module();
438     defer module.deinit();
439     const options = fixture.options();
440     const bounds = try lossWork(.{ .operation = module.choir_module, .state = &options });
441     const bytes = try allocator.alloc(u8, @intCast(bounds.workspace));
442     defer allocator.free(bytes);
443     var storage = @import("alloc_fixed").Tracked.init(bytes);
444     var cache = passes.AnalysisCache.init(allocator, null);
445     defer cache.deinit();
446     var context = passes.PassContext.init(module.choir_module, module.context(), allocator, &cache);
447     defer context.deinit();
448     context.allocator = storage.allocator();
449     defer context.allocator = allocator;
450     const result = runLossLoweringWithOptions(&context, options);
451     try testing.expect(!storage.exhausted);
452     try testing.expectEqual(null, module.context().exhaustedSegment());
453     try testing.expectEqual(.success, result);
454     try testing.expect(storage.status().high_water_bytes <= bounds.workspace);
455     try testing.expect(storage.status().high_water_bytes > 0);
456     try fixture.check(module, lowered);
457 }
458 
459 const loss_storage_cases = [_]LossCase{
460     .{ .rows = 1, .classes = 1, .threads = 1 },
461     .{},
462     .{ .rows = 129, .classes = 1000, .threads = 128 },
463     .{ .rows = std.math.maxInt(i32), .threads = std.math.maxInt(i32) },
464     .{ .classes = std.math.maxInt(i32) },
465 };
466 
467 test "loss lowering scratch covers family descriptors and repeated producers" {
468     for (loss_storage_cases) |case| {
469         for ([_]u32{ 1, 17 }) |copies| {
470             var fixture = case;
471             fixture.copies = copies;
472             try checkLossStorage(fixture, true);
473         }
474     }
475     try checkLossStorage(.{ .dtype = .f16 }, false);
476     try checkLossStorage(.{ .threads = 0 }, false);
477     try checkLossStorage(.{ .threads = 7 }, false);
478 }
479 
480 test "loss lowering descriptor storage covers owned specialization and family" {
481     const allocator = testing.allocator;
482     for (loss_storage_cases) |case| {
483         const bytes = try allocator.alloc(u8, @intCast(lossDescriptorStorage()));
484         defer allocator.free(bytes);
485         var storage = @import("alloc_fixed").Tracked.init(bytes);
486         var selected = (try kernel_library.selectOwned(storage.allocator(), .{
487             .row_sparse_cross_entropy = .{
488                 .dtype = .f32,
489                 .rows = case.rows,
490                 .classes = case.classes,
491                 .schedule = .{ .thread_blocks = case.threads },
492             },
493         })) orelse return error.TestExpectedDescriptor;
494         defer selected.deinit();
495         const instance = kernel_library.loss.rowSparseCrossEntropyInstanceFromSpecialization(
496             selected.descriptor.metadata.specialization,
497         ).?;
498         try testing.expectEqual(@as(u64, case.rows), instance.rows);
499         try testing.expectEqual(@as(u64, case.classes), instance.classes);
500         try testing.expectEqual(case.threads, instance.threads);
501         try testing.expect(!storage.exhausted);
502         try testing.expect(storage.status().high_water_bytes <= bytes.len);
503         try testing.expect(storage.status().high_water_bytes > 0);
504     }
505 }
506 
507 fn sparseCrossEntropyTestModule(allocator: std.mem.Allocator) !*semantic.SemanticModule {
508     var builder = try semantic.Builder.init(allocator, semantic.Builder.ContextLimits.standard);
509     defer builder.deinit();
510     const logits_ty = try builder.tensor(.f32, &.{ 6, 11 });
511     const targets_ty = try builder.tensor(.i32, &.{6});
512     const losses_ty = try builder.tensor(.f32, &.{6});
513     var fb = try builder.beginFunction("loss_lowering_sparse_cross_entropy", &.{ logits_ty, targets_ty }, &.{losses_ty});
514     const out = try fb.sparseCrossEntropy(fb.parameter(0), fb.parameter(1), losses_ty);
515     try fb.return_(&.{out});
516     try fb.finish();
517     return try builder.finish();
518 }
519 
520 test "loss lowering pass keeps sparse cross entropy generic without schedule" {
521     const allocator = testing.allocator;
522 
523     const module = try sparseCrossEntropyTestModule(allocator);
524     defer module.deinit();
525 
526     var options = Options{ .kernel_library = .enabled };
527     var pm = passes.PassManager.init(allocator);
528     defer pm.deinit();
529     try pm.addPass(lossLoweringPassWithOptions(&options));
530 
531     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
532     try module.verify();
533     try testing.expectEqual(@as(usize, 1), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.SparseCrossEntropyOp.operation_name));
534     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name));
535 }
536 
537 test "loss lowering pass selects scheduled row sparse cross entropy family" {
538     const allocator = testing.allocator;
539 
540     const module = try sparseCrossEntropyTestModule(allocator);
541     defer module.deinit();
542 
543     var options = Options{
544         .kernel_library = .enabled,
545         .row_sparse_cross_entropy_schedule = .{ .thread_blocks = 4 },
546     };
547     var pm = passes.PassManager.init(allocator);
548     defer pm.deinit();
549     try pm.addPass(lossLoweringPassWithOptions(&options));
550 
551     try testing.expectEqual(passes.PassResult.success, pm.run(module.choir_module, module.context()));
552     try module.verify();
553     try testing.expectEqual(@as(usize, 0), ir.inspection.countOperationsNamed(module.choir_module, dialect_mod.AccyDialect.SparseCrossEntropyOp.operation_name));
554     const kernel_call = ir.inspection.findOperationNamed(module.choir_module, dialect_mod.AccyDialect.KernelCallOp.operation_name) orelse {
555         return error.TestExpectedKernelCall;
556     };
557     const target_attr = kernel_call.getAttr("target") orelse return error.TestExpectedKernelCallTarget;
558     const target = target_attr.cast(ir.Attribute.DialectAttr) orelse return error.TestExpectedKernelCallTarget;
559     try testing.expectEqualStrings("accy.kernel.loss.row_sparse_cross_entropy_family_4_f32", target.payload);
560 }