tiny.accy.preparation.target

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

1const std = @import("std");2const gpu = @import("gpu");3const choir = @import("choir");4const accy_root = @import("../root.zig");56const ir = choir.ir;78pub const backend_kind_attr_name = "accy.backend.kind";9pub const artifact_format_attr_name = "accy.backend.artifact_format";10pub const math_tier_attr_name = "accy.backend.math_tier";11pub const dtype_bits_attr_name = "accy.backend.dtype_bits";12pub const feature_bits_attr_name = "accy.backend.feature_bits";13pub const generated_scan_schedule_attr_name = "accy.backend.scan_schedule";1415pub const GeneratedScanSchedule = struct {16    threads: u32,17    items: u32,1819    pub fn eql(self: GeneratedScanSchedule, other: GeneratedScanSchedule) bool {20        return self.threads == other.threads and self.items == other.items;21    }22};2324pub const GeneratedScanScheduleDecision = struct {25    total: ?u64 = null,26    schedule: GeneratedScanSchedule,27};2829pub fn setGeneratedScanSchedules(30    allocator: std.mem.Allocator,31    ctx: *ir.Context,32    module: *ir.Operation,33    decisions: []const GeneratedScanScheduleDecision,34) !void {35    const encoded = try encodeGeneratedScanSchedules(allocator, decisions);36    defer allocator.free(encoded);37    try module.setAttr(generated_scan_schedule_attr_name, try ctx.getStringAttr(encoded));38}3940/// The recipe calls this function to record caller decisions for one generated kernel, such as its41/// thread count, optionally tied to a problem size, as a request's scan schedule choices in the42/// same bytes the module stores. The recipe encodes a stage's options into its stage record. The43/// function writes the choices as comma-separated entries: `threads`x`items` or44/// `total`=`threads`x`items` when a choice is tied to a total length. The bytes match the module45/// attribute the same choices are stored in, and the caller owns them. The call returns46/// `error.InvalidArtifact` for an empty list.47pub fn encodeGeneratedScanSchedules(48    allocator: std.mem.Allocator,49    decisions: []const GeneratedScanScheduleDecision,50) ![]u8 {51    if (decisions.len == 0) return error.InvalidArtifact;52    var text = std.ArrayListUnmanaged(u8).empty;53    errdefer text.deinit(allocator);54    for (decisions, 0..) |decision, index| {55        if (index != 0) try text.append(allocator, ',');56        var entry_buffer: [64]u8 = undefined;57        const entry = if (decision.total) |total|58            try std.fmt.bufPrint(&entry_buffer, "{d}={d}x{d}", .{ total, decision.schedule.threads, decision.schedule.items })59        else60            try std.fmt.bufPrint(&entry_buffer, "{d}x{d}", .{ decision.schedule.threads, decision.schedule.items });61        try text.appendSlice(allocator, entry);62    }63    return text.toOwnedSlice(allocator);64}6566pub fn readGeneratedScanSchedules(module: *const ir.Operation) ?[]const u8 {67    const attr = module.getAttrAs(ir.Attribute.StringAttr, generated_scan_schedule_attr_name) orelse return null;68    return attr.getValue();69}7071pub fn resolveGeneratedScanSchedule(72    encoded: []const u8,73    total: u64,74) error{InvalidArtifact}!?GeneratedScanSchedule {75    var fallback: ?GeneratedScanSchedule = null;76    var entries = std.mem.splitScalar(u8, encoded, ',');77    while (entries.next()) |entry| {78        if (std.mem.indexOfScalar(u8, entry, '=')) |split_index| {79            const entry_total = std.fmt.parseInt(u64, entry[0..split_index], 10) catch return error.InvalidArtifact;80            const schedule = try parseGeneratedScanSchedule(entry[split_index + 1 ..]);81            if (entry_total == total) return schedule;82        } else {83            fallback = try parseGeneratedScanSchedule(entry);84        }85    }86    return fallback;87}8889fn parseGeneratedScanSchedule(text: []const u8) error{InvalidArtifact}!GeneratedScanSchedule {90    const split_index = std.mem.indexOfScalar(u8, text, 'x') orelse return error.InvalidArtifact;91    const threads = std.fmt.parseInt(u32, text[0..split_index], 10) catch return error.InvalidArtifact;92    const items = std.fmt.parseInt(u32, text[split_index + 1 ..], 10) catch return error.InvalidArtifact;93    if (threads == 0 or items == 0) return error.InvalidArtifact;94    return .{ .threads = threads, .items = items };95}9697pub const generated_row_pipeline_schedule_attr_name = "accy.backend.row_pipeline_schedule";9899pub const GeneratedRowPipelineSchedule = struct {100    threads: u32,101102    pub fn eql(self: GeneratedRowPipelineSchedule, other: GeneratedRowPipelineSchedule) bool {103        return self.threads == other.threads;104    }105};106107pub const GeneratedRowPipelineShape = struct {108    rows: u64,109    cols: u64,110};111112pub const GeneratedRowPipelineScheduleDecision = struct {113    shape: ?GeneratedRowPipelineShape = null,114    schedule: GeneratedRowPipelineSchedule,115};116117pub fn setGeneratedRowPipelineSchedules(118    allocator: std.mem.Allocator,119    ctx: *ir.Context,120    module: *ir.Operation,121    decisions: []const GeneratedRowPipelineScheduleDecision,122) !void {123    const encoded = try encodeGeneratedRowPipelineSchedules(allocator, decisions);124    defer allocator.free(encoded);125    try module.setAttr(generated_row_pipeline_schedule_attr_name, try ctx.getStringAttr(encoded));126}127128/// The recipe calls this function to record a request's row schedule choices in the same bytes the129/// module stores. The function writes the choices as comma-separated entries: `threads` or130/// `rows`x`cols`=`threads` when a choice is tied to a row shape. The bytes match the module131/// attribute the same choices are stored in, and the caller owns them. The call returns132/// `error.InvalidArtifact` for an empty list.133pub fn encodeGeneratedRowPipelineSchedules(134    allocator: std.mem.Allocator,135    decisions: []const GeneratedRowPipelineScheduleDecision,136) ![]u8 {137    if (decisions.len == 0) return error.InvalidArtifact;138    var text = std.ArrayListUnmanaged(u8).empty;139    errdefer text.deinit(allocator);140    for (decisions, 0..) |decision, index| {141        if (index != 0) try text.append(allocator, ',');142        var entry_buffer: [96]u8 = undefined;143        const entry = if (decision.shape) |shape|144            try std.fmt.bufPrint(&entry_buffer, "{d}x{d}={d}", .{ shape.rows, shape.cols, decision.schedule.threads })145        else146            try std.fmt.bufPrint(&entry_buffer, "{d}", .{decision.schedule.threads});147        try text.appendSlice(allocator, entry);148    }149    return text.toOwnedSlice(allocator);150}151152pub fn readGeneratedRowPipelineSchedules(module: *const ir.Operation) ?[]const u8 {153    const attr = module.getAttrAs(ir.Attribute.StringAttr, generated_row_pipeline_schedule_attr_name) orelse return null;154    return attr.getValue();155}156157pub fn resolveGeneratedRowPipelineSchedule(158    encoded: []const u8,159    rows: u64,160    cols: u64,161) error{InvalidArtifact}!?GeneratedRowPipelineSchedule {162    var fallback: ?GeneratedRowPipelineSchedule = null;163    var entries = std.mem.splitScalar(u8, encoded, ',');164    while (entries.next()) |entry| {165        if (std.mem.indexOfScalar(u8, entry, '=')) |split_index| {166            const shape_text = entry[0..split_index];167            const shape_split = std.mem.indexOfScalar(u8, shape_text, 'x') orelse return error.InvalidArtifact;168            const entry_rows = std.fmt.parseInt(u64, shape_text[0..shape_split], 10) catch return error.InvalidArtifact;169            const entry_cols = std.fmt.parseInt(u64, shape_text[shape_split + 1 ..], 10) catch return error.InvalidArtifact;170            const schedule = try parseGeneratedRowPipelineSchedule(entry[split_index + 1 ..]);171            if (entry_rows == rows and entry_cols == cols) return schedule;172        } else {173            fallback = try parseGeneratedRowPipelineSchedule(entry);174        }175    }176    return fallback;177}178179fn parseGeneratedRowPipelineSchedule(text: []const u8) error{InvalidArtifact}!GeneratedRowPipelineSchedule {180    const threads = std.fmt.parseInt(u32, text, 10) catch return error.InvalidArtifact;181    if (threads == 0) return error.InvalidArtifact;182    return .{ .threads = threads };183}184185pub const BackendTargetProfile = accy_root.choir.record.target.BackendTargetProfile;186187pub fn setBackendTargetProfile(188    ctx: *ir.Context,189    module: *ir.Operation,190    profile: BackendTargetProfile,191) !void {192    try module.setAttr(backend_kind_attr_name, try ctx.getStringAttr(@tagName(profile.backend_kind)));193    try module.setAttr(artifact_format_attr_name, try ctx.getStringAttr(@tagName(profile.artifact_format)));194    try module.setAttr(math_tier_attr_name, try ctx.getStringAttr(@tagName(profile.math_tier)));195    try module.setAttr(dtype_bits_attr_name, try ctx.getIntegerAttr(@bitCast(profile.dtype_bits), 64, false));196    try module.setAttr(feature_bits_attr_name, try ctx.getIntegerAttr(@bitCast(profile.feature_bits), 64, false));197}198199pub fn readBackendTargetProfile(module: *const ir.Operation) ?BackendTargetProfile {200    const backend_kind = readEnumAttr(gpu.BackendKind, module, backend_kind_attr_name) orelse return null;201    const artifact_format = readEnumAttr(gpu.ArtifactFormat, module, artifact_format_attr_name) orelse return null;202    const math_tier = readEnumAttr(gpu.BackendMathTier, module, math_tier_attr_name) orelse return null;203    const dtype_bits = readU64Attr(module, dtype_bits_attr_name) orelse return null;204    const feature_bits = readU64Attr(module, feature_bits_attr_name) orelse return null;205    return .{206        .backend_kind = backend_kind,207        .artifact_format = artifact_format,208        .math_tier = math_tier,209        .dtype_bits = dtype_bits,210        .feature_bits = feature_bits,211    };212}213214fn readEnumAttr(comptime E: type, module: *const ir.Operation, attr_name: []const u8) ?E {215    const string_attr = module.getAttrAs(ir.Attribute.StringAttr, attr_name) orelse return null;216    return std.meta.stringToEnum(E, string_attr.getValue());217}218219fn readU64Attr(module: *const ir.Operation, attr_name: []const u8) ?u64 {220    const int_attr = module.getAttrAs(ir.Attribute.IntegerAttr, attr_name) orelse return null;221    return int_attr.getUnsignedValue();222}223224const testing = std.testing;225226test "backend target profile round trips through prepared module attrs" {227    var arena = std.heap.ArenaAllocator.init(testing.allocator);228    defer arena.deinit();229    const allocator = arena.allocator();230231    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);232    defer ctx.deinit(allocator);233234    const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());235    const profile = try BackendTargetProfile.init(.{236        .identity = .{237            .backend = .vulkan,238            .family = .vulkan,239        },240        .dtypes = gpu.DTypeSet.init(&.{ .i32, .f32 }),241        .features = .{ .atomic_i32 = true },242        .artifact_formats = gpu.ArtifactFormatSet.init(&.{.vulkan_spirv}),243    }, .vulkan, .vulkan_spirv);244245    try setBackendTargetProfile(&ctx, module.op, profile);246    const read = readBackendTargetProfile(module.op) orelse return error.MissingTargetProfile;247248    try testing.expectEqual(profile.backend_kind, read.backend_kind);249    try testing.expectEqual(profile.artifact_format, read.artifact_format);250    try testing.expectEqual(profile.math_tier, read.math_tier);251    try testing.expectEqual(profile.dtype_bits, read.dtype_bits);252    try testing.expectEqual(profile.feature_bits, read.feature_bits);253    try testing.expect(read.supportsDType(.f32));254    try testing.expect(!read.supportsDType(.f64));255}256257test "generated scan schedule decisions round trip through module attrs" {258    var arena = std.heap.ArenaAllocator.init(testing.allocator);259    defer arena.deinit();260    const allocator = arena.allocator();261262    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);263    defer ctx.deinit(allocator);264265    const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());266    try setGeneratedScanSchedules(allocator, &ctx, module.op, &.{267        .{ .total = 16777216, .schedule = .{ .threads = 256, .items = 16 } },268        .{ .schedule = .{ .threads = 512, .items = 16 } },269    });270271    const encoded = readGeneratedScanSchedules(module.op) orelse return error.MissingScanSchedules;272    try testing.expectEqualStrings("16777216=256x16,512x16", encoded);273274    const exact = (try resolveGeneratedScanSchedule(encoded, 16777216)) orelse return error.MissingScanSchedule;275    try testing.expect(exact.eql(.{ .threads = 256, .items = 16 }));276277    const wildcard = (try resolveGeneratedScanSchedule(encoded, 8192)) orelse return error.MissingScanSchedule;278    try testing.expect(wildcard.eql(.{ .threads = 512, .items = 16 }));279}280281test "generated scan schedule resolution rejects malformed encodings" {282    try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("", 4096));283    try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("512", 4096));284    try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("4096=0x16", 4096));285    try testing.expectError(error.InvalidArtifact, resolveGeneratedScanSchedule("abc=512x16", 4096));286    try testing.expectEqual(@as(?GeneratedScanSchedule, null), try resolveGeneratedScanSchedule("8192=512x16", 4096));287}288289test "generated row pipeline schedule decisions round trip through module attrs" {290    var arena = std.heap.ArenaAllocator.init(testing.allocator);291    defer arena.deinit();292    const allocator = arena.allocator();293294    var ctx = try ir.Context.init(allocator, ir.Context.Limits.testing);295    defer ctx.deinit(allocator);296297    const module = try choir.dialects.builtin.BuiltinDialect.ModuleOp.create(&ctx, ir.Location.getUnknown());298    try setGeneratedRowPipelineSchedules(allocator, &ctx, module.op, &.{299        .{ .shape = .{ .rows = 4096, .cols = 4096 }, .schedule = .{ .threads = 512 } },300        .{ .schedule = .{ .threads = 128 } },301    });302303    const encoded = readGeneratedRowPipelineSchedules(module.op) orelse return error.MissingRowPipelineSchedules;304    try testing.expectEqualStrings("4096x4096=512,128", encoded);305306    const exact = (try resolveGeneratedRowPipelineSchedule(encoded, 4096, 4096)) orelse return error.MissingRowPipelineSchedule;307    try testing.expect(exact.eql(.{ .threads = 512 }));308309    const wildcard = (try resolveGeneratedRowPipelineSchedule(encoded, 8, 2048)) orelse return error.MissingRowPipelineSchedule;310    try testing.expect(wildcard.eql(.{ .threads = 128 }));311}312313test "generated row pipeline schedule resolution rejects malformed encodings" {314    try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("", 4, 1024));315    try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("4096=512", 4, 1024));316    try testing.expectError(error.InvalidArtifact, resolveGeneratedRowPipelineSchedule("4x1024=0", 4, 1024));317    try testing.expectEqual(@as(?GeneratedRowPipelineSchedule, null), try resolveGeneratedRowPipelineSchedule("8x2048=512", 4, 1024));318}319320test "backend target profile gates tf32 tensor math on cuda tensor cores" {321    const caps = gpu.BackendCapabilities{322        .identity = .{323            .backend = .cuda,324            .family = .nvidia_cuda,325        },326        .dtypes = gpu.DTypeSet.init(&.{.f32}),327        .features = .{ .tensor_cores = true },328        .artifact_formats = gpu.ArtifactFormatSet.init(&.{.cuda_ptx}),329    };330    const profile = try BackendTargetProfile.initWithMathTier(caps, .cuda, .cuda_ptx, .tf32_tensor);331    try testing.expectEqual(gpu.BackendMathTier.tf32_tensor, profile.math_tier);332    try testing.expect(profile.isSupportedBy(caps));333334    var no_tensor = caps;335    no_tensor.features.tensor_cores = false;336    try testing.expectError(error.CapabilityMismatch, BackendTargetProfile.initWithMathTier(no_tensor, .cuda, .cuda_ptx, .tf32_tensor));337}