lib/accy/src/executable/tuning.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const choir = @import("choir");
5 const accy_root = @import("../root.zig");
6 const artifact_product = @import("../artifact/root.zig");
7
8 pub const product_name = "accy.exec.tuning";
9
10 pub const LaunchCandidateMeasurement = struct {
11 kernel_id: usize,
12 candidate_index: usize,
13 median_ns: u64,
14 sample_count: u32 = 1,
15 };
16
17 pub const LaunchTuningSelection = struct {
18 kernel_id: usize,
19 candidate_index: usize,
20 median_ns: u64 = 0,
21 sample_count: u32 = 0,
22 };
23
24 pub const launch_tuning_cache_record_version: u32 = 6;
25 pub const launch_tuning_artifact_magic: u32 = 0x41515431;
26 pub const launch_tuning_artifact_version: u32 = 6;
27
28 pub fn artifactFingerprint(bytes: []const u8) choir.product.incremental.Fingerprint {
29 var builder = choir.product.incremental.FingerprintBuilder{};
30 builder.updateBytes(product_name);
31 builder.updateU32(launch_tuning_artifact_version);
32 builder.updateBytes(bytes);
33 return builder.finish();
34 }
35
36 pub fn artifactProductStamp(bytes: []const u8) choir.product.incremental.ProductStamp {
37 return choir.product.incremental.productStamp(product_name, artifactFingerprint(bytes));
38 }
39
40 pub const LaunchTuningCacheKey = struct {
41 backend: gpu.BackendKind,
42 family: gpu.DeviceFamily,
43 format: gpu.ArtifactFormat,
44 vendor_id: u32 = 0,
45 has_vendor_id: bool = false,
46 device_id: u32 = 0,
47 has_device_id: bool = false,
48 name_fingerprint: u64 = 0,
49 driver_version_fingerprint: u64 = 0,
50 has_driver_version: bool = false,
51 subgroup_supported: bool = false,
52 subgroup_size_min: u32 = 0,
53 subgroup_size_max: u32 = 0,
54 subgroup_shuffle: bool = false,
55 subgroup_ballot: bool = false,
56 subgroup_vote: bool = false,
57 subgroup_arithmetic: bool = false,
58 subgroup_scan: bool = false,
59 max_threads: u32 = 0,
60 max_threads_per_dim_x: u32 = 0,
61 max_threads_per_dim_y: u32 = 0,
62 max_threads_per_dim_z: u32 = 0,
63 max_grid_per_dim_x: u32 = 0,
64 max_grid_per_dim_y: u32 = 0,
65 max_grid_per_dim_z: u32 = 0,
66 shared_memory_bytes: u32 = 0,
67 dtype_bits: u64 = 0,
68 tensor_cores: bool = false,
69 cooperative_matrix: bool = false,
70 dynamic_shared_memory: bool = false,
71 output_layout_fingerprint: u64 = 0,
72 input_layout_fingerprint: u64 = 0,
73 element_count: u64 = 0,
74 resource_class: artifact_product.LaunchResourceClass = .unknown,
75 element_ops_per_kib: u64 = 0,
76 static_bytes_complete: bool = false,
77 tile_kind: artifact_product.LaunchTileKind = .none,
78 tile_m: u32 = 0,
79 tile_n: u32 = 0,
80 tile_k: u32 = 0,
81 tile_batch: u32 = 1,
82 tile_has_input_dtype: bool = false,
83 tile_input_dtype: choir_abi.DType = .f32,
84 tile_has_output_dtype: bool = false,
85 tile_output_dtype: choir_abi.DType = .f32,
86 tile_input_tile_bytes: u32 = 0,
87 tile_output_tile_bytes: u32 = 0,
88 tile_scratch_memory_bytes: u32 = 0,
89 tile_reduction_kind: artifact_product.LaunchReductionKind = .none,
90 tile_reduction_rank: u32 = 0,
91 tile_reduction_axis: u32 = 0,
92 tile_reduction_extent: u32 = 0,
93 candidate_count: u32 = 0,
94 launch_resource_fingerprint: u64 = 0,
95
96 pub fn init(
97 caps: gpu.BackendCapabilities,
98 planned: artifact_product.PlannedKernel,
99 ) LaunchTuningCacheKey {
100 const tile = planned.launch_resources.tile;
101 return .{
102 .backend = caps.identity.backend,
103 .family = caps.identity.family,
104 .format = planned.launch_resources.format,
105 .vendor_id = caps.identity.vendor_id orelse 0,
106 .has_vendor_id = caps.identity.vendor_id != null,
107 .device_id = caps.identity.device_id orelse 0,
108 .has_device_id = caps.identity.device_id != null,
109 .name_fingerprint = bytesFingerprint("accy.exec.tuning.device.name", caps.identity.name),
110 .driver_version_fingerprint = if (caps.identity.driver_version) |version| bytesFingerprint("accy.exec.tuning.driver.version", version) else 0,
111 .has_driver_version = caps.identity.driver_version != null,
112 .subgroup_supported = caps.subgroup.supported,
113 .subgroup_size_min = caps.subgroup.size_min,
114 .subgroup_size_max = caps.subgroup.size_max,
115 .subgroup_shuffle = caps.subgroup.shuffle,
116 .subgroup_ballot = caps.subgroup.ballot,
117 .subgroup_vote = caps.subgroup.vote,
118 .subgroup_arithmetic = caps.subgroup.arithmetic,
119 .subgroup_scan = caps.subgroup.scan,
120 .max_threads = caps.threadgroup.max_threads,
121 .max_threads_per_dim_x = caps.threadgroup.max_threads_per_dim[0],
122 .max_threads_per_dim_y = caps.threadgroup.max_threads_per_dim[1],
123 .max_threads_per_dim_z = caps.threadgroup.max_threads_per_dim[2],
124 .max_grid_per_dim_x = caps.threadgroup.max_grid_per_dim[0],
125 .max_grid_per_dim_y = caps.threadgroup.max_grid_per_dim[1],
126 .max_grid_per_dim_z = caps.threadgroup.max_grid_per_dim[2],
127 .shared_memory_bytes = caps.threadgroup.shared_memory_bytes,
128 .dtype_bits = caps.dtypes.bits,
129 .tensor_cores = caps.features.tensor_cores,
130 .cooperative_matrix = caps.features.cooperative_matrix,
131 .dynamic_shared_memory = caps.features.dynamic_shared_memory,
132 .output_layout_fingerprint = planned.output_layout_fingerprint,
133 .input_layout_fingerprint = planned.input_layout_fingerprint,
134 .element_count = planned.launch_resources.element_count,
135 .resource_class = planned.launch_resources.resource_class,
136 .element_ops_per_kib = planned.launch_resources.element_ops_per_kib,
137 .static_bytes_complete = planned.launch_resources.static_bytes_complete,
138 .tile_kind = tile.kind,
139 .tile_m = tile.m,
140 .tile_n = tile.n,
141 .tile_k = tile.k,
142 .tile_batch = tile.batch,
143 .tile_has_input_dtype = tile.input_dtype != null,
144 .tile_input_dtype = tile.input_dtype orelse .f32,
145 .tile_has_output_dtype = tile.output_dtype != null,
146 .tile_output_dtype = tile.output_dtype orelse .f32,
147 .tile_input_tile_bytes = tile.input_tile_bytes,
148 .tile_output_tile_bytes = tile.output_tile_bytes,
149 .tile_scratch_memory_bytes = tile.scratch_memory_bytes,
150 .tile_reduction_kind = tile.reduction_kind,
151 .tile_reduction_rank = tile.reduction_rank,
152 .tile_reduction_axis = tile.reduction_axis,
153 .tile_reduction_extent = tile.reduction_extent,
154 .candidate_count = @intCast(planned.launch_resources.candidate_count),
155 .launch_resource_fingerprint = artifact_product.launchResourcePlanFingerprint(planned.launch_resources),
156 };
157 }
158 };
159
160 pub const LaunchTuningCacheRecord = struct {
161 version: u32 = launch_tuning_cache_record_version,
162 key: LaunchTuningCacheKey,
163 selection: LaunchTuningSelection,
164 };
165
166 pub const LaunchTuningCache = struct {
167 selections: std.AutoHashMap(LaunchTuningCacheKey, LaunchTuningSelection),
168
169 pub fn init(allocator: std.mem.Allocator) LaunchTuningCache {
170 return .{
171 .selections = std.AutoHashMap(LaunchTuningCacheKey, LaunchTuningSelection).init(allocator),
172 };
173 }
174
175 pub fn deinit(self: *LaunchTuningCache) void {
176 self.selections.deinit();
177 self.* = undefined;
178 }
179
180 pub fn count(self: *const LaunchTuningCache) usize {
181 return self.selections.count();
182 }
183
184 pub fn exportRecords(
185 self: *const LaunchTuningCache,
186 result_allocator: std.mem.Allocator,
187 ) gpu.BackendError![]LaunchTuningCacheRecord {
188 const records = result_allocator.alloc(LaunchTuningCacheRecord, self.selections.count()) catch return error.OutOfMemory;
189 errdefer result_allocator.free(records);
190
191 var iterator = self.selections.iterator();
192 var index: usize = 0;
193 while (iterator.next()) |entry| {
194 records[index] = .{
195 .key = entry.key_ptr.*,
196 .selection = entry.value_ptr.*,
197 };
198 index += 1;
199 }
200 std.mem.sort(LaunchTuningCacheRecord, records, {}, launchTuningCacheRecordSortsBefore);
201 return records;
202 }
203
204 pub fn importRecords(
205 self: *LaunchTuningCache,
206 records: []const LaunchTuningCacheRecord,
207 ) gpu.BackendError!void {
208 for (records) |record| try validateLaunchTuningCacheRecord(record);
209 for (records) |record| {
210 if (self.selections.getPtr(record.key)) |existing| {
211 if (launchSelectionBeats(record.selection, existing.*)) existing.* = record.selection;
212 continue;
213 }
214 self.selections.put(record.key, record.selection) catch return error.OutOfMemory;
215 }
216 }
217
218 pub fn recordMeasurements(
219 self: *LaunchTuningCache,
220 caps: gpu.BackendCapabilities,
221 artifact_plan: *const artifact_product.BackendArtifactPlan,
222 measurements: []const LaunchCandidateMeasurement,
223 ) gpu.BackendError!void {
224 for (artifact_plan.kernels.items) |planned| {
225 if (try selectedLaunchMeasurement(planned, measurements)) |measurement| {
226 try self.recordMeasuredSelection(caps, planned, .{
227 .kernel_id = measurement.kernel_id,
228 .candidate_index = measurement.candidate_index,
229 .median_ns = measurement.median_ns,
230 .sample_count = measurement.sample_count,
231 });
232 }
233 }
234 }
235
236 pub fn recordMeasuredSelection(
237 self: *LaunchTuningCache,
238 caps: gpu.BackendCapabilities,
239 planned: artifact_product.PlannedKernel,
240 selection: LaunchTuningSelection,
241 ) gpu.BackendError!void {
242 if (selection.kernel_id != planned.kernel_id) return error.LaunchArgumentMismatch;
243 if (selection.candidate_index >= planned.launch_resources.candidate_count) return error.LaunchArgumentMismatch;
244 if (selection.sample_count == 0) return error.LaunchArgumentMismatch;
245 const key = LaunchTuningCacheKey.init(caps, planned);
246 if (self.selections.getPtr(key)) |existing| {
247 if (launchSelectionBeats(selection, existing.*)) existing.* = selection;
248 return;
249 }
250 self.selections.put(key, selection) catch return error.OutOfMemory;
251 }
252
253 pub fn selectionForKernel(
254 self: *const LaunchTuningCache,
255 caps: gpu.BackendCapabilities,
256 planned: artifact_product.PlannedKernel,
257 ) gpu.BackendError!?LaunchTuningSelection {
258 const key = LaunchTuningCacheKey.init(caps, planned);
259 if (self.selections.get(key)) |selection| {
260 if (selection.candidate_index >= planned.launch_resources.candidate_count) return error.LaunchArgumentMismatch;
261 return .{
262 .kernel_id = planned.kernel_id,
263 .candidate_index = selection.candidate_index,
264 .median_ns = selection.median_ns,
265 .sample_count = selection.sample_count,
266 };
267 }
268 return null;
269 }
270 };
271
272 pub fn encodeLaunchTuningArtifact(
273 result_allocator: std.mem.Allocator,
274 records: []const LaunchTuningCacheRecord,
275 ) gpu.BackendError![]u8 {
276 if (records.len > std.math.maxInt(u32)) return error.InvalidArtifact;
277 var writer = artifact_product.wire.ByteWriter{};
278 errdefer writer.deinit(result_allocator);
279
280 try writer.writeU32(result_allocator, launch_tuning_artifact_magic);
281 try writer.writeU32(result_allocator, launch_tuning_artifact_version);
282 try writer.writeU32(result_allocator, @intCast(records.len));
283 for (records) |record| {
284 try validateLaunchTuningCacheRecord(record);
285 try writer.writeU32(result_allocator, record.version);
286 try writer.writeEnum(result_allocator, gpu.BackendKind, record.key.backend);
287 try writer.writeEnum(result_allocator, gpu.DeviceFamily, record.key.family);
288 try writer.writeEnum(result_allocator, gpu.ArtifactFormat, record.key.format);
289 try writer.writeU32(result_allocator, record.key.vendor_id);
290 try writer.writeBool(result_allocator, record.key.has_vendor_id);
291 try writer.writeU32(result_allocator, record.key.device_id);
292 try writer.writeBool(result_allocator, record.key.has_device_id);
293 try writer.writeU64(result_allocator, record.key.name_fingerprint);
294 try writer.writeU64(result_allocator, record.key.driver_version_fingerprint);
295 try writer.writeBool(result_allocator, record.key.has_driver_version);
296 try writer.writeBool(result_allocator, record.key.subgroup_supported);
297 try writer.writeU32(result_allocator, record.key.subgroup_size_min);
298 try writer.writeU32(result_allocator, record.key.subgroup_size_max);
299 try writer.writeBool(result_allocator, record.key.subgroup_shuffle);
300 try writer.writeBool(result_allocator, record.key.subgroup_ballot);
301 try writer.writeBool(result_allocator, record.key.subgroup_vote);
302 try writer.writeBool(result_allocator, record.key.subgroup_arithmetic);
303 try writer.writeBool(result_allocator, record.key.subgroup_scan);
304 try writer.writeU32(result_allocator, record.key.max_threads);
305 try writer.writeU32(result_allocator, record.key.max_threads_per_dim_x);
306 try writer.writeU32(result_allocator, record.key.max_threads_per_dim_y);
307 try writer.writeU32(result_allocator, record.key.max_threads_per_dim_z);
308 try writer.writeU32(result_allocator, record.key.max_grid_per_dim_x);
309 try writer.writeU32(result_allocator, record.key.max_grid_per_dim_y);
310 try writer.writeU32(result_allocator, record.key.max_grid_per_dim_z);
311 try writer.writeU32(result_allocator, record.key.shared_memory_bytes);
312 try writer.writeU64(result_allocator, record.key.dtype_bits);
313 try writer.writeBool(result_allocator, record.key.tensor_cores);
314 try writer.writeBool(result_allocator, record.key.cooperative_matrix);
315 try writer.writeBool(result_allocator, record.key.dynamic_shared_memory);
316 try writer.writeU64(result_allocator, record.key.output_layout_fingerprint);
317 try writer.writeU64(result_allocator, record.key.input_layout_fingerprint);
318 try writer.writeU64(result_allocator, record.key.element_count);
319 try writer.writeEnum(result_allocator, artifact_product.LaunchResourceClass, record.key.resource_class);
320 try writer.writeU64(result_allocator, record.key.element_ops_per_kib);
321 try writer.writeBool(result_allocator, record.key.static_bytes_complete);
322 try writer.writeEnum(result_allocator, artifact_product.LaunchTileKind, record.key.tile_kind);
323 try writer.writeU32(result_allocator, record.key.tile_m);
324 try writer.writeU32(result_allocator, record.key.tile_n);
325 try writer.writeU32(result_allocator, record.key.tile_k);
326 try writer.writeU32(result_allocator, record.key.tile_batch);
327 try writer.writeBool(result_allocator, record.key.tile_has_input_dtype);
328 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.tile_input_dtype);
329 try writer.writeBool(result_allocator, record.key.tile_has_output_dtype);
330 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.tile_output_dtype);
331 try writer.writeU32(result_allocator, record.key.tile_input_tile_bytes);
332 try writer.writeU32(result_allocator, record.key.tile_output_tile_bytes);
333 try writer.writeU32(result_allocator, record.key.tile_scratch_memory_bytes);
334 try writer.writeEnum(result_allocator, artifact_product.LaunchReductionKind, record.key.tile_reduction_kind);
335 try writer.writeU32(result_allocator, record.key.tile_reduction_rank);
336 try writer.writeU32(result_allocator, record.key.tile_reduction_axis);
337 try writer.writeU32(result_allocator, record.key.tile_reduction_extent);
338 try writer.writeU32(result_allocator, record.key.candidate_count);
339 try writer.writeU64(result_allocator, record.key.launch_resource_fingerprint);
340 try writer.writeUsize(result_allocator, record.selection.kernel_id);
341 try writer.writeUsize(result_allocator, record.selection.candidate_index);
342 try writer.writeU64(result_allocator, record.selection.median_ns);
343 try writer.writeU32(result_allocator, record.selection.sample_count);
344 }
345
346 return writer.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
347 }
348
349 pub fn decodeLaunchTuningArtifact(
350 result_allocator: std.mem.Allocator,
351 bytes: []const u8,
352 ) gpu.BackendError![]LaunchTuningCacheRecord {
353 var reader = artifact_product.wire.ByteReader{ .bytes = bytes };
354 if ((try reader.readU32()) != launch_tuning_artifact_magic) return error.InvalidArtifact;
355 if ((try reader.readU32()) != launch_tuning_artifact_version) return error.InvalidArtifact;
356 const record_count = try reader.readU32();
357
358 var records = std.ArrayListUnmanaged(LaunchTuningCacheRecord).empty;
359 errdefer records.deinit(result_allocator);
360
361 var index: u32 = 0;
362 while (index < record_count) : (index += 1) {
363 const record = LaunchTuningCacheRecord{
364 .version = try reader.readU32(),
365 .key = .{
366 .backend = try reader.readEnum(gpu.BackendKind),
367 .family = try reader.readEnum(gpu.DeviceFamily),
368 .format = try reader.readEnum(gpu.ArtifactFormat),
369 .vendor_id = try reader.readU32(),
370 .has_vendor_id = try reader.readBool(),
371 .device_id = try reader.readU32(),
372 .has_device_id = try reader.readBool(),
373 .name_fingerprint = try reader.readU64(),
374 .driver_version_fingerprint = try reader.readU64(),
375 .has_driver_version = try reader.readBool(),
376 .subgroup_supported = try reader.readBool(),
377 .subgroup_size_min = try reader.readU32(),
378 .subgroup_size_max = try reader.readU32(),
379 .subgroup_shuffle = try reader.readBool(),
380 .subgroup_ballot = try reader.readBool(),
381 .subgroup_vote = try reader.readBool(),
382 .subgroup_arithmetic = try reader.readBool(),
383 .subgroup_scan = try reader.readBool(),
384 .max_threads = try reader.readU32(),
385 .max_threads_per_dim_x = try reader.readU32(),
386 .max_threads_per_dim_y = try reader.readU32(),
387 .max_threads_per_dim_z = try reader.readU32(),
388 .max_grid_per_dim_x = try reader.readU32(),
389 .max_grid_per_dim_y = try reader.readU32(),
390 .max_grid_per_dim_z = try reader.readU32(),
391 .shared_memory_bytes = try reader.readU32(),
392 .dtype_bits = try reader.readU64(),
393 .tensor_cores = try reader.readBool(),
394 .cooperative_matrix = try reader.readBool(),
395 .dynamic_shared_memory = try reader.readBool(),
396 .output_layout_fingerprint = try reader.readU64(),
397 .input_layout_fingerprint = try reader.readU64(),
398 .element_count = try reader.readU64(),
399 .resource_class = try reader.readEnum(artifact_product.LaunchResourceClass),
400 .element_ops_per_kib = try reader.readU64(),
401 .static_bytes_complete = try reader.readBool(),
402 .tile_kind = try reader.readEnum(artifact_product.LaunchTileKind),
403 .tile_m = try reader.readU32(),
404 .tile_n = try reader.readU32(),
405 .tile_k = try reader.readU32(),
406 .tile_batch = try reader.readU32(),
407 .tile_has_input_dtype = try reader.readBool(),
408 .tile_input_dtype = try reader.readEnum(choir_abi.DType),
409 .tile_has_output_dtype = try reader.readBool(),
410 .tile_output_dtype = try reader.readEnum(choir_abi.DType),
411 .tile_input_tile_bytes = try reader.readU32(),
412 .tile_output_tile_bytes = try reader.readU32(),
413 .tile_scratch_memory_bytes = try reader.readU32(),
414 .tile_reduction_kind = try reader.readEnum(artifact_product.LaunchReductionKind),
415 .tile_reduction_rank = try reader.readU32(),
416 .tile_reduction_axis = try reader.readU32(),
417 .tile_reduction_extent = try reader.readU32(),
418 .candidate_count = try reader.readU32(),
419 .launch_resource_fingerprint = try reader.readU64(),
420 },
421 .selection = .{
422 .kernel_id = try reader.readUsize(),
423 .candidate_index = try reader.readUsize(),
424 .median_ns = try reader.readU64(),
425 .sample_count = try reader.readU32(),
426 },
427 };
428 try validateLaunchTuningCacheRecord(record);
429 records.append(result_allocator, record) catch return error.OutOfMemory;
430 }
431 try reader.expectDone();
432 return records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
433 }
434 pub const LaunchTuning = struct {
435 selections: []const LaunchTuningSelection = &.{},
436 measurements: []const LaunchCandidateMeasurement = &.{},
437
438 pub fn selectedCandidateIndex(
439 self: LaunchTuning,
440 planned: artifact_product.PlannedKernel,
441 ) gpu.BackendError!usize {
442 if (try selectedLaunchSelection(planned, self.selections)) |selection| return selection.candidate_index;
443 if (try selectedLaunchMeasurement(planned, self.measurements)) |measurement| return measurement.candidate_index;
444 return 0;
445 }
446 };
447 pub fn launchTuningIsEmpty(tuning: LaunchTuning) bool {
448 return tuning.selections.len == 0 and tuning.measurements.len == 0;
449 }
450
451 fn selectedLaunchSelection(
452 planned: artifact_product.PlannedKernel,
453 selections: []const LaunchTuningSelection,
454 ) gpu.BackendError!?LaunchTuningSelection {
455 var selected: ?LaunchTuningSelection = null;
456 for (selections) |selection| {
457 if (selection.kernel_id != planned.kernel_id) continue;
458 if (selection.candidate_index >= planned.launch_resources.candidate_count) return error.LaunchArgumentMismatch;
459 if (selected != null) return error.LaunchArgumentMismatch;
460 selected = selection;
461 }
462 return selected;
463 }
464
465 pub fn selectedLaunchMeasurement(
466 planned: artifact_product.PlannedKernel,
467 measurements: []const LaunchCandidateMeasurement,
468 ) gpu.BackendError!?LaunchCandidateMeasurement {
469 var best: ?LaunchCandidateMeasurement = null;
470 for (measurements) |measurement| {
471 if (measurement.kernel_id != planned.kernel_id) continue;
472 if (measurement.candidate_index >= planned.launch_resources.candidate_count) return error.LaunchArgumentMismatch;
473 if (measurement.sample_count == 0) continue;
474 if (best == null or
475 measurement.median_ns < best.?.median_ns or
476 (measurement.median_ns == best.?.median_ns and measurement.sample_count > best.?.sample_count) or
477 (measurement.median_ns == best.?.median_ns and measurement.sample_count == best.?.sample_count and measurement.candidate_index < best.?.candidate_index))
478 {
479 best = measurement;
480 }
481 }
482 return best;
483 }
484
485 fn launchSelectionBeats(
486 lhs: LaunchTuningSelection,
487 rhs: LaunchTuningSelection,
488 ) bool {
489 if (lhs.median_ns != rhs.median_ns) return lhs.median_ns < rhs.median_ns;
490 if (lhs.sample_count != rhs.sample_count) return lhs.sample_count > rhs.sample_count;
491 return lhs.candidate_index < rhs.candidate_index;
492 }
493
494 fn validateLaunchTuningCacheRecord(record: LaunchTuningCacheRecord) gpu.BackendError!void {
495 if (record.version != launch_tuning_cache_record_version) return error.InvalidArtifact;
496 try validateLaunchTuningTileKey(record.key);
497 if (record.key.candidate_count == 0) return error.InvalidArtifact;
498 if (record.selection.sample_count == 0) return error.LaunchArgumentMismatch;
499 if (record.selection.candidate_index >= @as(usize, record.key.candidate_count)) return error.LaunchArgumentMismatch;
500 }
501
502 fn validateLaunchTuningTileKey(key: LaunchTuningCacheKey) gpu.BackendError!void {
503 switch (key.tile_kind) {
504 .none => {
505 if (key.tile_m != 0 or key.tile_n != 0 or key.tile_k != 0) return error.InvalidArtifact;
506 if (key.tile_batch != 1) return error.InvalidArtifact;
507 if (key.tile_has_input_dtype or key.tile_has_output_dtype) return error.InvalidArtifact;
508 if (key.tile_input_tile_bytes != 0 or key.tile_output_tile_bytes != 0) return error.InvalidArtifact;
509 if (key.tile_scratch_memory_bytes != 0) return error.InvalidArtifact;
510 if (key.tile_reduction_kind != .none) return error.InvalidArtifact;
511 if (key.tile_reduction_rank != 0 or key.tile_reduction_axis != 0) return error.InvalidArtifact;
512 if (key.tile_reduction_extent != 0) return error.InvalidArtifact;
513 },
514 .dot_general => {
515 if (key.tile_m == 0 or key.tile_n == 0 or key.tile_k == 0) return error.InvalidArtifact;
516 if (key.tile_batch == 0) return error.InvalidArtifact;
517 if (!key.tile_has_input_dtype or !key.tile_has_output_dtype) return error.InvalidArtifact;
518 if (key.tile_input_tile_bytes == 0 or key.tile_output_tile_bytes == 0) return error.InvalidArtifact;
519 if (key.tile_reduction_kind != .none) return error.InvalidArtifact;
520 if (key.tile_reduction_rank != 0 or key.tile_reduction_axis != 0) return error.InvalidArtifact;
521 if (key.tile_reduction_extent != 0) return error.InvalidArtifact;
522 },
523 .reduction => {
524 if (key.tile_m == 0 or key.tile_n == 0) return error.InvalidArtifact;
525 if (key.tile_batch != 1) return error.InvalidArtifact;
526 if (!key.tile_has_input_dtype or !key.tile_has_output_dtype) return error.InvalidArtifact;
527 if (key.tile_input_tile_bytes == 0 or key.tile_output_tile_bytes == 0) return error.InvalidArtifact;
528 if (key.tile_reduction_kind == .none) return error.InvalidArtifact;
529 if (key.tile_reduction_rank == 0 or key.tile_reduction_extent == 0) return error.InvalidArtifact;
530 if (key.tile_reduction_axis >= key.tile_reduction_rank) return error.InvalidArtifact;
531 if (key.tile_k != key.tile_reduction_axis) return error.InvalidArtifact;
532 if (key.tile_n != key.tile_reduction_extent) return error.InvalidArtifact;
533 },
534 .elementwise_rank2 => {
535 if (key.tile_m == 0 or key.tile_n == 0 or key.tile_k != 0) return error.InvalidArtifact;
536 if (key.tile_batch != 1) return error.InvalidArtifact;
537 if (key.tile_has_input_dtype or key.tile_has_output_dtype) return error.InvalidArtifact;
538 if (key.tile_input_tile_bytes != 0 or key.tile_output_tile_bytes != 0) return error.InvalidArtifact;
539 if (key.tile_scratch_memory_bytes != 0) return error.InvalidArtifact;
540 if (key.tile_reduction_kind != .none) return error.InvalidArtifact;
541 if (key.tile_reduction_rank != 0 or key.tile_reduction_axis != 0) return error.InvalidArtifact;
542 if (key.tile_reduction_extent != 0) return error.InvalidArtifact;
543 },
544 }
545 }
546
547 fn launchTuningCacheRecordSortsBefore(
548 _: void,
549 lhs: LaunchTuningCacheRecord,
550 rhs: LaunchTuningCacheRecord,
551 ) bool {
552 if (launchTuningCacheKeySortsBefore(lhs.key, rhs.key)) return true;
553 if (launchTuningCacheKeySortsBefore(rhs.key, lhs.key)) return false;
554 return launchTuningSelectionSortsBefore(lhs.selection, rhs.selection);
555 }
556
557 fn launchTuningCacheKeySortsBefore(
558 lhs: LaunchTuningCacheKey,
559 rhs: LaunchTuningCacheKey,
560 ) bool {
561 if (compareEnum(gpu.BackendKind, lhs.backend, rhs.backend)) |less| return less;
562 if (compareEnum(gpu.DeviceFamily, lhs.family, rhs.family)) |less| return less;
563 if (compareEnum(gpu.ArtifactFormat, lhs.format, rhs.format)) |less| return less;
564 if (compareU32(lhs.vendor_id, rhs.vendor_id)) |less| return less;
565 if (compareBool(lhs.has_vendor_id, rhs.has_vendor_id)) |less| return less;
566 if (compareU32(lhs.device_id, rhs.device_id)) |less| return less;
567 if (compareBool(lhs.has_device_id, rhs.has_device_id)) |less| return less;
568 if (compareU64(lhs.name_fingerprint, rhs.name_fingerprint)) |less| return less;
569 if (compareU64(lhs.driver_version_fingerprint, rhs.driver_version_fingerprint)) |less| return less;
570 if (compareBool(lhs.has_driver_version, rhs.has_driver_version)) |less| return less;
571 if (compareBool(lhs.subgroup_supported, rhs.subgroup_supported)) |less| return less;
572 if (compareU32(lhs.subgroup_size_min, rhs.subgroup_size_min)) |less| return less;
573 if (compareU32(lhs.subgroup_size_max, rhs.subgroup_size_max)) |less| return less;
574 if (compareBool(lhs.subgroup_shuffle, rhs.subgroup_shuffle)) |less| return less;
575 if (compareBool(lhs.subgroup_ballot, rhs.subgroup_ballot)) |less| return less;
576 if (compareBool(lhs.subgroup_vote, rhs.subgroup_vote)) |less| return less;
577 if (compareBool(lhs.subgroup_arithmetic, rhs.subgroup_arithmetic)) |less| return less;
578 if (compareBool(lhs.subgroup_scan, rhs.subgroup_scan)) |less| return less;
579 if (compareU32(lhs.max_threads, rhs.max_threads)) |less| return less;
580 if (compareU32(lhs.max_threads_per_dim_x, rhs.max_threads_per_dim_x)) |less| return less;
581 if (compareU32(lhs.max_threads_per_dim_y, rhs.max_threads_per_dim_y)) |less| return less;
582 if (compareU32(lhs.max_threads_per_dim_z, rhs.max_threads_per_dim_z)) |less| return less;
583 if (compareU32(lhs.max_grid_per_dim_x, rhs.max_grid_per_dim_x)) |less| return less;
584 if (compareU32(lhs.max_grid_per_dim_y, rhs.max_grid_per_dim_y)) |less| return less;
585 if (compareU32(lhs.max_grid_per_dim_z, rhs.max_grid_per_dim_z)) |less| return less;
586 if (compareU32(lhs.shared_memory_bytes, rhs.shared_memory_bytes)) |less| return less;
587 if (compareU64(lhs.dtype_bits, rhs.dtype_bits)) |less| return less;
588 if (compareBool(lhs.tensor_cores, rhs.tensor_cores)) |less| return less;
589 if (compareBool(lhs.cooperative_matrix, rhs.cooperative_matrix)) |less| return less;
590 if (compareBool(lhs.dynamic_shared_memory, rhs.dynamic_shared_memory)) |less| return less;
591 if (compareU64(lhs.output_layout_fingerprint, rhs.output_layout_fingerprint)) |less| return less;
592 if (compareU64(lhs.input_layout_fingerprint, rhs.input_layout_fingerprint)) |less| return less;
593 if (compareU64(lhs.element_count, rhs.element_count)) |less| return less;
594 if (compareEnum(artifact_product.LaunchResourceClass, lhs.resource_class, rhs.resource_class)) |less| return less;
595 if (compareU64(lhs.element_ops_per_kib, rhs.element_ops_per_kib)) |less| return less;
596 if (compareBool(lhs.static_bytes_complete, rhs.static_bytes_complete)) |less| return less;
597 if (compareEnum(artifact_product.LaunchTileKind, lhs.tile_kind, rhs.tile_kind)) |less| return less;
598 if (compareU32(lhs.tile_m, rhs.tile_m)) |less| return less;
599 if (compareU32(lhs.tile_n, rhs.tile_n)) |less| return less;
600 if (compareU32(lhs.tile_k, rhs.tile_k)) |less| return less;
601 if (compareU32(lhs.tile_batch, rhs.tile_batch)) |less| return less;
602 if (compareBool(lhs.tile_has_input_dtype, rhs.tile_has_input_dtype)) |less| return less;
603 if (compareEnum(choir_abi.DType, lhs.tile_input_dtype, rhs.tile_input_dtype)) |less| return less;
604 if (compareBool(lhs.tile_has_output_dtype, rhs.tile_has_output_dtype)) |less| return less;
605 if (compareEnum(choir_abi.DType, lhs.tile_output_dtype, rhs.tile_output_dtype)) |less| return less;
606 if (compareU32(lhs.tile_input_tile_bytes, rhs.tile_input_tile_bytes)) |less| return less;
607 if (compareU32(lhs.tile_output_tile_bytes, rhs.tile_output_tile_bytes)) |less| return less;
608 if (compareU32(lhs.tile_scratch_memory_bytes, rhs.tile_scratch_memory_bytes)) |less| return less;
609 if (compareEnum(artifact_product.LaunchReductionKind, lhs.tile_reduction_kind, rhs.tile_reduction_kind)) |less| return less;
610 if (compareU32(lhs.tile_reduction_rank, rhs.tile_reduction_rank)) |less| return less;
611 if (compareU32(lhs.tile_reduction_axis, rhs.tile_reduction_axis)) |less| return less;
612 if (compareU32(lhs.tile_reduction_extent, rhs.tile_reduction_extent)) |less| return less;
613 if (compareU32(lhs.candidate_count, rhs.candidate_count)) |less| return less;
614 if (compareU64(lhs.launch_resource_fingerprint, rhs.launch_resource_fingerprint)) |less| return less;
615 return false;
616 }
617
618 fn launchTuningSelectionSortsBefore(
619 lhs: LaunchTuningSelection,
620 rhs: LaunchTuningSelection,
621 ) bool {
622 if (compareUsize(lhs.kernel_id, rhs.kernel_id)) |less| return less;
623 if (compareUsize(lhs.candidate_index, rhs.candidate_index)) |less| return less;
624 if (compareU64(lhs.median_ns, rhs.median_ns)) |less| return less;
625 if (compareU32(lhs.sample_count, rhs.sample_count)) |less| return less;
626 return false;
627 }
628
629 fn compareEnum(comptime T: type, lhs: T, rhs: T) ?bool {
630 return compareU64(@backingInt(lhs), @backingInt(rhs));
631 }
632
633 fn compareBool(lhs: bool, rhs: bool) ?bool {
634 if (lhs == rhs) return null;
635 return !lhs and rhs;
636 }
637
638 fn compareU32(lhs: u32, rhs: u32) ?bool {
639 if (lhs == rhs) return null;
640 return lhs < rhs;
641 }
642
643 fn compareUsize(lhs: usize, rhs: usize) ?bool {
644 if (lhs == rhs) return null;
645 return lhs < rhs;
646 }
647
648 fn compareU64(lhs: u64, rhs: u64) ?bool {
649 if (lhs == rhs) return null;
650 return lhs < rhs;
651 }
652
653 fn bytesFingerprint(domain: []const u8, bytes: []const u8) choir.product.incremental.Fingerprint {
654 var builder = choir.product.incremental.FingerprintBuilder{};
655 builder.updateBytes(domain);
656 builder.updateBytes(bytes);
657 return builder.finish();
658 }
659
660 test "launch tuning artifact round-trips elementwise rank-2 tile keys" {
661 const allocator = std.testing.allocator;
662 const records = [_]LaunchTuningCacheRecord{.{
663 .key = .{
664 .backend = .cuda,
665 .family = .nvidia_cuda,
666 .format = .cuda_ptx,
667 .max_threads = 1024,
668 .max_threads_per_dim_x = 1024,
669 .max_threads_per_dim_y = 1024,
670 .max_threads_per_dim_z = 64,
671 .max_grid_per_dim_x = 2_147_483_647,
672 .max_grid_per_dim_y = 65_535,
673 .max_grid_per_dim_z = 65_535,
674 .element_count = 2048 * 2048,
675 .resource_class = .balanced,
676 .tile_kind = .elementwise_rank2,
677 .tile_m = 2048,
678 .tile_n = 2048,
679 .candidate_count = 8,
680 .launch_resource_fingerprint = 0x1234,
681 },
682 .selection = .{
683 .kernel_id = 3,
684 .candidate_index = 2,
685 .median_ns = 42,
686 .sample_count = 5,
687 },
688 }};
689
690 const encoded = try encodeLaunchTuningArtifact(allocator, records[0..]);
691 defer allocator.free(encoded);
692 const decoded = try decodeLaunchTuningArtifact(allocator, encoded);
693 defer allocator.free(decoded);
694
695 try std.testing.expectEqual(@as(usize, 1), decoded.len);
696 try std.testing.expectEqual(launch_tuning_cache_record_version, decoded[0].version);
697 try std.testing.expectEqual(records[0].key, decoded[0].key);
698 try std.testing.expectEqual(records[0].selection, decoded[0].selection);
699 }