lib/accy/src/tensor/lower.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const choir = @import("choir");
5 const accy = @import("../root.zig");
6 const interpret = @import("interpret/root.zig");
7 const program_mod = @import("program.zig");
8 const trace = @import("trace/root.zig");
9 const unroll = @import("unroll.zig");
10 const types = @import("type/root.zig");
11
12 const ir = choir.ir;
13
14 pub const SemanticModule = accy.choir.SemanticModule;
15 pub const BackendPreparedJob = accy.preparation.BackendPreparedJob;
16 pub const BackendPreparationRunOptions = accy.preparation.BackendPreparationRunOptions;
17 pub const GeneratedScheduleKind = accy.preparation.GeneratedScheduleKind;
18 pub const GeneratedSchedule = accy.preparation.GeneratedSchedule;
19 pub const GeneratedKernelProgram = accy.preparation.GeneratedKernelProgram;
20 pub const GeneratedKernelSummary = accy.preparation.GeneratedKernelSummary;
21 pub const GeneratedKernelSummaries = accy.preparation.GeneratedKernelSummaries;
22 pub const ArtifactJob = accy.artifact.ArtifactJob;
23 pub const BackendHandle = gpu.BackendHandle;
24 pub const CompiledFragment = accy.executable.CompiledFragment;
25 pub const LoadedFragment = accy.executable.LoadedFragment;
26 pub const FragmentCompilerOptions = accy.executable.FragmentCompilerOptions;
27 pub const ArtifactKernelSource = accy.artifact.KernelSource;
28 pub const ArtifactKernelSummary = accy.artifact.KernelSummary;
29 pub const ArtifactKernelSummaries = accy.artifact.KernelSummaries;
30
31 pub const ScatterAddLowering = enum {
32 expanded,
33 semantic_kernel,
34 };
35
36 pub const SparseCrossEntropyLowering = enum {
37 expanded,
38 semantic_kernel,
39 };
40
41 pub const SemanticLoweringOptions = struct {
42 context_limits: accy.choir.SemanticBuilder.ContextLimits = .standard,
43 scatter_add: ScatterAddLowering = .expanded,
44 sparse_cross_entropy: SparseCrossEntropyLowering = .expanded,
45
46 pub fn eql(self: SemanticLoweringOptions, other: SemanticLoweringOptions) bool {
47 return self.scatter_add == other.scatter_add and
48 self.sparse_cross_entropy == other.sparse_cross_entropy and
49 std.meta.eql(self.context_limits, other.context_limits);
50 }
51 };
52
53 pub const FragmentCompilerCacheUpdate = accy.executable.FragmentCompilerCacheUpdate;
54 pub const FragmentCompilationRequest = accy.executable.FragmentCompilationRequest;
55
56 /// A caller keeps one of these per changing program so each recompile reuses the stages that did
57 /// not change. Each refresh lowers the program into a new semantic module inside a bounded compiler
58 /// job, sized by the context limits in the options. The compile that follows reuses a stage only
59 /// when an exact check admits the earlier stage record, the sealed result of one compile stage, for
60 /// the new request. A failed refresh keeps the previous compiled result and the previous loaded
61 /// fragment, the compiled and loaded program ready to launch. `currentPrepared` and
62 /// `currentFragment` return the result of the last successful refresh.
63 pub const FragmentCompilerCache = struct {
64 allocator: std.mem.Allocator,
65 cache: accy.executable.FragmentCompilerCache,
66
67 pub fn init(allocator: std.mem.Allocator, handle: BackendHandle) FragmentCompilerCache {
68 return .{
69 .allocator = allocator,
70 .cache = accy.executable.FragmentCompilerCache.init(allocator, handle),
71 };
72 }
73
74 pub fn deinit(self: *FragmentCompilerCache) void {
75 self.cache.deinit();
76 self.* = undefined;
77 }
78
79 pub fn currentPrepared(
80 self: *const FragmentCompilerCache,
81 ) ?*const accy.preparation.pipeline.BackendPreparedModule {
82 return self.cache.currentPrepared();
83 }
84
85 pub fn currentFragment(self: *const FragmentCompilerCache) ?*LoadedFragment {
86 return self.cache.currentFragment();
87 }
88
89 pub fn refreshFromProgram(
90 self: *FragmentCompilerCache,
91 program: *const program_mod.Program,
92 options: FragmentCompilerOptions,
93 request: FragmentCompilationRequest,
94 report: *accy.preparation.publication.PreparationReport,
95 comptime configuration: choir.product.operation.Configuration,
96 ) !FragmentCompilerCacheUpdate {
97 const source = try toSemanticModuleWithOptions(
98 self.allocator,
99 program,
100 semanticLoweringOptionsFromFragmentCompilerOptions(options),
101 );
102 return self.cache.refreshFromSemanticModule(source, options, request, report, configuration);
103 }
104 };
105
106 pub fn toSemanticModule(
107 allocator: std.mem.Allocator,
108 program: *const program_mod.Program,
109 ) !*SemanticModule {
110 return toSemanticModuleWithOptions(allocator, program, .{});
111 }
112
113 pub fn toSemanticModuleWithOptions(
114 allocator: std.mem.Allocator,
115 program: *const program_mod.Program,
116 options: SemanticLoweringOptions,
117 ) !*SemanticModule {
118 if (program.containsScan()) {
119 if (!programScansLowerAsIterate(program)) {
120 var expanded = try unroll.apply(allocator, program);
121 defer expanded.deinit();
122 return toSemanticModuleWithOptions(allocator, &expanded, options);
123 }
124 }
125
126 var builder = try accy.choir.SemanticBuilder.init(allocator, options.context_limits);
127 errdefer builder.deinit();
128
129 const choir_types = try allocator.alloc(ir.Type, program.valueCount());
130 defer allocator.free(choir_types);
131 for (program.values, 0..) |ty, index| {
132 const shape = try types.extents(allocator, ty.dims);
133 defer allocator.free(@constCast(shape));
134 choir_types[index] = try builder.tensor(ty.dtype, shape);
135 }
136
137 const parameter_types = try allocator.alloc(ir.Type, program.parameters.len);
138 defer allocator.free(parameter_types);
139 for (program.parameters, 0..) |id, index| {
140 parameter_types[index] = choir_types[id.index];
141 }
142
143 const result_types = try allocator.alloc(ir.Type, program.outputs.len);
144 defer allocator.free(result_types);
145 for (program.outputs, 0..) |id, index| {
146 result_types[index] = choir_types[id.index];
147 }
148
149 var function = try builder.beginFunction(program.name, parameter_types, result_types);
150
151 var scan_state = ScanLoweringState{};
152 defer scan_state.deinit(allocator);
153
154 try interpret.run(allocator, program, semantics(allocator, &function, choir_types, &scan_state, options));
155
156 return try builder.finish();
157 }
158
159 fn semanticLoweringOptionsFromPreparationOptions(options: BackendPreparationRunOptions) SemanticLoweringOptions {
160 return .{
161 .scatter_add = if (options.tensor.indexing.kernel_library == .enabled and options.tensor.indexing.scatter_add_schedule != null)
162 .semantic_kernel
163 else
164 .expanded,
165 .sparse_cross_entropy = if (options.tensor.loss.kernel_library == .enabled and options.tensor.loss.row_sparse_cross_entropy_schedule != null)
166 .semantic_kernel
167 else
168 .expanded,
169 };
170 }
171
172 fn semanticLoweringOptionsFromFragmentCompilerOptions(options: FragmentCompilerOptions) SemanticLoweringOptions {
173 return .{
174 .context_limits = options.semantic_context_limits,
175 .scatter_add = if (options.kernel_call_registry != null and options.scatter_add_schedule != null)
176 .semantic_kernel
177 else
178 .expanded,
179 .sparse_cross_entropy = if (options.kernel_call_registry != null and options.row_sparse_cross_entropy_schedule != null)
180 .semantic_kernel
181 else
182 .expanded,
183 };
184 }
185
186 pub fn prepare(
187 allocator: std.mem.Allocator,
188 program: *const program_mod.Program,
189 ) !BackendPreparedJob {
190 return prepareWith(allocator, program, .{});
191 }
192
193 pub fn prepareWith(
194 allocator: std.mem.Allocator,
195 program: *const program_mod.Program,
196 options: BackendPreparationRunOptions,
197 ) !BackendPreparedJob {
198 const module_value = try toSemanticModuleWithOptions(
199 allocator,
200 program,
201 semanticLoweringOptionsFromPreparationOptions(options),
202 );
203 return accy.preparation.prepareBackendJobFromSemanticModule(allocator, module_value, options);
204 }
205
206 pub fn prepareFragment(
207 allocator: std.mem.Allocator,
208 handle: BackendHandle,
209 program: *const program_mod.Program,
210 options: FragmentCompilerOptions,
211 ) !BackendPreparedJob {
212 const module_value = try toSemanticModuleWithOptions(
213 allocator,
214 program,
215 semanticLoweringOptionsFromFragmentCompilerOptions(options),
216 );
217 return accy.executable.prepareFragmentFromSemanticModule(
218 allocator,
219 handle,
220 module_value,
221 options,
222 );
223 }
224
225 pub fn compileFragmentFromPreparedJob(
226 allocator: std.mem.Allocator,
227 handle: BackendHandle,
228 prepared: *BackendPreparedJob,
229 options: FragmentCompilerOptions,
230 ) !*CompiledFragment {
231 return try accy.executable.compileFragmentFromPreparedJob(
232 allocator,
233 handle,
234 prepared,
235 options,
236 );
237 }
238
239 pub fn createArtifactJobFromPreparedJob(
240 allocator: std.mem.Allocator,
241 handle: BackendHandle,
242 prepared: *BackendPreparedJob,
243 options: FragmentCompilerOptions,
244 ) !*ArtifactJob {
245 return accy.executable.createArtifactJobFromPreparedJob(
246 allocator,
247 handle,
248 prepared,
249 options,
250 );
251 }
252
253 pub fn compileFragmentFromArtifactJob(
254 allocator: std.mem.Allocator,
255 artifact_module: *ArtifactJob,
256 ) !*CompiledFragment {
257 return try accy.executable.compileFragmentFromArtifactJob(allocator, artifact_module);
258 }
259
260 pub fn createArtifactJob(
261 allocator: std.mem.Allocator,
262 handle: BackendHandle,
263 program: *const program_mod.Program,
264 options: FragmentCompilerOptions,
265 ) !*ArtifactJob {
266 var prepared = try prepareFragment(allocator, handle, program, options);
267 defer prepared.deinit();
268 return try createArtifactJobFromPreparedJob(allocator, handle, &prepared, options);
269 }
270
271 pub fn compileFragment(
272 allocator: std.mem.Allocator,
273 handle: BackendHandle,
274 program: *const program_mod.Program,
275 options: FragmentCompilerOptions,
276 ) !*CompiledFragment {
277 if (try compileHostLoopScanFragment(allocator, handle, program, options)) |fragment| {
278 return fragment;
279 }
280 const module_value = try toSemanticModuleWithOptions(
281 allocator,
282 program,
283 semanticLoweringOptionsFromFragmentCompilerOptions(options),
284 );
285 return try accy.executable.compileFragmentFromSemanticModule(allocator, handle, module_value, options);
286 }
287
288 fn compileHostLoopScanFragment(
289 allocator: std.mem.Allocator,
290 handle: BackendHandle,
291 program: *const program_mod.Program,
292 options: FragmentCompilerOptions,
293 ) !?*CompiledFragment {
294 const scan = hostLoopScanCandidate(program) orelse return null;
295
296 var body_program = try hostLoopScanBodyProgram(allocator, program.name, scan);
297 defer body_program.deinit();
298
299 const module_value = try toSemanticModuleWithOptions(
300 allocator,
301 &body_program,
302 semanticLoweringOptionsFromFragmentCompilerOptions(options),
303 );
304 var prepared = try accy.executable.prepareFragmentFromSemanticModule(
305 allocator,
306 handle,
307 module_value,
308 options,
309 );
310 defer prepared.deinit();
311
312 try accy.executable.fragment.recordBackendPreparationRun(options.instrumentation, prepared.run);
313
314 const artifact_module = try createArtifactJobFromPreparedJob(allocator, handle, &prepared, options);
315 defer artifact_module.deinit();
316
317 const launch_plan = try hostLoopScanLaunchPlan(
318 allocator,
319 artifact_module.artifactPlan(),
320 @intCast(scan.length),
321 );
322
323 return try accy.executable.compileFragmentFromArtifactJobWithLaunchPlan(
324 allocator,
325 artifact_module,
326 launch_plan,
327 );
328 }
329
330 fn hostLoopScanCandidate(program: *const program_mod.Program) ?program_mod.Scan {
331 if (program.outputs.len != 1) return null;
332
333 var scan_op: ?program_mod.Operation = null;
334 for (program.operations) |op| {
335 switch (op.kind) {
336 .scan => {
337 if (scan_op != null) return null;
338 scan_op = op;
339 },
340 .projection => return null,
341 else => {},
342 }
343 }
344
345 const op = scan_op orelse return null;
346 if (program.outputs[0].index != op.id.index) return null;
347 const scan = op.kind.scan;
348 if (scan.length <= 0) return null;
349 if (scan.inits.len != 1 or scan.body.outputs.len != 1) return null;
350 if (scanLowersAsIterate(program, scan)) return null;
351 if (!hostLoopScanBodySupported(scan)) return null;
352 return scan;
353 }
354
355 fn hostLoopScanBodySupported(scan: program_mod.Scan) bool {
356 var has_complex_body = false;
357 for (scan.body.operations) |op| {
358 switch (op.kind) {
359 .parameter,
360 .constant,
361 .unary,
362 .binary,
363 .iota,
364 .broadcast,
365 .broadcast_in_dim,
366 .reshape,
367 .transpose,
368 .compare,
369 .select,
370 => {},
371 .reduce,
372 .dot_general,
373 => has_complex_body = true,
374 else => return false,
375 }
376 }
377 return has_complex_body;
378 }
379
380 fn hostLoopScanBodyProgram(
381 allocator: std.mem.Allocator,
382 name: []const u8,
383 scan: program_mod.Scan,
384 ) !program_mod.Program {
385 var arena = std.heap.ArenaAllocator.init(allocator);
386 errdefer arena.deinit();
387
388 const arena_allocator = arena.allocator();
389 const body = try program_mod.cloneSubgraph(arena_allocator, scan.body);
390 const body_name = try std.fmt.allocPrint(arena_allocator, "{s}_scan_body", .{name});
391 const outputs = try arena_allocator.alloc(program_mod.Id, 1);
392 outputs[0] = if (@mod(scan.length, 2) == 0) body.parameters[0] else body.outputs[0];
393
394 return .{
395 .arena = arena,
396 .name = body_name,
397 .values = body.values,
398 .operations = body.operations,
399 .parameters = body.parameters,
400 .outputs = outputs,
401 };
402 }
403
404 fn hostLoopScanLaunchPlan(
405 allocator: std.mem.Allocator,
406 artifact_plan: *const accy.artifact.BackendArtifactPlan,
407 trip_count: u64,
408 ) !accy.executable.OwnedLaunchGraphPlan {
409 if (artifact_plan.input_slot_ids.len != 1) return error.UnsupportedOperation;
410 if (artifact_plan.output_slot_ids.len != 1) return error.UnsupportedOperation;
411 if (artifact_plan.kernels.items.len == 0) return error.UnsupportedOperation;
412
413 var launch_plan = try accy.executable.createDataflowLaunchGraphPlan(
414 allocator,
415 artifact_plan,
416 .{},
417 );
418 errdefer launch_plan.deinit();
419
420 const carries = allocator.alloc(accy.executable.LaunchGraphLoopCarry, 1) catch return error.OutOfMemory;
421 errdefer allocator.free(carries);
422
423 const initial_slot_id = artifact_plan.input_slot_ids[0];
424 const output_slot_id = artifact_plan.kernels.items[artifact_plan.kernels.items.len - 1].output_slot_id;
425 const final_slot_id = artifact_plan.output_slot_ids[0];
426 carries[0] = .{
427 .initial_slot_id = initial_slot_id,
428 .input_slot_id = initial_slot_id,
429 .output_slot_id = output_slot_id,
430 .final_slot_id = final_slot_id,
431 };
432 if (accy.executable.launchGraphLoopCarryFinalSlot(carries[0], trip_count) != final_slot_id) {
433 return error.InvalidArtifact;
434 }
435
436 const loops = allocator.alloc(accy.executable.LaunchGraphLoop, 1) catch return error.OutOfMemory;
437 errdefer allocator.free(loops);
438 loops[0] = .{
439 .first_node_index = 0,
440 .node_count = launch_plan.nodes.len,
441 .trip_count = trip_count,
442 .carries = carries,
443 };
444
445 launch_plan.loops = loops;
446 return launch_plan;
447 }
448
449 pub fn module(allocator: std.mem.Allocator) Module {
450 return .{ .allocator = allocator };
451 }
452
453 pub const Module = struct {
454 allocator: std.mem.Allocator,
455
456 pub fn attach(self: @This(), next: anytype) interpret.Layer(trace.Value, @TypeOf(next), ModuleLower) {
457 return interpret.layer(trace.Value, next, ModuleLower{ .allocator = self.allocator });
458 }
459 };
460
461 pub const ModuleLower = struct {
462 allocator: std.mem.Allocator,
463 options: SemanticLoweringOptions = .{},
464
465 pub const Result = *SemanticModule;
466
467 pub fn finish(self: *@This(), ctx: anytype, outputs: []const trace.Value) !Result {
468 var graph = try ctx.default(outputs);
469 defer graph.deinit();
470
471 return toSemanticModuleWithOptions(self.allocator, &graph, self.options);
472 }
473 };
474
475 pub fn semantics(
476 allocator: std.mem.Allocator,
477 function: *accy.choir.semantic.FunctionBuilder,
478 choir_types: []const ir.Type,
479 scan_state: *ScanLoweringState,
480 options: SemanticLoweringOptions,
481 ) Semantics {
482 return .{
483 .allocator = allocator,
484 .function = function,
485 .choir_types = choir_types,
486 .scan_state = scan_state,
487 .options = options,
488 };
489 }
490
491 pub const ScanLoweringState = struct {
492 results: std.AutoHashMapUnmanaged(u32, []const *ir.Value) = .empty,
493
494 fn deinit(self: *ScanLoweringState, allocator: std.mem.Allocator) void {
495 var iter = self.results.valueIterator();
496 while (iter.next()) |values| allocator.free(values.*);
497 self.results.deinit(allocator);
498 self.* = undefined;
499 }
500 };
501
502 pub const Semantics = struct {
503 allocator: std.mem.Allocator,
504 function: *accy.choir.semantic.FunctionBuilder,
505 choir_types: []const ir.Type,
506 scan_state: *ScanLoweringState,
507 options: SemanticLoweringOptions,
508
509 pub const Value = *ir.Value;
510 pub const Result = void;
511
512 pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {
513 switch (step.op.kind) {
514 .scan => |scan| return self.lowerScan(step.op, scan, step.values),
515 .projection => |projection| return self.scan_state.results.get(projection.source.index).?[projection.index],
516 else => {},
517 }
518 var buffer: [program_mod.max_operation_operands]Value = undefined;
519 return self.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));
520 }
521
522 pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {
523 return lowerOperation(self.allocator, self.function, self.choir_types[op.id.index], op, args, self.options);
524 }
525
526 fn lowerScan(self: *@This(), op: *const program_mod.Operation, scan: program_mod.Scan, values: []const Value) !Value {
527 var inits: [program_mod.max_scan_carries]Value = undefined;
528 for (scan.inits, 0..) |init, index| inits[index] = values[init.index];
529
530 var iterate = try self.function.beginIterate(inits[0..scan.inits.len], scan.length);
531 errdefer iterate.body.terminated = true;
532
533 const body_values = try self.allocator.alloc(Value, scan.body.values.len);
534 defer self.allocator.free(body_values);
535
536 for (scan.body.operations) |*body_op| {
537 body_values[body_op.id.index] = switch (body_op.kind) {
538 .parameter => |parameter| iterate.carry(parameter.index),
539 .scan, .projection => return error.ScanRequiresExpansion,
540 else => blk: {
541 var buffer: [program_mod.max_operation_operands]Value = undefined;
542 const args = interpret.arguments(Value, body_op, body_values, &buffer);
543 const result_type = try semanticType(self.allocator, iterate.body.ctx, body_op.result);
544 break :blk try lowerOperation(self.allocator, iterate.inner(), result_type, body_op, args, self.options);
545 },
546 };
547 }
548
549 const output_count = scan.body.outputs.len;
550 const yielded = try self.allocator.alloc(Value, output_count);
551 defer self.allocator.free(yielded);
552 for (scan.body.outputs, 0..) |id, index| yielded[index] = body_values[id.index];
553
554 const predicate = try truePredicate(self.allocator, iterate.inner(), scan.body.typeOf(scan.body.outputs[0]));
555 try iterate.yield_(predicate, yielded);
556
557 const results = try self.allocator.alloc(Value, output_count);
558 errdefer self.allocator.free(results);
559 for (results, 0..) |*result, index| result.* = iterate.result(index);
560 try self.scan_state.results.put(self.allocator, op.id.index, results);
561 return results[0];
562 }
563
564 pub fn finish(self: *@This(), outputs: []const Value) !Result {
565 try self.function.return_(outputs);
566 try self.function.finish();
567 }
568 };
569
570 fn lowerOperation(
571 allocator: std.mem.Allocator,
572 function: *accy.choir.semantic.FunctionBuilder,
573 result_type: ir.Type,
574 op: *const program_mod.Operation,
575 args: []const *ir.Value,
576 options: SemanticLoweringOptions,
577 ) !*ir.Value {
578 return switch (op.kind) {
579 .scan, .projection => error.ScanRequiresExpansion,
580 .parameter => |parameter| function.parameter(parameter.index),
581 .constant => |constant| function.constant(result_type, constant.payload),
582 .unary => |unary| lowerUnary(function, unary.op, args[0]),
583 .binary => |binary| lowerBinary(function, binary.op, args[0], args[1]),
584 .iota => |iota| function.iota(result_type, @intCast(iota.axis)),
585 .broadcast => |broadcast| function.broadcast(args[0], result_type, broadcast.sizes),
586 .broadcast_in_dim => |broadcast| blk: {
587 const result_shape = try types.extents(allocator, op.result.dims);
588 defer allocator.free(@constCast(result_shape));
589 break :blk function.broadcastInDim(
590 args[0],
591 result_type,
592 result_shape,
593 broadcast.broadcast_dims,
594 );
595 },
596 .reshape => |reshape| function.reshape(args[0], result_type, reshape.new_shape),
597 .transpose => |transpose| function.transpose(args[0], result_type, transpose.permutation),
598 .reduce => |reduce| function.reduce(
599 args[0],
600 args[1],
601 result_type,
602 reduce.reducer.name(),
603 reduce.dimensions,
604 ),
605 .gather => |gather| function.gather(args[0], args[1], result_type, gather.axis),
606 .scatter_add => |scatter_add| lowerScatterAdd(
607 allocator,
608 function,
609 result_type,
610 op,
611 scatter_add,
612 args[0],
613 args[1],
614 args[2],
615 options,
616 ),
617 .sparse_cross_entropy => |sparse_cross_entropy| lowerSparseCrossEntropy(
618 allocator,
619 function,
620 result_type,
621 op,
622 sparse_cross_entropy,
623 args[0],
624 args[1],
625 options,
626 ),
627 .compare => |compare| function.compare(args[0], args[1], result_type, switch (compare.direction) {
628 .lt => .lt,
629 .le => .le,
630 .gt => .gt,
631 .ge => .ge,
632 .eq => .eq,
633 .ne => .ne,
634 }),
635 .select => function.select(args[0], args[1], args[2]),
636 .custom_call => |custom| blk: {
637 const effects = @as([program_mod.max_custom_call_operands]accy.choir.semantic.KernelOperandEffect, @splat(.read));
638 const aliases = [_]?usize{null};
639 const call = try function.kernelCall(args, &.{result_type}, .{
640 .target = custom.target,
641 .version = custom.version,
642 .operand_effects = effects[0..args.len],
643 .result_aliases = aliases[0..],
644 });
645 break :blk call.getFirstResult();
646 },
647 .dot_general => |dot| function.dotGeneral(
648 args[0],
649 args[1],
650 result_type,
651 dot.lhs_contract,
652 dot.rhs_contract,
653 dot.lhs_batch,
654 dot.rhs_batch,
655 ),
656 };
657 }
658
659 fn lowerScatterAdd(
660 allocator: std.mem.Allocator,
661 function: *accy.choir.semantic.FunctionBuilder,
662 result_type: ir.Type,
663 op: *const program_mod.Operation,
664 scatter_add: program_mod.ScatterAdd,
665 input: *ir.Value,
666 indices: *ir.Value,
667 updates: *ir.Value,
668 options: SemanticLoweringOptions,
669 ) !*ir.Value {
670 if (options.scatter_add == .semantic_kernel and
671 try scatterAddSemanticKernelEligible(allocator, op, scatter_add, indices, updates))
672 {
673 return function.scatterAdd(input, indices, updates, result_type, scatter_add.axis);
674 }
675 return lowerScatterAddExpanded(allocator, function, result_type, op, scatter_add, input, indices, updates);
676 }
677
678 fn lowerScatterAddExpanded(
679 allocator: std.mem.Allocator,
680 function: *accy.choir.semantic.FunctionBuilder,
681 result_type: ir.Type,
682 op: *const program_mod.Operation,
683 scatter_add: program_mod.ScatterAdd,
684 input: *ir.Value,
685 indices: *ir.Value,
686 updates: *ir.Value,
687 ) !*ir.Value {
688 if (scatter_add.axis < 0) return error.AxisOutOfRange;
689 const axis = std.math.cast(usize, scatter_add.axis) orelse return error.AxisOutOfRange;
690 if (axis >= op.result.rank()) return error.AxisOutOfRange;
691 const input_shape = try types.extents(allocator, op.result.dims);
692 defer allocator.free(@constCast(input_shape));
693 const indices_type = try accy.choir.dialect.decodeTensorType(allocator, indices.type);
694 defer allocator.free(@constCast(indices_type.dims));
695 const updates_type = try accy.choir.dialect.decodeTensorType(allocator, updates.type);
696 defer allocator.free(@constCast(updates_type.dims));
697 const index_rank = indices_type.dims.len;
698 if (updates_type.dims.len != op.result.rank() - 1 + index_rank) return error.RankMismatch;
699
700 const expanded_shape = try allocator.alloc(i64, op.result.rank() + index_rank);
701 var expanded_out: usize = 0;
702 for (input_shape[0..axis]) |dim| {
703 expanded_shape[expanded_out] = dim;
704 expanded_out += 1;
705 }
706 expanded_shape[expanded_out] = input_shape[axis];
707 expanded_out += 1;
708 for (indices_type.dims) |dim| {
709 expanded_shape[expanded_out] = dim;
710 expanded_out += 1;
711 }
712 for (input_shape[axis + 1 ..]) |dim| {
713 expanded_shape[expanded_out] = dim;
714 expanded_out += 1;
715 }
716 defer allocator.free(expanded_shape);
717
718 const expanded_data_type = try accy.choir.dialect.accyTensorType(function.ctx, op.result.dtype, expanded_shape);
719 const expanded_index_type = try accy.choir.dialect.accyTensorType(function.ctx, .i32, expanded_shape);
720 const expanded_mask_type = try accy.choir.dialect.accyTensorType(function.ctx, .i1, expanded_shape);
721
722 const source_positions = try function.iota(expanded_index_type, @intCast(axis));
723 const indices_broadcast_dims = try allocator.alloc(i64, index_rank);
724 defer allocator.free(indices_broadcast_dims);
725 for (indices_broadcast_dims, 0..) |*slot, index| {
726 slot.* = @intCast(axis + 1 + index);
727 }
728 const broadcasted_indices = try function.broadcastInDim(indices, expanded_index_type, expanded_shape, indices_broadcast_dims);
729
730 const update_mapping = try allocator.alloc(i64, updates_type.dims.len);
731 defer allocator.free(update_mapping);
732 for (update_mapping, 0..) |*slot, index| {
733 slot.* = if (index < axis)
734 @intCast(index)
735 else if (index < axis + index_rank)
736 @intCast(index + 1)
737 else
738 @intCast(index + 1);
739 }
740 const broadcasted_updates = try function.broadcastInDim(updates, expanded_data_type, expanded_shape, update_mapping);
741 const mask = try function.compare(source_positions, broadcasted_indices, expanded_mask_type, .eq);
742 const zero_scalar_type = try accy.choir.dialect.accyTensorType(function.ctx, op.result.dtype, &.{});
743 var zero_bytes: [32]u8 = @as([32]u8, @splat(0));
744 const zero_scalar = try function.constant(zero_scalar_type, zero_bytes[0..op.result.dtype.sizeOf()]);
745 const zero_updates = try function.broadcastInDim(zero_scalar, expanded_data_type, expanded_shape, &.{});
746 const selected = try function.select(mask, broadcasted_updates, zero_updates);
747 const reduce_axes = try allocator.alloc(i64, index_rank);
748 defer allocator.free(reduce_axes);
749 for (reduce_axes, 0..) |*slot, index| {
750 slot.* = @intCast(axis + 1 + index);
751 }
752 const reduced = if (index_rank == 0)
753 selected
754 else
755 try function.reduce(selected, zero_scalar, result_type, "sum", reduce_axes);
756 return function.add(input, reduced);
757 }
758
759 fn lowerSparseCrossEntropy(
760 allocator: std.mem.Allocator,
761 function: *accy.choir.semantic.FunctionBuilder,
762 result_type: ir.Type,
763 op: *const program_mod.Operation,
764 sparse_cross_entropy: program_mod.SparseCrossEntropy,
765 logits: *ir.Value,
766 targets: *ir.Value,
767 options: SemanticLoweringOptions,
768 ) !*ir.Value {
769 if (options.sparse_cross_entropy == .semantic_kernel and
770 try sparseCrossEntropySemanticKernelEligible(allocator, op, sparse_cross_entropy, logits, targets))
771 {
772 return function.sparseCrossEntropy(logits, targets, result_type);
773 }
774 return lowerSparseCrossEntropyExpanded(allocator, function, result_type, sparse_cross_entropy, logits, targets);
775 }
776
777 fn lowerSparseCrossEntropyExpanded(
778 allocator: std.mem.Allocator,
779 function: *accy.choir.semantic.FunctionBuilder,
780 result_type: ir.Type,
781 sparse_cross_entropy: program_mod.SparseCrossEntropy,
782 logits: *ir.Value,
783 targets: *ir.Value,
784 ) !*ir.Value {
785 if (sparse_cross_entropy.axis < 0) return error.AxisOutOfRange;
786 const axis = std.math.cast(usize, sparse_cross_entropy.axis) orelse return error.AxisOutOfRange;
787 const logits_type = try accy.choir.dialect.decodeTensorType(allocator, logits.type);
788 defer allocator.free(@constCast(logits_type.dims));
789 if (axis >= logits_type.dims.len) return error.AxisOutOfRange;
790 const class_axes = [_]i64{@intCast(axis)};
791
792 const logits_index_type = try accy.choir.dialect.accyTensorType(function.ctx, .i32, logits_type.dims);
793 const logits_mask_type = try accy.choir.dialect.accyTensorType(function.ctx, .i1, logits_type.dims);
794 const logits_data_type = logits.type;
795
796 const back_mapping = try allocator.alloc(i64, logits_type.dims.len - 1);
797 defer allocator.free(back_mapping);
798 var mapping_out: usize = 0;
799 for (0..logits_type.dims.len) |position| {
800 if (position == axis) continue;
801 back_mapping[mapping_out] = @intCast(position);
802 mapping_out += 1;
803 }
804
805 const max_init = try lowerFloatLowestConstant(function, logits_type.dtype);
806 const row_max = try function.reduce(logits, max_init, result_type, "max", class_axes[0..]);
807 const row_max_full = try function.broadcastInDim(row_max, logits_data_type, logits_type.dims, back_mapping);
808 const shifted = try function.sub(logits, row_max_full);
809 const exponentials = try function.exp(shifted);
810
811 const zero_scalar_type = try accy.choir.dialect.accyTensorType(function.ctx, logits_type.dtype, &.{});
812 var zero_bytes: [32]u8 = @as([32]u8, @splat(0));
813 const zero_scalar = try function.constant(zero_scalar_type, zero_bytes[0..logits_type.dtype.sizeOf()]);
814 const denominator = try function.reduce(exponentials, zero_scalar, result_type, "sum", class_axes[0..]);
815 const log_denominator = try function.log(denominator);
816
817 const class_positions = try function.iota(logits_index_type, @intCast(axis));
818 const target_positions = try function.broadcastInDim(targets, logits_index_type, logits_type.dims, back_mapping);
819 const mask = try function.compare(class_positions, target_positions, logits_mask_type, .eq);
820 const zero_full = try function.broadcastInDim(zero_scalar, logits_data_type, logits_type.dims, &.{});
821 const selected = try function.select(mask, shifted, zero_full);
822 const target_shifted = try function.reduce(selected, zero_scalar, result_type, "sum", class_axes[0..]);
823 return function.sub(log_denominator, target_shifted);
824 }
825
826 fn lowerFloatLowestConstant(
827 function: *accy.choir.semantic.FunctionBuilder,
828 dtype: choir_abi.DType,
829 ) !*ir.Value {
830 const scalar_type = try accy.choir.dialect.accyTensorType(function.ctx, dtype, &.{});
831 return switch (dtype) {
832 .f16 => blk: {
833 const value: f16 = -std.math.floatMax(f16);
834 break :blk function.constant(scalar_type, std.mem.asBytes(&value));
835 },
836 .bf16 => blk: {
837 const value = choir_abi.DType.bf16.ZigType().fromF32(-std.math.floatMax(f32));
838 break :blk function.constant(scalar_type, std.mem.asBytes(&value));
839 },
840 .f32 => blk: {
841 const value: f32 = -std.math.floatMax(f32);
842 break :blk function.constant(scalar_type, std.mem.asBytes(&value));
843 },
844 .f64 => blk: {
845 const value: f64 = -std.math.floatMax(f64);
846 break :blk function.constant(scalar_type, std.mem.asBytes(&value));
847 },
848 else => error.NonFloatDType,
849 };
850 }
851
852 fn sparseCrossEntropySemanticKernelEligible(
853 allocator: std.mem.Allocator,
854 op: *const program_mod.Operation,
855 sparse_cross_entropy: program_mod.SparseCrossEntropy,
856 logits: *ir.Value,
857 targets: *ir.Value,
858 ) !bool {
859 if (!accy.kernel.library.loss.rowSparseCrossEntropyDTypeSupported(op.result.dtype)) return false;
860
861 const logits_type = try accy.choir.dialect.decodeTensorType(allocator, logits.type);
862 defer allocator.free(@constCast(logits_type.dims));
863 const targets_type = try accy.choir.dialect.decodeTensorType(allocator, targets.type);
864 defer allocator.free(@constCast(targets_type.dims));
865 if (logits_type.dims.len != 2 or targets_type.dims.len != 1) return false;
866 if (sparse_cross_entropy.axis != 1) return false;
867 if (targets_type.dtype != .i32) return false;
868 if (logits_type.dims[0] <= 0 or logits_type.dims[1] <= 0) return false;
869 if (targets_type.dims[0] != logits_type.dims[0]) return false;
870 return true;
871 }
872
873 fn scatterAddSemanticKernelEligible(
874 allocator: std.mem.Allocator,
875 op: *const program_mod.Operation,
876 scatter_add: program_mod.ScatterAdd,
877 indices: *ir.Value,
878 updates: *ir.Value,
879 ) !bool {
880 if (scatter_add.axis < 0) return false;
881 const axis = std.math.cast(usize, scatter_add.axis) orelse return false;
882 if (axis >= op.result.rank()) return false;
883 if (!accy.kernel.library.indexing.scatterAddDTypeSupported(op.result.dtype)) return false;
884
885 const input_shape = try types.extents(allocator, op.result.dims);
886 defer allocator.free(@constCast(input_shape));
887 for (input_shape) |dim| {
888 if (dim <= 0) return false;
889 }
890
891 const indices_type = try accy.choir.dialect.decodeTensorType(allocator, indices.type);
892 defer allocator.free(@constCast(indices_type.dims));
893 const updates_type = try accy.choir.dialect.decodeTensorType(allocator, updates.type);
894 defer allocator.free(@constCast(updates_type.dims));
895 if (indices_type.dtype != .i32) return false;
896 if (updates_type.dtype != op.result.dtype) return false;
897 if (indices_type.dims.len != 1) return false;
898 if (updates_type.dims.len != op.result.rank()) return false;
899 if (indices_type.dims[0] <= 0) return false;
900 for (updates_type.dims) |dim| {
901 if (dim <= 0) return false;
902 }
903
904 if (!std.mem.eql(i64, updates_type.dims[0..axis], input_shape[0..axis])) return false;
905 if (updates_type.dims[axis] != indices_type.dims[0]) return false;
906 if (!std.mem.eql(i64, updates_type.dims[axis + 1 ..], input_shape[axis + 1 ..])) return false;
907 return true;
908 }
909
910 fn semanticType(allocator: std.mem.Allocator, ctx: *ir.Context, ty: program_mod.Type) !ir.Type {
911 const shape = try types.extents(allocator, ty.dims);
912 defer allocator.free(@constCast(shape));
913 return accy.choir.dialect.accyTensorType(ctx, ty.dtype, shape);
914 }
915
916 fn truePredicate(
917 allocator: std.mem.Allocator,
918 function: *accy.choir.semantic.FunctionBuilder,
919 carry_ty: program_mod.Type,
920 ) !*ir.Value {
921 const scalar_ty = try accy.choir.dialect.accyTensorType(function.ctx, .i1, &.{});
922 var payload = [_]u8{1};
923 const scalar = try function.constant(scalar_ty, payload[0..]);
924 if (carry_ty.rank() == 0) return scalar;
925 const shape = try types.extents(allocator, carry_ty.dims);
926 defer allocator.free(@constCast(shape));
927 const pred_ty = try accy.choir.dialect.accyTensorType(function.ctx, .i1, shape);
928 return function.broadcast(scalar, pred_ty, shape);
929 }
930
931 fn programScansLowerAsIterate(program: *const program_mod.Program) bool {
932 for (program.operations) |op| {
933 switch (op.kind) {
934 .scan => |scan| if (!scanLowersAsIterate(program, scan)) return false,
935 else => {},
936 }
937 }
938 return true;
939 }
940
941 fn scanLowersAsIterate(source: anytype, scan: program_mod.Scan) bool {
942 if (scan.length <= 0) return false;
943 if (scan.inits.len == 0 or scan.inits.len > 8) return false;
944 const domain = scanElementDomain(source, scan) orelse return false;
945 if (scan.body.outputs.len != scan.inits.len) return false;
946 for (scan.body.outputs, scan.inits) |output, init| {
947 const output_ty = scan.body.typeOf(output);
948 const init_ty = source.typeOf(init);
949 if (!output_ty.eql(init_ty)) return false;
950 }
951 for (scan.body.operations) |op| {
952 if (!bodyOperationLowersAsIterate(scan.body, op, domain)) return false;
953 }
954 return true;
955 }
956
957 fn scanElementDomain(source: anytype, scan: program_mod.Scan) ?usize {
958 var domain: ?usize = null;
959 for (scan.inits) |init| {
960 const count = typeElementCount(source.typeOf(init)) orelse return null;
961 if (count == 1) continue;
962 if (domain) |existing| {
963 if (count != existing) return null;
964 } else {
965 domain = count;
966 }
967 }
968 return domain orelse 1;
969 }
970
971 fn typeElementCount(ty: program_mod.Type) ?usize {
972 var count: usize = 1;
973 for (ty.dims) |dim| {
974 if (dim.extent <= 0) return null;
975 count = std.math.mul(usize, count, @intCast(dim.extent)) catch return null;
976 }
977 return count;
978 }
979
980 fn bodyOperationLowersAsIterate(body: *const program_mod.Subgraph, op: program_mod.Operation, domain: usize) bool {
981 if (!domainCompatible(body.typeOf(op.id), domain)) return false;
982 return switch (op.kind) {
983 .parameter => true,
984 .constant => |constant| constantLowersAsSplat(body.typeOf(op.id), constant.payload),
985 .unary, .binary, .compare, .select => true,
986 .reshape => |reshape| sameElementCount(body.typeOf(reshape.input), body.typeOf(op.id)),
987 .broadcast => |broadcast| constantProducer(body, broadcast.input),
988 .broadcast_in_dim => |broadcast| constantProducer(body, broadcast.input),
989 .custom_call => |custom| customCallLowersAsIterate(body, op, custom, domain),
990 else => false,
991 };
992 }
993
994 fn domainCompatible(ty: program_mod.Type, domain: usize) bool {
995 const count = typeElementCount(ty) orelse return false;
996 return count == 1 or count == domain;
997 }
998
999 fn sameElementCount(lhs: program_mod.Type, rhs: program_mod.Type) bool {
1000 const lhs_count = typeElementCount(lhs) orelse return false;
1001 const rhs_count = typeElementCount(rhs) orelse return false;
1002 return lhs_count == rhs_count;
1003 }
1004
1005 fn customCallLowersAsIterate(
1006 body: *const program_mod.Subgraph,
1007 op: program_mod.Operation,
1008 custom: program_mod.CustomCall,
1009 domain: usize,
1010 ) bool {
1011 const spec = parsePhiloxKeyCounterUniformTarget(custom.target) orelse return false;
1012 if (custom.version != accy.kernel.library.random.philox_key_counter_uniform_family_version) return false;
1013 if (custom.operands.len != 2) return false;
1014 if (body.typeOf(custom.operands[0]).rank() != 0 or body.typeOf(custom.operands[0]).dtype != .key) return false;
1015 if (body.typeOf(custom.operands[1]).rank() != 0 or body.typeOf(custom.operands[1]).dtype != .i32) return false;
1016 if (op.result.dtype != spec.dtype) return false;
1017 return (typeElementCount(op.result) orelse return false) == domain;
1018 }
1019
1020 const PhiloxKeyCounterUniformTarget = struct {
1021 rounds: u32,
1022 dtype: choir_abi.DType,
1023 };
1024
1025 fn parsePhiloxKeyCounterUniformTarget(target: []const u8) ?PhiloxKeyCounterUniformTarget {
1026 const prefix = "accy.kernel.random.philox_key_counter_uniform_family_";
1027 if (!std.mem.startsWith(u8, target, prefix)) return null;
1028 const rest = target[prefix.len..];
1029 const rounds_end = std.mem.indexOf(u8, rest, "r_") orelse return null;
1030 const rounds = std.fmt.parseInt(u32, rest[0..rounds_end], 10) catch return null;
1031 const after_rounds = rest[rounds_end + 2 ..];
1032 const dtype_start = std.mem.lastIndexOfScalar(u8, after_rounds, '_') orelse return null;
1033 const dtype = choir_abi.DType.fromName(after_rounds[dtype_start + 1 ..]) orelse return null;
1034 if (!accy.kernel.library.random.randomDTypeSupported(dtype)) return null;
1035 return .{ .rounds = rounds, .dtype = dtype };
1036 }
1037
1038 fn constantProducer(body: *const program_mod.Subgraph, id: program_mod.Id) bool {
1039 const op = body.operation(id);
1040 return switch (op.kind) {
1041 .constant => |constant| constantLowersAsSplat(op.result, constant.payload),
1042 else => false,
1043 };
1044 }
1045
1046 fn constantLowersAsSplat(ty: program_mod.Type, payload: []const u8) bool {
1047 return switch (ty.dtype) {
1048 .f32, .i32 => blk: {
1049 if (payload.len < 4 or payload.len % 4 != 0) break :blk false;
1050 var first: u32 = undefined;
1051 @memcpy(std.mem.asBytes(&first), payload[0..4]);
1052 var offset: usize = 4;
1053 while (offset < payload.len) : (offset += 4) {
1054 var value: u32 = undefined;
1055 @memcpy(std.mem.asBytes(&value), payload[offset..][0..4]);
1056 if (value != first) break :blk false;
1057 }
1058 break :blk true;
1059 },
1060 .i1 => blk: {
1061 if (payload.len == 0) break :blk false;
1062 const first = payload[0];
1063 for (payload[1..]) |value| {
1064 if (value != first) break :blk false;
1065 }
1066 break :blk true;
1067 },
1068 .key => payload.len == @sizeOf(choir_abi.Key),
1069 else => false,
1070 };
1071 }
1072
1073 fn lowerUnary(function: *accy.choir.semantic.FunctionBuilder, op: program_mod.Unary, input: *ir.Value) !*ir.Value {
1074 return switch (op) {
1075 .neg => function.neg(input),
1076 .abs => function.abs(input),
1077 .exp => function.exp(input),
1078 .log => function.log(input),
1079 .sqrt => function.sqrt(input),
1080 .tanh => function.tanh(input),
1081 .sin => function.sin(input),
1082 .cos => function.cos(input),
1083 .tan => function.tan(input),
1084 };
1085 }
1086
1087 fn lowerBinary(
1088 function: *accy.choir.semantic.FunctionBuilder,
1089 op: program_mod.Binary,
1090 lhs: *ir.Value,
1091 rhs: *ir.Value,
1092 ) !*ir.Value {
1093 return switch (op) {
1094 .add => function.add(lhs, rhs),
1095 .sub => function.sub(lhs, rhs),
1096 .mul => function.mul(lhs, rhs),
1097 .div => function.div(lhs, rhs),
1098 .max => function.max(lhs, rhs),
1099 .min => function.min(lhs, rhs),
1100 .pow => function.pow(lhs, rhs),
1101 };
1102 }
1103
1104 test "tensor lowering produces verified semantic Choir" {
1105 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_dense");
1106 defer builder.deinit();
1107
1108 const x = try builder.input(.f32, .{ .m = 4, .k = 8 });
1109 const w = try builder.input(.f32, .{ .k = 8, .n = 3 });
1110 const b = try builder.input(.f32, .{ .n = 3 });
1111 const out = try (try (try x.contract(w, .k)).add(b)).tanh();
1112 var program = try builder.finish(&.{out});
1113 defer program.deinit();
1114
1115 const lowered = try toSemanticModule(std.testing.allocator, &program);
1116 defer lowered.deinit();
1117
1118 try lowered.verify();
1119 }
1120
1121 test "tensor lowering expands scatter add through semantic accumulation" {
1122 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_indexing");
1123 defer builder.deinit();
1124
1125 const table = try builder.input(.f32, .{ .vocab = 16, .channel = 4 });
1126 const ids = try builder.input(.i32, .{ .token = 3 });
1127 const gathered = try table.gather(ids, .vocab);
1128 const zero_table = try builder.full(.f32, .{ .vocab = 16, .channel = 4 }, 0.0);
1129 const out = try zero_table.scatterAdd(ids, gathered, .vocab);
1130 var program = try builder.finish(&.{out});
1131 defer program.deinit();
1132
1133 const lowered = try toSemanticModule(std.testing.allocator, &program);
1134 defer lowered.deinit();
1135
1136 try lowered.verify();
1137 try std.testing.expectEqual(
1138 @as(usize, 1),
1139 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.GatherOp.operation_name),
1140 );
1141 try std.testing.expect(ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ReduceOp.operation_name) >= 1);
1142 }
1143
1144 test "tensor lowering expands scatter add with shaped indices" {
1145 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_shaped_indexing");
1146 defer builder.deinit();
1147
1148 const table = try builder.input(.f32, .{ .vocab = 16, .channel = 4 });
1149 const ids = try builder.input(.i32, .{ .batch = 2, .token = 3 });
1150 const gathered = try table.gather(ids, .vocab);
1151 const zero_table = try builder.full(.f32, .{ .vocab = 16, .channel = 4 }, 0.0);
1152 const out = try zero_table.scatterAdd(ids, gathered, .vocab);
1153 var program = try builder.finish(&.{out});
1154 defer program.deinit();
1155
1156 const lowered = try toSemanticModule(std.testing.allocator, &program);
1157 defer lowered.deinit();
1158
1159 try lowered.verify();
1160 try std.testing.expectEqual(
1161 @as(usize, 1),
1162 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.GatherOp.operation_name),
1163 );
1164 try std.testing.expect(ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ReduceOp.operation_name) >= 1);
1165 }
1166
1167 test "tensor lowering emits semantic scatter add for scheduled kernel path" {
1168 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_semantic_scatter_add");
1169 defer builder.deinit();
1170
1171 const table = try builder.input(.f32, .{ .vocab = 16, .channel = 4 });
1172 const ids = try builder.input(.i32, .{ .token = 3 });
1173 const gathered = try table.gather(ids, .vocab);
1174 const zero_table = try builder.full(.f32, .{ .vocab = 16, .channel = 4 }, 0.0);
1175 const out = try zero_table.scatterAdd(ids, gathered, .vocab);
1176 var program = try builder.finish(&.{out});
1177 defer program.deinit();
1178
1179 const lowered = try toSemanticModuleWithOptions(std.testing.allocator, &program, .{
1180 .scatter_add = .semantic_kernel,
1181 });
1182 defer lowered.deinit();
1183
1184 try lowered.verify();
1185 try std.testing.expectEqual(
1186 @as(usize, 1),
1187 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.GatherOp.operation_name),
1188 );
1189 try std.testing.expectEqual(
1190 @as(usize, 1),
1191 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ScatterAddOp.operation_name),
1192 );
1193 try std.testing.expectEqual(
1194 @as(usize, 0),
1195 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ReduceOp.operation_name),
1196 );
1197 }
1198
1199 test "tensor lowering expands sparse cross entropy through semantic reduction" {
1200 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_sparse_cross_entropy");
1201 defer builder.deinit();
1202
1203 const logits = try builder.input(.f32, .{ .sample = 6, .vocab = 11 });
1204 const targets = try builder.input(.i32, .{ .sample = 6 });
1205 const out = try builder.sparseCrossEntropyLoss(logits, targets, .vocab);
1206 var program = try builder.finish(&.{out});
1207 defer program.deinit();
1208
1209 const lowered = try toSemanticModule(std.testing.allocator, &program);
1210 defer lowered.deinit();
1211
1212 try lowered.verify();
1213 try std.testing.expectEqual(
1214 @as(usize, 0),
1215 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.SparseCrossEntropyOp.operation_name),
1216 );
1217 try std.testing.expect(ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ReduceOp.operation_name) >= 2);
1218 }
1219
1220 test "tensor lowering emits semantic sparse cross entropy for scheduled kernel path" {
1221 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_semantic_sparse_cross_entropy");
1222 defer builder.deinit();
1223
1224 const logits = try builder.input(.f32, .{ .sample = 6, .vocab = 11 });
1225 const targets = try builder.input(.i32, .{ .sample = 6 });
1226 const out = try builder.sparseCrossEntropyLoss(logits, targets, .vocab);
1227 var program = try builder.finish(&.{out});
1228 defer program.deinit();
1229
1230 const lowered = try toSemanticModuleWithOptions(std.testing.allocator, &program, .{
1231 .sparse_cross_entropy = .semantic_kernel,
1232 });
1233 defer lowered.deinit();
1234
1235 try lowered.verify();
1236 try std.testing.expectEqual(
1237 @as(usize, 1),
1238 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.SparseCrossEntropyOp.operation_name),
1239 );
1240 try std.testing.expectEqual(
1241 @as(usize, 0),
1242 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.ReduceOp.operation_name),
1243 );
1244 }
1245
1246 test "tensor executable preparation lowers scheduled sparse cross entropy to catalog call" {
1247 const allocator = std.testing.allocator;
1248 var builder = try accy.tensor.Builder.init(allocator, "prepare_tensor_sparse_cross_entropy_catalog");
1249 defer builder.deinit();
1250
1251 const logits = try builder.input(.f32, .{ .sample = 6, .vocab = 11 });
1252 const targets = try builder.input(.i32, .{ .sample = 6 });
1253 const out = try builder.sparseCrossEntropyLoss(logits, targets, .vocab);
1254 var program = try builder.finish(&.{out});
1255 defer program.deinit();
1256
1257 var state = gpu.recording.BackendState{
1258 .allocator = allocator,
1259 .kind = .cuda,
1260 .format = .cuda_ptx,
1261 };
1262 const registry = accy.artifact.KernelCallRegistry{ .entries = &.{} };
1263 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1264 .artifact_format = .cuda_ptx,
1265 .kernel_call_registry = ®istry,
1266 .row_sparse_cross_entropy_schedule = .{ .thread_blocks = 4 },
1267 });
1268 defer prepared.deinit();
1269
1270 try std.testing.expectEqual(
1271 @as(usize, 0),
1272 ir.inspection.countOperationsNamed(prepared.choir_module, accy.choir.dialect.AccyDialect.SparseCrossEntropyOp.operation_name),
1273 );
1274 try std.testing.expectEqual(
1275 @as(usize, 1),
1276 ir.inspection.countOperationsNamed(prepared.choir_module, accy.choir.dialect.AccyDialect.KernelCallOp.operation_name),
1277 );
1278 }
1279
1280 test "tensor executable preparation lowers scheduled scatter add to catalog call" {
1281 const allocator = std.testing.allocator;
1282 var builder = try accy.tensor.Builder.init(allocator, "prepare_tensor_scatter_add_catalog");
1283 defer builder.deinit();
1284
1285 const table = try builder.input(.f32, .{ .vocab = 16, .channel = 4 });
1286 const ids = try builder.input(.i32, .{ .token = 3 });
1287 const gathered = try table.gather(ids, .vocab);
1288 const zero_table = try builder.full(.f32, .{ .vocab = 16, .channel = 4 }, 0.0);
1289 const out = try zero_table.scatterAdd(ids, gathered, .vocab);
1290 var program = try builder.finish(&.{out});
1291 defer program.deinit();
1292
1293 var state = gpu.recording.BackendState{
1294 .allocator = allocator,
1295 .kind = .cuda,
1296 .format = .cuda_ptx,
1297 };
1298 const registry = accy.artifact.KernelCallRegistry{ .entries = &.{} };
1299 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1300 .artifact_format = .cuda_ptx,
1301 .kernel_call_registry = ®istry,
1302 .scatter_add_schedule = .{ .thread_blocks = 4 },
1303 });
1304 defer prepared.deinit();
1305
1306 try std.testing.expectEqual(
1307 @as(usize, 0),
1308 ir.inspection.countOperationsNamed(prepared.choir_module, accy.choir.dialect.AccyDialect.ScatterAddOp.operation_name),
1309 );
1310 try std.testing.expectEqual(
1311 @as(usize, 1),
1312 ir.inspection.countOperationsNamed(prepared.choir_module, accy.choir.dialect.AccyDialect.KernelCallOp.operation_name),
1313 );
1314 }
1315
1316 fn lowerScanStep(scan_builder: *accy.tensor.Builder, carry: anytype) !@TypeOf(carry) {
1317 const half = try scan_builder.full(.f32, .{ .lane = 4 }, 0.5);
1318 return .{
1319 .x = try (try (try carry.x.mul(carry.x)).mul(half)).add(carry.c),
1320 .c = carry.c,
1321 };
1322 }
1323
1324 test "tensor lowering preserves eligible scan as iterate" {
1325 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_scan");
1326 defer builder.deinit();
1327
1328 const x0 = try builder.input(.f32, .{ .lane = 4 });
1329 const c = try builder.input(.f32, .{ .lane = 4 });
1330 const walked = try builder.scan(.{
1331 .length = 3,
1332 .init = .{ .x = x0, .c = c },
1333 .body = lowerScanStep,
1334 });
1335 var program = try builder.finish(&.{ walked.x, walked.c });
1336 defer program.deinit();
1337
1338 const lowered = try toSemanticModule(std.testing.allocator, &program);
1339 defer lowered.deinit();
1340
1341 try lowered.verify();
1342 try std.testing.expectEqual(
1343 @as(usize, 1),
1344 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.IterateOp.operation_name),
1345 );
1346 }
1347
1348 fn iterateIncrementStep(scan_builder: *accy.tensor.Builder, carry: accy.tensor.Value) !accy.tensor.Value {
1349 const one = try scan_builder.full(.f32, .{ .lane = 8 }, 1.0);
1350 return carry.add(one);
1351 }
1352
1353 test "tensor eligible scan prepares as iterate kernel" {
1354 const allocator = std.testing.allocator;
1355 var builder = try accy.tensor.Builder.init(allocator, "prepare_scan_iterate");
1356 defer builder.deinit();
1357
1358 const x0 = try builder.input(.f32, .{ .lane = 8 });
1359 const walked = try builder.scan(.{
1360 .length = 5,
1361 .init = x0,
1362 .body = iterateIncrementStep,
1363 });
1364 var program = try builder.finish(&.{walked});
1365 defer program.deinit();
1366
1367 var state = gpu.recording.BackendState{
1368 .allocator = allocator,
1369 .kind = .cuda,
1370 .format = .cuda_ptx,
1371 };
1372 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1373 .artifact_format = .cuda_ptx,
1374 });
1375 defer prepared.deinit();
1376
1377 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1378 const summary = try prepared.generatedKernelSummary(0);
1379 try std.testing.expect(std.mem.startsWith(u8, summary.entry_name, "accy_choir_iterate1_cap5"));
1380 }
1381
1382 fn iterateScalarCarryStep(scan_builder: *accy.tensor.Builder, carry: anytype) !@TypeOf(carry) {
1383 const one = try scan_builder.full(.f32, .{ .lane = 8 }, 1.0);
1384 const one_step = try scan_builder.scalar(.i32, 1);
1385 return .{
1386 .x = try carry.x.add(one),
1387 .step = try carry.step.add(one_step),
1388 };
1389 }
1390
1391 test "tensor eligible scan carries scalar loop state in iterate kernel" {
1392 const allocator = std.testing.allocator;
1393 var builder = try accy.tensor.Builder.init(allocator, "prepare_scan_iterate_scalar_carry");
1394 defer builder.deinit();
1395
1396 const x0 = try builder.input(.f32, .{ .lane = 8 });
1397 const step0 = try builder.scalar(.i32, 0);
1398 const walked = try builder.scan(.{
1399 .length = 5,
1400 .init = .{ .x = x0, .step = step0 },
1401 .body = iterateScalarCarryStep,
1402 });
1403 var program = try builder.finish(&.{ walked.x, walked.step });
1404 defer program.deinit();
1405
1406 const lowered = try toSemanticModule(allocator, &program);
1407 defer lowered.deinit();
1408 try lowered.verify();
1409 try std.testing.expectEqual(
1410 @as(usize, 1),
1411 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.IterateOp.operation_name),
1412 );
1413
1414 var state = gpu.recording.BackendState{
1415 .allocator = allocator,
1416 .kind = .cuda,
1417 .format = .cuda_ptx,
1418 };
1419 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1420 .artifact_format = .cuda_ptx,
1421 });
1422 defer prepared.deinit();
1423
1424 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1425 const summary = try prepared.generatedKernelSummary(0);
1426 try std.testing.expect(std.mem.startsWith(u8, summary.entry_name, "accy_choir_iterate2_cap5"));
1427 }
1428
1429 fn iterateRandomStep(scan_builder: *accy.tensor.Builder, carry: anytype) !@TypeOf(carry) {
1430 const seed = try accy.tensor.random.seed(scan_builder, 0x01234567_89abcdef, .{ .threads = 8 });
1431 const unit = try seed.counterUniform(carry.step, .{ .lane = 8 }, .f32);
1432 const one_step = try scan_builder.scalar(.i32, 1);
1433 return .{
1434 .x = try carry.x.add(unit),
1435 .step = try carry.step.add(one_step),
1436 };
1437 }
1438
1439 test "tensor eligible scan inlines counter uniform random in iterate kernel" {
1440 const allocator = std.testing.allocator;
1441 var builder = try accy.tensor.Builder.init(allocator, "prepare_scan_iterate_random");
1442 defer builder.deinit();
1443
1444 const x0 = try builder.input(.f32, .{ .lane = 8 });
1445 const step0 = try builder.scalar(.i32, 0);
1446 const walked = try builder.scan(.{
1447 .length = 5,
1448 .init = .{ .x = x0, .step = step0 },
1449 .body = iterateRandomStep,
1450 });
1451 var program = try builder.finish(&.{walked.x});
1452 defer program.deinit();
1453
1454 const lowered = try toSemanticModule(allocator, &program);
1455 defer lowered.deinit();
1456 try lowered.verify();
1457 try std.testing.expectEqual(
1458 @as(usize, 1),
1459 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.IterateOp.operation_name),
1460 );
1461 try std.testing.expectEqual(
1462 @as(usize, 1),
1463 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.KernelCallOp.operation_name),
1464 );
1465
1466 var state = gpu.recording.BackendState{
1467 .allocator = allocator,
1468 .kind = .cuda,
1469 .format = .cuda_ptx,
1470 };
1471 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1472 .artifact_format = .cuda_ptx,
1473 });
1474 defer prepared.deinit();
1475
1476 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1477 const summary = try prepared.generatedKernelSummary(0);
1478 try std.testing.expect(std.mem.startsWith(u8, summary.entry_name, "accy_choir_iterate2_cap5"));
1479 }
1480
1481 test "tensor eligible scan with counter uniform random prepares for CPU object" {
1482 const allocator = std.testing.allocator;
1483 var builder = try accy.tensor.Builder.init(allocator, "prepare_scan_iterate_random_cpu");
1484 defer builder.deinit();
1485
1486 const x0 = try builder.input(.f32, .{ .lane = 8 });
1487 const step0 = try builder.scalar(.i32, 0);
1488 const walked = try builder.scan(.{
1489 .length = 5,
1490 .init = .{ .x = x0, .step = step0 },
1491 .body = iterateRandomStep,
1492 });
1493 var program = try builder.finish(&.{walked.x});
1494 defer program.deinit();
1495
1496 var state = gpu.cpu.State.init(allocator);
1497 defer state.deinit();
1498 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1499 .artifact_format = .cpu_object,
1500 });
1501 defer prepared.deinit();
1502
1503 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1504 }
1505
1506 fn iterateRandomTrigStep(scan_builder: *accy.tensor.Builder, carry: anytype) !@TypeOf(carry) {
1507 const seed = try accy.tensor.random.seed(scan_builder, 0x01234567_89abcdef, .{ .threads = 8 });
1508 const unit = try seed.counterUniform(carry.step, .{ .lane = 8 }, .f32);
1509 const tau = try scan_builder.full(.f32, .{ .lane = 8 }, 6.283185307179586);
1510 const one = try scan_builder.full(.f32, .{ .lane = 8 }, 1.0);
1511 const zero = try scan_builder.full(.f32, .{ .lane = 8 }, 0.0);
1512 const cutoff = try scan_builder.full(.f32, .{ .lane = 8 }, 0.001);
1513 const radius2 = try (try carry.x.mul(carry.x)).add(try carry.y.mul(carry.y));
1514 const radius = try (try one.sub(try radius2.sqrt())).max(zero);
1515 const can_move = try radius.compare(.gt, cutoff);
1516 const inactive = try scan_builder.full(.i1, .{ .lane = 8 }, false);
1517 const moving = try carry.active.select(can_move, inactive);
1518 const angle = try unit.mul(tau);
1519 const next_x = try carry.x.add(try (try angle.cos()).mul(radius));
1520 const next_y = try carry.y.add(try (try angle.sin()).mul(radius));
1521 const one_step = try scan_builder.scalar(.i32, 1);
1522 return .{
1523 .x = try moving.select(next_x, carry.x),
1524 .y = try moving.select(next_y, carry.y),
1525 .active = moving,
1526 .step = try carry.step.add(one_step),
1527 };
1528 }
1529
1530 test "tensor eligible scan with random trig mask prepares for CPU object" {
1531 const allocator = std.testing.allocator;
1532 var builder = try accy.tensor.Builder.init(allocator, "prepare_scan_iterate_random_trig_cpu");
1533 defer builder.deinit();
1534
1535 const x0 = try builder.input(.f32, .{ .lane = 8 });
1536 const y0 = try builder.input(.f32, .{ .lane = 8 });
1537 const active0 = try builder.full(.i1, .{ .lane = 8 }, true);
1538 const step0 = try builder.scalar(.i32, 0);
1539 const walked = try builder.scan(.{
1540 .length = 5,
1541 .init = .{ .x = x0, .y = y0, .active = active0, .step = step0 },
1542 .body = iterateRandomTrigStep,
1543 });
1544 var program = try builder.finish(&.{ walked.x, walked.y });
1545 defer program.deinit();
1546
1547 var state = gpu.cpu.State.init(allocator);
1548 defer state.deinit();
1549 var prepared = try prepareFragment(allocator, state.handle(), &program, .{
1550 .artifact_format = .cpu_object,
1551 });
1552 defer prepared.deinit();
1553
1554 try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());
1555 }
1556
1557 fn ineligibleScanStep(scan_builder: *accy.tensor.Builder, carry: accy.tensor.Value) !accy.tensor.Value {
1558 const generated = try scan_builder.customCall("accy.custom.generate", 1, &.{}, carry.ty);
1559 return carry.add(generated);
1560 }
1561
1562 test "tensor lowering unrolls ineligible scan operations" {
1563 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_scan_custom_call");
1564 defer builder.deinit();
1565
1566 const x0 = try builder.input(.f32, .{ .lane = 4 });
1567 const walked = try builder.scan(.{
1568 .length = 3,
1569 .init = x0,
1570 .body = ineligibleScanStep,
1571 });
1572 var program = try builder.finish(&.{walked});
1573 defer program.deinit();
1574
1575 const lowered = try toSemanticModule(std.testing.allocator, &program);
1576 defer lowered.deinit();
1577
1578 try lowered.verify();
1579 try std.testing.expectEqual(
1580 @as(usize, 0),
1581 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.IterateOp.operation_name),
1582 );
1583 try std.testing.expectEqual(
1584 @as(usize, 3),
1585 ir.inspection.countOperationsNamed(lowered.choir_module, accy.choir.dialect.AccyDialect.KernelCallOp.operation_name),
1586 );
1587 }
1588
1589 test "tensor complex scan executable uses bounded host loop launch plan" {
1590 const allocator = std.testing.allocator;
1591
1592 var short_program = try reductionBodyScanProgram(allocator, "host_loop_scan_short", 3);
1593 defer short_program.deinit();
1594 var long_program = try reductionBodyScanProgram(allocator, "host_loop_scan_long", 7);
1595 defer long_program.deinit();
1596
1597 var state = gpu.recording.BackendState{
1598 .allocator = allocator,
1599 .kind = .cuda,
1600 .format = .cuda_ptx,
1601 };
1602 const handle = state.handle();
1603 const options = FragmentCompilerOptions{ .artifact_format = .cuda_ptx };
1604
1605 const short_compiled = try compileFragment(allocator, handle, &short_program, options);
1606 var short_fragment = try accy.executable.loadFragment(allocator, handle, short_compiled, options);
1607 defer short_fragment.deinit();
1608 const long_compiled = try compileFragment(allocator, handle, &long_program, options);
1609 var long_fragment = try accy.executable.loadFragment(allocator, handle, long_compiled, options);
1610 defer long_fragment.deinit();
1611
1612 var short_graph = try short_fragment.createLaunchGraphPlan(allocator, .{});
1613 defer short_graph.deinit();
1614 var long_graph = try long_fragment.createLaunchGraphPlan(allocator, .{});
1615 defer long_graph.deinit();
1616
1617 try std.testing.expectEqual(@as(usize, 1), short_graph.loops.len);
1618 try std.testing.expectEqual(@as(usize, 1), long_graph.loops.len);
1619 try std.testing.expectEqual(@as(u64, 3), short_graph.loops[0].trip_count);
1620 try std.testing.expectEqual(@as(u64, 7), long_graph.loops[0].trip_count);
1621 try std.testing.expectEqual(short_fragment.kernelCount(), long_fragment.kernelCount());
1622 try std.testing.expectEqual(short_graph.nodes.len, long_graph.nodes.len);
1623 try std.testing.expect(short_fragment.kernelCount() > 0);
1624 try std.testing.expect(short_fragment.kernelCount() < 7);
1625
1626 const input = @as([8]f32, @splat(1.0));
1627 const bindings = try accy.executable.prepareInvocation(long_fragment, allocator, &.{std.mem.asBytes(&input)});
1628 defer bindings.deinit();
1629 try bindings.launch(allocator);
1630
1631 try std.testing.expectEqual(long_fragment.kernelCount() * 7, state.launch_count);
1632 }
1633
1634 test "tensor lowering accepts zero operand custom calls" {
1635 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_zero_operand_custom_call");
1636 defer builder.deinit();
1637
1638 const ty = try accy.tensor.Type.init(builder.arena.allocator(), .f32, &.{
1639 .{ .name = "lane", .extent = 4 },
1640 });
1641 const out = try builder.customCall("accy.custom.generate", 1, &.{}, ty);
1642 var program = try builder.finish(&.{out});
1643 defer program.deinit();
1644
1645 const lowered = try toSemanticModule(std.testing.allocator, &program);
1646 defer lowered.deinit();
1647
1648 try lowered.verify();
1649 }
1650
1651 fn reductionBodyScanStep(scan_builder: *accy.tensor.Builder, carry: accy.tensor.Value) !accy.tensor.Value {
1652 _ = scan_builder;
1653 const total = try carry.sum(.lane);
1654 const lifted = try total.broadcast(.{ .lane = 8 });
1655 return carry.add(lifted);
1656 }
1657
1658 fn reductionBodyScanProgram(
1659 allocator: std.mem.Allocator,
1660 name: []const u8,
1661 length: i64,
1662 ) !program_mod.Program {
1663 var builder = try accy.tensor.Builder.init(allocator, name);
1664 defer builder.deinit();
1665
1666 const x0 = try builder.input(.f32, .{ .lane = 8 });
1667 const walked = try builder.scan(.{
1668 .length = length,
1669 .init = x0,
1670 .body = reductionBodyScanStep,
1671 });
1672 return try builder.finish(&.{walked});
1673 }
1674
1675 const CountLowerAdd = struct {
1676 count: *usize,
1677
1678 pub fn bind(self: *@This(), ctx: anytype) !*ir.Value {
1679 switch (ctx.op.kind) {
1680 .binary => |binary| {
1681 if (binary.op == .add) self.count.* += 1;
1682 },
1683 else => {},
1684 }
1685 return ctx.default();
1686 }
1687 };
1688
1689 fn toSemanticModuleCountingAdds(
1690 allocator: std.mem.Allocator,
1691 program: *const program_mod.Program,
1692 count: *usize,
1693 ) !*accy.choir.SemanticModule {
1694 var builder = try accy.choir.SemanticBuilder.init(allocator, accy.choir.SemanticBuilder.ContextLimits.standard);
1695 errdefer builder.deinit();
1696
1697 const choir_types = try allocator.alloc(ir.Type, program.valueCount());
1698 defer allocator.free(choir_types);
1699 for (program.values, 0..) |ty, index| {
1700 const shape = try types.extents(allocator, ty.dims);
1701 defer allocator.free(@constCast(shape));
1702 choir_types[index] = try builder.tensor(ty.dtype, shape);
1703 }
1704
1705 const parameter_types = try allocator.alloc(ir.Type, program.parameters.len);
1706 defer allocator.free(parameter_types);
1707 for (program.parameters, 0..) |id, index| {
1708 parameter_types[index] = choir_types[id.index];
1709 }
1710
1711 const result_types = try allocator.alloc(ir.Type, program.outputs.len);
1712 defer allocator.free(result_types);
1713 for (program.outputs, 0..) |id, index| {
1714 result_types[index] = choir_types[id.index];
1715 }
1716
1717 var function = try builder.beginFunction(program.name, parameter_types, result_types);
1718 var scan_state = ScanLoweringState{};
1719 defer scan_state.deinit(allocator);
1720 try interpret.run(allocator, program, interpret.layer(*ir.Value, semantics(allocator, &function, choir_types, &scan_state, .{}), CountLowerAdd{ .count = count }));
1721
1722 return try builder.finish();
1723 }
1724
1725 test "tensor lowering semantics composes with interpreter layers" {
1726 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_layer");
1727 defer builder.deinit();
1728
1729 const x = try builder.input(.f32, .{ .lane = 4 });
1730 const y = try builder.input(.f32, .{ .lane = 4 });
1731 const out = try x.add(y);
1732 var program = try builder.finish(&.{out});
1733 defer program.deinit();
1734
1735 var add_count: usize = 0;
1736 const lowered = try toSemanticModuleCountingAdds(std.testing.allocator, &program, &add_count);
1737 defer lowered.deinit();
1738
1739 try std.testing.expectEqual(@as(usize, 1), add_count);
1740 try lowered.verify();
1741 }
1742
1743 test "tensor lowering module attaches to graph interpretation" {
1744 var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_module");
1745 defer builder.deinit();
1746
1747 const x = try builder.input(.f32, .{ .lane = 4 });
1748 const y = try builder.input(.f32, .{ .lane = 4 });
1749 const out = try (try x.add(y)).tanh();
1750 var program = try builder.finish(&.{out});
1751 defer program.deinit();
1752
1753 var graph_builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_module");
1754 errdefer graph_builder.deinit();
1755
1756 const graph = interpret.Graph{ .builder = &graph_builder };
1757 const lowered = try interpret.run(std.testing.allocator, &program, module(std.testing.allocator).attach(graph));
1758 defer lowered.deinit();
1759
1760 try lowered.verify();
1761 }