tiny.accy.preparation.recipe

1const std = @import("std");2const choir = @import("choir");3const stage_owner = @import("stage.zig");4const run = @import("run.zig");5const target = @import("target.zig");6const accy_choir = @import("../choir/root.zig");7const records = accy_choir.record;8const Stage = accy_choir.publication.Stage;9const Options = run.BackendPreparationRunOptions;10const library = @import("../kernel/library/root.zig");1112pub const identity = choir.product.revision.record.Version{13    .name = "accy-stage-recipe",14    .version = 1,15};1617pub const Policy = struct {18    max_threads: usize = 1,19    worker_allocator: ?void = null,20    verification: choir.ir.VerifyOptions = choir.ir.verify.default_options,21};2223pub const policy = Policy{};2425pub fn runOptions() choir.passes.PassManagerRunOptions {26    return .{ .max_threads = policy.max_threads, .worker_allocator = null };27}2829pub fn pipelineName(comptime stage: Stage) []const u8 {30    return @field(stage_owner, @tagName(stage) ++ "_pipeline_name");31}3233/// The compile chain records this list with each stage so a change to the pass pipeline changes the34/// stage record, the sealed result of one compile stage. The function returns the ordered35/// identities of the passes in the stage's pipeline: each pass's declared name, version and cost36/// estimate, or its name at version 1 when it has none. The list is derived at compile time from37/// the same stage declarations that build the pipeline, so the recorded list and the pipeline that38/// runs cannot drift apart.39pub fn pipeline(comptime stage: Stage) []const choir.product.revision.record.Version {40    const declarations = comptime stages(stage);41    const identities = comptime blk: {42        var result: [declarations.len]choir.product.revision.record.Version = undefined;43        for (declarations, 0..) |entry, index| {44            result[index] = if (entry.pass.work_contract) |contract|45                contract.identity46            else47                .{ .name = entry.name, .version = 1 };48        }49        break :blk result;50    };51    return &identities;52}5354fn stages(comptime stage: Stage) []const stage_owner.BackendPreparationStage {55    return switch (stage) {56        .semantic => &.{},57        .contract => &stage_owner.contract_stages,58        .tensor => &stage_owner.tensor_stages,59        .dispatch => &stage_owner.dispatch_stages,60        .memory => &stage_owner.memory_stages,61        .kernel => &stage_owner.kernel_stages,62        .target => &stage_owner.target_stages,63    };64}6566/// The compile chain calls this to fill a pass manager with one stage's pipeline. The call turns on67/// verification and adds the stage's passes to `manager`. The tensor stage's passes borrow68/// `options.tensor`, so `options` must outlive `manager`.69pub fn configure(70    manager: *choir.passes.PassManager,71    comptime stage: Stage,72    options: *const Options,73) !void {74    manager.enableVerifierWithOptions(policy.verification);75    switch (stage) {76        .semantic => {},77        .contract => try stage_owner.addContractPipeline(manager),78        .tensor => try stage_owner.addTensorPipelineWithOptions(manager, &options.tensor),79        .dispatch => try stage_owner.addDispatchPipeline(manager),80        .memory => try stage_owner.addMemoryPipeline(manager),81        .kernel => try stage_owner.addKernelPipeline(manager),82        .target => try stage_owner.addTargetPipeline(manager),83    }84}8586/// The compile chain calls this to write a request's target settings, the device profile and87/// per-kernel schedule choices stored on the module, onto a module in a mutable compiler job. The88/// call writes the device profile when one is set, and the scan and row schedule choices when their89/// lists are non-empty. An unset profile or an empty list leaves the setting inherited from the90/// source exactly as it was.91pub fn applyTargetOptions(92    allocator: std.mem.Allocator,93    root: *choir.ir.Operation,94    options: Options,95) !void {96    if (options.target_profile) |profile| {97        try target.setBackendTargetProfile(root.context, root, profile);98    }99    if (options.generated_scan_schedules.len != 0) {100        try target.setGeneratedScanSchedules(101            allocator,102            root.context,103            root,104            options.generated_scan_schedules,105        );106    }107    if (options.generated_row_pipeline_schedules.len != 0) {108        try target.setGeneratedRowPipelineSchedules(109            allocator,110            root.context,111            root,112            options.generated_row_pipeline_schedules,113        );114    }115}116117const Einsum = struct {118    strategy: accy_choir.einsum.Strategy,119    exact_state_limit: usize,120    beam_width: usize,121    auto_beam_width: usize,122    kernel_library: @import("library.zig").KernelLibraryLowering,123    matrix_product_schedule: ?library.MatrixProductSchedule,124    matrix_product_tuning: ?library.linalg.MatrixProductScheduleReader = null,125    family_tuning: ?library.tuning.FamilyTuningReader = null,126};127128const Indexing = struct {129    kernel_library: @import("library.zig").KernelLibraryLowering,130    gather_schedule: ?library.GatherSchedule,131    scatter_schedule: ?library.ScatterSchedule,132    scatter_add_schedule: ?library.ScatterAddSchedule,133    family_tuning: ?library.tuning.FamilyTuningReader = null,134};135136const Tensor = struct {137    activation: stage_owner.ActivationLoweringOptions,138    einsum: Einsum,139    indexing: Indexing,140    loss: stage_owner.LossLoweringOptions,141};142143pub const TargetInputs = struct {144    profile: ?target.BackendTargetProfile,145    generated_scan_schedules: ?[]const u8,146    generated_row_pipeline_schedules: ?[]const u8,147};148149fn Lowering(comptime stage: Stage) type {150    return switch (stage) {151        .tensor => Tensor,152        .kernel, .target => TargetInputs,153        else => struct {},154    };155}156157pub fn Record(comptime stage: Stage) type {158    return struct { execution: Policy, options: Lowering(stage) };159}160161/// The compile chain calls this to encode the options a stage ran with, for its stage record. The162/// call reads the target settings already on `root` without changing the module, applies the163/// request's overrides, and returns bytes the caller owns. An empty override keeps the inherited164/// setting. A tuning table given in the options is recorded with its device facts and its records165/// in order, so a table that lacks a key is recorded as lacking it. Timing, the failure destination166/// and the clock are administrative and are left out, so they never change a record.167pub fn encode(168    allocator: std.mem.Allocator,169    comptime stage: Stage,170    root: *choir.ir.Operation,171    options: Options,172) ![]u8 {173    return encodeWithTarget(allocator, stage, .{174        .profile = target.readBackendTargetProfile(root),175        .generated_scan_schedules = target.readGeneratedScanSchedules(root),176        .generated_row_pipeline_schedules = target.readGeneratedRowPipelineSchedules(root),177    }, options);178}179180/// The compile chain calls this to encode a stage's options from a retained source without a181/// mutable module at hand. The call starts from the target settings captured from the original182/// source, applies this request's overrides, and returns bytes the caller owns. The call decodes183/// its own output and compares it with the input before returning, so a value that does not survive184/// encoding is caught at once.185pub fn encodeWithTarget(186    allocator: std.mem.Allocator,187    comptime stage: Stage,188    inherited: TargetInputs,189    options: Options,190) ![]u8 {191    comptime classify();192    const scan = if ((stage == .kernel or stage == .target) and193        options.generated_scan_schedules.len != 0)194        try target.encodeGeneratedScanSchedules(allocator, options.generated_scan_schedules)195    else196        null;197    defer if (scan) |bytes| allocator.free(bytes);198    const row = if ((stage == .kernel or stage == .target) and199        options.generated_row_pipeline_schedules.len != 0)200        try target.encodeGeneratedRowPipelineSchedules(201            allocator,202            options.generated_row_pipeline_schedules,203        )204    else205        null;206    defer if (row) |bytes| allocator.free(bytes);207    const source = switch (stage) {208        .tensor => options.tensor,209        .kernel, .target => TargetInputs{210            .profile = options.target_profile orelse inherited.profile,211            .generated_scan_schedules = scan orelse inherited.generated_scan_schedules,212            .generated_row_pipeline_schedules = row orelse213                inherited.generated_row_pipeline_schedules,214        },215        else => Lowering(stage){},216    };217    var references = records.reference.Index{ .allocator = allocator, .limit = 0 };218    defer references.deinit();219    const effective = .{ .execution = policy, .options = source };220    const bytes = try records.codec.encode(allocator, Record(stage), stage, effective, &references);221    errdefer allocator.free(bytes);222    var decoded = try records.codec.decode(allocator, Record(stage), stage, bytes);223    defer decoded.deinit();224    try records.codec.compare(decoded.value, effective, &references);225    return bytes;226}227228fn classify() void {229    const fields = choir.product.revision.record.requireFields;230    fields(Policy, &.{ "max_threads", "worker_allocator", "verification" });231    fields(choir.passes.PassManagerRunOptions, &.{ "max_threads", "worker_allocator" });232    fields(choir.ir.VerifyOptions, &.{233        "check_terminators",     "require_terminators", "recursive", "check_use_def",234        "check_local_dominance", "check_cfg",           "max_depth",235    });236    fields(Options, &.{237        "timing",                           "failure", "target_profile", "generated_scan_schedules",238        "generated_row_pipeline_schedules", "tensor",  "now",239    });240    records.codec.coverage(stage_owner.TensorLoweringOptions, Tensor, &.{});241    records.codec.coverage(stage_owner.EinsumLoweringOptions, Einsum, &.{});242    records.codec.coverage(stage_owner.IndexingLoweringOptions, Indexing, &.{});243    fields(stage_owner.ActivationLoweringOptions, &.{"kernel_library"});244    fields(stage_owner.LossLoweringOptions, &.{245        "kernel_library", "row_sparse_cross_entropy_schedule",246    });247    fields(target.BackendTargetProfile, &.{248        "backend_kind", "artifact_format", "math_tier", "dtype_bits", "feature_bits",249    });250}

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