lib/accy/src/artifact/summary.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const artifact = @import("root.zig");
5
6 pub const KernelSource = enum {
7 tensor,
8 kernel_call,
9 choir_kernel,
10 };
11
12 pub const KernelSummary = struct {
13 source: KernelSource,
14 kernel_id: usize,
15 work_item_id: usize,
16 output_layout_fingerprint: u64,
17 input_layout_fingerprint: u64,
18 element_count: u64,
19 op_count: usize,
20 entry_name: []const u8,
21 artifact_format: gpu.ArtifactFormat,
22 artifact_payload_bytes: usize,
23 compile_argument_count: u32,
24 compile_required_dtype_bits: u64,
25 compile_required_features: choir_abi.Features,
26 compile_required_subgroup: choir_abi.SubgroupRequirements,
27 compile_payload: artifact.PlannedKernelCompilePayload,
28 compile_payload_bytes: usize,
29 compile_launch: artifact.PlannedKernelCompileLaunch,
30 runtime_scalar_argument_count: u32,
31 launch_geometry: choir_abi.LaunchGeometry,
32 launch_candidate_count: usize,
33 launch_resource_class: []const u8,
34 fixed_threadgroup: bool,
35 subgroup_aligned: bool,
36 subgroup_size: ?u32,
37 static_bytes_complete: bool,
38 };
39
40 pub const KernelSummaries = struct {
41 allocator: std.mem.Allocator,
42 items: []KernelSummary,
43
44 pub fn deinit(self: *KernelSummaries) void {
45 for (self.items) |item| {
46 self.allocator.free(item.entry_name);
47 }
48 self.allocator.free(self.items);
49 self.* = undefined;
50 }
51
52 pub fn len(self: *const KernelSummaries) usize {
53 return self.items.len;
54 }
55
56 pub fn summary(self: *const KernelSummaries, kernel_index: usize) gpu.BackendError!KernelSummary {
57 if (kernel_index >= self.items.len) return error.InvalidArtifact;
58 return self.items[kernel_index];
59 }
60
61 pub fn summaryForWork(self: *const KernelSummaries, work_item_id: usize) gpu.BackendError!KernelSummary {
62 for (self.items) |item| {
63 if (item.work_item_id == work_item_id) return item;
64 }
65 return error.InvalidArtifact;
66 }
67 };
68
69 pub fn copyKernelSummary(allocator: std.mem.Allocator, summary: KernelSummary) gpu.BackendError!KernelSummary {
70 var copied = summary;
71 copied.entry_name = allocator.dupe(u8, summary.entry_name) catch return error.OutOfMemory;
72 return copied;
73 }
74
75 pub fn kernelSource(source: artifact.PlannedKernelSource) KernelSource {
76 return switch (source) {
77 .tensor => .tensor,
78 .kernel_call => .kernel_call,
79 .choir_kernel => .choir_kernel,
80 };
81 }
82
83 pub fn artifactPayloadByteCount(kernel_artifact: gpu.KernelArtifact) gpu.BackendError!usize {
84 return switch (kernel_artifact.payload) {
85 .bytes => |bytes| bytes.len,
86 .words_u32 => |words| words.len * @sizeOf(u32),
87 .text => |text| text.len,
88 .none, .external => error.InvalidArtifact,
89 };
90 }
91
92 pub fn launchResourceClassName(class: artifact.LaunchResourceClass) []const u8 {
93 return switch (class) {
94 .unknown => "unknown",
95 .memory_bound => "memory_bound",
96 .balanced => "balanced",
97 .compute_weighted => "compute_weighted",
98 };
99 }
100
101 pub fn summarizePlannedKernel(planned: artifact.PlannedKernel) gpu.BackendError!KernelSummary {
102 return .{
103 .source = kernelSource(planned.compile.source),
104 .kernel_id = planned.kernel_id,
105 .work_item_id = planned.work_item_id,
106 .output_layout_fingerprint = planned.output_layout_fingerprint,
107 .input_layout_fingerprint = planned.input_layout_fingerprint,
108 .element_count = planned.element_count,
109 .op_count = planned.op_count,
110 .entry_name = planned.artifact.entry_name,
111 .artifact_format = planned.artifact.format,
112 .artifact_payload_bytes = try artifactPayloadByteCount(planned.artifact),
113 .compile_argument_count = planned.compile.argument_count,
114 .compile_required_dtype_bits = planned.compile.required_dtypes.bits,
115 .compile_required_features = planned.compile.required_features,
116 .compile_required_subgroup = planned.compile.required_subgroup,
117 .compile_payload = planned.compile.payload,
118 .compile_payload_bytes = planned.compile.payload_byte_count,
119 .compile_launch = planned.compile.launch,
120 .runtime_scalar_argument_count = planned.runtime_scalar_argument_count,
121 .launch_geometry = planned.launch_resources.geometry,
122 .launch_candidate_count = planned.launch_resources.candidate_count,
123 .launch_resource_class = launchResourceClassName(planned.launch_resources.resource_class),
124 .fixed_threadgroup = planned.launch_resources.fixed_threadgroup,
125 .subgroup_aligned = planned.launch_resources.subgroup_aligned,
126 .subgroup_size = planned.launch_resources.subgroup_size,
127 .static_bytes_complete = planned.launch_resources.static_bytes_complete,
128 };
129 }
130
131 pub fn kernelSummariesEqual(lhs: KernelSummary, rhs: KernelSummary) bool {
132 return lhs.source == rhs.source and
133 lhs.kernel_id == rhs.kernel_id and
134 lhs.work_item_id == rhs.work_item_id and
135 lhs.output_layout_fingerprint == rhs.output_layout_fingerprint and
136 lhs.input_layout_fingerprint == rhs.input_layout_fingerprint and
137 lhs.element_count == rhs.element_count and
138 lhs.op_count == rhs.op_count and
139 std.mem.eql(u8, lhs.entry_name, rhs.entry_name) and
140 lhs.artifact_format == rhs.artifact_format and
141 lhs.artifact_payload_bytes == rhs.artifact_payload_bytes and
142 lhs.compile_argument_count == rhs.compile_argument_count and
143 lhs.compile_required_dtype_bits == rhs.compile_required_dtype_bits and
144 std.meta.eql(lhs.compile_required_features, rhs.compile_required_features) and
145 std.meta.eql(lhs.compile_required_subgroup, rhs.compile_required_subgroup) and
146 lhs.compile_payload == rhs.compile_payload and
147 lhs.compile_payload_bytes == rhs.compile_payload_bytes and
148 lhs.compile_launch == rhs.compile_launch and
149 lhs.runtime_scalar_argument_count == rhs.runtime_scalar_argument_count and
150 launchGeometriesEqual(lhs.launch_geometry, rhs.launch_geometry) and
151 lhs.launch_candidate_count == rhs.launch_candidate_count and
152 std.mem.eql(u8, lhs.launch_resource_class, rhs.launch_resource_class) and
153 lhs.fixed_threadgroup == rhs.fixed_threadgroup and
154 lhs.subgroup_aligned == rhs.subgroup_aligned and
155 lhs.subgroup_size == rhs.subgroup_size and
156 lhs.static_bytes_complete == rhs.static_bytes_complete;
157 }
158
159 pub fn launchGeometriesEqual(lhs: choir_abi.LaunchGeometry, rhs: choir_abi.LaunchGeometry) bool {
160 return lhs.grid[0] == rhs.grid[0] and
161 lhs.grid[1] == rhs.grid[1] and
162 lhs.grid[2] == rhs.grid[2] and
163 lhs.threadgroup[0] == rhs.threadgroup[0] and
164 lhs.threadgroup[1] == rhs.threadgroup[1] and
165 lhs.threadgroup[2] == rhs.threadgroup[2] and
166 lhs.dynamic_shared_memory_bytes == rhs.dynamic_shared_memory_bytes;
167 }
168
169 fn testKernelSummary() KernelSummary {
170 return .{
171 .source = .tensor,
172 .kernel_id = 7,
173 .work_item_id = 3,
174 .output_layout_fingerprint = 11,
175 .input_layout_fingerprint = 13,
176 .element_count = 1024,
177 .op_count = 4,
178 .entry_name = "accy_test_kernel",
179 .artifact_format = .cuda_ptx,
180 .artifact_payload_bytes = 128,
181 .compile_argument_count = 3,
182 .compile_required_dtype_bits = gpu.DTypeSet.init(&.{ .f32, .i32 }).bits,
183 .compile_required_features = .{},
184 .compile_required_subgroup = .{},
185 .compile_payload = .text,
186 .compile_payload_bytes = 96,
187 .compile_launch = .generic,
188 .runtime_scalar_argument_count = 0,
189 .launch_geometry = .{
190 .grid = .{ 8, 2, 1 },
191 .threadgroup = .{ 128, 1, 1 },
192 .dynamic_shared_memory_bytes = 256,
193 },
194 .launch_candidate_count = 2,
195 .launch_resource_class = "balanced",
196 .fixed_threadgroup = false,
197 .subgroup_aligned = true,
198 .subgroup_size = 32,
199 .static_bytes_complete = true,
200 };
201 }
202
203 test "artifact kernel summary equality includes compile contract" {
204 const base = testKernelSummary();
205 try std.testing.expect(kernelSummariesEqual(base, testKernelSummary()));
206
207 var changed = testKernelSummary();
208 changed.compile_argument_count += 1;
209 try std.testing.expect(!kernelSummariesEqual(base, changed));
210
211 changed = testKernelSummary();
212 changed.compile_required_dtype_bits ^= gpu.DTypeSet.init(&.{.f16}).bits;
213 try std.testing.expect(!kernelSummariesEqual(base, changed));
214
215 changed = testKernelSummary();
216 changed.compile_required_features = .{ .tensor_cores = true };
217 try std.testing.expect(!kernelSummariesEqual(base, changed));
218
219 changed = testKernelSummary();
220 changed.compile_required_subgroup = .{ .supported = true, .shuffle = true };
221 try std.testing.expect(!kernelSummariesEqual(base, changed));
222
223 changed = testKernelSummary();
224 changed.compile_payload = .bytes;
225 try std.testing.expect(!kernelSummariesEqual(base, changed));
226
227 changed = testKernelSummary();
228 changed.compile_payload_bytes += 1;
229 try std.testing.expect(!kernelSummariesEqual(base, changed));
230
231 changed = testKernelSummary();
232 changed.compile_launch = .dot_general;
233 try std.testing.expect(!kernelSummariesEqual(base, changed));
234
235 changed = testKernelSummary();
236 changed.runtime_scalar_argument_count += 1;
237 try std.testing.expect(!kernelSummariesEqual(base, changed));
238
239 changed = testKernelSummary();
240 changed.launch_geometry.dynamic_shared_memory_bytes += 1;
241 try std.testing.expect(!kernelSummariesEqual(base, changed));
242 }