tiny.accy.choir.record.kernel

1const std = @import("std");2const records = @import("root.zig");3const schedule = @import("../../kernel/model/root.zig").core.schedule;45pub const DotGeneralMmaTile = struct {6    bm: u32 = 128,7    bn: u32 = 64,8    bk: u32 = 16,9    warps_m: u32 = 2,10    warps_n: u32 = 2,11    stages: u32 = 1,12    splits: u32 = 1,1314    pub fn warps(self: DotGeneralMmaTile) u32 {15        return self.warps_m * self.warps_n;16    }1718    pub fn warpTileM(self: DotGeneralMmaTile) u32 {19        return self.bm / self.warps_m;20    }2122    pub fn warpTileN(self: DotGeneralMmaTile) u32 {23        return self.bn / self.warps_n;24    }2526    pub fn threads(self: DotGeneralMmaTile) u32 {27        return self.warps() * 32;28    }2930    pub fn exact(self: DotGeneralMmaTile, dims: DotGeneralStaticDims) bool {31        return dims.m % self.bm == 0 and dims.n % self.bn == 0 and dims.k % self.bk == 0;32    }33};3435pub const ReductionAtomicPlan = struct {36    threads: u32,37    blocks: u32,3839    pub fn stride(self: ReductionAtomicPlan) u64 {40        return @as(u64, self.threads) * self.blocks;41    }42};4344pub const DotGeneralBlockTile = struct {45    bm: u32 = 64,46    bn: u32 = 64,47    bk: u32 = 16,48    tm: u32 = 4,49    tn: u32 = 4,50    splits: u32 = 1,51    stages: u32 = 2,5253    pub fn threadsX(self: DotGeneralBlockTile) u32 {54        return self.bn / self.tn;55    }5657    pub fn threadsY(self: DotGeneralBlockTile) u32 {58        return self.bm / self.tm;59    }6061    pub fn exact(self: DotGeneralBlockTile, dims: DotGeneralStaticDims) bool {62        return dims.m % self.bm == 0 and dims.n % self.bn == 0 and dims.k % self.bk == 0;63    }64};6566pub const DotGeneralStaticDims = struct {67    m: u32,68    n: u32,69    k: u32,70    batch: u32 = 1,71};7273pub const ReductionWarpRowsPlan = struct {74    threads: u32,75    rows: u32,76};7778pub const ElementwiseVectorPlan = struct {79    quads: u64,80};8182pub const LoweredKernelBody = union(enum) {83    generic,84    dot_block_tile: DotGeneralBlockTile,85    dot_mma_tile: DotGeneralMmaTile,86    reduction_single_block: u32,87    reduction_atomic: ReductionAtomicPlan,88    reduction_warp_rows: ReductionWarpRowsPlan,89    row_pipeline: RowPipelinePlan,90    flash_attention: FlashAttentionPlan,91    scan: ScanPlan,92    elementwise_rank2: ElementwiseRank2Plan,93    elementwise_vector: ElementwiseVectorPlan,94};9596pub const ElementwiseRank2Plan = struct {97    rows: u32,98    cols: u32,99    threads_x: u32,100    threads_y: u32,101};102103pub const ScanPlan = struct {104    blocks: u32,105    threads: u32,106    items: u32,107};108109pub const FlashAttentionPlan = struct {110    seq: u32,111    dim: u32,112    br: u32,113    bc: u32,114    threads_x: u32,115    threads_y: u32,116};117118pub const RowPipelinePlan = struct {119    threads: u32,120    rows: u32,121    cols: u32,122    warps: u32,123};124125pub const GeneratedSchedule = struct {126    kind: GeneratedScheduleKind,127    threads: @import("../../kernel/model/root.zig").logical.schedule.Threads,128};129130pub const GeneratedScheduleKind = enum {131    flat,132    matrix,133};134135pub const KernelOutlineKind = enum {136    elementwise,137    shape,138    dot_general,139    reduction,140    kernel_call,141    row_pipeline,142    iterate,143    flash_attention,144    scan,145};146147pub const Outline = struct {148    id: usize,149    name: []const u8,150    kind: KernelOutlineKind,151    work_item_id: usize,152    root: records.reference.Operation,153    input_slot_ids: []const usize,154    output_slot_id: usize,155    element_count: u64,156    op_count: usize,157};158159pub const Outlines = struct {160    kernels: []const Outline,161    total_input_slots: usize,162    total_scheduled_ops: usize,163};164165pub const Lowered = struct {166    work_item_id: usize,167    entry_name: []const u8,168    program: records.program.Record,169    argument_count: u32,170    body_fingerprint: u64,171    dynamic_shared_memory_bytes: u32,172    schedule: GeneratedSchedule,173    launch: ?schedule.Launch,174    output_fill_pattern: ?u32,175    scratch_fill_pattern: ?u32,176    body: LoweredKernelBody,177};178179pub const Generated = struct {180    kernels: []const Lowered,181};182183pub const Record = struct {184    outlines: Outlines,185    generated: Generated,186};187188pub fn validate(189    allocator: std.mem.Allocator,190    value: Record,191    memory: records.memory.Record,192) !void {193    var work = std.AutoHashMapUnmanaged(usize, bool).empty;194    defer work.deinit(allocator);195    const count = std.math.cast(u32, value.outlines.kernels.len) orelse196        return error.InvalidStageRecord;197    try work.ensureTotalCapacity(allocator, count);198    var inputs: usize = 0;199    var operations: usize = 0;200    for (value.outlines.kernels, 0..) |outline, index| {201        if (outline.id != index or outline.output_slot_id >= memory.buffers.slots.len) {202            return error.InvalidStageRecord;203        }204        for (outline.input_slot_ids) |slot| {205            if (slot >= memory.buffers.slots.len) return error.InvalidStageRecord;206        }207        const entry = work.getOrPutAssumeCapacity(outline.work_item_id);208        if (entry.found_existing) return error.InvalidStageRecord;209        entry.value_ptr.* = false;210        inputs = std.math.add(usize, inputs, outline.input_slot_ids.len) catch211            return error.InvalidStageRecord;212        operations = std.math.add(usize, operations, outline.op_count) catch213            return error.InvalidStageRecord;214    }215    if (inputs != value.outlines.total_input_slots or216        operations != value.outlines.total_scheduled_ops) return error.InvalidStageRecord;217    for (value.generated.kernels) |kernel| {218        const seen = work.getPtr(kernel.work_item_id) orelse return error.InvalidStageRecord;219        if (seen.*) return error.InvalidStageRecord;220        seen.* = true;221        if (kernel.argument_count != kernel.program.params.len) return error.InvalidStageRecord;222    }223}

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