tiny.accy.preparation.capture

1const std = @import("std");2const choir = @import("choir");3const gpu = choir.backends.gpu;4const preparation = @import("root.zig");5const records = @import("../choir/root.zig").record;6const codec = records.codec;7const dispatch_record = records.dispatch;8const memory_record = records.memory;910pub fn dispatch(11    allocator: std.mem.Allocator,12    root: *choir.ir.Operation,13    fusion: *const preparation.fusion.FusionPlanAnalysis,14    schedule: *const preparation.schedule.SchedulePlanAnalysis,15    entity_limit: u32,16) ![]u8 {17    comptime {18        codec.coverage(preparation.fusion.FusionCluster, dispatch_record.Cluster, &.{});19        codec.coverage(20            preparation.fusion.FusionPlanAnalysis,21            dispatch_record.Fusion,22            &.{"allocator"},23        );24        codec.coverage(preparation.schedule.ScheduleWorkItem, dispatch_record.WorkItem, &.{});25        codec.coverage(preparation.schedule.SchedulePlanAnalysis, dispatch_record.Schedule, &.{26            "allocator", "root_to_item",27        });28    }29    if (schedule.root_to_item.count() != schedule.work_items.items.len) {30        return error.UnencodableProduct;31    }32    for (schedule.work_items.items, 0..) |item, index| {33        if (schedule.root_to_item.get(item.root) != index) return error.UnencodableProduct;34    }35    var references = try records.reference.Index.init(allocator, root, entity_limit);36    defer references.deinit();37    return qualified(allocator, dispatch_record.Record, .dispatch, .{38        .fusion = fusion,39        .schedule = schedule,40    }, &references);41}4243pub fn memory(44    allocator: std.mem.Allocator,45    root: *choir.ir.Operation,46    buffers: *const preparation.bufferization.BufferPlanAnalysis,47    spaces: *const preparation.memory.MemorySpacePlanAnalysis,48    layouts: *const preparation.layout.LayoutPlanAnalysis,49    entity_limit: u32,50) ![]u8 {51    comptime memoryCoverage();52    if (buffers.value_to_slot.count() < buffers.slots.items.len or53        buffers.value_to_elision.count() != buffers.elisions.items.len or54        spaces.slot_to_assignment.count() != spaces.assignments.items.len or55        layouts.slot_to_assignment.count() != layouts.assignments.items.len)56    {57        return error.UnencodableProduct;58    }59    for (buffers.slots.items, 0..) |slot, index| {60        if (buffers.value_to_slot.get(slot.value) != index) return error.UnencodableProduct;61    }62    for (buffers.elisions.items, 0..) |elision, index| {63        if (buffers.value_to_elision.get(elision.value) != index) return error.UnencodableProduct;64    }65    for (spaces.assignments.items, 0..) |assignment, index| {66        if (spaces.slot_to_assignment.get(assignment.slot_id) != index) {67            return error.UnencodableProduct;68        }69    }70    for (layouts.assignments.items, 0..) |assignment, index| {71        if (layouts.slot_to_assignment.get(assignment.slot_id) != index) {72            return error.UnencodableProduct;73        }74    }75    var references = try records.reference.Index.init(allocator, root, entity_limit);76    defer references.deinit();77    const bindings = try memoryBindings(allocator, buffers, &references);78    defer allocator.free(bindings);79    const bytes = try qualified(allocator, memory_record.Record, .memory, .{80        .buffers = buffers,81        .bindings = bindings,82        .spaces = spaces,83        .layouts = layouts,84    }, &references);85    errdefer allocator.free(bytes);86    try compareMemoryBindings(allocator, bytes, buffers, &references);87    return bytes;88}8990fn compareMemoryBindings(91    allocator: std.mem.Allocator,92    bytes: []const u8,93    buffers: *const preparation.bufferization.BufferPlanAnalysis,94    references: *const records.reference.Index,95) !void {96    var decoded = try codec.decode(allocator, memory_record.Record, .memory, bytes);97    defer decoded.deinit();98    if (decoded.value.bindings.len != buffers.value_to_slot.count()) return error.UnencodableProduct;99    var entries = buffers.value_to_slot.iterator();100    while (entries.next()) |entry| {101        const expected = try references.value(entry.key_ptr.*);102        const found = for (decoded.value.bindings) |binding| {103            if (std.meta.eql(expected, binding.value)) break binding.slot_id;104        } else return error.UnencodableProduct;105        if (found != entry.value_ptr.*) return error.UnencodableProduct;106    }107}108109fn memoryBindings(110    allocator: std.mem.Allocator,111    buffers: *const preparation.bufferization.BufferPlanAnalysis,112    references: *const records.reference.Index,113) ![]memory_record.Binding {114    const bindings = try allocator.alloc(memory_record.Binding, buffers.value_to_slot.count());115    errdefer allocator.free(bindings);116    var entries = buffers.value_to_slot.iterator();117    var count: usize = 0;118    while (entries.next()) |entry| {119        if (entry.value_ptr.* >= buffers.slots.items.len) return error.UnencodableProduct;120        bindings[count] = .{121            .value = try references.value(entry.key_ptr.*),122            .slot_id = entry.value_ptr.*,123        };124        count += 1;125    }126    std.debug.assert(count == bindings.len);127    std.mem.sort(memory_record.Binding, bindings, {}, memory_record.Binding.lessThan);128    return bindings;129}130131fn memoryCoverage() void {132    codec.coverage(preparation.bufferization.BufferSlot, memory_record.Slot, &.{});133    codec.coverage(preparation.bufferization.FusionElision, memory_record.Elision, &.{});134    codec.coverage(preparation.bufferization.BufferPlanAnalysis, memory_record.Buffers, &.{135        "allocator", "value_to_slot", "value_to_elision",136    });137    codec.coverage(preparation.memory.MemorySpaceAssignment, memory_record.Assignment, &.{});138    codec.coverage(preparation.memory.MemorySpacePlanAnalysis, memory_record.Spaces, &.{139        "allocator", "slot_to_assignment",140    });141    codec.coverage(preparation.layout.LayoutAssignment, memory_record.Layout, &.{});142    codec.coverage(preparation.layout.LayoutPlanAnalysis, memory_record.Layouts, &.{143        "allocator", "slot_to_assignment",144    });145}146147pub fn kernel(148    allocator: std.mem.Allocator,149    root: *choir.ir.Operation,150    outlines: *const preparation.kernelization.product.KernelOutlinePlanAnalysis,151    generated: *const preparation.kernelization.KernelizationAnalysis,152    comptime configuration: choir.product.operation.Configuration,153) ![]u8 {154    comptime kernelCoverage();155    try requireWorkMap(outlines);156    try requireWorkMap(generated);157    var programs = try Generated.init(allocator, generated.kernels.items, configuration);158    defer programs.deinit();159    var references = try records.reference.Index.init(allocator, root, configuration.image.entities);160    defer references.deinit();161    return qualified(allocator, records.kernel.Record, .kernel, .{162        .outlines = outlines,163        .generated = .{ .kernels = programs.lowered },164    }, &references);165}166167const Generated = struct {168    allocator: std.mem.Allocator,169    programs: []codec.Decoded(records.program.Record),170    lowered: []records.kernel.Lowered,171172    fn init(173        allocator: std.mem.Allocator,174        source: []const preparation.kernelization.LoweredKernel,175        comptime configuration: choir.product.operation.Configuration,176    ) !Generated {177        comptime kernelCoverage();178        if (source.len > configuration.codec.fields) return error.RecordLimit;179        const programs = try allocator.alloc(180            codec.Decoded(records.program.Record),181            source.len,182        );183        errdefer allocator.free(programs);184        var initialized: usize = 0;185        errdefer for (programs[0..initialized]) |*program| program.deinit();186        const lowered = try allocator.alloc(records.kernel.Lowered, programs.len);187        errdefer allocator.free(lowered);188        for (source, 0..) |*kernel_value, index| {189            programs[index] = try records.program.capture(190                allocator,191                &kernel_value.program,192                configuration,193            );194            initialized += 1;195            inline for (@typeInfo(records.kernel.Lowered).@"struct".field_names) |name| {196                @field(lowered[index], name) = if (comptime std.mem.eql(u8, name, "program"))197                    programs[index].value198                else199                    @field(kernel_value, name);200            }201        }202        return .{ .allocator = allocator, .programs = programs, .lowered = lowered };203    }204205    fn deinit(self: *Generated) void {206        for (self.programs) |*program| program.deinit();207        self.allocator.free(self.programs);208        self.allocator.free(self.lowered);209    }210};211212pub fn target(213    allocator: std.mem.Allocator,214    root: *choir.ir.Operation,215    schedule: *const preparation.schedule.SchedulePlanAnalysis,216    generated: []const preparation.kernelization.LoweredKernel,217    comptime configuration: choir.product.operation.Configuration,218) ![]u8 {219    var programs = try Generated.init(allocator, generated, configuration);220    defer programs.deinit();221    const profile = preparation.readBackendTargetProfile(root);222    const kernels = try allocator.alloc(records.target.Kernel, programs.lowered.len);223    defer allocator.free(kernels);224    var initialized: usize = 0;225    defer for (kernels[0..initialized]) |kernel_value| {226        if (kernel_value.abi) |abi| allocator.free(abi.static_arguments);227    };228    for (programs.lowered, generated, 0..) |lowered, *source, index| {229        const work = for (schedule.work_items.items) |item| {230            if (item.id == lowered.work_item_id) break item;231        } else return error.UnencodableProduct;232        kernels[index] = try targetKernel(allocator, lowered, source, work, profile);233        initialized += 1;234    }235    var references = records.reference.Index{ .allocator = allocator, .limit = 0 };236    defer references.deinit();237    const row_schedules = preparation.target.readGeneratedRowPipelineSchedules(root);238    return qualified(allocator, records.target.Record, .target, records.target.Record{239        .profile = profile,240        .math_tier = if (profile) |value| value.math_tier else .exact,241        .generated_scan_schedules = preparation.target.readGeneratedScanSchedules(root),242        .generated_row_pipeline_schedules = row_schedules,243        .kernels = kernels,244    }, &references);245}246247fn targetKernel(248    allocator: std.mem.Allocator,249    lowered: records.kernel.Lowered,250    source: *const preparation.kernelization.LoweredKernel,251    work: preparation.schedule.ScheduleWorkItem,252    profile: ?preparation.BackendTargetProfile,253) !records.target.Kernel {254    var dtypes = source.requiredDTypes();255    dtypes.insert(work.dtype);256    return .{257        .lowered = lowered,258        .work_dtype = work.dtype,259        .element_count = work.element_count,260        .required_dtype_bits = dtypes.bits,261        .required_features = gpu.featureRequirementsForModule(source.program.kernelModule()),262        .required_subgroup = gpu.subgroupRequirementsForModule(source.program.kernelModule()),263        .runtime_scalar_argument_count = source.runtimeScalarArgumentCount(),264        .abi = if (profile) |value| try targetAbi(allocator, source, work, value) else null,265    };266}267268fn targetAbi(269    allocator: std.mem.Allocator,270    source: *const preparation.kernelization.LoweredKernel,271    work: preparation.schedule.ScheduleWorkItem,272    profile: preparation.BackendTargetProfile,273) !records.target.Abi {274    const owner = @import("../target/root.zig");275    const format = profile.artifact_format;276    const count = try owner.abi.argumentCount(format, source.argument_count);277    const launch = try owner.abi.launchGeometry(278        format,279        work.element_count,280        source.launchGeometry(),281    );282    return .{283        .argument_count = count,284        .static_arguments = try owner.abi.staticArguments(285            allocator,286            format,287            work.element_count,288            launch,289        ),290        .launch = launch,291        .compile_options = owner.compileOptionsForArtifactFormat(format, work.element_count),292    };293}294295fn requireWorkMap(analysis: anytype) !void {296    if (analysis.work_to_kernel.count() != analysis.kernels.items.len) {297        return error.UnencodableProduct;298    }299    for (analysis.kernels.items, 0..) |item, index| {300        if (analysis.work_to_kernel.get(item.work_item_id) != index) {301            return error.UnencodableProduct;302        }303    }304}305306fn kernelCoverage() void {307    const product = preparation.kernelization.product;308    codec.coverage(product.KernelOutline, records.kernel.Outline, &.{});309    codec.coverage(product.KernelOutlinePlanAnalysis, records.kernel.Outlines, &.{310        "allocator", "work_to_kernel",311    });312    codec.coverage(product.LoweredKernel, records.kernel.Lowered, &.{});313    codec.coverage(product.KernelizationAnalysis, records.kernel.Generated, &.{314        "allocator", "context", "work_to_kernel",315    });316}317318fn qualified(319    allocator: std.mem.Allocator,320    comptime T: type,321    stage: @import("../choir/root.zig").publication.Stage,322    source: anytype,323    references: *const records.reference.Index,324) ![]u8 {325    const bytes = try codec.encode(allocator, T, stage, source, references);326    errdefer allocator.free(bytes);327    var decoded = try codec.decode(allocator, T, stage, bytes);328    defer decoded.deinit();329    try codec.compare(decoded.value, source, references);330    const compared = try codec.encode(allocator, T, stage, decoded.value, references);331    defer allocator.free(compared);332    if (!std.mem.eql(u8, bytes, compared)) return error.UnencodableProduct;333    return bytes;334}

1pub const capture = @import("capture.zig");