tiny.accy.tensor.lower

1const std = @import("std");2const gpu = @import("gpu");3const choir_abi = @import("choir_abi");4const choir = @import("choir");5const accy = @import("../root.zig");6const interpret = @import("interpret/root.zig");7const program_mod = @import("program.zig");8const trace = @import("trace/root.zig");9const unroll = @import("unroll.zig");10const types = @import("type/root.zig");1112const ir = choir.ir;1314pub const SemanticModule = accy.choir.SemanticModule;15pub const BackendPreparedJob = accy.preparation.BackendPreparedJob;16pub const BackendPreparationRunOptions = accy.preparation.BackendPreparationRunOptions;17pub const GeneratedScheduleKind = accy.preparation.GeneratedScheduleKind;18pub const GeneratedSchedule = accy.preparation.GeneratedSchedule;19pub const GeneratedKernelProgram = accy.preparation.GeneratedKernelProgram;20pub const GeneratedKernelSummary = accy.preparation.GeneratedKernelSummary;21pub const GeneratedKernelSummaries = accy.preparation.GeneratedKernelSummaries;22pub const ArtifactJob = accy.artifact.ArtifactJob;23pub const BackendHandle = gpu.BackendHandle;24pub const CompiledFragment = accy.executable.CompiledFragment;25pub const LoadedFragment = accy.executable.LoadedFragment;26pub const FragmentCompilerOptions = accy.executable.FragmentCompilerOptions;27pub const ArtifactKernelSource = accy.artifact.KernelSource;28pub const ArtifactKernelSummary = accy.artifact.KernelSummary;29pub const ArtifactKernelSummaries = accy.artifact.KernelSummaries;3031pub const ScatterAddLowering = enum {32    expanded,33    semantic_kernel,34};3536pub const SparseCrossEntropyLowering = enum {37    expanded,38    semantic_kernel,39};4041pub const SemanticLoweringOptions = struct {42    context_limits: accy.choir.SemanticBuilder.ContextLimits = .standard,43    scatter_add: ScatterAddLowering = .expanded,44    sparse_cross_entropy: SparseCrossEntropyLowering = .expanded,4546    pub fn eql(self: SemanticLoweringOptions, other: SemanticLoweringOptions) bool {47        return self.scatter_add == other.scatter_add and48            self.sparse_cross_entropy == other.sparse_cross_entropy and49            std.meta.eql(self.context_limits, other.context_limits);50    }51};5253pub const FragmentCompilerCacheUpdate = accy.executable.FragmentCompilerCacheUpdate;54pub const FragmentCompilationRequest = accy.executable.FragmentCompilationRequest;5556/// A caller keeps one of these per changing program so each recompile reuses the stages that did57/// not change. Each refresh lowers the program into a new semantic module inside a bounded compiler58/// job, sized by the context limits in the options. The compile that follows reuses a stage only59/// when an exact check admits the earlier stage record, the sealed result of one compile stage, for60/// the new request. A failed refresh keeps the previous compiled result and the previous loaded61/// fragment, the compiled and loaded program ready to launch. `currentPrepared` and62/// `currentFragment` return the result of the last successful refresh.63pub const FragmentCompilerCache = struct {64    allocator: std.mem.Allocator,65    cache: accy.executable.FragmentCompilerCache,6667    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    }7374    pub fn deinit(self: *FragmentCompilerCache) void {75        self.cache.deinit();76        self.* = undefined;77    }7879    pub fn currentPrepared(80        self: *const FragmentCompilerCache,81    ) ?*const accy.preparation.pipeline.BackendPreparedModule {82        return self.cache.currentPrepared();83    }8485    pub fn currentFragment(self: *const FragmentCompilerCache) ?*LoadedFragment {86        return self.cache.currentFragment();87    }8889    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};105106pub fn toSemanticModule(107    allocator: std.mem.Allocator,108    program: *const program_mod.Program,109) !*SemanticModule {110    return toSemanticModuleWithOptions(allocator, program, .{});111}112113pub 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    }125126    var builder = try accy.choir.SemanticBuilder.init(allocator, options.context_limits);127    errdefer builder.deinit();128129    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    }136137    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    }142143    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    }148149    var function = try builder.beginFunction(program.name, parameter_types, result_types);150151    var scan_state = ScanLoweringState{};152    defer scan_state.deinit(allocator);153154    try interpret.run(allocator, program, semantics(allocator, &function, choir_types, &scan_state, options));155156    return try builder.finish();157}158159fn 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_kernel163        else164            .expanded,165        .sparse_cross_entropy = if (options.tensor.loss.kernel_library == .enabled and options.tensor.loss.row_sparse_cross_entropy_schedule != null)166            .semantic_kernel167        else168            .expanded,169    };170}171172fn 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_kernel177        else178            .expanded,179        .sparse_cross_entropy = if (options.kernel_call_registry != null and options.row_sparse_cross_entropy_schedule != null)180            .semantic_kernel181        else182            .expanded,183    };184}185186pub fn prepare(187    allocator: std.mem.Allocator,188    program: *const program_mod.Program,189) !BackendPreparedJob {190    return prepareWith(allocator, program, .{});191}192193pub 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}205206pub 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}224225pub 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}238239pub 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}252253pub fn compileFragmentFromArtifactJob(254    allocator: std.mem.Allocator,255    artifact_module: *ArtifactJob,256) !*CompiledFragment {257    return try accy.executable.compileFragmentFromArtifactJob(allocator, artifact_module);258}259260pub 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}270271pub 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}287288fn 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;295296    var body_program = try hostLoopScanBodyProgram(allocator, program.name, scan);297    defer body_program.deinit();298299    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();311312    try accy.executable.fragment.recordBackendPreparationRun(options.instrumentation, prepared.run);313314    const artifact_module = try createArtifactJobFromPreparedJob(allocator, handle, &prepared, options);315    defer artifact_module.deinit();316317    const launch_plan = try hostLoopScanLaunchPlan(318        allocator,319        artifact_module.artifactPlan(),320        @intCast(scan.length),321    );322323    return try accy.executable.compileFragmentFromArtifactJobWithLaunchPlan(324        allocator,325        artifact_module,326        launch_plan,327    );328}329330fn hostLoopScanCandidate(program: *const program_mod.Program) ?program_mod.Scan {331    if (program.outputs.len != 1) return null;332333    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    }344345    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}354355fn 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}379380fn 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();387388    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];393394    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}403404fn 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;412413    var launch_plan = try accy.executable.createDataflowLaunchGraphPlan(414        allocator,415        artifact_plan,416        .{},417    );418    errdefer launch_plan.deinit();419420    const carries = allocator.alloc(accy.executable.LaunchGraphLoopCarry, 1) catch return error.OutOfMemory;421    errdefer allocator.free(carries);422423    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    }435436    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    };444445    launch_plan.loops = loops;446    return launch_plan;447}448449pub fn module(allocator: std.mem.Allocator) Module {450    return .{ .allocator = allocator };451}452453pub const Module = struct {454    allocator: std.mem.Allocator,455456    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};460461pub const ModuleLower = struct {462    allocator: std.mem.Allocator,463    options: SemanticLoweringOptions = .{},464465    pub const Result = *SemanticModule;466467    pub fn finish(self: *@This(), ctx: anytype, outputs: []const trace.Value) !Result {468        var graph = try ctx.default(outputs);469        defer graph.deinit();470471        return toSemanticModuleWithOptions(self.allocator, &graph, self.options);472    }473};474475pub 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}490491pub const ScanLoweringState = struct {492    results: std.AutoHashMapUnmanaged(u32, []const *ir.Value) = .empty,493494    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};501502pub 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,508509    pub const Value = *ir.Value;510    pub const Result = void;511512    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    }521522    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    }525526    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];529530        var iterate = try self.function.beginIterate(inits[0..scan.inits.len], scan.length);531        errdefer iterate.body.terminated = true;532533        const body_values = try self.allocator.alloc(Value, scan.body.values.len);534        defer self.allocator.free(body_values);535536        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        }548549        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];553554        const predicate = try truePredicate(self.allocator, iterate.inner(), scan.body.typeOf(scan.body.outputs[0]));555        try iterate.yield_(predicate, yielded);556557        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    }563564    pub fn finish(self: *@This(), outputs: []const Value) !Result {565        try self.function.return_(outputs);566        try self.function.finish();567    }568};569570fn 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}658659fn 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 and671        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}677678fn 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;699700    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);717718    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);721722    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);729730    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        else738            @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        selected754    else755        try function.reduce(selected, zero_scalar, result_type, "sum", reduce_axes);756    return function.add(input, reduced);757}758759fn 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 and770        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}776777fn 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)};791792    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;795796    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    }804805    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);810811    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);816817    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}825826fn 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}851852fn 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;860861    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}872873fn 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;884885    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    }890891    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    }903904    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}909910fn 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}915916fn 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}930931fn 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}940941fn 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}956957fn 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}970971fn 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}979980fn 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}993994fn domainCompatible(ty: program_mod.Type, domain: usize) bool {995    const count = typeElementCount(ty) orelse return false;996    return count == 1 or count == domain;997}998999fn 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}10041005fn 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}10191020const PhiloxKeyCounterUniformTarget = struct {1021    rounds: u32,1022    dtype: choir_abi.DType,1023};10241025fn 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}10371038fn 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}10451046fn 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}10721073fn 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}10861087fn 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}11031104test "tensor lowering produces verified semantic Choir" {1105    var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_dense");1106    defer builder.deinit();11071108    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();11141115    const lowered = try toSemanticModule(std.testing.allocator, &program);1116    defer lowered.deinit();11171118    try lowered.verify();1119}11201121test "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();11241125    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();11321133    const lowered = try toSemanticModule(std.testing.allocator, &program);1134    defer lowered.deinit();11351136    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}11431144test "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();11471148    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();11551156    const lowered = try toSemanticModule(std.testing.allocator, &program);1157    defer lowered.deinit();11581159    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}11661167test "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();11701171    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();11781179    const lowered = try toSemanticModuleWithOptions(std.testing.allocator, &program, .{1180        .scatter_add = .semantic_kernel,1181    });1182    defer lowered.deinit();11831184    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}11981199test "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();12021203    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();12081209    const lowered = try toSemanticModule(std.testing.allocator, &program);1210    defer lowered.deinit();12111212    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}12191220test "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();12231224    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();12291230    const lowered = try toSemanticModuleWithOptions(std.testing.allocator, &program, .{1231        .sparse_cross_entropy = .semantic_kernel,1232    });1233    defer lowered.deinit();12341235    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}12451246test "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();12501251    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();12561257    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 = &registry,1266        .row_sparse_cross_entropy_schedule = .{ .thread_blocks = 4 },1267    });1268    defer prepared.deinit();12691270    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}12791280test "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();12841285    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();12921293    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 = &registry,1302        .scatter_add_schedule = .{ .thread_blocks = 4 },1303    });1304    defer prepared.deinit();13051306    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}13151316fn 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}13231324test "tensor lowering preserves eligible scan as iterate" {1325    var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_scan");1326    defer builder.deinit();13271328    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();13371338    const lowered = try toSemanticModule(std.testing.allocator, &program);1339    defer lowered.deinit();13401341    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}13471348fn 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}13521353test "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();13571358    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();13661367    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();13761377    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}13811382fn 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}13901391test "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();13951396    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();14051406    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    );14131414    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();14231424    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}14281429fn 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}14381439test "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();14431444    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();14531454    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    );14651466    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();14751476    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}14801481test "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();14851486    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();14951496    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();15021503    try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());1504}15051506fn 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}15291530test "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();15341535    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();15461547    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();15531554    try std.testing.expectEqual(@as(usize, 1), try prepared.generatedKernelCount());1555}15561557fn 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}15611562test "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();15651566    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();15741575    const lowered = try toSemanticModule(std.testing.allocator, &program);1576    defer lowered.deinit();15771578    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}15881589test "tensor complex scan executable uses bounded host loop launch plan" {1590    const allocator = std.testing.allocator;15911592    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();15961597    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 };16041605    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();16111612    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();16161617    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);16251626    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);16301631    try std.testing.expectEqual(long_fragment.kernelCount() * 7, state.launch_count);1632}16331634test "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();16371638    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();16441645    const lowered = try toSemanticModule(std.testing.allocator, &program);1646    defer lowered.deinit();16471648    try lowered.verify();1649}16501651fn 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}16571658fn 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();16651666    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}16741675const CountLowerAdd = struct {1676    count: *usize,16771678    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};16881689fn 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();16961697    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    }17041705    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    }17101711    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    }17161717    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 }));17211722    return try builder.finish();1723}17241725test "tensor lowering semantics composes with interpreter layers" {1726    var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_layer");1727    defer builder.deinit();17281729    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();17341735    var add_count: usize = 0;1736    const lowered = try toSemanticModuleCountingAdds(std.testing.allocator, &program, &add_count);1737    defer lowered.deinit();17381739    try std.testing.expectEqual(@as(usize, 1), add_count);1740    try lowered.verify();1741}17421743test "tensor lowering module attaches to graph interpretation" {1744    var builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_module");1745    defer builder.deinit();17461747    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();17521753    var graph_builder = try accy.tensor.Builder.init(std.testing.allocator, "lower_module");1754    errdefer graph_builder.deinit();17551756    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();17591760    try lowered.verify();1761}