lib/accy/src/kernel/library/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
6 const artifact = @import("../../artifact/model/root.zig");
7
8 pub const product_name = "accy.kernel.family_tuning";
9
10 pub const family_tuning_artifact_magic: u32 = 0x41515432;
11 pub const family_tuning_artifact_version: u32 = 1;
12 pub const family_tuning_max_extents: usize = 8;
13 pub const family_tuning_default_margin_percent: u32 = 2;
14
15 /// A caller builds this key to look up or record the measured schedule, the thread and block
16 /// structure a kernel runs with, for one problem on one device. The key identifies a problem by
17 /// device, the shape of a family, which is a group of related hand-written kernels such as sorts or
18 /// matrix products, operation, element type, family version, and up to eight problem sizes, so
19 /// `init` returns null when more than eight sizes are given. Two keys are equal only when every
20 /// field and every size is equal.
21 pub const FamilyTuningKey = struct {
22 device_fingerprint: u64,
23 family_fingerprint: u64,
24 operation_fingerprint: u64,
25 dtype: choir_abi.DType,
26 family_version: u32,
27 extent_count: u32 = 0,
28 extents: [family_tuning_max_extents]u64 = @as([family_tuning_max_extents]u64, @splat(0)),
29
30 pub fn init(
31 device_fingerprint: u64,
32 family_fingerprint: u64,
33 operation_fingerprint: u64,
34 key_dtype: choir_abi.DType,
35 family_version: u32,
36 extents: []const u64,
37 ) ?FamilyTuningKey {
38 if (extents.len > family_tuning_max_extents) return null;
39 var key = FamilyTuningKey{
40 .device_fingerprint = device_fingerprint,
41 .family_fingerprint = family_fingerprint,
42 .operation_fingerprint = operation_fingerprint,
43 .dtype = key_dtype,
44 .family_version = family_version,
45 .extent_count = @intCast(extents.len),
46 };
47 @memcpy(key.extents[0..extents.len], extents);
48 return key;
49 }
50
51 pub fn extentSlice(self: *const FamilyTuningKey) []const u64 {
52 return self.extents[0..self.extent_count];
53 }
54
55 pub fn eql(self: FamilyTuningKey, other: FamilyTuningKey) bool {
56 return self.device_fingerprint == other.device_fingerprint and
57 self.family_fingerprint == other.family_fingerprint and
58 self.operation_fingerprint == other.operation_fingerprint and
59 self.dtype == other.dtype and
60 self.family_version == other.family_version and
61 self.extent_count == other.extent_count and
62 std.mem.eql(u64, self.extents[0..self.extent_count], other.extents[0..other.extent_count]);
63 }
64 };
65
66 pub fn deviceFingerprint(caps: gpu.BackendCapabilities) u64 {
67 var builder = choir.product.incremental.FingerprintBuilder{};
68 builder.updateBytes(product_name);
69 builder.updateU32(@backingInt(caps.identity.backend));
70 builder.updateU32(@backingInt(caps.identity.family));
71 builder.updateU32(caps.identity.vendor_id orelse 0);
72 builder.updateU32(@intFromBool(caps.identity.vendor_id != null));
73 builder.updateU32(caps.identity.device_id orelse 0);
74 builder.updateU32(@intFromBool(caps.identity.device_id != null));
75 builder.updateBytes(caps.identity.name);
76 return builder.finish();
77 }
78
79 pub const FamilyCandidateMeasurement = struct {
80 key: FamilyTuningKey,
81 target: []const u8,
82 median_ns: u64,
83 sample_count: u32 = 1,
84 };
85
86 pub const FamilyMeasurementAccumulator = struct {
87 allocator: std.mem.Allocator,
88 measurements: std.ArrayListUnmanaged(FamilyCandidateMeasurement) = .empty,
89
90 pub fn init(measurement_allocator: std.mem.Allocator) FamilyMeasurementAccumulator {
91 return .{ .allocator = measurement_allocator };
92 }
93
94 pub fn count(self: *const FamilyMeasurementAccumulator) usize {
95 return self.measurements.items.len;
96 }
97
98 pub fn append(
99 self: *FamilyMeasurementAccumulator,
100 key: FamilyTuningKey,
101 target: []const u8,
102 median_ns: u64,
103 sample_count: u32,
104 ) gpu.BackendError!void {
105 if (target.len == 0 or median_ns == 0) return error.InvalidArtifact;
106 const owned_target = self.allocator.dupe(u8, target) catch return error.OutOfMemory;
107 errdefer self.allocator.free(owned_target);
108 self.measurements.append(self.allocator, .{
109 .key = key,
110 .target = owned_target,
111 .median_ns = median_ns,
112 .sample_count = sample_count,
113 }) catch return error.OutOfMemory;
114 }
115
116 pub fn selectWinners(
117 self: *const FamilyMeasurementAccumulator,
118 result_allocator: std.mem.Allocator,
119 margin_percent: u32,
120 ) gpu.BackendError!OwnedFamilyTuningRecords {
121 return selectFamilyWinners(result_allocator, self.measurements.items, margin_percent);
122 }
123
124 pub fn deinit(self: *FamilyMeasurementAccumulator) void {
125 for (self.measurements.items) |measurement| self.allocator.free(measurement.target);
126 self.measurements.deinit(self.allocator);
127 self.* = undefined;
128 }
129 };
130
131 pub const FamilyTuningRecord = struct {
132 key: FamilyTuningKey,
133 target: []const u8,
134 winner_median_ns: u64,
135 runner_up_median_ns: u64,
136 sample_count: u32,
137 };
138
139 pub const FamilyTuningTable = struct {
140 records: []const FamilyTuningRecord = &.{},
141
142 pub fn find(self: FamilyTuningTable, key: FamilyTuningKey) ?FamilyTuningRecord {
143 for (self.records) |record| {
144 if (record.key.eql(key)) return record;
145 }
146 return null;
147 }
148 };
149
150 pub const OwnedFamilyTuningRecords = struct {
151 allocator: std.mem.Allocator,
152 records: []FamilyTuningRecord = &.{},
153
154 pub fn table(self: *const OwnedFamilyTuningRecords) FamilyTuningTable {
155 return .{ .records = self.records };
156 }
157
158 pub fn deinit(self: *OwnedFamilyTuningRecords) void {
159 for (self.records) |record| self.allocator.free(record.target);
160 if (self.records.len != 0) self.allocator.free(self.records);
161 self.* = undefined;
162 }
163 };
164
165 /// A caller hands this reader to a kernel family so its schedule choice follows measurements taken
166 /// on the same device. The reader holds one device fingerprint and a table of measured records. The
167 /// device fingerprint covers the backend kind, device family, vendor and device ids, and the device
168 /// name, and leaves out the driver version. A family keeps the schedule it would choose with no
169 /// table when it finds no record with an exact key.
170 pub const FamilyTuningReader = struct {
171 device_fingerprint: u64,
172 table: FamilyTuningTable,
173
174 pub fn init(caps: gpu.BackendCapabilities, table: FamilyTuningTable) FamilyTuningReader {
175 return .{ .device_fingerprint = deviceFingerprint(caps), .table = table };
176 }
177 };
178
179 /// A caller uses this function to turn benchmark measurements into tuning records that change code
180 /// generation only when the evidence is clear. The function groups measurements by key and emits
181 /// one record per group, naming the fastest candidate by median time and keeping the runner-up
182 /// time. The function emits nothing for a group with one candidate and nothing for a group whose
183 /// runner-up is within `margin_percent` of the winner, so a single or near-tied measurement cannot
184 /// change generated code. The call returns `error.InvalidArtifact` when any measurement has a
185 /// median of zero. The returned records and their target strings are allocated with
186 /// `result_allocator` and owned by the caller.
187 pub fn selectFamilyWinners(
188 result_allocator: std.mem.Allocator,
189 measurements: []const FamilyCandidateMeasurement,
190 margin_percent: u32,
191 ) gpu.BackendError!OwnedFamilyTuningRecords {
192 var records = std.ArrayListUnmanaged(FamilyTuningRecord).empty;
193 errdefer {
194 for (records.items) |record| result_allocator.free(record.target);
195 records.deinit(result_allocator);
196 }
197
198 for (measurements, 0..) |measurement, index| {
199 if (measurement.median_ns == 0) return error.InvalidArtifact;
200 var seen_before = false;
201 for (measurements[0..index]) |previous| {
202 if (previous.key.eql(measurement.key)) {
203 seen_before = true;
204 break;
205 }
206 }
207 if (seen_before) continue;
208
209 var winner = measurement;
210 var runner_up_ns: u64 = 0;
211 var group_count: usize = 1;
212 for (measurements[index + 1 ..]) |candidate| {
213 if (!candidate.key.eql(measurement.key)) continue;
214 group_count += 1;
215 if (candidate.median_ns < winner.median_ns) {
216 runner_up_ns = winner.median_ns;
217 winner = candidate;
218 } else if (runner_up_ns == 0 or candidate.median_ns < runner_up_ns) {
219 runner_up_ns = candidate.median_ns;
220 }
221 }
222 if (group_count < 2) continue;
223 const margin_floor = winner.median_ns + winner.median_ns * margin_percent / 100;
224 if (runner_up_ns < margin_floor) continue;
225
226 const target = result_allocator.dupe(u8, winner.target) catch return error.OutOfMemory;
227 errdefer result_allocator.free(target);
228 records.append(result_allocator, .{
229 .key = winner.key,
230 .target = target,
231 .winner_median_ns = winner.median_ns,
232 .runner_up_median_ns = runner_up_ns,
233 .sample_count = winner.sample_count,
234 }) catch return error.OutOfMemory;
235 }
236
237 return .{
238 .allocator = result_allocator,
239 .records = records.toOwnedSlice(result_allocator) catch return error.OutOfMemory,
240 };
241 }
242
243 pub fn encodeFamilyTuningArtifact(
244 result_allocator: std.mem.Allocator,
245 records: []const FamilyTuningRecord,
246 ) gpu.BackendError![]u8 {
247 if (records.len > std.math.maxInt(u32)) return error.InvalidArtifact;
248 var writer = artifact.wire.ByteWriter{};
249 errdefer writer.deinit(result_allocator);
250
251 try writer.writeU32(result_allocator, family_tuning_artifact_magic);
252 try writer.writeU32(result_allocator, family_tuning_artifact_version);
253 try writer.writeU32(result_allocator, @intCast(records.len));
254 for (records) |record| {
255 if (record.target.len == 0) return error.InvalidArtifact;
256 if (record.key.extent_count > family_tuning_max_extents) return error.InvalidArtifact;
257 try writer.writeU64(result_allocator, record.key.device_fingerprint);
258 try writer.writeU64(result_allocator, record.key.family_fingerprint);
259 try writer.writeU64(result_allocator, record.key.operation_fingerprint);
260 try writer.writeEnum(result_allocator, choir_abi.DType, record.key.dtype);
261 try writer.writeU32(result_allocator, record.key.family_version);
262 try writer.writeU32(result_allocator, record.key.extent_count);
263 for (record.key.extentSlice()) |extent| {
264 try writer.writeU64(result_allocator, extent);
265 }
266 try writer.writeLengthPrefixedBytes(result_allocator, record.target);
267 try writer.writeU64(result_allocator, record.winner_median_ns);
268 try writer.writeU64(result_allocator, record.runner_up_median_ns);
269 try writer.writeU32(result_allocator, record.sample_count);
270 }
271
272 return writer.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
273 }
274
275 pub fn decodeFamilyTuningArtifact(
276 result_allocator: std.mem.Allocator,
277 bytes: []const u8,
278 ) gpu.BackendError!OwnedFamilyTuningRecords {
279 var reader = artifact.wire.ByteReader{ .bytes = bytes };
280 if ((try reader.readU32()) != family_tuning_artifact_magic) return error.InvalidArtifact;
281 if ((try reader.readU32()) != family_tuning_artifact_version) return error.InvalidArtifact;
282 const record_count = try reader.readU32();
283
284 var owned = OwnedFamilyTuningRecords{ .allocator = result_allocator };
285 var records = std.ArrayListUnmanaged(FamilyTuningRecord).empty;
286 errdefer {
287 for (records.items) |record| result_allocator.free(record.target);
288 records.deinit(result_allocator);
289 }
290
291 var index: u32 = 0;
292 while (index < record_count) : (index += 1) {
293 var key = FamilyTuningKey{
294 .device_fingerprint = try reader.readU64(),
295 .family_fingerprint = try reader.readU64(),
296 .operation_fingerprint = try reader.readU64(),
297 .dtype = try reader.readEnum(choir_abi.DType),
298 .family_version = try reader.readU32(),
299 };
300 key.extent_count = try reader.readU32();
301 if (key.extent_count > family_tuning_max_extents) return error.InvalidArtifact;
302 for (key.extents[0..key.extent_count]) |*extent| {
303 extent.* = try reader.readU64();
304 }
305 const target_bytes = try reader.readLengthPrefixedBytes();
306 if (target_bytes.len == 0) return error.InvalidArtifact;
307 const target = result_allocator.dupe(u8, target_bytes) catch return error.OutOfMemory;
308 errdefer result_allocator.free(target);
309 const winner_median_ns = try reader.readU64();
310 const runner_up_median_ns = try reader.readU64();
311 const sample_count = try reader.readU32();
312 records.append(result_allocator, .{
313 .key = key,
314 .target = target,
315 .winner_median_ns = winner_median_ns,
316 .runner_up_median_ns = runner_up_median_ns,
317 .sample_count = sample_count,
318 }) catch return error.OutOfMemory;
319 }
320 try reader.expectDone();
321
322 owned.records = records.toOwnedSlice(result_allocator) catch return error.OutOfMemory;
323 return owned;
324 }
325
326 pub fn artifactFingerprint(bytes: []const u8) choir.product.incremental.Fingerprint {
327 var builder = choir.product.incremental.FingerprintBuilder{};
328 builder.updateBytes(product_name);
329 builder.updateU32(family_tuning_artifact_version);
330 builder.updateBytes(bytes);
331 return builder.finish();
332 }
333
334 pub fn artifactProductStamp(bytes: []const u8) choir.product.incremental.ProductStamp {
335 return choir.product.incremental.productStamp(product_name, artifactFingerprint(bytes));
336 }
337
338 const testing = std.testing;
339
340 fn testKey(device: u64, family: u64, extents: []const u64) FamilyTuningKey {
341 return FamilyTuningKey.init(device, family, 77, .f32, 1, extents) orelse unreachable;
342 }
343
344 fn testCapabilities(device_id: u32) gpu.BackendCapabilities {
345 return .{ .identity = .{
346 .backend = .cuda,
347 .family = .nvidia_cuda,
348 .name = "test-cuda-device",
349 .vendor_id = 0x10de,
350 .device_id = device_id,
351 } };
352 }
353
354 test "family tuning winner selection requires a margin over the runner up" {
355 const allocator = testing.allocator;
356 const key = testKey(11, 22, &.{ 64, 64 });
357 const flappy_key = testKey(11, 33, &.{128});
358 const lone_key = testKey(11, 44, &.{256});
359
360 const measurements = [_]FamilyCandidateMeasurement{
361 .{ .key = key, .target = "family_a_16", .median_ns = 1200, .sample_count = 30 },
362 .{ .key = key, .target = "family_a_32", .median_ns = 1000, .sample_count = 30 },
363 .{ .key = key, .target = "family_a_64", .median_ns = 1500, .sample_count = 30 },
364 .{ .key = flappy_key, .target = "family_b_32", .median_ns = 1000, .sample_count = 30 },
365 .{ .key = flappy_key, .target = "family_b_64", .median_ns = 1010, .sample_count = 30 },
366 .{ .key = lone_key, .target = "family_c_32", .median_ns = 900, .sample_count = 30 },
367 };
368
369 var winners = try selectFamilyWinners(allocator, measurements[0..], family_tuning_default_margin_percent);
370 defer winners.deinit();
371
372 try testing.expectEqual(@as(usize, 1), winners.records.len);
373 try testing.expectEqualStrings("family_a_32", winners.records[0].target);
374 try testing.expectEqual(@as(u64, 1000), winners.records[0].winner_median_ns);
375 try testing.expectEqual(@as(u64, 1200), winners.records[0].runner_up_median_ns);
376
377 const table = winners.table();
378 const found = table.find(key) orelse return error.TestExpectedTuningRecord;
379 try testing.expectEqualStrings("family_a_32", found.target);
380 try testing.expectEqual(@as(?FamilyTuningRecord, null), table.find(flappy_key));
381 try testing.expectEqual(@as(?FamilyTuningRecord, null), table.find(lone_key));
382 }
383
384 test "family tuning artifact round-trips through the wire" {
385 const allocator = testing.allocator;
386 const records = [_]FamilyTuningRecord{
387 .{
388 .key = testKey(7, 99, &.{ 1024, 500, 1 }),
389 .target = "accy.kernel.indexing.gather_family_512_f32",
390 .winner_median_ns = 10350,
391 .runner_up_median_ns = 10419,
392 .sample_count = 30,
393 },
394 .{
395 .key = testKey(7, 55, &.{4096}),
396 .target = "accy.kernel.segmented.segment_sum_family_warp_512_f32",
397 .winner_median_ns = 7905,
398 .runner_up_median_ns = 16842,
399 .sample_count = 30,
400 },
401 };
402
403 const encoded = try encodeFamilyTuningArtifact(allocator, records[0..]);
404 defer allocator.free(encoded);
405
406 var decoded = try decodeFamilyTuningArtifact(allocator, encoded);
407 defer decoded.deinit();
408
409 try testing.expectEqual(records.len, decoded.records.len);
410 for (records, decoded.records) |want, got| {
411 try testing.expect(want.key.eql(got.key));
412 try testing.expectEqualStrings(want.target, got.target);
413 try testing.expectEqual(want.winner_median_ns, got.winner_median_ns);
414 try testing.expectEqual(want.runner_up_median_ns, got.runner_up_median_ns);
415 try testing.expectEqual(want.sample_count, got.sample_count);
416 }
417 }
418
419 test "family tuning artifact rejects truncated and oversized payloads" {
420 const allocator = testing.allocator;
421 const records = [_]FamilyTuningRecord{.{
422 .key = testKey(1, 2, &.{8}),
423 .target = "accy.kernel.scan.prefix_sum_family_32_f32",
424 .winner_median_ns = 100,
425 .runner_up_median_ns = 200,
426 .sample_count = 10,
427 }};
428 const encoded = try encodeFamilyTuningArtifact(allocator, records[0..]);
429 defer allocator.free(encoded);
430
431 try testing.expectError(error.InvalidArtifact, decodeFamilyTuningArtifact(allocator, encoded[0 .. encoded.len - 1]));
432
433 var oversized_key = testKey(1, 2, &.{8});
434 oversized_key.extent_count = family_tuning_max_extents + 1;
435 const invalid = [_]FamilyTuningRecord{.{
436 .key = oversized_key,
437 .target = "x",
438 .winner_median_ns = 1,
439 .runner_up_median_ns = 2,
440 .sample_count = 1,
441 }};
442 try testing.expectError(error.InvalidArtifact, encodeFamilyTuningArtifact(allocator, invalid[0..]));
443 }
444
445 test "family measurement accumulator owns targets and folds winners" {
446 const allocator = testing.allocator;
447 var accumulator = FamilyMeasurementAccumulator.init(allocator);
448 defer accumulator.deinit();
449
450 const key = testKey(11, 22, &.{ 64, 64 });
451 var target_buffer: [10]u8 = undefined;
452 @memcpy(target_buffer[0..], "family_a_8");
453 try accumulator.append(key, target_buffer[0..], 1500, 30);
454 @memcpy(target_buffer[0..], "XXXXXXXXXX");
455 try accumulator.append(key, "family_a_32", 1000, 30);
456 try accumulator.append(key, "family_a_64", 1200, 30);
457 try testing.expectEqual(@as(usize, 3), accumulator.count());
458 try testing.expectEqualStrings("family_a_8", accumulator.measurements.items[0].target);
459
460 var winners = try accumulator.selectWinners(allocator, family_tuning_default_margin_percent);
461 defer winners.deinit();
462 try testing.expectEqual(@as(usize, 1), winners.records.len);
463 try testing.expectEqualStrings("family_a_32", winners.records[0].target);
464 try testing.expectEqual(@as(u64, 1000), winners.records[0].winner_median_ns);
465 try testing.expectEqual(@as(u64, 1200), winners.records[0].runner_up_median_ns);
466
467 try testing.expectError(error.InvalidArtifact, accumulator.append(key, "", 100, 1));
468 try testing.expectError(error.InvalidArtifact, accumulator.append(key, "x", 0, 1));
469 try testing.expectEqual(@as(usize, 3), accumulator.count());
470 }
471
472 test "family tuning artifact product identity tracks bytes" {
473 const allocator = testing.allocator;
474 const records = [_]FamilyTuningRecord{.{
475 .key = testKey(1, 2, &.{8}),
476 .target = "accy.kernel.scan.prefix_sum_family_32_f32",
477 .winner_median_ns = 100,
478 .runner_up_median_ns = 200,
479 .sample_count = 10,
480 }};
481 const encoded = try encodeFamilyTuningArtifact(allocator, records[0..]);
482 defer allocator.free(encoded);
483
484 const fingerprint = artifactFingerprint(encoded);
485 try testing.expectEqual(fingerprint, artifactFingerprint(encoded));
486
487 const mutated = try allocator.dupe(u8, encoded);
488 defer allocator.free(mutated);
489 mutated[mutated.len - 1] +%= 1;
490 try testing.expect(artifactFingerprint(mutated) != fingerprint);
491
492 const stamp = artifactProductStamp(encoded);
493 try testing.expectEqualStrings(product_name, stamp.name);
494 try testing.expectEqual(fingerprint, stamp.fingerprint);
495 }
496
497 test "family tuning reader resolves winners by instance and device" {
498 const caps = testCapabilities(0x2684);
499 const device = deviceFingerprint(caps);
500
501 const key = testKey(device, 64, &.{64});
502 const records = [_]FamilyTuningRecord{.{
503 .key = key,
504 .target = "family_64",
505 .winner_median_ns = 900,
506 .runner_up_median_ns = 1500,
507 .sample_count = 30,
508 }};
509
510 const reader = FamilyTuningReader.init(caps, .{ .records = records[0..] });
511 try testing.expectEqual(device, reader.device_fingerprint);
512
513 const hit = reader.table.find(key) orelse return error.TestExpectedTuningRecord;
514 try testing.expectEqualStrings("family_64", hit.target);
515
516 const other_extent = testKey(device, 64, &.{128});
517 try testing.expectEqual(@as(?FamilyTuningRecord, null), reader.table.find(other_extent));
518
519 const other_family = testKey(device, 65, &.{64});
520 try testing.expectEqual(@as(?FamilyTuningRecord, null), reader.table.find(other_family));
521
522 const other_device = FamilyTuningReader.init(testCapabilities(0x1b80), .{ .records = records[0..] });
523 const other_device_key = testKey(other_device.device_fingerprint, 64, &.{64});
524 try testing.expectEqual(@as(?FamilyTuningRecord, null), other_device.table.find(other_device_key));
525 }
526
527 pub const matrix_product_family_schedule_tuning_product_name =
528 "accy.exec.matrix_product_family_schedule_tuning";
529 pub const matrix_product_family_schedule_tuning_record_version: u32 = 1;
530 pub const matrix_product_family_schedule_tuning_max_candidates: usize = 8;
531
532 pub const MatrixProductFamilyScheduleThreads = struct {
533 x: u32,
534 y: u32,
535
536 pub fn lessThan(
537 _: void,
538 lhs: MatrixProductFamilyScheduleThreads,
539 rhs: MatrixProductFamilyScheduleThreads,
540 ) bool {
541 return if (lhs.x != rhs.x) lhs.x < rhs.x else lhs.y < rhs.y;
542 }
543
544 pub fn eql(
545 self: MatrixProductFamilyScheduleThreads,
546 other: MatrixProductFamilyScheduleThreads,
547 ) bool {
548 return self.x == other.x and self.y == other.y;
549 }
550 };
551
552 pub const MatrixProductFamilyScheduleTuningProblem = struct {
553 format: gpu.ArtifactFormat,
554 m: u64,
555 n: u64,
556 k: u64,
557 dtype: choir_abi.DType = .f32,
558 accumulation_dtype: choir_abi.DType = .f32,
559 family_version: u32,
560 candidates: []const MatrixProductFamilyScheduleThreads,
561 };
562
563 pub const MatrixProductFamilyScheduleTuningKey = struct {
564 backend: gpu.BackendKind,
565 family: gpu.DeviceFamily,
566 format: gpu.ArtifactFormat,
567 vendor_id: u32 = 0,
568 has_vendor_id: bool = false,
569 device_id: u32 = 0,
570 has_device_id: bool = false,
571 name_fingerprint: u64 = 0,
572 driver_version_fingerprint: u64 = 0,
573 has_driver_version: bool = false,
574 m: u64,
575 n: u64,
576 k: u64,
577 dtype: choir_abi.DType,
578 accumulation_dtype: choir_abi.DType,
579 family_version: u32,
580 candidate_count: u32,
581 candidate_set_fingerprint: u64,
582
583 pub fn init(
584 device: gpu.DeviceIdentity,
585 problem: MatrixProductFamilyScheduleTuningProblem,
586 ) gpu.BackendError!MatrixProductFamilyScheduleTuningKey {
587 if (problem.m == 0 or problem.n == 0 or problem.k == 0) return error.InvalidArtifact;
588 if (problem.family_version == 0) return error.InvalidArtifact;
589 const candidate_set_fingerprint = try matrixProductFamilyScheduleCandidateSetFingerprint(
590 problem.candidates,
591 );
592 return .{
593 .backend = device.backend,
594 .family = device.family,
595 .format = problem.format,
596 .vendor_id = device.vendor_id orelse 0,
597 .has_vendor_id = device.vendor_id != null,
598 .device_id = device.device_id orelse 0,
599 .has_device_id = device.device_id != null,
600 .name_fingerprint = matrixScheduleBytesFingerprint(
601 matrix_product_family_schedule_tuning_product_name ++ ".device.name",
602 device.name,
603 ),
604 .driver_version_fingerprint = if (device.driver_version) |version|
605 matrixScheduleBytesFingerprint(
606 matrix_product_family_schedule_tuning_product_name ++ ".driver.version",
607 version,
608 )
609 else
610 0,
611 .has_driver_version = device.driver_version != null,
612 .m = problem.m,
613 .n = problem.n,
614 .k = problem.k,
615 .dtype = problem.dtype,
616 .accumulation_dtype = problem.accumulation_dtype,
617 .family_version = problem.family_version,
618 .candidate_count = @intCast(problem.candidates.len),
619 .candidate_set_fingerprint = candidate_set_fingerprint,
620 };
621 }
622
623 pub fn eql(
624 self: MatrixProductFamilyScheduleTuningKey,
625 other: MatrixProductFamilyScheduleTuningKey,
626 ) bool {
627 return self.backend == other.backend and
628 self.family == other.family and
629 self.format == other.format and
630 self.vendor_id == other.vendor_id and
631 self.has_vendor_id == other.has_vendor_id and
632 self.device_id == other.device_id and
633 self.has_device_id == other.has_device_id and
634 self.name_fingerprint == other.name_fingerprint and
635 self.driver_version_fingerprint == other.driver_version_fingerprint and
636 self.has_driver_version == other.has_driver_version and
637 self.m == other.m and
638 self.n == other.n and
639 self.k == other.k and
640 self.dtype == other.dtype and
641 self.accumulation_dtype == other.accumulation_dtype and
642 self.family_version == other.family_version and
643 self.candidate_count == other.candidate_count and
644 self.candidate_set_fingerprint == other.candidate_set_fingerprint;
645 }
646 };
647
648 pub const MatrixProductFamilyScheduleTuningSelection = struct {
649 threads: MatrixProductFamilyScheduleThreads,
650 winner_median_ns: u64,
651 runner_up_median_ns: u64,
652 sample_count: u32,
653 };
654
655 pub const MatrixProductFamilyScheduleTuningRecord = struct {
656 version: u32 = matrix_product_family_schedule_tuning_record_version,
657 key: MatrixProductFamilyScheduleTuningKey,
658 selection: MatrixProductFamilyScheduleTuningSelection,
659 };
660
661 pub fn matrixProductFamilyScheduleCandidateSetFingerprint(
662 candidates: []const MatrixProductFamilyScheduleThreads,
663 ) gpu.BackendError!choir.product.incremental.Fingerprint {
664 const set = try sortedMatrixProductFamilyScheduleCandidates(candidates);
665 var builder = choir.product.incremental.FingerprintBuilder{};
666 builder.updateBytes(matrix_product_family_schedule_tuning_product_name);
667 builder.updateU32(@intCast(set.count));
668 for (set.slice()) |candidate| {
669 builder.updateU32(candidate.x);
670 builder.updateU32(candidate.y);
671 }
672 return builder.finish();
673 }
674
675 const MatrixProductFamilyScheduleCandidateSet = struct {
676 const capacity = matrix_product_family_schedule_tuning_max_candidates;
677 count: usize = 0,
678 items: [capacity]MatrixProductFamilyScheduleThreads = @splat(.{ .x = 1, .y = 1 }),
679
680 fn slice(
681 self: *const MatrixProductFamilyScheduleCandidateSet,
682 ) []const MatrixProductFamilyScheduleThreads {
683 return self.items[0..self.count];
684 }
685 };
686
687 fn sortedMatrixProductFamilyScheduleCandidates(
688 candidates: []const MatrixProductFamilyScheduleThreads,
689 ) gpu.BackendError!MatrixProductFamilyScheduleCandidateSet {
690 if (candidates.len < 2) return error.InvalidArtifact;
691 if (candidates.len > matrix_product_family_schedule_tuning_max_candidates) {
692 return error.InvalidArtifact;
693 }
694 var set = MatrixProductFamilyScheduleCandidateSet{ .count = candidates.len };
695 @memcpy(set.items[0..candidates.len], candidates);
696 for (set.slice()) |candidate| {
697 if (candidate.x == 0 or candidate.y == 0) return error.InvalidArtifact;
698 }
699 std.mem.sort(
700 MatrixProductFamilyScheduleThreads,
701 set.items[0..set.count],
702 {},
703 MatrixProductFamilyScheduleThreads.lessThan,
704 );
705 for (set.slice()[1..], 1..) |candidate, index| {
706 if (candidate.eql(set.items[index - 1])) return error.InvalidArtifact;
707 }
708 return set;
709 }
710
711 fn matrixScheduleBytesFingerprint(domain: []const u8, bytes: []const u8) u64 {
712 var builder = choir.product.incremental.FingerprintBuilder{};
713 builder.updateBytes(domain);
714 builder.updateBytes(bytes);
715 return builder.finish();
716 }
717
718 test "matrix product family schedule candidate set fingerprint canonicalizes candidates" {
719 const candidates = [_]MatrixProductFamilyScheduleThreads{
720 .{ .x = 17, .y = 9 },
721 .{ .x = 16, .y = 16 },
722 .{ .x = 8, .y = 8 },
723 };
724 const reordered = [_]MatrixProductFamilyScheduleThreads{
725 candidates[2],
726 candidates[0],
727 candidates[1],
728 };
729 const changed = [_]MatrixProductFamilyScheduleThreads{
730 candidates[0],
731 candidates[1],
732 .{ .x = 4, .y = 4 },
733 };
734
735 const fingerprint = try matrixProductFamilyScheduleCandidateSetFingerprint(candidates[0..]);
736 try testing.expectEqual(
737 fingerprint,
738 try matrixProductFamilyScheduleCandidateSetFingerprint(reordered[0..]),
739 );
740 const changed_fingerprint = try matrixProductFamilyScheduleCandidateSetFingerprint(&changed);
741 try testing.expect(fingerprint != changed_fingerprint);
742 }
743
744 test "matrix product family schedule candidate set rejects invalid candidates" {
745 const candidates = [_]MatrixProductFamilyScheduleThreads{
746 .{ .x = 17, .y = 9 },
747 .{ .x = 16, .y = 16 },
748 .{ .x = 8, .y = 8 },
749 };
750 const single = [_]MatrixProductFamilyScheduleThreads{candidates[0]};
751 try testing.expectError(
752 error.InvalidArtifact,
753 matrixProductFamilyScheduleCandidateSetFingerprint(&single),
754 );
755
756 const duplicate = [_]MatrixProductFamilyScheduleThreads{
757 candidates[0],
758 candidates[1],
759 candidates[0],
760 };
761 try testing.expectError(
762 error.InvalidArtifact,
763 matrixProductFamilyScheduleCandidateSetFingerprint(&duplicate),
764 );
765
766 const zero = [_]MatrixProductFamilyScheduleThreads{
767 candidates[0],
768 .{ .x = 0, .y = 16 },
769 };
770 try testing.expectError(
771 error.InvalidArtifact,
772 matrixProductFamilyScheduleCandidateSetFingerprint(&zero),
773 );
774
775 const too_many = [_]MatrixProductFamilyScheduleThreads{
776 .{ .x = 1, .y = 1 },
777 .{ .x = 2, .y = 1 },
778 .{ .x = 3, .y = 1 },
779 .{ .x = 4, .y = 1 },
780 .{ .x = 5, .y = 1 },
781 .{ .x = 6, .y = 1 },
782 .{ .x = 7, .y = 1 },
783 .{ .x = 8, .y = 1 },
784 .{ .x = 9, .y = 1 },
785 };
786 try testing.expectError(
787 error.InvalidArtifact,
788 matrixProductFamilyScheduleCandidateSetFingerprint(&too_many),
789 );
790 }