tiny.accy.preparation.kernelization

1pub fn dotGeneralDescriptionForWork(2    outline: kernelization_model.KernelOutline,3    work: schedule_planning.ScheduleWorkItem,4    buffer_plan: *const bufferization.BufferPlanAnalysis,5) gpu.BackendError!DotGeneralDescription {6    if (work.kind != .dot_general) return error.UnsupportedOperation;7    if (work.ops.len == 0) return error.UnsupportedOperation;8    if (!isName(work.ops[0].name.name, dialect_mod.AccyDialect.DotGeneralOp.operation_name)) return error.UnsupportedOperation;9    if (work.ops.len == 1 and outline.inputCount() != 2) return error.UnsupportedOperation;10    if (outline.inputCount() < 2) return error.UnsupportedOperation;1112    const lhs_slot = bufferSlotById(buffer_plan, outline.input_slot_ids[0]) orelse return error.InvalidArtifact;13    const rhs_slot = bufferSlotById(buffer_plan, outline.input_slot_ids[1]) orelse return error.InvalidArtifact;14    const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;15    try validateCanonicalDotLayout(work.ops[0], lhs_slot.dims.len);16    const dims = try dotGeneralStaticDims(lhs_slot.*, rhs_slot.*, output_slot.*);17    try validateDotGeneralDTypes(lhs_slot.*, rhs_slot.*, output_slot.*, work);1819    return .{20        .input_dtype = lhs_slot.dtype,21        .output_dtype = output_slot.dtype,22        .dims = dims,23    };24}

1pub fn expectedInputCount(kind: ElementwiseKernel) usize {2    return switch (kind) {3        .add, .sub, .mul, .div, .min, .max, .pow, .atan2, .compare => 2,4        .neg, .abs, .sqrt, .exp, .log, .tanh, .sin, .cos, .tan, .floor, .round, .trunc, .convert => 1,5        .select => 3,6    };7}

1pub fn kernelForOperation(op: *ir.Operation) ?ElementwiseKernel {2    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ConvertOp.operation_name)) return .convert;3    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.CompareOp.operation_name)) return .compare;4    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SelectOp.operation_name)) return .select;5    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.AddOp.operation_name)) return .add;6    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SubOp.operation_name)) return .sub;7    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MulOp.operation_name)) return .mul;8    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.DivOp.operation_name)) return .div;9    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MinOp.operation_name)) return .min;10    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.MaxOp.operation_name)) return .max;11    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.NegOp.operation_name)) return .neg;12    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.AbsOp.operation_name)) return .abs;13    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SqrtOp.operation_name)) return .sqrt;14    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ExpOp.operation_name)) return .exp;15    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.LogOp.operation_name)) return .log;16    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TanhOp.operation_name)) return .tanh;17    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SinOp.operation_name)) return .sin;18    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.CosOp.operation_name)) return .cos;19    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TanOp.operation_name)) return .tan;20    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.FloorOp.operation_name)) return .floor;21    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.RoundOp.operation_name)) return .round;22    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TruncOp.operation_name)) return .trunc;23    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.PowOp.operation_name)) return .pow;24    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.Atan2Op.operation_name)) return .atan2;25    return null;26}

1pub fn getKernelizationAnalysis(2    pass_ctx: *passes.PassContext,3    op: *ir.Operation,4) !*KernelizationAnalysis {5    const ptr = try pass_ctx.getAnalysis(6        op,7        &kernelization_analysis_descriptor,8        computeKernelizationAnalysis,9        cleanupKernelizationAnalysis,10    );11    return @ptrCast(@alignCast(ptr));12}

1pub fn kernelizationPass() passes.Pass {2    return .{3        .name = kernelization_pass_name,4        .description = kernelization_pass_description,5        .run_fn = runKernelizationPass,6        .work_contract = .{7            .identity = .{ .name = kernelization_pass_name, .version = 1 },8            .estimate = lowering.work.pass,9        },10    };11}

1pub const kernelization_analysis_descriptor = passes.AnalysisDescriptor{2    .id = passes.analysisId(kernelization_analysis_name),3    .name = kernelization_analysis_name,4    .work_contract = .{5        .identity = .{ .name = kernelization_analysis_name, .version = 1 },6        .estimate = lowering.work.analysis,7    },8};

1pub const kernelization_analysis_name = "accy-choir-kernelization-plan";

1pub const kernelization_pass_description =2    "Lower scheduled Accy tensor work to kernel-language programs";

1pub const kernelization_pass_name = "accy-choir-plan-kernelization";

1pub fn reductionDescriptionForWork(2    outline: kernelization_model.KernelOutline,3    work: schedule_planning.ScheduleWorkItem,4    buffer_plan: *const bufferization.BufferPlanAnalysis,5) gpu.BackendError!ReductionDescription {6    if (work.kind != .reduction) return error.UnsupportedOperation;7    if (work.ops.len == 0) return error.UnsupportedOperation;8    const root = work.ops[work.ops.len - 1];9    if (isName(root.name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name)) {10        return concatReductionDescription(root, outline, work, buffer_plan);11    }12    const op = root;13    if (!isName(op.name.name, dialect_mod.AccyDialect.ReduceOp.operation_name)) return error.UnsupportedOperation;14    const operands = op.getOperandValues();15    if (operands.len != 2) return error.UnsupportedOperation;1617    const fused = reductionWorkIsFused(work, buffer_plan);18    if (!fused) {19        const input_slot = buffer_plan.getSlot(operands[0]) orelse return error.InvalidArtifact;20        if (externalInputIndex(outline, input_slot.id) != 0) return error.UnsupportedOperation;21    }2223    var input_arena_buffer: [256]u8 = undefined;24    var input_arena = alloc_fixed.FixedBuffer.init(input_arena_buffer[0..]);25    const input_type = dialect_mod.decodeTensorType(input_arena.allocator(), operands[0].type) catch return error.InvalidArtifact;2627    const init_slot = buffer_plan.getSlot(operands[1]) orelse return error.InvalidArtifact;28    const output_slot = bufferSlotById(buffer_plan, outline.output_slot_id) orelse return error.InvalidArtifact;29    var axes_buffer: [i64_attr_list_stack_capacity]i64 = undefined;30    const axes = try readI64ListAttrBounded(op, "dimensions", dialect_mod.AccyDialect.ReduceOp.dialectAttrName("dimensions"), &axes_buffer);3132    const dims = try reductionStaticDims(input_type.dims, output_slot.dims, axes);33    try validateReductionDTypes(input_type.dtype, output_slot.dtype, work);34    const kind = try reductionKindForOp(op);35    const init = try reductionInitValue(init_slot.*, input_type.dtype, outline);3637    return .{38        .kind = kind,39        .input_dtype = input_type.dtype,40        .output_dtype = output_slot.dtype,41        .dims = dims,42        .init = init,43    };44}

1pub fn scheduleCandidates(2    cols: u64,3    format: ?gpu.ArtifactFormat,4    buffer: *[max_schedule_candidates]Schedule,5) []const Schedule {6    if (format != .cuda_ptx) return buffer[0..0];7    var count: usize = 0;8    for (block_schedules) |candidate_threads| {9        if (!scheduleViable(candidate_threads, cols)) continue;10        buffer[count] = .{ .threads = candidate_threads };11        count += 1;12    }13    return buffer[0..count];14}

1pub const schedule_version: u32 = 1;

1pub fn scheduleCandidates(2    total: u64,3    format: ?gpu.ArtifactFormat,4    buffer: *[max_schedule_candidates]Schedule,5) []const Schedule {6    if (format != .cuda_ptx) return buffer[0..0];7    var count: usize = 0;8    for (block_schedules) |schedule| {9        if (!scheduleViable(schedule, total)) continue;10        buffer[count] = schedule;11        count += 1;12    }13    return buffer[0..count];14}

1pub const schedule_version: u32 = 1;

1pub fn kernelForOperation(op: *ir.Operation) ?ShapeKernel {2    if (isName(op.name.name, dialect_mod.AccyDialect.BroadcastInDimOp.operation_name)) {3        return .broadcast_in_dim;4    }5    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.IotaOp.operation_name)) return .iota;6    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ReshapeOp.operation_name)) return .reshape;7    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.TransposeOp.operation_name)) return .transpose;8    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.SliceOp.operation_name)) return .slice;9    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.PadOp.operation_name)) return .pad;10    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ConcatenateOp.operation_name)) return .concatenate;11    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.GatherOp.operation_name)) return .gather;12    if (std.mem.eql(u8, op.name.name, dialect_mod.AccyDialect.ScatterOp.operation_name)) return .scatter;13    return null;14}

1pub const max_schedule_candidates = block_schedules.len;

1pub const max_schedule_candidates = block_schedules.len;

1pub const product = @import("model/root.zig");2const lowering = @import("lowering/root.zig");34pub const ElementwiseKernel = product.ElementwiseKernel;5pub const ShapeKernel = product.ShapeKernel;6pub const GeneratedScheduleKind = product.GeneratedScheduleKind;7pub const GeneratedSchedule = product.GeneratedSchedule;8pub const GeneratedKernelProgram = product.GeneratedKernelProgram;9pub const GeneratedKernelSummary = product.GeneratedKernelSummary;10pub const GeneratedKernelSummaries = product.GeneratedKernelSummaries;11pub const LoweredKernel = product.LoweredKernel;12pub const KernelizationAnalysis = product.KernelizationAnalysis;13pub const DotGeneralStaticDims = product.DotGeneralStaticDims;14pub const DotGeneralDescription = product.DotGeneralDescription;15pub const DotGeneralBlockTile = product.DotGeneralBlockTile;16pub const DotGeneralMmaTile = product.DotGeneralMmaTile;17pub const dotGeneralBlockTileFor = product.dotGeneralBlockTileFor;18pub const reductionSingleBlockThreads = product.reductionSingleBlockThreads;19pub const LoweredKernelBody = product.LoweredKernelBody;20pub const ReductionAtomicPlan = product.ReductionAtomicPlan;21pub const reductionAtomicPlanFor = product.reductionAtomicPlanFor;22pub const ReductionStaticDims = product.ReductionStaticDims;23pub const ReductionKind = product.ReductionKind;24pub const ReductionDescription = product.ReductionDescription;2526pub const kernelization_analysis_name = lowering.kernelization_analysis_name;27pub const kernelization_pass_name = lowering.kernelization_pass_name;28pub const kernelization_pass_description = lowering.kernelization_pass_description;29pub const kernelization_analysis_descriptor = lowering.kernelization_analysis_descriptor;30pub const getKernelizationAnalysis = lowering.getKernelizationAnalysis;31pub const kernelizationPass = lowering.kernelizationPass;32pub const dotGeneralDescriptionForWork = lowering.dotGeneralDescriptionForWork;33pub const reductionDescriptionForWork = lowering.reductionDescriptionForWork;34pub const elementwiseKernelForOperation = lowering.elementwiseKernelForOperation;35pub const shapeKernelForOperation = lowering.shapeKernelForOperation;36pub const expectedInputCount = lowering.expectedInputCount;37pub const generated_scan_schedule_version = lowering.generated_scan_schedule_version;38pub const max_scan_schedule_candidates = lowering.max_scan_schedule_candidates;39pub const scanScheduleCandidates = lowering.scanScheduleCandidates;40pub const generated_row_pipeline_schedule_version = lowering.generated_row_pipeline_schedule_version;41pub const max_row_pipeline_schedule_candidates = lowering.max_row_pipeline_schedule_candidates;42pub const rowPipelineScheduleCandidates = lowering.rowPipelineScheduleCandidates;

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