tiny.accy.choir.record.target

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

1const std = @import("std");2const gpu = @import("gpu");3const choir_abi = @import("choir_abi");4const records = @import("root.zig");56pub const CompileOptions = struct {7    cpu_vector_width: ?u32 = null,8};910pub const BackendTargetProfile = struct {11    backend_kind: gpu.BackendKind,12    artifact_format: gpu.ArtifactFormat,13    math_tier: gpu.BackendMathTier = .exact,14    dtype_bits: u64 = 0,15    feature_bits: u64 = 0,1617    pub fn init(18        caps: gpu.BackendCapabilities,19        backend_kind: gpu.BackendKind,20        artifact_format: gpu.ArtifactFormat,21    ) gpu.BackendError!BackendTargetProfile {22        return initWithMathTier(caps, backend_kind, artifact_format, .exact);23    }2425    pub fn initWithMathTier(26        caps: gpu.BackendCapabilities,27        backend_kind: gpu.BackendKind,28        artifact_format: gpu.ArtifactFormat,29        math_tier: gpu.BackendMathTier,30    ) gpu.BackendError!BackendTargetProfile {31        if (!caps.supportsArtifactFormat(artifact_format)) return error.UnsupportedArtifactFormat;32        if (!mathTierSupported(caps, backend_kind, artifact_format, math_tier)) return error.CapabilityMismatch;33        return .{34            .backend_kind = backend_kind,35            .artifact_format = artifact_format,36            .math_tier = math_tier,37            .dtype_bits = caps.dtypes.bits,38            .feature_bits = featureBits(caps.features),39        };40    }4142    pub fn dtypes(self: BackendTargetProfile) gpu.DTypeSet {43        return .{ .bits = self.dtype_bits };44    }4546    pub fn supportsDType(self: BackendTargetProfile, element_type: choir_abi.DType) bool {47        return self.dtypes().contains(element_type);48    }4950    pub fn supportsFeaturesFrom(self: BackendTargetProfile, caps: gpu.BackendCapabilities) bool {51        return (featureBits(caps.features) & self.feature_bits) == self.feature_bits;52    }5354    pub fn isSupportedBy(self: BackendTargetProfile, caps: gpu.BackendCapabilities) bool {55        return caps.supportsArtifactFormat(self.artifact_format) and56            caps.dtypes.containsAll(self.dtypes()) and57            self.supportsFeaturesFrom(caps) and58            mathTierSupported(caps, self.backend_kind, self.artifact_format, self.math_tier);59    }6061    pub fn eql(self: BackendTargetProfile, other: BackendTargetProfile) bool {62        return self.backend_kind == other.backend_kind and63            self.artifact_format == other.artifact_format and64            self.math_tier == other.math_tier and65            self.dtype_bits == other.dtype_bits and66            self.feature_bits == other.feature_bits;67    }68};6970pub const Abi = struct {71    argument_count: u32,72    static_arguments: []const choir_abi.ScalarArgument,73    launch: ?choir_abi.LaunchGeometry,74    compile_options: CompileOptions,75};7677pub const Kernel = struct {78    lowered: records.kernel.Lowered,79    work_dtype: choir_abi.DType,80    element_count: u64,81    required_dtype_bits: u64,82    required_features: choir_abi.Features,83    required_subgroup: choir_abi.SubgroupRequirements,84    runtime_scalar_argument_count: u32,85    abi: ?Abi,86};8788pub const Record = struct {89    profile: ?BackendTargetProfile,90    math_tier: gpu.BackendMathTier,91    generated_scan_schedules: ?[]const u8,92    generated_row_pipeline_schedules: ?[]const u8,93    kernels: []const Kernel,94};9596pub fn validate(allocator: std.mem.Allocator, value: Record, parent: records.kernel.Record) !void {97    if (value.kernels.len != parent.generated.kernels.len) return error.InvalidStageRecord;98    const tier = if (value.profile) |profile| profile.math_tier else .exact;99    if (value.math_tier != tier) return error.InvalidStageRecord;100    var work = std.AutoHashMapUnmanaged(usize, bool).empty;101    defer work.deinit(allocator);102    const count = std.math.cast(u32, parent.generated.kernels.len) orelse103        return error.InvalidStageRecord;104    try work.ensureTotalCapacity(allocator, count);105    for (parent.generated.kernels) |kernel| work.putAssumeCapacity(kernel.work_item_id, false);106    for (value.kernels) |kernel| {107        const seen = work.getPtr(kernel.lowered.work_item_id) orelse108            return error.InvalidStageRecord;109        if (seen.*) return error.InvalidStageRecord;110        seen.* = true;111        const arguments = kernel.lowered.program.params.len;112        if (kernel.lowered.argument_count != arguments or113            kernel.runtime_scalar_argument_count != try runtimeScalarCount(kernel.lowered.program)) return error.InvalidStageRecord;114        if ((value.profile == null) != (kernel.abi == null)) return error.InvalidStageRecord;115        if (kernel.abi) |abi| {116            if (abi.argument_count < arguments or117                abi.static_arguments.len != abi.argument_count - arguments)118            {119                return error.InvalidStageRecord;120            }121        }122    }123}124125fn runtimeScalarCount(program: records.program.Record) !u32 {126    var count: u32 = 0;127    for (program.params) |param| {128        switch (param) {129            .scalar => count = std.math.add(u32, count, 1) catch130                return error.InvalidStageRecord,131            .buffer => {},132        }133    }134    return count;135}136137fn featureBits(features: choir_abi.Features) u64 {138    var bits: u64 = 0;139    if (features.atomic_i32) bits |= 1 << 0;140    if (features.atomic_u32) bits |= 1 << 1;141    if (features.atomic_index) bits |= 1 << 2;142    if (features.atomic_f32_add_device) bits |= 1 << 3;143    if (features.atomic_f32_add_shared) bits |= 1 << 4;144    if (features.unsupported_atomic) bits |= 1 << 5;145    if (features.async_copy) bits |= 1 << 6;146    if (features.tensor_cores) bits |= 1 << 7;147    if (features.cooperative_matrix) bits |= 1 << 8;148    if (features.dynamic_shared_memory) bits |= 1 << 9;149    if (features.indirect_launch) bits |= 1 << 10;150    return bits;151}152153fn mathTierSupported(154    caps: gpu.BackendCapabilities,155    backend_kind: gpu.BackendKind,156    artifact_format: gpu.ArtifactFormat,157    math_tier: gpu.BackendMathTier,158) bool {159    return switch (math_tier) {160        .exact => true,161        .tf32_tensor => backend_kind == .cuda and artifact_format == .cuda_ptx and caps.features.tensor_cores,162    };163}