tiny.accy.preparation.kernelization.product

1pub const KernelOutline = struct {2    id: usize,3    name: []u8,4    kind: KernelOutlineKind,5    work_item_id: usize,6    root: *ir.Operation,7    input_slot_ids: []usize,8    output_slot_id: usize,9    element_count: u64,10    op_count: usize,1112    pub fn inputCount(self: KernelOutline) usize {13        return self.input_slot_ids.len;14    }1516    fn deinit(self: *KernelOutline, allocator: std.mem.Allocator) void {17        allocator.free(self.name);18        allocator.free(self.input_slot_ids);19        self.* = undefined;20    }21};

1pub const KernelOutlinePlanAnalysis = struct {2    allocator: std.mem.Allocator,3    kernels: std.ArrayListUnmanaged(KernelOutline),4    work_to_kernel: std.AutoHashMap(usize, usize),5    total_input_slots: usize = 0,6    total_scheduled_ops: usize = 0,78    pub fn init(allocator: std.mem.Allocator) KernelOutlinePlanAnalysis {9        return .{10            .allocator = allocator,11            .kernels = .empty,12            .work_to_kernel = std.AutoHashMap(usize, usize).init(allocator),13        };14    }1516    pub fn deinit(self: *KernelOutlinePlanAnalysis) void {17        for (self.kernels.items) |*kernel| {18            kernel.deinit(self.allocator);19        }20        self.kernels.deinit(self.allocator);21        self.work_to_kernel.deinit();22        self.* = undefined;23    }2425    pub fn kernelCount(self: KernelOutlinePlanAnalysis) usize {26        return self.kernels.items.len;27    }2829    pub fn getKernelForWork(30        self: *const KernelOutlinePlanAnalysis,31        work_item_id: usize,32    ) ?*const KernelOutline {33        const index = self.work_to_kernel.get(work_item_id) orelse return null;34        return &self.kernels.items[index];35    }36};

1pub const DotGeneralDescription = struct {2    input_dtype: choir_abi.DType,3    output_dtype: choir_abi.DType,4    dims: DotGeneralStaticDims,5};

1pub const ElementwiseKernel = enum {2    add,3    sub,4    mul,5    div,6    min,7    max,8    neg,9    abs,10    sqrt,11    exp,12    log,13    tanh,14    sin,15    cos,16    tan,17    floor,18    round,19    trunc,20    pow,21    atan2,22    convert,23    compare,24    select,25};

1pub const GeneratedKernelSummaries = struct {2    allocator: std.mem.Allocator,3    items: []GeneratedKernelSummary,45    pub fn deinit(self: *GeneratedKernelSummaries) void {6        for (self.items) |item| {7            self.allocator.free(item.entry_name);8        }9        self.allocator.free(self.items);10        self.* = undefined;11    }1213    pub fn len(self: *const GeneratedKernelSummaries) usize {14        return self.items.len;15    }1617    pub fn summary(18        self: *const GeneratedKernelSummaries,19        kernel_index: usize,20    ) !GeneratedKernelSummary {21        if (kernel_index >= self.items.len) return error.InvalidIndex;22        return self.items[kernel_index];23    }2425    pub fn summaryForWork(26        self: *const GeneratedKernelSummaries,27        work_item_id: usize,28    ) !GeneratedKernelSummary {29        for (self.items) |item| {30            if (item.work_item_id == work_item_id) return item;31        }32        return error.MissingKernelization;33    }34};

1pub const GeneratedKernelSummary = struct {2    work_item_id: usize,3    entry_name: []const u8,4    argument_count: u32,5    body_fingerprint: u64,6    dynamic_shared_memory_bytes: u32 = 0,7    schedule: GeneratedSchedule,8    launch_geometry: ?choir_abi.LaunchGeometry,9};

1pub const KernelizationAnalysis = struct {2    allocator: std.mem.Allocator,3    context: *ir.Context,4    kernels: std.ArrayListUnmanaged(LoweredKernel),5    work_to_kernel: std.AutoHashMap(usize, usize),67    /// The kernel stage calls this for its work bound to charge the result object before any kernel8    /// is generated, so compilation can refuse a pass that would exceed the caller's limits. That9    /// bound is the costs a pass declares before it runs. The bound covers the result object, its10    /// shared compiler context, and the list of kernels and the map from a scheduled unit of work11    /// to kernel, sized for `kernel_count` kernels. That context is the object that owns the12    /// operations and values of the generated kernels. Generated programs, names, generation13    /// scratch and the analyses it depends on are charged separately.14    pub fn baseStorageBound(limits: ir.Context.Limits, kernel_count: u64) !u64 {15        const context = ir.Context.Capacity.derive(limits) catch return error.WorkOverflow;16        var bytes: u64 = @sizeOf(KernelizationAnalysis) + @alignOf(KernelizationAnalysis);17        bytes = try accounting.add(bytes, @sizeOf(ir.Context) + @alignOf(ir.Context));18        bytes = try accounting.add(bytes, context.storage_bytes);19        bytes = try accounting.add(bytes, context.storage_alignment.toByteUnits());20        const kernels = try accounting.arrayListGrowth(LoweredKernel, kernel_count);21        bytes = try accounting.add(bytes, kernels);22        bytes = try accounting.add(bytes, try accounting.hashMapGrowth(usize, usize, kernel_count));23        if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;24        return bytes;25    }2627    pub fn reserveKernelCapacity(self: *KernelizationAnalysis, count: usize) !void {28        const map_count = std.math.cast(u32, count) orelse return error.InvalidArtifact;29        try self.kernels.ensureTotalCapacity(self.allocator, count);30        try self.work_to_kernel.ensureTotalCapacity(map_count);31    }3233    pub fn init(allocator: std.mem.Allocator, limits: ir.Context.Limits) !KernelizationAnalysis {34        const context = try ir.Context.create(allocator, limits);35        errdefer {36            context.deinit(allocator);37            allocator.destroy(context);38        }39        dialects.registerChoirDialect(context) catch |err| {40            return if (err == error.OutOfMemory) error.WorkExhausted else err;41        };42        return .{43            .allocator = allocator,44            .context = context,45            .kernels = .empty,46            .work_to_kernel = std.AutoHashMap(usize, usize).init(allocator),47        };48    }4950    pub fn deinit(self: *KernelizationAnalysis) void {51        for (self.kernels.items) |*kernel| {52            kernel.deinit(self.allocator);53        }54        self.kernels.deinit(self.allocator);55        self.work_to_kernel.deinit();56        self.context.deinit(self.allocator);57        self.allocator.destroy(self.context);58        self.* = undefined;59    }6061    pub fn kernelCount(self: KernelizationAnalysis) usize {62        return self.kernels.items.len;63    }6465    pub fn getForWork(66        self: *const KernelizationAnalysis,67        work_item_id: usize,68    ) ?*const LoweredKernel {69        const index = self.work_to_kernel.get(work_item_id) orelse return null;70        return &self.kernels.items[index];71    }7273    pub fn kernelSummary(74        self: *const KernelizationAnalysis,75        kernel_index: usize,76    ) !GeneratedKernelSummary {77        const kernel_value = try self.kernelAt(kernel_index);78        return kernel_value.summary();79    }8081    pub fn kernelProgram(82        self: *const KernelizationAnalysis,83        kernel_index: usize,84    ) !*const GeneratedKernelProgram {85        const kernel_value = try self.kernelAt(kernel_index);86        return &kernel_value.program;87    }8889    pub fn kernelSummaryForWork(90        self: *const KernelizationAnalysis,91        work_item_id: usize,92    ) !GeneratedKernelSummary {93        const kernel_value = self.getForWork(work_item_id) orelse return error.MissingKernelization;94        return kernel_value.summary();95    }9697    pub fn kernelProgramForWork(98        self: *const KernelizationAnalysis,99        work_item_id: usize,100    ) !*const GeneratedKernelProgram {101        const kernel_value = self.getForWork(work_item_id) orelse return error.MissingKernelization;102        return &kernel_value.program;103    }104105    pub fn copyKernelSummaries(106        self: *const KernelizationAnalysis,107        result_allocator: std.mem.Allocator,108    ) !GeneratedKernelSummaries {109        const items = try result_allocator.alloc(GeneratedKernelSummary, self.kernels.items.len);110        var copied: usize = 0;111        errdefer {112            for (items[0..copied]) |summary| {113                result_allocator.free(summary.entry_name);114            }115            result_allocator.free(items);116        }117118        for (items, self.kernels.items) |*item, *kernel_value| {119            item.* = kernel_value.summary();120            item.entry_name = try result_allocator.dupe(u8, kernel_value.entry_name);121            copied += 1;122        }123124        return .{125            .allocator = result_allocator,126            .items = items,127        };128    }129130    fn kernelAt(131        self: *const KernelizationAnalysis,132        kernel_index: usize,133    ) !*const LoweredKernel {134        if (kernel_index >= self.kernels.items.len) return error.InvalidIndex;135        return &self.kernels.items[kernel_index];136    }137};

1pub const LoweredKernel = struct {2    work_item_id: usize,3    entry_name: []u8,4    program: kernel_program.Program,5    argument_count: u32,6    body_fingerprint: u64,7    dynamic_shared_memory_bytes: u32 = 0,8    schedule: GeneratedSchedule,9    launch: ?kernel_schedule.Launch = null,10    output_fill_pattern: ?u32 = null,11    scratch_fill_pattern: ?u32 = null,12    body: LoweredKernelBody = .generic,1314    pub fn deinit(self: *LoweredKernel, allocator: std.mem.Allocator) void {15        self.program.deinit();16        allocator.free(self.entry_name);17        self.* = undefined;18    }1920    pub fn runtimeScalarArgumentCount(self: *const LoweredKernel) u32 {21        var count: u32 = 0;22        for (self.program.params()) |param| {23            switch (param) {24                .scalar => count += 1,25                .buffer => {},26            }27        }28        return count;29    }3031    pub fn requiredDTypes(self: *const LoweredKernel) gpu.DTypeSet {32        var dtypes: gpu.DTypeSet = .{};33        for (self.program.params()) |param| {34            switch (param) {35                .scalar => |dtype| dtypes.insert(dtype),36                .buffer => |buffer| dtypes.insert(buffer.dtype),37            }38        }39        return dtypes;40    }4142    pub fn launchGeometry(self: *const LoweredKernel) ?choir_abi.LaunchGeometry {43        const launch_value = self.launch orelse return null;44        return .{45            .grid = launch_value.grid,46            .threadgroup = launch_value.block,47            .dynamic_shared_memory_bytes = self.dynamic_shared_memory_bytes,48        };49    }5051    pub fn summary(self: *const LoweredKernel) GeneratedKernelSummary {52        return .{53            .work_item_id = self.work_item_id,54            .entry_name = self.entry_name,55            .argument_count = self.argument_count,56            .body_fingerprint = self.body_fingerprint,57            .dynamic_shared_memory_bytes = self.dynamic_shared_memory_bytes,58            .schedule = self.schedule,59            .launch_geometry = self.launchGeometry(),60        };61    }62};

1pub const ReductionDescription = struct {2    kind: ReductionKind,3    input_dtype: choir_abi.DType,4    output_dtype: choir_abi.DType,5    dims: ReductionStaticDims,6    init: ReductionInit,7};

1pub const ReductionInit = union(enum) {2    constant: ReductionInitValue,3    input_buffer,4};

1pub const ReductionInitValue = union(enum) {2    f32: f32,3    i32: i32,4    u32: u32,5    f16: f16,6};

1pub const ReductionKind = enum {2    sum,3    max,4    min,5};

1pub const ReductionStaticDims = struct {2    input_rank: u8,3    axis: u8,4    input_element_count: u32,5    output_element_count: u32,6    rows: u32 = 1,7    cols: u32 = 1,8    inner: u32 = 1,9};

1pub const ShapeKernel = enum {2    broadcast_in_dim,3    iota,4    reshape,5    transpose,6    slice,7    pad,8    concatenate,9    gather,10    scatter,11};

1pub const kernel_outline_plan_analysis_name = "accy-choir-kernel-outline-plan";

1pub fn addKernel(analysis: *KernelizationAnalysis, kernel: LoweredKernel) !void {2    const index = analysis.kernels.items.len;3    try analysis.work_to_kernel.put(kernel.work_item_id, index);4    errdefer _ = analysis.work_to_kernel.remove(kernel.work_item_id);5    try analysis.kernels.append(analysis.allocator, kernel);6}

1pub fn dotGeneralBlockTileFor(desc: DotGeneralDescription) ?DotGeneralBlockTile {2    if (desc.input_dtype != .f32 or desc.output_dtype != .f32) return null;3    if (desc.dims.k >= 128 and blockTileFits(wide_block_tile, desc.dims, wide_block_tile_min_blocks)) return wide_block_tile;4    var tile = DotGeneralBlockTile{ .stages = 4 };5    if (!blockTileFits(tile, desc.dims, block_tile_min_blocks)) return null;6    if (desc.dims.batch == 1) {7        const tiles = (desc.dims.m / tile.bm) * (desc.dims.n / tile.bn);8        while (tile.splits < split_k_max) {9            const doubled = tile.splits * 2;10            if (tiles * tile.splits >= split_k_target_blocks) break;11            if (desc.dims.k % doubled != 0) break;12            const chunk = desc.dims.k / doubled;13            if (chunk % tile.bk != 0 or chunk < 2 * tile.bk) break;14            tile.splits = doubled;15        }16    }17    return tile;18}

1pub fn dotGeneralMmaTileFor(desc: DotGeneralDescription) ?DotGeneralMmaTile {2    if (desc.input_dtype != .f32 or desc.output_dtype != .f32) return null;3    var tile = DotGeneralMmaTile{};4    if (!tile.exact(desc.dims)) return null;5    const blocks = @as(u64, desc.dims.batch) * (desc.dims.m / tile.bm) * (desc.dims.n / tile.bn);6    if (blocks < mma_pipeline_max_blocks) {7        tile.stages = 2;8    }9    if (desc.dims.batch == 1 and blocks < 32) {10        while (tile.splits < split_k_max) {11            const doubled = tile.splits * 2;12            if (blocks * tile.splits >= 128) break;13            if (desc.dims.k % doubled != 0) break;14            const chunk = desc.dims.k / doubled;15            if (chunk % (2 * @as(u64, tile.bk)) != 0 or chunk < 4 * @as(u64, tile.bk)) break;16            tile.splits = doubled;17        }18    }19    return tile;20}

1pub fn dynamicSharedMemoryBytes(module: *ir.Operation) gpu.BackendError!u32 {2    var state = DynamicSharedMemoryState{};3    _ = module.walk(.{ .order = .pre_order }, &state, DynamicSharedMemoryState.visit) catch |err| switch (err) {4        error.OutOfMemory => return error.OutOfMemory,5        error.InvalidArtifact => return error.InvalidArtifact,6        else => return error.InvalidArtifact,7    };8    return state.bytes;9}

1pub fn reductionAtomicPlanFor(2    dims: ReductionStaticDims,3    kind: ReductionKind,4    input_dtype: choir_abi.DType,5    init_is_constant: bool,6) ?ReductionAtomicPlan {7    if (dims.output_element_count != 1) return null;8    if (dims.input_element_count < reduction_atomic_min_extent) return null;9    if (kind != .sum) return null;10    if (!init_is_constant) return null;11    switch (input_dtype) {12        .f32, .i32, .u32 => {},13        else => return null,14    }15    const threads = reduction_single_block_threads;16    const per_block = @as(u64, threads) * reduction_atomic_target_chain;17    const wanted = (dims.input_element_count + per_block - 1) / per_block;18    const blocks: u32 = @intCast(@min(@max(wanted, 1), reduction_atomic_max_blocks));19    return .{ .threads = threads, .blocks = blocks };20}

1pub fn reductionSingleBlockThreads(dims: ReductionStaticDims) ?u32 {2    if (dims.output_element_count != 1) return null;3    if (dims.input_element_count < reduction_single_block_min_extent) return null;4    return reduction_single_block_threads;5}

1pub fn reductionWarpRowsThreads(dims: ReductionStaticDims) ?u32 {2    if (dims.input_rank != 2 or dims.axis != 1) return null;3    if (dims.cols < reduction_warp_rows_min_cols) return null;4    return reduction_warp_rows_threads;5}

1pub const reduction_single_block_threads: u32 = 256;

1pub const reduction_warp_rows_threads: u32 = 256;

1pub const KernelOutlineKind = records.kernel.KernelOutlineKind;

1pub const DotGeneralBlockTile = records.kernel.DotGeneralBlockTile;

1pub const DotGeneralMmaTile = records.kernel.DotGeneralMmaTile;

1pub const DotGeneralStaticDims = records.kernel.DotGeneralStaticDims;

1pub const ElementwiseRank2Plan = records.kernel.ElementwiseRank2Plan;

1pub const ElementwiseVectorPlan = records.kernel.ElementwiseVectorPlan;

1pub const FlashAttentionPlan = records.kernel.FlashAttentionPlan;

1pub const GeneratedSchedule = records.kernel.GeneratedSchedule;

1pub const GeneratedScheduleKind = records.kernel.GeneratedScheduleKind;

1pub const LoweredKernelBody = records.kernel.LoweredKernelBody;

1pub const ReductionAtomicPlan = records.kernel.ReductionAtomicPlan;

1pub const ReductionWarpRowsPlan = records.kernel.ReductionWarpRowsPlan;

1pub const RowPipelinePlan = records.kernel.RowPipelinePlan;

1pub const ScanPlan = records.kernel.ScanPlan;

1const product = @import("product.zig");2const outline = @import("outline.zig");34pub const kernel_outline_plan_analysis_name = outline.kernel_outline_plan_analysis_name;5pub const KernelOutlineKind = outline.KernelOutlineKind;6pub const KernelOutline = outline.KernelOutline;7pub const KernelOutlinePlanAnalysis = outline.KernelOutlinePlanAnalysis;89pub const ElementwiseKernel = product.ElementwiseKernel;10pub const ShapeKernel = product.ShapeKernel;11pub const GeneratedScheduleKind = product.GeneratedScheduleKind;12pub const GeneratedSchedule = product.GeneratedSchedule;13pub const GeneratedKernelProgram = product.GeneratedKernelProgram;14pub const GeneratedKernelSummary = product.GeneratedKernelSummary;15pub const GeneratedKernelSummaries = product.GeneratedKernelSummaries;16pub const RowPipelinePlan = product.RowPipelinePlan;17pub const FlashAttentionPlan = product.FlashAttentionPlan;18pub const ScanPlan = product.ScanPlan;19pub const ElementwiseRank2Plan = product.ElementwiseRank2Plan;20pub const LoweredKernelBody = product.LoweredKernelBody;21pub const ElementwiseVectorPlan = product.ElementwiseVectorPlan;22pub const ReductionWarpRowsPlan = product.ReductionWarpRowsPlan;23pub const LoweredKernel = product.LoweredKernel;24pub const KernelizationAnalysis = product.KernelizationAnalysis;25pub const DotGeneralStaticDims = product.DotGeneralStaticDims;26pub const DotGeneralDescription = product.DotGeneralDescription;27pub const DotGeneralBlockTile = product.DotGeneralBlockTile;28pub const reduction_single_block_threads = product.reduction_single_block_threads;29pub const reduction_warp_rows_threads = product.reduction_warp_rows_threads;30pub const reductionWarpRowsThreads = product.reductionWarpRowsThreads;31pub const reductionSingleBlockThreads = product.reductionSingleBlockThreads;32pub const ReductionAtomicPlan = product.ReductionAtomicPlan;33pub const reductionAtomicPlanFor = product.reductionAtomicPlanFor;34pub const DotGeneralMmaTile = product.DotGeneralMmaTile;35pub const dotGeneralMmaTileFor = product.dotGeneralMmaTileFor;36pub const dotGeneralBlockTileFor = product.dotGeneralBlockTileFor;37pub const ReductionStaticDims = product.ReductionStaticDims;38pub const ReductionKind = product.ReductionKind;39pub const ReductionInitValue = product.ReductionInitValue;40pub const ReductionInit = product.ReductionInit;41pub const ReductionDescription = product.ReductionDescription;42pub const addKernel = product.addKernel;43pub const dynamicSharedMemoryBytes = product.dynamicSharedMemoryBytes;

1pub const product = @import("model/root.zig");