tiny.accy.choir.record.program

1const std = @import("std");2const choir = @import("choir");3const model = @import("../../kernel/model/root.zig");4const records = @import("root.zig");5const schedule = model.core.schedule;6const Configuration = choir.product.operation.Configuration;78/// The stored form of one generated kernel program: its own compiler bytecode,9/// the number of its function within that bytecode, its parameters, its loop10/// and thread structure, its launch shape and its arithmetic policy. A kernel11/// stage record carries one of these for each kernel it generated. The function12/// number counts operations in this program's own bytecode, which is separate13/// from the bytecode of the whole stage.14pub const Record = struct {15    image: []const u8,16    function: u32,17    params: []const model.Param,18    schedule: schedule.Record,19    launch: schedule.Launch,20    arithmetic: choir.product.recipe.ArithmeticPolicy,21};2223/// Copies each generated program into its record for the kernel stage: verifies24/// the program, encodes its bytecode, and fails with `error.RecordLimit` when25/// that bytecode exceeds the configured image size. The call numbers the26/// program's operations, snapshots its schedule, encodes the whole record,27/// decodes it again and compares the result with the live program. The caller28/// receives the decoded copy, which owns its own memory apart from `source`.29pub fn capture(30    allocator: std.mem.Allocator,31    source: *const model.Program,32    comptime configuration: Configuration,33) !records.codec.Decoded(Record) {34    comptime coverage();35    const root = source.kernelModule();36    var context = try choir.ir.Context.init(allocator, configuration.context);37    defer context.deinit(allocator);38    context.arithmetic_policy = root.context.arithmetic_policy;39    try configuration.register(&context);40    try choir.ir.verifyOperation(root, configuration.verify);41    const image = try choir.bytecode.qualification.encode(42        allocator,43        root,44        &.{},45        &context,46        configuration.codec,47    );48    defer allocator.free(image);49    if (image.len > configuration.image.bytes) return error.RecordLimit;50    const entity_limit = configuration.image.entities;51    var references = try records.reference.Index.init(allocator, root, entity_limit);52    defer references.deinit();53    var snapshot = try source.scheduleSnapshot(allocator);54    defer snapshot.deinit(allocator);55    const projected = Record{56        .image = image,57        .function = (try references.operation(source.storage.kernel.func().op)).ordinal,58        .params = source.params(),59        .schedule = try snapshot.record(),60        .launch = try source.launch(),61        .arithmetic = root.context.arithmetic_policy,62    };63    const bytes = try records.codec.encode(allocator, Record, .kernel, projected, &references);64    defer allocator.free(bytes);65    var decoded = try records.codec.decode(allocator, Record, .kernel, bytes);66    errdefer decoded.deinit();67    try records.codec.compare(decoded.value, projected, &references);68    return decoded;69}7071/// Validates each stored program for the checker before a record is accepted:72/// decodes the bytecode into a scratch context, verifies it, and checks that73/// the function number is in range (`error.UnboundProductInput`). The function74/// must be a kernel with a body, named `entry_name`, whose arguments match the75/// stored parameters in count and type. The stored launch shape must match the76/// result of replaying the schedule, and any mismatch is77/// `error.InvalidStageRecord`.78pub fn validate(79    allocator: std.mem.Allocator,80    value: Record,81    entry_name: []const u8,82    comptime configuration: Configuration,83) !void {84    const image = try choir.bytecode.image.Index.create(85        allocator,86        value.image,87        configuration.image,88    );89    defer image.destroy();90    if (value.function >= image.view().operations.len) return error.UnboundProductInput;91    var context = try choir.ir.Context.init(allocator, configuration.context);92    defer context.deinit(allocator);93    context.arithmetic_policy = value.arithmetic;94    try configuration.register(&context);95    var decoded = try choir.bytecode.decodeModule(allocator, &context, value.image);96    defer decoded.deinit();97    try choir.ir.verifyOperation(decoded.module, configuration.verify);98    var references = try records.reference.Index.init(99        allocator,100        decoded.module,101        configuration.image.entities,102    );103    defer references.deinit();104    const function = try selectedFunction(&references, value.function);105    const name = function.getName() orelse return error.InvalidStageRecord;106    if (!std.mem.eql(u8, name, entry_name)) return error.InvalidStageRecord;107    if (!function.isKernel() or !function.hasBody()) return error.InvalidStageRecord;108    if (function.getNumArguments() != value.params.len) return error.InvalidStageRecord;109    for (value.params, function.getArguments()) |param, argument| {110        if (!argument.type.eql(try param.getType(&context))) return error.InvalidStageRecord;111    }112    var replayed = try schedule.Schedule.init(allocator, value.schedule.replay_limits);113    defer replayed.deinit(allocator);114    try replayed.replay(value.schedule);115    if (!std.meta.eql(try replayed.launch(), value.launch)) return error.InvalidStageRecord;116}117118/// Restores a working kernel program from a stored record for a later compile:119/// rebuilds the whole program in new storage sized from the record, then120/// verifies it, checks its name, and replays its schedule. Storage that runs121/// out gives `error.WorkExhausted`. The returned program copies the bytecode,122/// the parameters and the schedule, and keeps no pointer into `value`.123pub fn restore(124    allocator: std.mem.Allocator,125    value: Record,126    entry_name: []const u8,127    comptime configuration: Configuration,128) !model.Program {129    const limits = try restorationLimits(value, entry_name, configuration);130    const capacity = model.program.Capacity.derive(limits) catch return error.WorkOverflow;131    var storage = try model.core.builder.Storage.init(allocator, limits.raw());132    var transferred = false;133    defer if (!transferred) storage.deinit(allocator);134    const program = restoreIn(135        allocator,136        value,137        entry_name,138        &storage,139        capacity,140        configuration,141    ) catch |err| {142        return if (storage.context.exhaustedSegment() != null) error.WorkExhausted else err;143    };144    transferred = true;145    return program;146}147148fn restorationLimits(149    value: Record,150    entry_name: []const u8,151    comptime configuration: Configuration,152) !model.program.Limits {153    var names: usize = 0;154    for (value.schedule.axes) |axis| {155        names = std.math.add(usize, names, axis.name.len) catch return error.WorkOverflow;156    }157    for (value.schedule.steps) |step| {158        if (step == .axis) {159            names = std.math.add(usize, names, step.axis.name.len) catch return error.WorkOverflow;160        }161    }162    return .{163        .context = configuration.context,164        .parameters = value.params.len,165        .kernel_name_bytes = entry_name.len,166        .temporary_values = 0,167        .temporary_types = 0,168        .schedule = value.schedule.replay_limits,169        .snapshot = .{170            .axes = value.schedule.axes.len,171            .steps = value.schedule.steps.len,172            .name_bytes = names,173        },174    };175}176177fn restoreIn(178    allocator: std.mem.Allocator,179    value: Record,180    entry_name: []const u8,181    storage: *model.core.builder.Storage,182    capacity: model.program.Capacity,183    comptime configuration: Configuration,184) !model.Program {185    const image = try choir.bytecode.image.Index.create(186        allocator,187        value.image,188        configuration.image,189    );190    defer image.destroy();191    if (value.function >= image.view().operations.len) return error.UnboundProductInput;192    storage.context.arithmetic_policy = value.arithmetic;193    try configuration.register(storage.context);194    var decoded = try choir.bytecode.decodeModule(195        choir.ir.context.operationAllocator(storage.context),196        storage.context,197        value.image,198    );199    errdefer decoded.deinit();200    errdefer decoded.module.erase();201    try choir.ir.verifyOperation(decoded.module, configuration.verify);202    var references = try records.reference.Index.init(203        allocator,204        decoded.module,205        configuration.image.entities,206    );207    defer references.deinit();208    const function = try selectedFunction(&references, value.function);209    const name = function.getName() orelse return error.InvalidStageRecord;210    if (!std.mem.eql(u8, name, entry_name)) return error.InvalidStageRecord;211    var replayed = try schedule.Schedule.init(allocator, value.schedule.replay_limits);212    errdefer replayed.deinit(allocator);213    try replayed.replay(value.schedule);214    if (!std.meta.eql(try replayed.launch(), value.launch)) return error.InvalidStageRecord;215    const kernel = try model.Kernel.fromDecoded(216        allocator,217        storage.*,218        decoded,219        function,220        value.params,221    );222    return model.Program.init(kernel, replayed, capacity);223}224225fn selectedFunction(226    references: *const records.reference.Index,227    ordinal: u32,228) !choir.dialects.FuncDialect.FuncOp {229    var selected: ?*choir.ir.Operation = null;230    var entries = references.operations.iterator();231    for (0..references.operations.count()) |_| {232        const entry = entries.next().?;233        if (entry.value_ptr.ordinal == ordinal) selected = entry.key_ptr.*;234    }235    const operation = selected orelse return error.UnboundProductInput;236    const FuncOp = choir.dialects.FuncDialect.FuncOp;237    if (!std.mem.eql(u8, operation.name.name, FuncOp.operation_name)) {238        return error.InvalidStageRecord;239    }240    return .{ .op = operation };241}242243fn coverage() void {244    const require = choir.product.revision.record.requireFields;245    require(model.Program, &.{ "capacity", "storage" });246    require(@FieldType(model.Program, "storage"), &.{ "kernel", "schedule" });247    require(model.Kernel, &.{ "allocator", "capacity", "storage", "state" });248    const State = @typeInfo(@FieldType(model.Kernel, "state")).pointer.child;249    require(State, &.{ "ctx", "module", "func", "params", "decoded" });250    require(schedule.Snapshot, &.{251        "phase",252        "capacity",253        "storage",254        "version",255        "axes_storage",256        "axes_len",257        "steps_storage",258        "steps_len",259        "names_storage",260        "names_len",261        "replay_limits",262        "captured",263    });264}

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