lib/accy/src/choir/record/target.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const records = @import("root.zig");
5
6 pub const CompileOptions = struct {
7 cpu_vector_width: ?u32 = null,
8 };
9
10 pub 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,
16
17 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 }
24
25 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 }
41
42 pub fn dtypes(self: BackendTargetProfile) gpu.DTypeSet {
43 return .{ .bits = self.dtype_bits };
44 }
45
46 pub fn supportsDType(self: BackendTargetProfile, element_type: choir_abi.DType) bool {
47 return self.dtypes().contains(element_type);
48 }
49
50 pub fn supportsFeaturesFrom(self: BackendTargetProfile, caps: gpu.BackendCapabilities) bool {
51 return (featureBits(caps.features) & self.feature_bits) == self.feature_bits;
52 }
53
54 pub fn isSupportedBy(self: BackendTargetProfile, caps: gpu.BackendCapabilities) bool {
55 return caps.supportsArtifactFormat(self.artifact_format) and
56 caps.dtypes.containsAll(self.dtypes()) and
57 self.supportsFeaturesFrom(caps) and
58 mathTierSupported(caps, self.backend_kind, self.artifact_format, self.math_tier);
59 }
60
61 pub fn eql(self: BackendTargetProfile, other: BackendTargetProfile) bool {
62 return self.backend_kind == other.backend_kind and
63 self.artifact_format == other.artifact_format and
64 self.math_tier == other.math_tier and
65 self.dtype_bits == other.dtype_bits and
66 self.feature_bits == other.feature_bits;
67 }
68 };
69
70 pub const Abi = struct {
71 argument_count: u32,
72 static_arguments: []const choir_abi.ScalarArgument,
73 launch: ?choir_abi.LaunchGeometry,
74 compile_options: CompileOptions,
75 };
76
77 pub 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 };
87
88 pub 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 };
95
96 pub 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) orelse
103 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) orelse
108 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 or
113 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 or
117 abi.static_arguments.len != abi.argument_count - arguments)
118 {
119 return error.InvalidStageRecord;
120 }
121 }
122 }
123 }
124
125 fn 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) catch
130 return error.InvalidStageRecord,
131 .buffer => {},
132 }
133 }
134 return count;
135 }
136
137 fn 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 }
152
153 fn 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 }