lib/accy/src/executable/schedule.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/model/root.zig");
7 const preparation = @import("../preparation/root.zig");
8
9 const GeneratedScanSchedule = preparation.target.GeneratedScanSchedule;
10 const GeneratedRowPipelineSchedule = preparation.target.GeneratedRowPipelineSchedule;
11
12 const matrix_tuning = @import("../kernel/library/root.zig").tuning;
13 const matrix_product_family_schedule_tuning_product_name =
14 matrix_tuning.matrix_product_family_schedule_tuning_product_name;
15 const matrix_product_family_schedule_tuning_record_version =
16 matrix_tuning.matrix_product_family_schedule_tuning_record_version;
17 const matrix_product_family_schedule_tuning_max_candidates =
18 matrix_tuning.matrix_product_family_schedule_tuning_max_candidates;
19 const MatrixProductFamilyScheduleThreads = matrix_tuning.MatrixProductFamilyScheduleThreads;
20 const MatrixProductFamilyScheduleTuningProblem =
21 matrix_tuning.MatrixProductFamilyScheduleTuningProblem;
22 const MatrixProductFamilyScheduleTuningKey = matrix_tuning.MatrixProductFamilyScheduleTuningKey;
23 const MatrixProductFamilyScheduleTuningSelection =
24 matrix_tuning.MatrixProductFamilyScheduleTuningSelection;
25 const MatrixProductFamilyScheduleTuningRecord =
26 matrix_tuning.MatrixProductFamilyScheduleTuningRecord;
27
28 pub const matrix_product_family_schedule_tuning_artifact_magic: u32 = 0x41515433;
29 pub const matrix_product_family_schedule_tuning_artifact_version: u32 = 1;
30
31 pub const MatrixProductFamilyScheduleTuningCache = struct {
32 selections: std.AutoHashMap(MatrixProductFamilyScheduleTuningKey, MatrixProductFamilyScheduleTuningSelection),
33
34 pub fn init(allocator: std.mem.Allocator) MatrixProductFamilyScheduleTuningCache {
35 return .{
36 .selections = std.AutoHashMap(MatrixProductFamilyScheduleTuningKey, MatrixProductFamilyScheduleTuningSelection).init(allocator),
37 };
38 }
39
40 pub fn deinit(self: *MatrixProductFamilyScheduleTuningCache) void {
41 self.selections.deinit();
42 self.* = undefined;
43 }
44
45 pub fn count(self: *const MatrixProductFamilyScheduleTuningCache) usize {
46 return self.selections.count();
47 }
48
49 pub fn exportRecords(
50 self: *const MatrixProductFamilyScheduleTuningCache,
51 result_allocator: std.mem.Allocator,
52 ) gpu.BackendError![]MatrixProductFamilyScheduleTuningRecord {
53 const records = result_allocator.alloc(MatrixProductFamilyScheduleTuningRecord, self.selections.count()) catch return error.OutOfMemory;
54 errdefer result_allocator.free(records);
55
56 var iterator = self.selections.iterator();
57 var index: usize = 0;
58 while (iterator.next()) |entry| {
59 records[index] = .{
60 .key = entry.key_ptr.*,
61 .selection = entry.value_ptr.*,
62 };
63 index += 1;
64 }
65 std.mem.sort(MatrixProductFamilyScheduleTuningRecord, records, {}, matrixProductFamilyScheduleTuningRecordSortsBefore);
66 return records;
67 }
68
69 pub fn importRecords(
70 self: *MatrixProductFamilyScheduleTuningCache,
71 records: []const MatrixProductFamilyScheduleTuningRecord,
72 ) gpu.BackendError!void {
73 for (records) |record| try validateMatrixProductFamilyScheduleTuningRecord(record);
74 for (records) |record| {
75 if (self.selections.getPtr(record.key)) |existing| {
76 if (matrixProductFamilyScheduleSelectionBeats(record.selection, existing.*)) existing.* = record.selection;
77 continue;
78 }
79 self.selections.put(record.key, record.selection) catch return error.OutOfMemory;
80 }
81 }
82
83 pub fn recordSelection(
84 self: *MatrixProductFamilyScheduleTuningCache,
85 caps: gpu.BackendCapabilities,
86 problem: MatrixProductFamilyScheduleTuningProblem,
87 selection: MatrixProductFamilyScheduleTuningSelection,
88 ) gpu.BackendError!void {
89 try validateMatrixProductFamilyScheduleTuningSelection(selection);
90 if (!matrixProductFamilyScheduleCandidatesContain(problem.candidates, selection.threads)) return error.LaunchArgumentMismatch;
91 const key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, problem);
92 if (self.selections.getPtr(key)) |existing| {
93 if (matrixProductFamilyScheduleSelectionBeats(selection, existing.*)) existing.* = selection;
94 return;
95 }
96 self.selections.put(key, selection) catch return error.OutOfMemory;
97 }
98
99 pub fn selectionForProblem(
100 self: *const MatrixProductFamilyScheduleTuningCache,
101 caps: gpu.BackendCapabilities,
102 problem: MatrixProductFamilyScheduleTuningProblem,
103 ) gpu.BackendError!?MatrixProductFamilyScheduleTuningSelection {
104 const key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, problem);
105 return self.selections.get(key);
106 }
107 };
108
109 pub fn encodeMatrixProductFamilyScheduleTuningArtifact(
110 result_allocator: std.mem.Allocator,
111 records: []const MatrixProductFamilyScheduleTuningRecord,
112 ) gpu.BackendError![]u8 {
113 if (records.len > std.math.maxInt(u32)) return error.InvalidArtifact;
114 var writer = artifact_product.wire.ByteWriter{};
115 errdefer writer.deinit(result_allocator);
116
117 try writer.writeU32(result_allocator, matrix_product_family_schedule_tuning_artifact_magic);
118 try writer.writeU32(result_allocator, matrix_product_family_schedule_tuning_artifact_version);
119 try writer.writeU32(result_allocator, @intCast(records.len));
120 for (records) |record| {
121 try validateMatrixProductFamilyScheduleTuningRecord(record);
122 try writer.writeU32(result_allocator, record.version);
123 try writer.writeEnum(result_allocator, gpu.BackendKind, record.key.backend);
124 try writer.writeEnum(result_allocator, gpu.DeviceFamily, record.key.family);
125 try writer.writeEnum(result_allocator, gpu.ArtifactFormat, record.key.format);
126 try writer.writeU32(result_allocator, record.key.vendor_id);
127 try writer.writeBool(result_allocator, record.key.has_vendor_id);
128 try writer.writeU32(result_allocator, record.key.device_id);
129 try writer.writeBool(result_allocator, record.key.has_device_id);
130 try writer.writeU64(result_allocator, record.key.name_fingerprint);
131 try writer.writeU64(result_allocator, record.key.driver_version_fingerprint);
132 try writer.writeBool(result_allocator, record.key.has_driver_version);
133 try writer.writeU64(result_allocator, record.key.m);
134 try writer.writeU64(result_allocator, record.key.n);
135 try writer.writeU64(result_allocator, record.key.k);
136 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.dtype);
137 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.accumulation_dtype);
138 try writer.writeU32(result_allocator, record.key.family_version);
139 try writer.writeU32(result_allocator, record.key.candidate_count);
140 try writer.writeU64(result_allocator, record.key.candidate_set_fingerprint);
141 try writer.writeU32(result_allocator, record.selection.threads.x);
142 try writer.writeU32(result_allocator, record.selection.threads.y);
143 try writer.writeU64(result_allocator, record.selection.winner_median_ns);
144 try writer.writeU64(result_allocator, record.selection.runner_up_median_ns);
145 try writer.writeU32(result_allocator, record.selection.sample_count);
146 }
147
148 return writer.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
149 }
150
151 pub fn decodeMatrixProductFamilyScheduleTuningArtifact(
152 result_allocator: std.mem.Allocator,
153 bytes: []const u8,
154 ) gpu.BackendError![]MatrixProductFamilyScheduleTuningRecord {
155 var reader = artifact_product.wire.ByteReader{ .bytes = bytes };
156 if ((try reader.readU32()) != matrix_product_family_schedule_tuning_artifact_magic) return error.InvalidArtifact;
157 if ((try reader.readU32()) != matrix_product_family_schedule_tuning_artifact_version) return error.InvalidArtifact;
158 const record_count = try reader.readU32();
159
160 var records = std.ArrayListUnmanaged(MatrixProductFamilyScheduleTuningRecord).empty;
161 errdefer records.deinit(result_allocator);
162
163 var index: u32 = 0;
164 while (index < record_count) : (index += 1) {
165 const record = MatrixProductFamilyScheduleTuningRecord{
166 .version = try reader.readU32(),
167 .key = .{
168 .backend = try reader.readEnum(gpu.BackendKind),
169 .family = try reader.readEnum(gpu.DeviceFamily),
170 .format = try reader.readEnum(gpu.ArtifactFormat),
171 .vendor_id = try reader.readU32(),
172 .has_vendor_id = try reader.readBool(),
173 .device_id = try reader.readU32(),
174 .has_device_id = try reader.readBool(),
175 .name_fingerprint = try reader.readU64(),
176 .driver_version_fingerprint = try reader.readU64(),
177 .has_driver_version = try reader.readBool(),
178 .m = try reader.readU64(),
179 .n = try reader.readU64(),
180 .k = try reader.readU64(),
181 .dtype = try reader.readEnum(choir_abi.DType),
182 .accumulation_dtype = try reader.readEnum(choir_abi.DType),
183 .family_version = try reader.readU32(),
184 .candidate_count = try reader.readU32(),
185 .candidate_set_fingerprint = try reader.readU64(),
186 },
187 .selection = .{
188 .threads = .{
189 .x = try reader.readU32(),
190 .y = try reader.readU32(),
191 },
192 .winner_median_ns = try reader.readU64(),
193 .runner_up_median_ns = try reader.readU64(),
194 .sample_count = try reader.readU32(),
195 },
196 };
197 try validateMatrixProductFamilyScheduleTuningRecord(record);
198 records.append(result_allocator, record) catch return error.OutOfMemory;
199 }
200 try reader.expectDone();
201 return records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
202 }
203
204 pub fn matrixProductFamilyScheduleTuningArtifactFingerprint(bytes: []const u8) choir.product.incremental.Fingerprint {
205 var builder = choir.product.incremental.FingerprintBuilder{};
206 builder.updateBytes(matrix_product_family_schedule_tuning_product_name);
207 builder.updateU32(matrix_product_family_schedule_tuning_artifact_version);
208 builder.updateBytes(bytes);
209 return builder.finish();
210 }
211
212 fn validateMatrixProductFamilyScheduleTuningRecord(record: MatrixProductFamilyScheduleTuningRecord) gpu.BackendError!void {
213 if (record.version != matrix_product_family_schedule_tuning_record_version) return error.InvalidArtifact;
214 if (record.key.m == 0 or record.key.n == 0 or record.key.k == 0) return error.InvalidArtifact;
215 if (record.key.family_version == 0) return error.InvalidArtifact;
216 if (record.key.candidate_count < 2) return error.InvalidArtifact;
217 if (record.key.candidate_count > matrix_product_family_schedule_tuning_max_candidates) return error.InvalidArtifact;
218 if (record.key.candidate_set_fingerprint == 0) return error.InvalidArtifact;
219 try validateMatrixProductFamilyScheduleTuningSelection(record.selection);
220 }
221
222 fn validateMatrixProductFamilyScheduleTuningSelection(selection: MatrixProductFamilyScheduleTuningSelection) gpu.BackendError!void {
223 try validateMatrixProductFamilyScheduleThreads(selection.threads);
224 if (selection.winner_median_ns == 0) return error.InvalidArtifact;
225 if (selection.runner_up_median_ns == 0) return error.InvalidArtifact;
226 if (selection.runner_up_median_ns < selection.winner_median_ns) return error.InvalidArtifact;
227 if (selection.sample_count == 0) return error.InvalidArtifact;
228 }
229
230 fn validateMatrixProductFamilyScheduleThreads(threads: MatrixProductFamilyScheduleThreads) gpu.BackendError!void {
231 if (threads.x == 0 or threads.y == 0) return error.InvalidArtifact;
232 }
233
234 fn matrixProductFamilyScheduleCandidatesContain(
235 candidates: []const MatrixProductFamilyScheduleThreads,
236 threads: MatrixProductFamilyScheduleThreads,
237 ) bool {
238 for (candidates) |candidate| {
239 if (candidate.eql(threads)) return true;
240 }
241 return false;
242 }
243
244 fn matrixProductFamilyScheduleSelectionBeats(
245 lhs: MatrixProductFamilyScheduleTuningSelection,
246 rhs: MatrixProductFamilyScheduleTuningSelection,
247 ) bool {
248 if (lhs.winner_median_ns != rhs.winner_median_ns) return lhs.winner_median_ns < rhs.winner_median_ns;
249 if (lhs.runner_up_median_ns != rhs.runner_up_median_ns) return lhs.runner_up_median_ns > rhs.runner_up_median_ns;
250 if (lhs.sample_count != rhs.sample_count) return lhs.sample_count > rhs.sample_count;
251 return MatrixProductFamilyScheduleThreads.lessThan({}, lhs.threads, rhs.threads);
252 }
253
254 fn matrixProductFamilyScheduleTuningRecordSortsBefore(
255 _: void,
256 lhs: MatrixProductFamilyScheduleTuningRecord,
257 rhs: MatrixProductFamilyScheduleTuningRecord,
258 ) bool {
259 if (matrixProductFamilyScheduleTuningKeySortsBefore(lhs.key, rhs.key)) return true;
260 if (matrixProductFamilyScheduleTuningKeySortsBefore(rhs.key, lhs.key)) return false;
261 return matrixProductFamilyScheduleTuningSelectionSortsBefore(lhs.selection, rhs.selection);
262 }
263
264 fn matrixProductFamilyScheduleTuningKeySortsBefore(
265 lhs: MatrixProductFamilyScheduleTuningKey,
266 rhs: MatrixProductFamilyScheduleTuningKey,
267 ) bool {
268 if (compareEnum(gpu.BackendKind, lhs.backend, rhs.backend)) |less| return less;
269 if (compareEnum(gpu.DeviceFamily, lhs.family, rhs.family)) |less| return less;
270 if (compareEnum(gpu.ArtifactFormat, lhs.format, rhs.format)) |less| return less;
271 if (compareU32(lhs.vendor_id, rhs.vendor_id)) |less| return less;
272 if (compareBool(lhs.has_vendor_id, rhs.has_vendor_id)) |less| return less;
273 if (compareU32(lhs.device_id, rhs.device_id)) |less| return less;
274 if (compareBool(lhs.has_device_id, rhs.has_device_id)) |less| return less;
275 if (compareU64(lhs.name_fingerprint, rhs.name_fingerprint)) |less| return less;
276 if (compareU64(lhs.driver_version_fingerprint, rhs.driver_version_fingerprint)) |less| return less;
277 if (compareBool(lhs.has_driver_version, rhs.has_driver_version)) |less| return less;
278 if (compareU64(lhs.m, rhs.m)) |less| return less;
279 if (compareU64(lhs.n, rhs.n)) |less| return less;
280 if (compareU64(lhs.k, rhs.k)) |less| return less;
281 if (compareEnum(choir_abi.DType, lhs.dtype, rhs.dtype)) |less| return less;
282 if (compareEnum(choir_abi.DType, lhs.accumulation_dtype, rhs.accumulation_dtype)) |less| return less;
283 if (compareU32(lhs.family_version, rhs.family_version)) |less| return less;
284 if (compareU32(lhs.candidate_count, rhs.candidate_count)) |less| return less;
285 if (compareU64(lhs.candidate_set_fingerprint, rhs.candidate_set_fingerprint)) |less| return less;
286 return false;
287 }
288
289 fn matrixProductFamilyScheduleTuningSelectionSortsBefore(
290 lhs: MatrixProductFamilyScheduleTuningSelection,
291 rhs: MatrixProductFamilyScheduleTuningSelection,
292 ) bool {
293 if (MatrixProductFamilyScheduleThreads.lessThan({}, lhs.threads, rhs.threads)) return true;
294 if (MatrixProductFamilyScheduleThreads.lessThan({}, rhs.threads, lhs.threads)) return false;
295 if (compareU64(lhs.winner_median_ns, rhs.winner_median_ns)) |less| return less;
296 if (compareU64(lhs.runner_up_median_ns, rhs.runner_up_median_ns)) |less| return less;
297 if (compareU32(lhs.sample_count, rhs.sample_count)) |less| return less;
298 return false;
299 }
300
301 pub const generated_scan_schedule_tuning_product_name = "accy.exec.generated_scan_schedule_tuning";
302 pub const generated_scan_schedule_tuning_record_version: u32 = 1;
303 pub const generated_scan_schedule_tuning_artifact_magic: u32 = 0x41515434;
304 pub const generated_scan_schedule_tuning_artifact_version: u32 = 1;
305 pub const generated_scan_schedule_tuning_max_candidates: usize = 4;
306
307 pub const GeneratedScanScheduleTuningProblem = struct {
308 format: gpu.ArtifactFormat,
309 total: u64,
310 dtype: choir_abi.DType = .f32,
311 schedule_version: u32,
312 candidates: []const GeneratedScanSchedule,
313 };
314
315 /// The key records the identity a measured scan schedule choice is stored
316 /// under: backend, device family, artifact format, vendor and device ids when
317 /// known, fingerprints of the device name and driver version, the element
318 /// count, the element type, the schedule version, and the candidate set. A
319 /// tuning run stores its winning scan schedule under this key, and a later
320 /// compile rebuilds the key to find it. A scan schedule is a thread count and
321 /// an item count per thread. `init` returns `error.InvalidArtifact` for a zero
322 /// element count, a zero schedule version, or a candidate set outside two to
323 /// four distinct valid schedules. `eql` compares every field, so a different
324 /// driver version gives a different key. A compile reading a tuning file skips
325 /// any record whose key does not match the key it rebuilds for the current
326 /// device and problem, so a stale or foreign record changes nothing and the
327 /// compile proceeds as if no record existed.
328 pub const GeneratedScanScheduleTuningKey = struct {
329 backend: gpu.BackendKind,
330 family: gpu.DeviceFamily,
331 format: gpu.ArtifactFormat,
332 vendor_id: u32 = 0,
333 has_vendor_id: bool = false,
334 device_id: u32 = 0,
335 has_device_id: bool = false,
336 name_fingerprint: u64 = 0,
337 driver_version_fingerprint: u64 = 0,
338 has_driver_version: bool = false,
339 total: u64,
340 dtype: choir_abi.DType,
341 schedule_version: u32,
342 candidate_count: u32,
343 candidate_set_fingerprint: u64,
344
345 pub fn init(
346 caps: gpu.BackendCapabilities,
347 problem: GeneratedScanScheduleTuningProblem,
348 ) gpu.BackendError!GeneratedScanScheduleTuningKey {
349 if (problem.total == 0) return error.InvalidArtifact;
350 if (problem.schedule_version == 0) return error.InvalidArtifact;
351 const candidate_set_fingerprint = try generatedScanScheduleCandidateSetFingerprint(problem.candidates);
352 return .{
353 .backend = caps.identity.backend,
354 .family = caps.identity.family,
355 .format = problem.format,
356 .vendor_id = caps.identity.vendor_id orelse 0,
357 .has_vendor_id = caps.identity.vendor_id != null,
358 .device_id = caps.identity.device_id orelse 0,
359 .has_device_id = caps.identity.device_id != null,
360 .name_fingerprint = bytesFingerprint(generated_scan_schedule_tuning_product_name ++ ".device.name", caps.identity.name),
361 .driver_version_fingerprint = if (caps.identity.driver_version) |version| bytesFingerprint(generated_scan_schedule_tuning_product_name ++ ".driver.version", version) else 0,
362 .has_driver_version = caps.identity.driver_version != null,
363 .total = problem.total,
364 .dtype = problem.dtype,
365 .schedule_version = problem.schedule_version,
366 .candidate_count = @intCast(problem.candidates.len),
367 .candidate_set_fingerprint = candidate_set_fingerprint,
368 };
369 }
370
371 pub fn eql(self: GeneratedScanScheduleTuningKey, other: GeneratedScanScheduleTuningKey) bool {
372 return self.backend == other.backend and
373 self.family == other.family and
374 self.format == other.format and
375 self.vendor_id == other.vendor_id and
376 self.has_vendor_id == other.has_vendor_id and
377 self.device_id == other.device_id and
378 self.has_device_id == other.has_device_id and
379 self.name_fingerprint == other.name_fingerprint and
380 self.driver_version_fingerprint == other.driver_version_fingerprint and
381 self.has_driver_version == other.has_driver_version and
382 self.total == other.total and
383 self.dtype == other.dtype and
384 self.schedule_version == other.schedule_version and
385 self.candidate_count == other.candidate_count and
386 self.candidate_set_fingerprint == other.candidate_set_fingerprint;
387 }
388 };
389
390 pub const GeneratedScanScheduleTuningSelection = struct {
391 schedule: GeneratedScanSchedule,
392 winner_median_ns: u64,
393 runner_up_median_ns: u64,
394 sample_count: u32,
395 };
396
397 pub const GeneratedScanScheduleTuningRecord = struct {
398 version: u32 = generated_scan_schedule_tuning_record_version,
399 key: GeneratedScanScheduleTuningKey,
400 selection: GeneratedScanScheduleTuningSelection,
401 };
402
403 pub const GeneratedScanScheduleTuningCache = struct {
404 selections: std.AutoHashMap(GeneratedScanScheduleTuningKey, GeneratedScanScheduleTuningSelection),
405
406 pub fn init(allocator: std.mem.Allocator) GeneratedScanScheduleTuningCache {
407 return .{
408 .selections = std.AutoHashMap(GeneratedScanScheduleTuningKey, GeneratedScanScheduleTuningSelection).init(allocator),
409 };
410 }
411
412 pub fn deinit(self: *GeneratedScanScheduleTuningCache) void {
413 self.selections.deinit();
414 self.* = undefined;
415 }
416
417 pub fn count(self: *const GeneratedScanScheduleTuningCache) usize {
418 return self.selections.count();
419 }
420
421 pub fn exportRecords(
422 self: *const GeneratedScanScheduleTuningCache,
423 result_allocator: std.mem.Allocator,
424 ) gpu.BackendError![]GeneratedScanScheduleTuningRecord {
425 const records = result_allocator.alloc(GeneratedScanScheduleTuningRecord, self.selections.count()) catch return error.OutOfMemory;
426 errdefer result_allocator.free(records);
427
428 var iterator = self.selections.iterator();
429 var index: usize = 0;
430 while (iterator.next()) |entry| {
431 records[index] = .{
432 .key = entry.key_ptr.*,
433 .selection = entry.value_ptr.*,
434 };
435 index += 1;
436 }
437 std.mem.sort(GeneratedScanScheduleTuningRecord, records, {}, generatedScanScheduleTuningRecordSortsBefore);
438 return records;
439 }
440
441 pub fn importRecords(
442 self: *GeneratedScanScheduleTuningCache,
443 records: []const GeneratedScanScheduleTuningRecord,
444 ) gpu.BackendError!void {
445 for (records) |record| try validateGeneratedScanScheduleTuningRecord(record);
446 for (records) |record| {
447 if (self.selections.getPtr(record.key)) |existing| {
448 if (generatedScanScheduleSelectionBeats(record.selection, existing.*)) existing.* = record.selection;
449 continue;
450 }
451 self.selections.put(record.key, record.selection) catch return error.OutOfMemory;
452 }
453 }
454
455 pub fn recordSelection(
456 self: *GeneratedScanScheduleTuningCache,
457 caps: gpu.BackendCapabilities,
458 problem: GeneratedScanScheduleTuningProblem,
459 selection: GeneratedScanScheduleTuningSelection,
460 ) gpu.BackendError!void {
461 try validateGeneratedScanScheduleTuningSelection(selection);
462 if (!generatedScanScheduleCandidatesContain(problem.candidates, selection.schedule)) return error.LaunchArgumentMismatch;
463 const key = try GeneratedScanScheduleTuningKey.init(caps, problem);
464 if (self.selections.getPtr(key)) |existing| {
465 if (generatedScanScheduleSelectionBeats(selection, existing.*)) existing.* = selection;
466 return;
467 }
468 self.selections.put(key, selection) catch return error.OutOfMemory;
469 }
470
471 pub fn selectionForProblem(
472 self: *const GeneratedScanScheduleTuningCache,
473 caps: gpu.BackendCapabilities,
474 problem: GeneratedScanScheduleTuningProblem,
475 ) gpu.BackendError!?GeneratedScanScheduleTuningSelection {
476 const key = try GeneratedScanScheduleTuningKey.init(caps, problem);
477 return self.selections.get(key);
478 }
479 };
480
481 pub fn generatedScanScheduleCandidateSetFingerprint(
482 candidates: []const GeneratedScanSchedule,
483 ) gpu.BackendError!choir.product.incremental.Fingerprint {
484 const set = try sortedGeneratedScanScheduleCandidates(candidates);
485 var builder = choir.product.incremental.FingerprintBuilder{};
486 builder.updateBytes(generated_scan_schedule_tuning_product_name);
487 builder.updateU32(@intCast(set.count));
488 for (set.slice()) |candidate| {
489 builder.updateU32(candidate.threads);
490 builder.updateU32(candidate.items);
491 }
492 return builder.finish();
493 }
494
495 pub fn encodeGeneratedScanScheduleTuningArtifact(
496 result_allocator: std.mem.Allocator,
497 records: []const GeneratedScanScheduleTuningRecord,
498 ) gpu.BackendError![]u8 {
499 if (records.len > std.math.maxInt(u32)) return error.InvalidArtifact;
500 var writer = artifact_product.wire.ByteWriter{};
501 errdefer writer.deinit(result_allocator);
502
503 try writer.writeU32(result_allocator, generated_scan_schedule_tuning_artifact_magic);
504 try writer.writeU32(result_allocator, generated_scan_schedule_tuning_artifact_version);
505 try writer.writeU32(result_allocator, @intCast(records.len));
506 for (records) |record| {
507 try validateGeneratedScanScheduleTuningRecord(record);
508 try writer.writeU32(result_allocator, record.version);
509 try writer.writeEnum(result_allocator, gpu.BackendKind, record.key.backend);
510 try writer.writeEnum(result_allocator, gpu.DeviceFamily, record.key.family);
511 try writer.writeEnum(result_allocator, gpu.ArtifactFormat, record.key.format);
512 try writer.writeU32(result_allocator, record.key.vendor_id);
513 try writer.writeBool(result_allocator, record.key.has_vendor_id);
514 try writer.writeU32(result_allocator, record.key.device_id);
515 try writer.writeBool(result_allocator, record.key.has_device_id);
516 try writer.writeU64(result_allocator, record.key.name_fingerprint);
517 try writer.writeU64(result_allocator, record.key.driver_version_fingerprint);
518 try writer.writeBool(result_allocator, record.key.has_driver_version);
519 try writer.writeU64(result_allocator, record.key.total);
520 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.dtype);
521 try writer.writeU32(result_allocator, record.key.schedule_version);
522 try writer.writeU32(result_allocator, record.key.candidate_count);
523 try writer.writeU64(result_allocator, record.key.candidate_set_fingerprint);
524 try writer.writeU32(result_allocator, record.selection.schedule.threads);
525 try writer.writeU32(result_allocator, record.selection.schedule.items);
526 try writer.writeU64(result_allocator, record.selection.winner_median_ns);
527 try writer.writeU64(result_allocator, record.selection.runner_up_median_ns);
528 try writer.writeU32(result_allocator, record.selection.sample_count);
529 }
530
531 return writer.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
532 }
533
534 pub fn decodeGeneratedScanScheduleTuningArtifact(
535 result_allocator: std.mem.Allocator,
536 bytes: []const u8,
537 ) gpu.BackendError![]GeneratedScanScheduleTuningRecord {
538 var reader = artifact_product.wire.ByteReader{ .bytes = bytes };
539 if ((try reader.readU32()) != generated_scan_schedule_tuning_artifact_magic) return error.InvalidArtifact;
540 if ((try reader.readU32()) != generated_scan_schedule_tuning_artifact_version) return error.InvalidArtifact;
541 const record_count = try reader.readU32();
542
543 var records = std.ArrayListUnmanaged(GeneratedScanScheduleTuningRecord).empty;
544 errdefer records.deinit(result_allocator);
545
546 var index: u32 = 0;
547 while (index < record_count) : (index += 1) {
548 const record = GeneratedScanScheduleTuningRecord{
549 .version = try reader.readU32(),
550 .key = .{
551 .backend = try reader.readEnum(gpu.BackendKind),
552 .family = try reader.readEnum(gpu.DeviceFamily),
553 .format = try reader.readEnum(gpu.ArtifactFormat),
554 .vendor_id = try reader.readU32(),
555 .has_vendor_id = try reader.readBool(),
556 .device_id = try reader.readU32(),
557 .has_device_id = try reader.readBool(),
558 .name_fingerprint = try reader.readU64(),
559 .driver_version_fingerprint = try reader.readU64(),
560 .has_driver_version = try reader.readBool(),
561 .total = try reader.readU64(),
562 .dtype = try reader.readEnum(choir_abi.DType),
563 .schedule_version = try reader.readU32(),
564 .candidate_count = try reader.readU32(),
565 .candidate_set_fingerprint = try reader.readU64(),
566 },
567 .selection = .{
568 .schedule = .{
569 .threads = try reader.readU32(),
570 .items = try reader.readU32(),
571 },
572 .winner_median_ns = try reader.readU64(),
573 .runner_up_median_ns = try reader.readU64(),
574 .sample_count = try reader.readU32(),
575 },
576 };
577 try validateGeneratedScanScheduleTuningRecord(record);
578 records.append(result_allocator, record) catch return error.OutOfMemory;
579 }
580 try reader.expectDone();
581 return records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
582 }
583
584 pub fn generatedScanScheduleTuningArtifactFingerprint(bytes: []const u8) choir.product.incremental.Fingerprint {
585 var builder = choir.product.incremental.FingerprintBuilder{};
586 builder.updateBytes(generated_scan_schedule_tuning_product_name);
587 builder.updateU32(generated_scan_schedule_tuning_artifact_version);
588 builder.updateBytes(bytes);
589 return builder.finish();
590 }
591
592 pub fn generatedScanScheduleTuningArtifactProductStamp(bytes: []const u8) choir.product.incremental.ProductStamp {
593 return choir.product.incremental.productStamp(generated_scan_schedule_tuning_product_name, generatedScanScheduleTuningArtifactFingerprint(bytes));
594 }
595
596 fn validateGeneratedScanScheduleTuningRecord(record: GeneratedScanScheduleTuningRecord) gpu.BackendError!void {
597 if (record.version != generated_scan_schedule_tuning_record_version) return error.InvalidArtifact;
598 if (record.key.total == 0) return error.InvalidArtifact;
599 if (record.key.schedule_version == 0) return error.InvalidArtifact;
600 if (record.key.candidate_count < 2) return error.InvalidArtifact;
601 if (record.key.candidate_count > generated_scan_schedule_tuning_max_candidates) return error.InvalidArtifact;
602 if (record.key.candidate_set_fingerprint == 0) return error.InvalidArtifact;
603 try validateGeneratedScanScheduleTuningSelection(record.selection);
604 }
605
606 fn validateGeneratedScanScheduleTuningSelection(selection: GeneratedScanScheduleTuningSelection) gpu.BackendError!void {
607 try validateGeneratedScanSchedule(selection.schedule);
608 if (selection.winner_median_ns == 0) return error.InvalidArtifact;
609 if (selection.runner_up_median_ns == 0) return error.InvalidArtifact;
610 if (selection.runner_up_median_ns < selection.winner_median_ns) return error.InvalidArtifact;
611 if (selection.sample_count == 0) return error.InvalidArtifact;
612 }
613
614 fn validateGeneratedScanSchedule(schedule: GeneratedScanSchedule) gpu.BackendError!void {
615 if (schedule.threads == 0 or schedule.items == 0) return error.InvalidArtifact;
616 }
617
618 const GeneratedScanScheduleCandidateSet = struct {
619 count: usize = 0,
620 items: [generated_scan_schedule_tuning_max_candidates]GeneratedScanSchedule = @as([generated_scan_schedule_tuning_max_candidates]GeneratedScanSchedule, @splat(.{ .threads = 1, .items = 1 })),
621
622 fn slice(self: *const GeneratedScanScheduleCandidateSet) []const GeneratedScanSchedule {
623 return self.items[0..self.count];
624 }
625 };
626
627 fn sortedGeneratedScanScheduleCandidates(
628 candidates: []const GeneratedScanSchedule,
629 ) gpu.BackendError!GeneratedScanScheduleCandidateSet {
630 if (candidates.len < 2) return error.InvalidArtifact;
631 if (candidates.len > generated_scan_schedule_tuning_max_candidates) return error.InvalidArtifact;
632 var set = GeneratedScanScheduleCandidateSet{ .count = candidates.len };
633 @memcpy(set.items[0..candidates.len], candidates);
634 for (set.slice()) |candidate| try validateGeneratedScanSchedule(candidate);
635 std.mem.sort(GeneratedScanSchedule, set.items[0..set.count], {}, generatedScanScheduleSortsBefore);
636 for (set.slice()[1..], 1..) |candidate, index| {
637 if (candidate.eql(set.items[index - 1])) return error.InvalidArtifact;
638 }
639 return set;
640 }
641
642 fn generatedScanScheduleCandidatesContain(
643 candidates: []const GeneratedScanSchedule,
644 schedule: GeneratedScanSchedule,
645 ) bool {
646 for (candidates) |candidate| {
647 if (candidate.eql(schedule)) return true;
648 }
649 return false;
650 }
651
652 fn generatedScanScheduleSelectionBeats(
653 lhs: GeneratedScanScheduleTuningSelection,
654 rhs: GeneratedScanScheduleTuningSelection,
655 ) bool {
656 if (lhs.winner_median_ns != rhs.winner_median_ns) return lhs.winner_median_ns < rhs.winner_median_ns;
657 if (lhs.runner_up_median_ns != rhs.runner_up_median_ns) return lhs.runner_up_median_ns > rhs.runner_up_median_ns;
658 if (lhs.sample_count != rhs.sample_count) return lhs.sample_count > rhs.sample_count;
659 return generatedScanScheduleSortsBefore({}, lhs.schedule, rhs.schedule);
660 }
661
662 fn generatedScanScheduleTuningRecordSortsBefore(
663 _: void,
664 lhs: GeneratedScanScheduleTuningRecord,
665 rhs: GeneratedScanScheduleTuningRecord,
666 ) bool {
667 if (generatedScanScheduleTuningKeySortsBefore(lhs.key, rhs.key)) return true;
668 if (generatedScanScheduleTuningKeySortsBefore(rhs.key, lhs.key)) return false;
669 return generatedScanScheduleTuningSelectionSortsBefore(lhs.selection, rhs.selection);
670 }
671
672 fn generatedScanScheduleTuningKeySortsBefore(
673 lhs: GeneratedScanScheduleTuningKey,
674 rhs: GeneratedScanScheduleTuningKey,
675 ) bool {
676 if (compareEnum(gpu.BackendKind, lhs.backend, rhs.backend)) |less| return less;
677 if (compareEnum(gpu.DeviceFamily, lhs.family, rhs.family)) |less| return less;
678 if (compareEnum(gpu.ArtifactFormat, lhs.format, rhs.format)) |less| return less;
679 if (compareU32(lhs.vendor_id, rhs.vendor_id)) |less| return less;
680 if (compareBool(lhs.has_vendor_id, rhs.has_vendor_id)) |less| return less;
681 if (compareU32(lhs.device_id, rhs.device_id)) |less| return less;
682 if (compareBool(lhs.has_device_id, rhs.has_device_id)) |less| return less;
683 if (compareU64(lhs.name_fingerprint, rhs.name_fingerprint)) |less| return less;
684 if (compareU64(lhs.driver_version_fingerprint, rhs.driver_version_fingerprint)) |less| return less;
685 if (compareBool(lhs.has_driver_version, rhs.has_driver_version)) |less| return less;
686 if (compareU64(lhs.total, rhs.total)) |less| return less;
687 if (compareEnum(choir_abi.DType, lhs.dtype, rhs.dtype)) |less| return less;
688 if (compareU32(lhs.schedule_version, rhs.schedule_version)) |less| return less;
689 if (compareU32(lhs.candidate_count, rhs.candidate_count)) |less| return less;
690 if (compareU64(lhs.candidate_set_fingerprint, rhs.candidate_set_fingerprint)) |less| return less;
691 return false;
692 }
693
694 fn generatedScanScheduleTuningSelectionSortsBefore(
695 lhs: GeneratedScanScheduleTuningSelection,
696 rhs: GeneratedScanScheduleTuningSelection,
697 ) bool {
698 if (generatedScanScheduleSortsBefore({}, lhs.schedule, rhs.schedule)) return true;
699 if (generatedScanScheduleSortsBefore({}, rhs.schedule, lhs.schedule)) return false;
700 if (compareU64(lhs.winner_median_ns, rhs.winner_median_ns)) |less| return less;
701 if (compareU64(lhs.runner_up_median_ns, rhs.runner_up_median_ns)) |less| return less;
702 if (compareU32(lhs.sample_count, rhs.sample_count)) |less| return less;
703 return false;
704 }
705
706 fn generatedScanScheduleSortsBefore(
707 _: void,
708 lhs: GeneratedScanSchedule,
709 rhs: GeneratedScanSchedule,
710 ) bool {
711 if (compareU32(lhs.threads, rhs.threads)) |less| return less;
712 if (compareU32(lhs.items, rhs.items)) |less| return less;
713 return false;
714 }
715
716 fn compareEnum(comptime T: type, lhs: T, rhs: T) ?bool {
717 return compareU64(@backingInt(lhs), @backingInt(rhs));
718 }
719
720 fn compareBool(lhs: bool, rhs: bool) ?bool {
721 if (lhs == rhs) return null;
722 return !lhs and rhs;
723 }
724
725 fn compareU32(lhs: u32, rhs: u32) ?bool {
726 if (lhs == rhs) return null;
727 return lhs < rhs;
728 }
729
730 fn compareU64(lhs: u64, rhs: u64) ?bool {
731 if (lhs == rhs) return null;
732 return lhs < rhs;
733 }
734
735 fn bytesFingerprint(domain: []const u8, bytes: []const u8) choir.product.incremental.Fingerprint {
736 var builder = choir.product.incremental.FingerprintBuilder{};
737 builder.updateBytes(domain);
738 builder.updateBytes(bytes);
739 return builder.finish();
740 }
741
742 pub const generated_row_pipeline_schedule_tuning_product_name = "accy.exec.generated_row_pipeline_schedule_tuning";
743 pub const generated_row_pipeline_schedule_tuning_record_version: u32 = 1;
744 pub const generated_row_pipeline_schedule_tuning_artifact_magic: u32 = 0x41515435;
745 pub const generated_row_pipeline_schedule_tuning_artifact_version: u32 = 1;
746 pub const generated_row_pipeline_schedule_tuning_max_candidates: usize = 4;
747
748 pub const GeneratedRowPipelineScheduleTuningProblem = struct {
749 format: gpu.ArtifactFormat,
750 rows: u64,
751 cols: u64,
752 dtype: choir_abi.DType = .f32,
753 schedule_version: u32,
754 candidates: []const GeneratedRowPipelineSchedule,
755 };
756
757 /// The key records the identity a measured row-pipeline schedule choice is
758 /// stored under: the same device, format, element type, schedule version and
759 /// candidate-set fields as the scan key, with the row and column counts as the
760 /// problem size. A tuning run stores its winning row-pipeline schedule under
761 /// this key, and a later compile rebuilds the key to find it. A row-pipeline
762 /// schedule is a thread count. `init` returns `error.InvalidArtifact` for zero
763 /// rows, zero columns, a zero schedule version, or an invalid candidate set. A
764 /// compile reading a tuning file skips any record whose key does not match the
765 /// key it rebuilds, so a stale or foreign record changes nothing.
766 pub const GeneratedRowPipelineScheduleTuningKey = struct {
767 backend: gpu.BackendKind,
768 family: gpu.DeviceFamily,
769 format: gpu.ArtifactFormat,
770 vendor_id: u32 = 0,
771 has_vendor_id: bool = false,
772 device_id: u32 = 0,
773 has_device_id: bool = false,
774 name_fingerprint: u64 = 0,
775 driver_version_fingerprint: u64 = 0,
776 has_driver_version: bool = false,
777 rows: u64,
778 cols: u64,
779 dtype: choir_abi.DType,
780 schedule_version: u32,
781 candidate_count: u32,
782 candidate_set_fingerprint: u64,
783
784 pub fn init(
785 caps: gpu.BackendCapabilities,
786 problem: GeneratedRowPipelineScheduleTuningProblem,
787 ) gpu.BackendError!GeneratedRowPipelineScheduleTuningKey {
788 if (problem.rows == 0 or problem.cols == 0) return error.InvalidArtifact;
789 if (problem.schedule_version == 0) return error.InvalidArtifact;
790 const candidate_set_fingerprint = try generatedRowPipelineScheduleCandidateSetFingerprint(problem.candidates);
791 return .{
792 .backend = caps.identity.backend,
793 .family = caps.identity.family,
794 .format = problem.format,
795 .vendor_id = caps.identity.vendor_id orelse 0,
796 .has_vendor_id = caps.identity.vendor_id != null,
797 .device_id = caps.identity.device_id orelse 0,
798 .has_device_id = caps.identity.device_id != null,
799 .name_fingerprint = bytesFingerprint(generated_row_pipeline_schedule_tuning_product_name ++ ".device.name", caps.identity.name),
800 .driver_version_fingerprint = if (caps.identity.driver_version) |version| bytesFingerprint(generated_row_pipeline_schedule_tuning_product_name ++ ".driver.version", version) else 0,
801 .has_driver_version = caps.identity.driver_version != null,
802 .rows = problem.rows,
803 .cols = problem.cols,
804 .dtype = problem.dtype,
805 .schedule_version = problem.schedule_version,
806 .candidate_count = @intCast(problem.candidates.len),
807 .candidate_set_fingerprint = candidate_set_fingerprint,
808 };
809 }
810
811 pub fn eql(self: GeneratedRowPipelineScheduleTuningKey, other: GeneratedRowPipelineScheduleTuningKey) bool {
812 return self.backend == other.backend and
813 self.family == other.family and
814 self.format == other.format and
815 self.vendor_id == other.vendor_id and
816 self.has_vendor_id == other.has_vendor_id and
817 self.device_id == other.device_id and
818 self.has_device_id == other.has_device_id and
819 self.name_fingerprint == other.name_fingerprint and
820 self.driver_version_fingerprint == other.driver_version_fingerprint and
821 self.has_driver_version == other.has_driver_version and
822 self.rows == other.rows and
823 self.cols == other.cols and
824 self.dtype == other.dtype and
825 self.schedule_version == other.schedule_version and
826 self.candidate_count == other.candidate_count and
827 self.candidate_set_fingerprint == other.candidate_set_fingerprint;
828 }
829 };
830
831 pub const GeneratedRowPipelineScheduleTuningSelection = struct {
832 schedule: GeneratedRowPipelineSchedule,
833 winner_median_ns: u64,
834 runner_up_median_ns: u64,
835 sample_count: u32,
836 };
837
838 pub const GeneratedRowPipelineScheduleTuningRecord = struct {
839 version: u32 = generated_row_pipeline_schedule_tuning_record_version,
840 key: GeneratedRowPipelineScheduleTuningKey,
841 selection: GeneratedRowPipelineScheduleTuningSelection,
842 };
843
844 pub const GeneratedRowPipelineScheduleTuningCache = struct {
845 selections: std.AutoHashMap(GeneratedRowPipelineScheduleTuningKey, GeneratedRowPipelineScheduleTuningSelection),
846
847 pub fn init(allocator: std.mem.Allocator) GeneratedRowPipelineScheduleTuningCache {
848 return .{
849 .selections = std.AutoHashMap(GeneratedRowPipelineScheduleTuningKey, GeneratedRowPipelineScheduleTuningSelection).init(allocator),
850 };
851 }
852
853 pub fn deinit(self: *GeneratedRowPipelineScheduleTuningCache) void {
854 self.selections.deinit();
855 self.* = undefined;
856 }
857
858 pub fn count(self: *const GeneratedRowPipelineScheduleTuningCache) usize {
859 return self.selections.count();
860 }
861
862 pub fn exportRecords(
863 self: *const GeneratedRowPipelineScheduleTuningCache,
864 result_allocator: std.mem.Allocator,
865 ) gpu.BackendError![]GeneratedRowPipelineScheduleTuningRecord {
866 const records = result_allocator.alloc(GeneratedRowPipelineScheduleTuningRecord, self.selections.count()) catch return error.OutOfMemory;
867 errdefer result_allocator.free(records);
868
869 var iterator = self.selections.iterator();
870 var index: usize = 0;
871 while (iterator.next()) |entry| {
872 records[index] = .{
873 .key = entry.key_ptr.*,
874 .selection = entry.value_ptr.*,
875 };
876 index += 1;
877 }
878 std.mem.sort(GeneratedRowPipelineScheduleTuningRecord, records, {}, generatedRowPipelineScheduleTuningRecordSortsBefore);
879 return records;
880 }
881
882 pub fn importRecords(
883 self: *GeneratedRowPipelineScheduleTuningCache,
884 records: []const GeneratedRowPipelineScheduleTuningRecord,
885 ) gpu.BackendError!void {
886 for (records) |record| try validateGeneratedRowPipelineScheduleTuningRecord(record);
887 for (records) |record| {
888 if (self.selections.getPtr(record.key)) |existing| {
889 if (generatedRowPipelineScheduleSelectionBeats(record.selection, existing.*)) existing.* = record.selection;
890 continue;
891 }
892 self.selections.put(record.key, record.selection) catch return error.OutOfMemory;
893 }
894 }
895
896 pub fn recordSelection(
897 self: *GeneratedRowPipelineScheduleTuningCache,
898 caps: gpu.BackendCapabilities,
899 problem: GeneratedRowPipelineScheduleTuningProblem,
900 selection: GeneratedRowPipelineScheduleTuningSelection,
901 ) gpu.BackendError!void {
902 try validateGeneratedRowPipelineScheduleTuningSelection(selection);
903 if (!generatedRowPipelineScheduleCandidatesContain(problem.candidates, selection.schedule)) return error.LaunchArgumentMismatch;
904 const key = try GeneratedRowPipelineScheduleTuningKey.init(caps, problem);
905 if (self.selections.getPtr(key)) |existing| {
906 if (generatedRowPipelineScheduleSelectionBeats(selection, existing.*)) existing.* = selection;
907 return;
908 }
909 self.selections.put(key, selection) catch return error.OutOfMemory;
910 }
911
912 pub fn selectionForProblem(
913 self: *const GeneratedRowPipelineScheduleTuningCache,
914 caps: gpu.BackendCapabilities,
915 problem: GeneratedRowPipelineScheduleTuningProblem,
916 ) gpu.BackendError!?GeneratedRowPipelineScheduleTuningSelection {
917 const key = try GeneratedRowPipelineScheduleTuningKey.init(caps, problem);
918 return self.selections.get(key);
919 }
920 };
921
922 pub fn generatedRowPipelineScheduleCandidateSetFingerprint(
923 candidates: []const GeneratedRowPipelineSchedule,
924 ) gpu.BackendError!choir.product.incremental.Fingerprint {
925 const set = try sortedGeneratedRowPipelineScheduleCandidates(candidates);
926 var builder = choir.product.incremental.FingerprintBuilder{};
927 builder.updateBytes(generated_row_pipeline_schedule_tuning_product_name);
928 builder.updateU32(@intCast(set.count));
929 for (set.slice()) |candidate| {
930 builder.updateU32(candidate.threads);
931 }
932 return builder.finish();
933 }
934
935 pub fn encodeGeneratedRowPipelineScheduleTuningArtifact(
936 result_allocator: std.mem.Allocator,
937 records: []const GeneratedRowPipelineScheduleTuningRecord,
938 ) gpu.BackendError![]u8 {
939 if (records.len > std.math.maxInt(u32)) return error.InvalidArtifact;
940 var writer = artifact_product.wire.ByteWriter{};
941 errdefer writer.deinit(result_allocator);
942
943 try writer.writeU32(result_allocator, generated_row_pipeline_schedule_tuning_artifact_magic);
944 try writer.writeU32(result_allocator, generated_row_pipeline_schedule_tuning_artifact_version);
945 try writer.writeU32(result_allocator, @intCast(records.len));
946 for (records) |record| {
947 try validateGeneratedRowPipelineScheduleTuningRecord(record);
948 try writer.writeU32(result_allocator, record.version);
949 try writer.writeEnum(result_allocator, gpu.BackendKind, record.key.backend);
950 try writer.writeEnum(result_allocator, gpu.DeviceFamily, record.key.family);
951 try writer.writeEnum(result_allocator, gpu.ArtifactFormat, record.key.format);
952 try writer.writeU32(result_allocator, record.key.vendor_id);
953 try writer.writeBool(result_allocator, record.key.has_vendor_id);
954 try writer.writeU32(result_allocator, record.key.device_id);
955 try writer.writeBool(result_allocator, record.key.has_device_id);
956 try writer.writeU64(result_allocator, record.key.name_fingerprint);
957 try writer.writeU64(result_allocator, record.key.driver_version_fingerprint);
958 try writer.writeBool(result_allocator, record.key.has_driver_version);
959 try writer.writeU64(result_allocator, record.key.rows);
960 try writer.writeU64(result_allocator, record.key.cols);
961 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.dtype);
962 try writer.writeU32(result_allocator, record.key.schedule_version);
963 try writer.writeU32(result_allocator, record.key.candidate_count);
964 try writer.writeU64(result_allocator, record.key.candidate_set_fingerprint);
965 try writer.writeU32(result_allocator, record.selection.schedule.threads);
966 try writer.writeU64(result_allocator, record.selection.winner_median_ns);
967 try writer.writeU64(result_allocator, record.selection.runner_up_median_ns);
968 try writer.writeU32(result_allocator, record.selection.sample_count);
969 }
970
971 return writer.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
972 }
973
974 pub fn decodeGeneratedRowPipelineScheduleTuningArtifact(
975 result_allocator: std.mem.Allocator,
976 bytes: []const u8,
977 ) gpu.BackendError![]GeneratedRowPipelineScheduleTuningRecord {
978 var reader = artifact_product.wire.ByteReader{ .bytes = bytes };
979 if ((try reader.readU32()) != generated_row_pipeline_schedule_tuning_artifact_magic) return error.InvalidArtifact;
980 if ((try reader.readU32()) != generated_row_pipeline_schedule_tuning_artifact_version) return error.InvalidArtifact;
981 const record_count = try reader.readU32();
982
983 var records = std.ArrayListUnmanaged(GeneratedRowPipelineScheduleTuningRecord).empty;
984 errdefer records.deinit(result_allocator);
985
986 var index: u32 = 0;
987 while (index < record_count) : (index += 1) {
988 const record = GeneratedRowPipelineScheduleTuningRecord{
989 .version = try reader.readU32(),
990 .key = .{
991 .backend = try reader.readEnum(gpu.BackendKind),
992 .family = try reader.readEnum(gpu.DeviceFamily),
993 .format = try reader.readEnum(gpu.ArtifactFormat),
994 .vendor_id = try reader.readU32(),
995 .has_vendor_id = try reader.readBool(),
996 .device_id = try reader.readU32(),
997 .has_device_id = try reader.readBool(),
998 .name_fingerprint = try reader.readU64(),
999 .driver_version_fingerprint = try reader.readU64(),
1000 .has_driver_version = try reader.readBool(),
1001 .rows = try reader.readU64(),
1002 .cols = try reader.readU64(),
1003 .dtype = try reader.readEnum(choir_abi.DType),
1004 .schedule_version = try reader.readU32(),
1005 .candidate_count = try reader.readU32(),
1006 .candidate_set_fingerprint = try reader.readU64(),
1007 },
1008 .selection = .{
1009 .schedule = .{
1010 .threads = try reader.readU32(),
1011 },
1012 .winner_median_ns = try reader.readU64(),
1013 .runner_up_median_ns = try reader.readU64(),
1014 .sample_count = try reader.readU32(),
1015 },
1016 };
1017 try validateGeneratedRowPipelineScheduleTuningRecord(record);
1018 records.append(result_allocator, record) catch return error.OutOfMemory;
1019 }
1020 try reader.expectDone();
1021 return records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
1022 }
1023
1024 pub fn generatedRowPipelineScheduleTuningArtifactFingerprint(bytes: []const u8) choir.product.incremental.Fingerprint {
1025 var builder = choir.product.incremental.FingerprintBuilder{};
1026 builder.updateBytes(generated_row_pipeline_schedule_tuning_product_name);
1027 builder.updateU32(generated_row_pipeline_schedule_tuning_artifact_version);
1028 builder.updateBytes(bytes);
1029 return builder.finish();
1030 }
1031
1032 pub fn generatedRowPipelineScheduleTuningArtifactProductStamp(bytes: []const u8) choir.product.incremental.ProductStamp {
1033 return choir.product.incremental.productStamp(generated_row_pipeline_schedule_tuning_product_name, generatedRowPipelineScheduleTuningArtifactFingerprint(bytes));
1034 }
1035
1036 fn validateGeneratedRowPipelineScheduleTuningRecord(record: GeneratedRowPipelineScheduleTuningRecord) gpu.BackendError!void {
1037 if (record.version != generated_row_pipeline_schedule_tuning_record_version) return error.InvalidArtifact;
1038 if (record.key.rows == 0 or record.key.cols == 0) return error.InvalidArtifact;
1039 if (record.key.schedule_version == 0) return error.InvalidArtifact;
1040 if (record.key.candidate_count < 2) return error.InvalidArtifact;
1041 if (record.key.candidate_count > generated_row_pipeline_schedule_tuning_max_candidates) return error.InvalidArtifact;
1042 if (record.key.candidate_set_fingerprint == 0) return error.InvalidArtifact;
1043 try validateGeneratedRowPipelineScheduleTuningSelection(record.selection);
1044 }
1045
1046 fn validateGeneratedRowPipelineScheduleTuningSelection(selection: GeneratedRowPipelineScheduleTuningSelection) gpu.BackendError!void {
1047 try validateGeneratedRowPipelineSchedule(selection.schedule);
1048 if (selection.winner_median_ns == 0) return error.InvalidArtifact;
1049 if (selection.runner_up_median_ns == 0) return error.InvalidArtifact;
1050 if (selection.runner_up_median_ns < selection.winner_median_ns) return error.InvalidArtifact;
1051 if (selection.sample_count == 0) return error.InvalidArtifact;
1052 }
1053
1054 fn validateGeneratedRowPipelineSchedule(schedule: GeneratedRowPipelineSchedule) gpu.BackendError!void {
1055 if (schedule.threads == 0) return error.InvalidArtifact;
1056 }
1057
1058 const GeneratedRowPipelineScheduleCandidateSet = struct {
1059 count: usize = 0,
1060 items: [generated_row_pipeline_schedule_tuning_max_candidates]GeneratedRowPipelineSchedule = @as([generated_row_pipeline_schedule_tuning_max_candidates]GeneratedRowPipelineSchedule, @splat(.{ .threads = 1 })),
1061
1062 fn slice(self: *const GeneratedRowPipelineScheduleCandidateSet) []const GeneratedRowPipelineSchedule {
1063 return self.items[0..self.count];
1064 }
1065 };
1066
1067 fn sortedGeneratedRowPipelineScheduleCandidates(
1068 candidates: []const GeneratedRowPipelineSchedule,
1069 ) gpu.BackendError!GeneratedRowPipelineScheduleCandidateSet {
1070 if (candidates.len < 2) return error.InvalidArtifact;
1071 if (candidates.len > generated_row_pipeline_schedule_tuning_max_candidates) return error.InvalidArtifact;
1072 var set = GeneratedRowPipelineScheduleCandidateSet{ .count = candidates.len };
1073 @memcpy(set.items[0..candidates.len], candidates);
1074 for (set.slice()) |candidate| try validateGeneratedRowPipelineSchedule(candidate);
1075 std.mem.sort(GeneratedRowPipelineSchedule, set.items[0..set.count], {}, generatedRowPipelineScheduleSortsBefore);
1076 for (set.slice()[1..], 1..) |candidate, index| {
1077 if (candidate.eql(set.items[index - 1])) return error.InvalidArtifact;
1078 }
1079 return set;
1080 }
1081
1082 fn generatedRowPipelineScheduleCandidatesContain(
1083 candidates: []const GeneratedRowPipelineSchedule,
1084 schedule: GeneratedRowPipelineSchedule,
1085 ) bool {
1086 for (candidates) |candidate| {
1087 if (candidate.eql(schedule)) return true;
1088 }
1089 return false;
1090 }
1091
1092 fn generatedRowPipelineScheduleSelectionBeats(
1093 lhs: GeneratedRowPipelineScheduleTuningSelection,
1094 rhs: GeneratedRowPipelineScheduleTuningSelection,
1095 ) bool {
1096 if (lhs.winner_median_ns != rhs.winner_median_ns) return lhs.winner_median_ns < rhs.winner_median_ns;
1097 if (lhs.runner_up_median_ns != rhs.runner_up_median_ns) return lhs.runner_up_median_ns > rhs.runner_up_median_ns;
1098 if (lhs.sample_count != rhs.sample_count) return lhs.sample_count > rhs.sample_count;
1099 return generatedRowPipelineScheduleSortsBefore({}, lhs.schedule, rhs.schedule);
1100 }
1101
1102 fn generatedRowPipelineScheduleTuningRecordSortsBefore(
1103 _: void,
1104 lhs: GeneratedRowPipelineScheduleTuningRecord,
1105 rhs: GeneratedRowPipelineScheduleTuningRecord,
1106 ) bool {
1107 if (generatedRowPipelineScheduleTuningKeySortsBefore(lhs.key, rhs.key)) return true;
1108 if (generatedRowPipelineScheduleTuningKeySortsBefore(rhs.key, lhs.key)) return false;
1109 return generatedRowPipelineScheduleTuningSelectionSortsBefore(lhs.selection, rhs.selection);
1110 }
1111
1112 fn generatedRowPipelineScheduleTuningKeySortsBefore(
1113 lhs: GeneratedRowPipelineScheduleTuningKey,
1114 rhs: GeneratedRowPipelineScheduleTuningKey,
1115 ) bool {
1116 if (compareEnum(gpu.BackendKind, lhs.backend, rhs.backend)) |less| return less;
1117 if (compareEnum(gpu.DeviceFamily, lhs.family, rhs.family)) |less| return less;
1118 if (compareEnum(gpu.ArtifactFormat, lhs.format, rhs.format)) |less| return less;
1119 if (compareU32(lhs.vendor_id, rhs.vendor_id)) |less| return less;
1120 if (compareBool(lhs.has_vendor_id, rhs.has_vendor_id)) |less| return less;
1121 if (compareU32(lhs.device_id, rhs.device_id)) |less| return less;
1122 if (compareBool(lhs.has_device_id, rhs.has_device_id)) |less| return less;
1123 if (compareU64(lhs.name_fingerprint, rhs.name_fingerprint)) |less| return less;
1124 if (compareU64(lhs.driver_version_fingerprint, rhs.driver_version_fingerprint)) |less| return less;
1125 if (compareBool(lhs.has_driver_version, rhs.has_driver_version)) |less| return less;
1126 if (compareU64(lhs.rows, rhs.rows)) |less| return less;
1127 if (compareU64(lhs.cols, rhs.cols)) |less| return less;
1128 if (compareEnum(choir_abi.DType, lhs.dtype, rhs.dtype)) |less| return less;
1129 if (compareU32(lhs.schedule_version, rhs.schedule_version)) |less| return less;
1130 if (compareU32(lhs.candidate_count, rhs.candidate_count)) |less| return less;
1131 if (compareU64(lhs.candidate_set_fingerprint, rhs.candidate_set_fingerprint)) |less| return less;
1132 return false;
1133 }
1134
1135 fn generatedRowPipelineScheduleTuningSelectionSortsBefore(
1136 lhs: GeneratedRowPipelineScheduleTuningSelection,
1137 rhs: GeneratedRowPipelineScheduleTuningSelection,
1138 ) bool {
1139 if (generatedRowPipelineScheduleSortsBefore({}, lhs.schedule, rhs.schedule)) return true;
1140 if (generatedRowPipelineScheduleSortsBefore({}, rhs.schedule, lhs.schedule)) return false;
1141 if (compareU64(lhs.winner_median_ns, rhs.winner_median_ns)) |less| return less;
1142 if (compareU64(lhs.runner_up_median_ns, rhs.runner_up_median_ns)) |less| return less;
1143 if (compareU32(lhs.sample_count, rhs.sample_count)) |less| return less;
1144 return false;
1145 }
1146
1147 fn generatedRowPipelineScheduleSortsBefore(
1148 _: void,
1149 lhs: GeneratedRowPipelineSchedule,
1150 rhs: GeneratedRowPipelineSchedule,
1151 ) bool {
1152 if (compareU32(lhs.threads, rhs.threads)) |less| return less;
1153 return false;
1154 }
1155
1156 const testing = std.testing;
1157
1158 fn matrixProductFamilyScheduleTestCapabilities(device_id: u32) gpu.BackendCapabilities {
1159 return .{
1160 .identity = .{
1161 .backend = .cuda,
1162 .family = .nvidia_cuda,
1163 .name = "test-cuda-device",
1164 .vendor_id = 0x10de,
1165 .device_id = device_id,
1166 .driver_version = "550.54",
1167 },
1168 .artifact_formats = gpu.ArtifactFormatSet.init(&.{.cuda_ptx}),
1169 };
1170 }
1171
1172 fn matrixProductFamilyScheduleTestProblem(candidates: []const MatrixProductFamilyScheduleThreads) MatrixProductFamilyScheduleTuningProblem {
1173 return .{
1174 .format = .cuda_ptx,
1175 .m = 17,
1176 .n = 17,
1177 .k = 13,
1178 .dtype = .f32,
1179 .accumulation_dtype = .f32,
1180 .family_version = 1,
1181 .candidates = candidates,
1182 };
1183 }
1184
1185 fn expectMatrixProductFamilyScheduleSelection(
1186 expected: MatrixProductFamilyScheduleTuningSelection,
1187 actual: MatrixProductFamilyScheduleTuningSelection,
1188 ) !void {
1189 try testing.expect(expected.threads.eql(actual.threads));
1190 try testing.expectEqual(expected.winner_median_ns, actual.winner_median_ns);
1191 try testing.expectEqual(expected.runner_up_median_ns, actual.runner_up_median_ns);
1192 try testing.expectEqual(expected.sample_count, actual.sample_count);
1193 }
1194
1195 test "matrix product family schedule tuning cache records device keyed selections" {
1196 const allocator = testing.allocator;
1197 const candidates = [_]MatrixProductFamilyScheduleThreads{
1198 .{ .x = 17, .y = 9 },
1199 .{ .x = 16, .y = 16 },
1200 .{ .x = 8, .y = 8 },
1201 };
1202 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1203 const problem = matrixProductFamilyScheduleTestProblem(candidates[0..]);
1204
1205 var cache = MatrixProductFamilyScheduleTuningCache.init(allocator);
1206 defer cache.deinit();
1207
1208 const selected = MatrixProductFamilyScheduleTuningSelection{
1209 .threads = candidates[0],
1210 .winner_median_ns = 1000,
1211 .runner_up_median_ns = 1300,
1212 .sample_count = 30,
1213 };
1214 try cache.recordSelection(caps, problem, selected);
1215 try testing.expectEqual(@as(usize, 1), cache.count());
1216
1217 const hit = (try cache.selectionForProblem(caps, problem)) orelse return error.TestExpectedTuningRecord;
1218 try expectMatrixProductFamilyScheduleSelection(selected, hit);
1219
1220 try cache.recordSelection(caps, problem, .{
1221 .threads = candidates[1],
1222 .winner_median_ns = 1200,
1223 .runner_up_median_ns = 1400,
1224 .sample_count = 40,
1225 });
1226 try expectMatrixProductFamilyScheduleSelection(selected, (try cache.selectionForProblem(caps, problem)).?);
1227
1228 const faster = MatrixProductFamilyScheduleTuningSelection{
1229 .threads = candidates[1],
1230 .winner_median_ns = 900,
1231 .runner_up_median_ns = 1300,
1232 .sample_count = 20,
1233 };
1234 try cache.recordSelection(caps, problem, faster);
1235 try expectMatrixProductFamilyScheduleSelection(faster, (try cache.selectionForProblem(caps, problem)).?);
1236
1237 try testing.expectEqual(@as(?MatrixProductFamilyScheduleTuningSelection, null), try cache.selectionForProblem(matrixProductFamilyScheduleTestCapabilities(0x1b80), problem));
1238
1239 var other_problem = problem;
1240 other_problem.m = 32;
1241 try testing.expectEqual(@as(?MatrixProductFamilyScheduleTuningSelection, null), try cache.selectionForProblem(caps, other_problem));
1242
1243 try testing.expectError(error.LaunchArgumentMismatch, cache.recordSelection(caps, problem, .{
1244 .threads = .{ .x = 3, .y = 3 },
1245 .winner_median_ns = 850,
1246 .runner_up_median_ns = 1300,
1247 .sample_count = 30,
1248 }));
1249
1250 const records = try cache.exportRecords(allocator);
1251 defer allocator.free(records);
1252 try testing.expectEqual(@as(usize, 1), records.len);
1253 try expectMatrixProductFamilyScheduleSelection(faster, records[0].selection);
1254
1255 var imported = MatrixProductFamilyScheduleTuningCache.init(allocator);
1256 defer imported.deinit();
1257 try imported.importRecords(records);
1258 try expectMatrixProductFamilyScheduleSelection(faster, (try imported.selectionForProblem(caps, problem)).?);
1259 }
1260
1261 test "matrix product family schedule tuning artifact round trips through the wire" {
1262 const allocator = testing.allocator;
1263 const candidates = [_]MatrixProductFamilyScheduleThreads{
1264 .{ .x = 17, .y = 9 },
1265 .{ .x = 16, .y = 16 },
1266 .{ .x = 8, .y = 8 },
1267 };
1268 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1269 const problem = matrixProductFamilyScheduleTestProblem(candidates[0..]);
1270 const other_problem = MatrixProductFamilyScheduleTuningProblem{
1271 .format = .cuda_ptx,
1272 .m = 64,
1273 .n = 32,
1274 .k = 16,
1275 .dtype = .f16,
1276 .accumulation_dtype = .f32,
1277 .family_version = 1,
1278 .candidates = candidates[0..],
1279 };
1280 const records = [_]MatrixProductFamilyScheduleTuningRecord{
1281 .{
1282 .key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, problem),
1283 .selection = .{
1284 .threads = candidates[0],
1285 .winner_median_ns = 1000,
1286 .runner_up_median_ns = 1300,
1287 .sample_count = 30,
1288 },
1289 },
1290 .{
1291 .key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, other_problem),
1292 .selection = .{
1293 .threads = candidates[1],
1294 .winner_median_ns = 700,
1295 .runner_up_median_ns = 900,
1296 .sample_count = 30,
1297 },
1298 },
1299 };
1300
1301 const encoded = try encodeMatrixProductFamilyScheduleTuningArtifact(allocator, records[0..]);
1302 defer allocator.free(encoded);
1303
1304 const decoded = try decodeMatrixProductFamilyScheduleTuningArtifact(allocator, encoded);
1305 defer allocator.free(decoded);
1306
1307 try testing.expectEqual(records.len, decoded.len);
1308 for (records, decoded) |want, got| {
1309 try testing.expect(want.key.eql(got.key));
1310 try expectMatrixProductFamilyScheduleSelection(want.selection, got.selection);
1311 }
1312
1313 const encoded_again = try encodeMatrixProductFamilyScheduleTuningArtifact(allocator, decoded);
1314 defer allocator.free(encoded_again);
1315 try testing.expectEqualSlices(u8, encoded, encoded_again);
1316 }
1317
1318 test "matrix product family schedule tuning artifact rejects invalid records" {
1319 const allocator = testing.allocator;
1320 const candidates = [_]MatrixProductFamilyScheduleThreads{
1321 .{ .x = 17, .y = 9 },
1322 .{ .x = 16, .y = 16 },
1323 };
1324 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1325 const problem = matrixProductFamilyScheduleTestProblem(candidates[0..]);
1326 const valid = MatrixProductFamilyScheduleTuningRecord{
1327 .key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, problem),
1328 .selection = .{
1329 .threads = candidates[0],
1330 .winner_median_ns = 1000,
1331 .runner_up_median_ns = 1300,
1332 .sample_count = 30,
1333 },
1334 };
1335 const encoded = try encodeMatrixProductFamilyScheduleTuningArtifact(allocator, &.{valid});
1336 defer allocator.free(encoded);
1337
1338 try testing.expectError(error.InvalidArtifact, decodeMatrixProductFamilyScheduleTuningArtifact(allocator, encoded[0 .. encoded.len - 1]));
1339
1340 var bad_count = valid;
1341 bad_count.key.candidate_count = 1;
1342 try testing.expectError(error.InvalidArtifact, encodeMatrixProductFamilyScheduleTuningArtifact(allocator, &.{bad_count}));
1343
1344 var bad_selection = valid;
1345 bad_selection.selection.runner_up_median_ns = 999;
1346 try testing.expectError(error.InvalidArtifact, encodeMatrixProductFamilyScheduleTuningArtifact(allocator, &.{bad_selection}));
1347
1348 var extra = std.ArrayListUnmanaged(u8).empty;
1349 defer extra.deinit(allocator);
1350 try extra.appendSlice(allocator, encoded);
1351 try extra.append(allocator, 0);
1352 try testing.expectError(error.InvalidArtifact, decodeMatrixProductFamilyScheduleTuningArtifact(allocator, extra.items));
1353 }
1354
1355 test "matrix product family schedule tuning artifact fingerprint tracks bytes" {
1356 const allocator = testing.allocator;
1357 const candidates = [_]MatrixProductFamilyScheduleThreads{
1358 .{ .x = 17, .y = 9 },
1359 .{ .x = 16, .y = 16 },
1360 };
1361 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1362 const problem = matrixProductFamilyScheduleTestProblem(candidates[0..]);
1363 const record = MatrixProductFamilyScheduleTuningRecord{
1364 .key = try MatrixProductFamilyScheduleTuningKey.init(caps.identity, problem),
1365 .selection = .{
1366 .threads = candidates[0],
1367 .winner_median_ns = 1000,
1368 .runner_up_median_ns = 1300,
1369 .sample_count = 30,
1370 },
1371 };
1372 const encoded = try encodeMatrixProductFamilyScheduleTuningArtifact(allocator, &.{record});
1373 defer allocator.free(encoded);
1374
1375 const fingerprint = matrixProductFamilyScheduleTuningArtifactFingerprint(encoded);
1376 try testing.expectEqual(fingerprint, matrixProductFamilyScheduleTuningArtifactFingerprint(encoded));
1377
1378 const mutated = try allocator.dupe(u8, encoded);
1379 defer allocator.free(mutated);
1380 mutated[mutated.len - 1] +%= 1;
1381 try testing.expect(matrixProductFamilyScheduleTuningArtifactFingerprint(mutated) != fingerprint);
1382 }
1383
1384 fn generatedScanScheduleTestProblem(candidates: []const GeneratedScanSchedule) GeneratedScanScheduleTuningProblem {
1385 return .{
1386 .format = .cuda_ptx,
1387 .total = 16777216,
1388 .dtype = .f32,
1389 .schedule_version = 1,
1390 .candidates = candidates,
1391 };
1392 }
1393
1394 fn expectGeneratedScanScheduleSelection(
1395 want: GeneratedScanScheduleTuningSelection,
1396 got: GeneratedScanScheduleTuningSelection,
1397 ) !void {
1398 try testing.expect(want.schedule.eql(got.schedule));
1399 try testing.expectEqual(want.winner_median_ns, got.winner_median_ns);
1400 try testing.expectEqual(want.runner_up_median_ns, got.runner_up_median_ns);
1401 try testing.expectEqual(want.sample_count, got.sample_count);
1402 }
1403
1404 test "generated scan schedule candidate set fingerprint canonicalizes candidates" {
1405 const forward = [_]GeneratedScanSchedule{
1406 .{ .threads = 512, .items = 16 },
1407 .{ .threads = 256, .items = 16 },
1408 };
1409 const reversed = [_]GeneratedScanSchedule{
1410 .{ .threads = 256, .items = 16 },
1411 .{ .threads = 512, .items = 16 },
1412 };
1413 const forward_fingerprint = try generatedScanScheduleCandidateSetFingerprint(forward[0..]);
1414 try testing.expectEqual(forward_fingerprint, try generatedScanScheduleCandidateSetFingerprint(reversed[0..]));
1415
1416 const widened = [_]GeneratedScanSchedule{
1417 .{ .threads = 512, .items = 16 },
1418 .{ .threads = 256, .items = 16 },
1419 .{ .threads = 128, .items = 16 },
1420 };
1421 try testing.expect(forward_fingerprint != try generatedScanScheduleCandidateSetFingerprint(widened[0..]));
1422
1423 const duplicated = [_]GeneratedScanSchedule{
1424 .{ .threads = 512, .items = 16 },
1425 .{ .threads = 512, .items = 16 },
1426 };
1427 try testing.expectError(error.InvalidArtifact, generatedScanScheduleCandidateSetFingerprint(duplicated[0..]));
1428 try testing.expectError(error.InvalidArtifact, generatedScanScheduleCandidateSetFingerprint(forward[0..1]));
1429 }
1430
1431 test "generated scan schedule tuning cache records device keyed selections" {
1432 const allocator = testing.allocator;
1433 const candidates = [_]GeneratedScanSchedule{
1434 .{ .threads = 512, .items = 16 },
1435 .{ .threads = 256, .items = 16 },
1436 };
1437 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1438 const problem = generatedScanScheduleTestProblem(candidates[0..]);
1439
1440 var cache = GeneratedScanScheduleTuningCache.init(allocator);
1441 defer cache.deinit();
1442
1443 const selected = GeneratedScanScheduleTuningSelection{
1444 .schedule = candidates[1],
1445 .winner_median_ns = 160000,
1446 .runner_up_median_ns = 164000,
1447 .sample_count = 10,
1448 };
1449 try cache.recordSelection(caps, problem, selected);
1450 try testing.expectEqual(@as(usize, 1), cache.count());
1451
1452 const hit = (try cache.selectionForProblem(caps, problem)) orelse return error.TestExpectedTuningRecord;
1453 try expectGeneratedScanScheduleSelection(selected, hit);
1454
1455 try cache.recordSelection(caps, problem, .{
1456 .schedule = candidates[0],
1457 .winner_median_ns = 170000,
1458 .runner_up_median_ns = 175000,
1459 .sample_count = 10,
1460 });
1461 try expectGeneratedScanScheduleSelection(selected, (try cache.selectionForProblem(caps, problem)).?);
1462
1463 const faster = GeneratedScanScheduleTuningSelection{
1464 .schedule = candidates[0],
1465 .winner_median_ns = 150000,
1466 .runner_up_median_ns = 160000,
1467 .sample_count = 10,
1468 };
1469 try cache.recordSelection(caps, problem, faster);
1470 try expectGeneratedScanScheduleSelection(faster, (try cache.selectionForProblem(caps, problem)).?);
1471
1472 try testing.expectEqual(@as(?GeneratedScanScheduleTuningSelection, null), try cache.selectionForProblem(matrixProductFamilyScheduleTestCapabilities(0x1b80), problem));
1473
1474 var other_problem = problem;
1475 other_problem.total = 8388608;
1476 try testing.expectEqual(@as(?GeneratedScanScheduleTuningSelection, null), try cache.selectionForProblem(caps, other_problem));
1477
1478 try testing.expectError(error.LaunchArgumentMismatch, cache.recordSelection(caps, problem, .{
1479 .schedule = .{ .threads = 128, .items = 16 },
1480 .winner_median_ns = 140000,
1481 .runner_up_median_ns = 150000,
1482 .sample_count = 10,
1483 }));
1484
1485 const records = try cache.exportRecords(allocator);
1486 defer allocator.free(records);
1487 try testing.expectEqual(@as(usize, 1), records.len);
1488 try expectGeneratedScanScheduleSelection(faster, records[0].selection);
1489
1490 var imported = GeneratedScanScheduleTuningCache.init(allocator);
1491 defer imported.deinit();
1492 try imported.importRecords(records);
1493 try expectGeneratedScanScheduleSelection(faster, (try imported.selectionForProblem(caps, problem)).?);
1494 }
1495
1496 test "generated scan schedule tuning artifact round trips through the wire" {
1497 const allocator = testing.allocator;
1498 const candidates = [_]GeneratedScanSchedule{
1499 .{ .threads = 512, .items = 16 },
1500 .{ .threads = 256, .items = 16 },
1501 };
1502 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1503 const problem = generatedScanScheduleTestProblem(candidates[0..]);
1504 var other_problem = problem;
1505 other_problem.total = 8388608;
1506 const records = [_]GeneratedScanScheduleTuningRecord{
1507 .{
1508 .key = try GeneratedScanScheduleTuningKey.init(caps, problem),
1509 .selection = .{
1510 .schedule = candidates[0],
1511 .winner_median_ns = 160000,
1512 .runner_up_median_ns = 164000,
1513 .sample_count = 10,
1514 },
1515 },
1516 .{
1517 .key = try GeneratedScanScheduleTuningKey.init(caps, other_problem),
1518 .selection = .{
1519 .schedule = candidates[1],
1520 .winner_median_ns = 80000,
1521 .runner_up_median_ns = 90000,
1522 .sample_count = 10,
1523 },
1524 },
1525 };
1526
1527 const encoded = try encodeGeneratedScanScheduleTuningArtifact(allocator, records[0..]);
1528 defer allocator.free(encoded);
1529
1530 const decoded = try decodeGeneratedScanScheduleTuningArtifact(allocator, encoded);
1531 defer allocator.free(decoded);
1532
1533 try testing.expectEqual(records.len, decoded.len);
1534 for (records, decoded) |want, got| {
1535 try testing.expect(want.key.eql(got.key));
1536 try expectGeneratedScanScheduleSelection(want.selection, got.selection);
1537 }
1538
1539 const encoded_again = try encodeGeneratedScanScheduleTuningArtifact(allocator, decoded);
1540 defer allocator.free(encoded_again);
1541 try testing.expectEqualSlices(u8, encoded, encoded_again);
1542 }
1543
1544 test "generated scan schedule tuning artifact rejects invalid records" {
1545 const allocator = testing.allocator;
1546 const candidates = [_]GeneratedScanSchedule{
1547 .{ .threads = 512, .items = 16 },
1548 .{ .threads = 256, .items = 16 },
1549 };
1550 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1551 const problem = generatedScanScheduleTestProblem(candidates[0..]);
1552
1553 var stale = GeneratedScanScheduleTuningRecord{
1554 .key = try GeneratedScanScheduleTuningKey.init(caps, problem),
1555 .selection = .{
1556 .schedule = candidates[0],
1557 .winner_median_ns = 160000,
1558 .runner_up_median_ns = 164000,
1559 .sample_count = 10,
1560 },
1561 };
1562 stale.version = generated_scan_schedule_tuning_record_version + 1;
1563 try testing.expectError(error.InvalidArtifact, encodeGeneratedScanScheduleTuningArtifact(allocator, &.{stale}));
1564
1565 var inverted = stale;
1566 inverted.version = generated_scan_schedule_tuning_record_version;
1567 inverted.selection.runner_up_median_ns = inverted.selection.winner_median_ns - 1;
1568 try testing.expectError(error.InvalidArtifact, encodeGeneratedScanScheduleTuningArtifact(allocator, &.{inverted}));
1569
1570 try testing.expectError(error.InvalidArtifact, decodeGeneratedScanScheduleTuningArtifact(allocator, &.{ 1, 2, 3 }));
1571 }
1572
1573 fn generatedRowPipelineScheduleTestProblem(candidates: []const GeneratedRowPipelineSchedule) GeneratedRowPipelineScheduleTuningProblem {
1574 return .{
1575 .format = .cuda_ptx,
1576 .rows = 4096,
1577 .cols = 4096,
1578 .dtype = .f32,
1579 .schedule_version = 1,
1580 .candidates = candidates,
1581 };
1582 }
1583
1584 fn expectGeneratedRowPipelineScheduleSelection(
1585 want: GeneratedRowPipelineScheduleTuningSelection,
1586 got: GeneratedRowPipelineScheduleTuningSelection,
1587 ) !void {
1588 try testing.expect(want.schedule.eql(got.schedule));
1589 try testing.expectEqual(want.winner_median_ns, got.winner_median_ns);
1590 try testing.expectEqual(want.runner_up_median_ns, got.runner_up_median_ns);
1591 try testing.expectEqual(want.sample_count, got.sample_count);
1592 }
1593
1594 test "generated row pipeline schedule tuning cache records device keyed selections" {
1595 const allocator = testing.allocator;
1596 const candidates = [_]GeneratedRowPipelineSchedule{
1597 .{ .threads = 256 },
1598 .{ .threads = 512 },
1599 .{ .threads = 128 },
1600 };
1601 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1602 const problem = generatedRowPipelineScheduleTestProblem(candidates[0..]);
1603
1604 var cache = GeneratedRowPipelineScheduleTuningCache.init(allocator);
1605 defer cache.deinit();
1606
1607 const selected = GeneratedRowPipelineScheduleTuningSelection{
1608 .schedule = candidates[1],
1609 .winner_median_ns = 145000,
1610 .runner_up_median_ns = 150000,
1611 .sample_count = 10,
1612 };
1613 try cache.recordSelection(caps, problem, selected);
1614 try expectGeneratedRowPipelineScheduleSelection(selected, (try cache.selectionForProblem(caps, problem)).?);
1615
1616 try testing.expectEqual(@as(?GeneratedRowPipelineScheduleTuningSelection, null), try cache.selectionForProblem(matrixProductFamilyScheduleTestCapabilities(0x1b80), problem));
1617
1618 var other_problem = problem;
1619 other_problem.cols = 2048;
1620 try testing.expectEqual(@as(?GeneratedRowPipelineScheduleTuningSelection, null), try cache.selectionForProblem(caps, other_problem));
1621
1622 try testing.expectError(error.LaunchArgumentMismatch, cache.recordSelection(caps, problem, .{
1623 .schedule = .{ .threads = 64 },
1624 .winner_median_ns = 140000,
1625 .runner_up_median_ns = 150000,
1626 .sample_count = 10,
1627 }));
1628
1629 const records = try cache.exportRecords(allocator);
1630 defer allocator.free(records);
1631 try testing.expectEqual(@as(usize, 1), records.len);
1632
1633 var imported = GeneratedRowPipelineScheduleTuningCache.init(allocator);
1634 defer imported.deinit();
1635 try imported.importRecords(records);
1636 try expectGeneratedRowPipelineScheduleSelection(selected, (try imported.selectionForProblem(caps, problem)).?);
1637 }
1638
1639 test "generated row pipeline schedule tuning artifact round trips through the wire" {
1640 const allocator = testing.allocator;
1641 const candidates = [_]GeneratedRowPipelineSchedule{
1642 .{ .threads = 256 },
1643 .{ .threads = 512 },
1644 };
1645 const caps = matrixProductFamilyScheduleTestCapabilities(0x2684);
1646 const problem = generatedRowPipelineScheduleTestProblem(candidates[0..]);
1647 var other_problem = problem;
1648 other_problem.rows = 8;
1649 other_problem.cols = 2048;
1650 const records = [_]GeneratedRowPipelineScheduleTuningRecord{
1651 .{
1652 .key = try GeneratedRowPipelineScheduleTuningKey.init(caps, problem),
1653 .selection = .{
1654 .schedule = candidates[1],
1655 .winner_median_ns = 145000,
1656 .runner_up_median_ns = 150000,
1657 .sample_count = 10,
1658 },
1659 },
1660 .{
1661 .key = try GeneratedRowPipelineScheduleTuningKey.init(caps, other_problem),
1662 .selection = .{
1663 .schedule = candidates[0],
1664 .winner_median_ns = 9000,
1665 .runner_up_median_ns = 9500,
1666 .sample_count = 10,
1667 },
1668 },
1669 };
1670
1671 const encoded = try encodeGeneratedRowPipelineScheduleTuningArtifact(allocator, records[0..]);
1672 defer allocator.free(encoded);
1673
1674 const decoded = try decodeGeneratedRowPipelineScheduleTuningArtifact(allocator, encoded);
1675 defer allocator.free(decoded);
1676
1677 try testing.expectEqual(records.len, decoded.len);
1678 for (records, decoded) |want, got| {
1679 try testing.expect(want.key.eql(got.key));
1680 try expectGeneratedRowPipelineScheduleSelection(want.selection, got.selection);
1681 }
1682
1683 const encoded_again = try encodeGeneratedRowPipelineScheduleTuningArtifact(allocator, decoded);
1684 defer allocator.free(encoded_again);
1685 try testing.expectEqualSlices(u8, encoded, encoded_again);
1686
1687 try testing.expectError(error.InvalidArtifact, decodeGeneratedRowPipelineScheduleTuningArtifact(allocator, &.{ 9, 9 }));
1688 }
1689
1690 const MatrixScheduleCandidates =
1691 [matrix_product_family_schedule_tuning_max_candidates]MatrixProductFamilyScheduleThreads;
1692
1693 fn matrixScheduleReaderProblem(
1694 instance: @import("../kernel/library/root.zig").linalg.MatrixProduct,
1695 storage: *MatrixScheduleCandidates,
1696 ) MatrixProductFamilyScheduleTuningProblem {
1697 const linalg = @import("../kernel/library/root.zig").linalg;
1698 const candidates = linalg.matrixProductThreadCandidatesForExtents(instance.m, instance.n);
1699 for (candidates.slice(), 0..) |candidate, index| {
1700 storage[index] = .{ .x = candidate.x, .y = candidate.y };
1701 }
1702 return .{
1703 .format = .cuda_ptx,
1704 .m = instance.m,
1705 .n = instance.n,
1706 .k = instance.k,
1707 .dtype = instance.dtype,
1708 .accumulation_dtype = instance.accumulation_dtype,
1709 .family_version = linalg.matrix_product_family_version,
1710 .candidates = storage[0..candidates.slice().len],
1711 };
1712 }
1713
1714 test "matrix product schedule reader agrees with cache hits and missing keys" {
1715 const linalg = @import("../kernel/library/root.zig").linalg;
1716 const allocator = std.testing.allocator;
1717 const caps = matrixProductFamilyScheduleTestCapabilities(42);
1718 const probes = [_]linalg.MatrixProduct{
1719 .{ .m = 5, .n = 7, .k = 3 },
1720 .{ .m = 64, .n = 32, .k = 16 },
1721 .{ .m = 64, .n = 32, .k = 16, .dtype = .f16 },
1722 };
1723 var cache = MatrixProductFamilyScheduleTuningCache.init(allocator);
1724 defer cache.deinit();
1725 for (probes) |probe| {
1726 var storage: MatrixScheduleCandidates = undefined;
1727 const problem = matrixScheduleReaderProblem(probe, &storage);
1728 try cache.recordSelection(caps, problem, .{
1729 .threads = problem.candidates[problem.candidates.len - 1],
1730 .winner_median_ns = 3,
1731 .runner_up_median_ns = 5,
1732 .sample_count = 7,
1733 });
1734 }
1735 const entries = try cache.exportRecords(allocator);
1736 defer allocator.free(entries);
1737 const reader: linalg.MatrixProductScheduleReader = .{
1738 .device = caps.identity,
1739 .format = .cuda_ptx,
1740 .records = entries,
1741 };
1742 for (probes) |probe| {
1743 var storage: MatrixScheduleCandidates = undefined;
1744 const problem = matrixScheduleReaderProblem(probe, &storage);
1745 const expected = (try cache.selectionForProblem(caps, problem)).?;
1746 const actual = (try reader.resolve(probe)).?;
1747 try testing.expectEqual(expected.threads.x, actual.x);
1748 try testing.expectEqual(expected.threads.y, actual.y);
1749 var missing = probe;
1750 missing.k += 1;
1751 try testing.expectEqual(null, try cache.selectionForProblem(
1752 caps,
1753 matrixScheduleReaderProblem(missing, &storage),
1754 ));
1755 try testing.expectEqual(null, try reader.resolve(missing));
1756 var other = reader;
1757 other.device.driver_version = "new driver";
1758 var other_caps = caps;
1759 other_caps.identity = other.device;
1760 try testing.expectEqual(null, try cache.selectionForProblem(other_caps, problem));
1761 try testing.expectEqual(null, try other.resolve(probe));
1762 }
1763 }