lib/accy/src/artifact/model/registry.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir_abi = @import("choir_abi");
4 const accy_choir = @import("../../choir/root.zig");
5
6 const testing = std.testing;
7
8 pub const ArtifactPlanOptions = struct {
9 format: ?gpu.ArtifactFormat = null,
10 kernel_call_registry: ?*const KernelCallRegistry = null,
11 };
12
13 pub const ElementCountArgument = enum {
14 none,
15 scalar_u32,
16 device_buffer_u32,
17 };
18 pub const KernelCallLaunch = union(enum) {
19 derived: KernelCallDerivedLaunch,
20 fixed: choir_abi.LaunchGeometry,
21 };
22
23 pub const KernelCallDerivedLaunchAxis = union(enum) {
24 fixed: u32,
25 runtime_u32_ceil_div: RuntimeU32CeilDiv,
26
27 pub const RuntimeU32CeilDiv = struct {
28 argument_index: u32,
29 divisor: u32 = 1,
30 };
31
32 pub fn extent(
33 self: KernelCallDerivedLaunchAxis,
34 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
35 ) gpu.BackendError!u32 {
36 return switch (self) {
37 .fixed => |value| if (value == 0) error.InvalidArtifact else value,
38 .runtime_u32_ceil_div => |axis| runtimeU32CeilDiv(runtime_scalar_arguments, axis),
39 };
40 }
41 };
42
43 pub const KernelCallDerivedLaunch = struct {
44 grid: [3]KernelCallDerivedLaunchAxis = .{
45 .{ .fixed = 1 },
46 .{ .fixed = 1 },
47 .{ .fixed = 1 },
48 },
49 threadgroup: [3]u32 = .{ 1, 1, 1 },
50 dynamic_shared_memory_bytes: u32 = 0,
51
52 pub fn geometry(
53 self: KernelCallDerivedLaunch,
54 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
55 ) gpu.BackendError!choir_abi.LaunchGeometry {
56 return .{
57 .grid = .{
58 try self.grid[0].extent(runtime_scalar_arguments),
59 try self.grid[1].extent(runtime_scalar_arguments),
60 try self.grid[2].extent(runtime_scalar_arguments),
61 },
62 .threadgroup = try validateThreadgroup(self.threadgroup),
63 .dynamic_shared_memory_bytes = self.dynamic_shared_memory_bytes,
64 };
65 }
66 };
67
68 pub const KernelCallShapeProfileBounds = accy_choir.shape.Bounds;
69
70 pub const KernelCallShapeProfileDimension = struct {
71 name: []const u8,
72 runtime_scalar_argument_index: u32,
73 bounds: KernelCallShapeProfileBounds,
74 };
75
76 pub const KernelCallShapeProfile = struct {
77 name: []const u8,
78 fingerprint: u64,
79 dimensions: []const KernelCallShapeProfileDimension = &.{},
80
81 pub fn validate(self: KernelCallShapeProfile, runtime_scalar_argument_count: u32) gpu.BackendError!void {
82 if (self.name.len == 0) return error.InvalidArtifact;
83 if (self.dimensions.len == 0) return error.InvalidArtifact;
84 for (self.dimensions, 0..) |item, index| {
85 if (item.name.len == 0) return error.InvalidArtifact;
86 if (item.runtime_scalar_argument_index >= runtime_scalar_argument_count) return error.InvalidArtifact;
87 if (!item.bounds.valid()) return error.InvalidArtifact;
88 for (self.dimensions[0..index]) |previous| {
89 if (previous.runtime_scalar_argument_index == item.runtime_scalar_argument_index) return error.InvalidArtifact;
90 if (std.mem.eql(u8, previous.name, item.name)) return error.InvalidArtifact;
91 }
92 }
93 }
94
95 pub fn dimension(self: KernelCallShapeProfile, name: []const u8) ?KernelCallShapeProfileDimension {
96 for (self.dimensions) |item| {
97 if (std.mem.eql(u8, item.name, name)) return item;
98 }
99 return null;
100 }
101
102 pub fn runtimeScalarDimension(self: KernelCallShapeProfile, index: u32) ?KernelCallShapeProfileDimension {
103 for (self.dimensions) |item| {
104 if (item.runtime_scalar_argument_index == index) return item;
105 }
106 return null;
107 }
108 };
109
110 /// Records one prebuilt kernel in a registry, found by the triple of target
111 /// name, for example "accy.kernel.linalg.matmul5x7x3_4x2_f32", its `version`,
112 /// and its artifact format, so compiled programs can call it by name. The entry
113 /// also records its entry point name, argument count, required element types,
114 /// device features, code payload, launch rule, static arguments, and optional
115 /// named dimensions. The `version` field stands for what the kernel computes
116 /// and is separate from the byte-layout version of the registry file. A
117 /// maintainer raises `version` before an existing entry changes its math, the
118 /// meaning of its element types or memory layout, its argument list, its static
119 /// arguments, how its launch is derived, what its shape profile means, or its
120 /// schedule. A change that touches only the registry file's byte layout, or
121 /// that keeps behavior unchanged, leaves `version` as it is, because the
122 /// registry file carries its own layout version. A family target names one
123 /// reusable schedule, and sizes that leave that schedule unchanged arrive at
124 /// launch as bounded runtime scalar arguments.
125 pub const KernelCallArtifact = struct {
126 target: []const u8,
127 version: u32,
128 format: gpu.ArtifactFormat,
129 entry_name: []const u8,
130 argument_count: u32,
131 shape_family_fingerprint: ?u64 = null,
132 shape_profile: ?KernelCallShapeProfile = null,
133 required_dtypes: gpu.DTypeSet = .{},
134 required_features: choir_abi.Features = .{},
135 required_subgroup: choir_abi.SubgroupRequirements = .{},
136 /// Layout the emitter gave the entry's push-constant block, if any.
137 push_constants: choir_abi.PushConstants = .{},
138 payload: gpu.CompilePayload,
139 launch: KernelCallLaunch = .{ .derived = .{} },
140 element_count_argument: ElementCountArgument = .none,
141 runtime_scalar_argument_count: u32 = 0,
142 static_arguments: []const choir_abi.ScalarArgument = &.{},
143
144 /// Scalar arguments the entry's kernel takes, runtime and static
145 /// together. The rest of `argument_count` are buffers.
146 pub fn scalarArgumentCount(self: KernelCallArtifact) gpu.BackendError!u32 {
147 const static_count = std.math.cast(u32, self.static_arguments.len) orelse
148 return error.InvalidArtifact;
149 const count = std.math.add(u32, self.runtime_scalar_argument_count, static_count) catch
150 return error.InvalidArtifact;
151 if (count > self.argument_count) return error.InvalidArtifact;
152 return count;
153 }
154 };
155
156 pub const StandaloneKernelOptions = struct {
157 runtime_scalar_argument_count: u32 = 0,
158 static_arguments: []const choir_abi.ScalarArgument = &.{},
159 };
160
161 pub const KernelCallRegistryIndexSlot = struct {
162 occupied: bool = false,
163 target: []const u8 = &.{},
164 version: u32 = 0,
165 format: gpu.ArtifactFormat = .cuda_ptx,
166 entry_index: usize = 0,
167 };
168
169 pub const KernelCallRegistryIndex = struct {
170 slots: []const KernelCallRegistryIndexSlot = &.{},
171
172 pub fn findEntryIndex(
173 self: KernelCallRegistryIndex,
174 target: []const u8,
175 version: u32,
176 format: gpu.ArtifactFormat,
177 ) ?usize {
178 if (self.slots.len == 0) return null;
179 var slot_index = kernelCallRegistrySlotIndex(self.slots.len, target, version, format);
180 var probe_count: usize = 0;
181 while (probe_count < self.slots.len) : (probe_count += 1) {
182 const slot = self.slots[slot_index];
183 if (!slot.occupied) return null;
184 if (kernelCallRegistryKeyMatches(slot, target, version, format)) return slot.entry_index;
185 slot_index = (slot_index + 1) & (self.slots.len - 1);
186 }
187 return null;
188 }
189 };
190
191 pub const KernelCallRegistry = struct {
192 entries: []const KernelCallArtifact = &.{},
193 index: ?KernelCallRegistryIndex = null,
194
195 pub fn find(
196 self: KernelCallRegistry,
197 target: []const u8,
198 version: u32,
199 format: gpu.ArtifactFormat,
200 ) ?KernelCallArtifact {
201 if (self.index) |index| {
202 const entry_index = index.findEntryIndex(target, version, format) orelse return null;
203 if (entry_index >= self.entries.len) return null;
204 return self.entries[entry_index];
205 }
206 for (self.entries) |entry| {
207 if (entry.version != version) continue;
208 if (entry.format != format) continue;
209 if (!std.mem.eql(u8, entry.target, target)) continue;
210 return entry;
211 }
212 return null;
213 }
214 };
215
216 pub const KernelCallRegistryIndexError = error{
217 DuplicateKernelCallArtifact,
218 } || std.mem.Allocator.Error;
219
220 pub fn buildKernelCallRegistryIndex(
221 allocator: std.mem.Allocator,
222 entries: []const KernelCallArtifact,
223 ) KernelCallRegistryIndexError!KernelCallRegistryIndex {
224 if (entries.len == 0) return .{};
225 const slots = try allocator.alloc(KernelCallRegistryIndexSlot, try kernelCallRegistryIndexCapacity(entries.len));
226 @memset(slots, .{});
227 const index = KernelCallRegistryIndex{ .slots = slots };
228 errdefer deinitKernelCallRegistryIndex(allocator, index);
229 for (entries, 0..) |entry, entry_index| try insertKernelCallRegistryIndexEntry(slots, entry, entry_index);
230 return index;
231 }
232
233 pub fn deinitKernelCallRegistryIndex(
234 allocator: std.mem.Allocator,
235 index: KernelCallRegistryIndex,
236 ) void {
237 if (index.slots.len != 0) allocator.free(@constCast(index.slots));
238 }
239
240 fn insertKernelCallRegistryIndexEntry(
241 slots: []KernelCallRegistryIndexSlot,
242 entry: KernelCallArtifact,
243 entry_index: usize,
244 ) KernelCallRegistryIndexError!void {
245 var slot_index = kernelCallRegistrySlotIndex(slots.len, entry.target, entry.version, entry.format);
246 while (true) {
247 const slot = &slots[slot_index];
248 if (!slot.occupied) {
249 slot.* = .{
250 .occupied = true,
251 .target = entry.target,
252 .version = entry.version,
253 .format = entry.format,
254 .entry_index = entry_index,
255 };
256 return;
257 }
258 if (kernelCallRegistryKeyMatches(slot.*, entry.target, entry.version, entry.format)) {
259 return error.DuplicateKernelCallArtifact;
260 }
261 slot_index = (slot_index + 1) & (slots.len - 1);
262 }
263 }
264
265 fn kernelCallRegistryIndexCapacity(entry_count: usize) KernelCallRegistryIndexError!usize {
266 const wanted = std.math.mul(usize, entry_count, 2) catch return error.OutOfMemory;
267 var capacity: usize = 1;
268 while (capacity < wanted) {
269 capacity = std.math.mul(usize, capacity, 2) catch return error.OutOfMemory;
270 }
271 return capacity;
272 }
273
274 fn kernelCallRegistrySlotIndex(
275 slot_count: usize,
276 target: []const u8,
277 version: u32,
278 format: gpu.ArtifactFormat,
279 ) usize {
280 return @intCast(kernelCallRegistryKeyHash(target, version, format) & @as(u64, @intCast(slot_count - 1)));
281 }
282
283 fn kernelCallRegistryKeyHash(target: []const u8, version: u32, format: gpu.ArtifactFormat) u64 {
284 var hasher = std.hash.Wyhash.init(0);
285 hasher.update(target);
286 hasher.update(std.mem.asBytes(&version));
287 const format_tag: u8 = @backingInt(format);
288 hasher.update(std.mem.asBytes(&format_tag));
289 return hasher.final();
290 }
291
292 fn kernelCallRegistryKeyMatches(
293 slot: KernelCallRegistryIndexSlot,
294 target: []const u8,
295 version: u32,
296 format: gpu.ArtifactFormat,
297 ) bool {
298 return slot.version == version and
299 slot.format == format and
300 std.mem.eql(u8, slot.target, target);
301 }
302
303 test "kernel call registry index accelerates lookup without changing literal fallback" {
304 const entries = [_]KernelCallArtifact{
305 .{
306 .target = "accy.kernel.test.a",
307 .version = 1,
308 .format = .cuda_ptx,
309 .entry_name = "a",
310 .argument_count = 1,
311 .payload = .{ .text = "// a" },
312 },
313 .{
314 .target = "accy.kernel.test.b",
315 .version = 2,
316 .format = .vulkan_spirv,
317 .entry_name = "b",
318 .argument_count = 2,
319 .payload = .{ .words_u32 = &.{1} },
320 },
321 };
322
323 const literal = KernelCallRegistry{ .entries = entries[0..] };
324 try testing.expect(literal.index == null);
325 try testing.expect(literal.find("accy.kernel.test.b", 2, .vulkan_spirv) != null);
326
327 const index = try buildKernelCallRegistryIndex(testing.allocator, entries[0..]);
328 defer deinitKernelCallRegistryIndex(testing.allocator, index);
329 const indexed = KernelCallRegistry{ .entries = entries[0..], .index = index };
330 try testing.expect(indexed.find("accy.kernel.test.b", 2, .vulkan_spirv) != null);
331 try testing.expect(indexed.find("accy.kernel.test.b", 2, .cuda_ptx) == null);
332 try testing.expect(indexed.find("missing", 1, .cuda_ptx) == null);
333 }
334
335 test "kernel call registry index rejects duplicate keys" {
336 const entries = [_]KernelCallArtifact{
337 .{
338 .target = "accy.kernel.test.dup",
339 .version = 1,
340 .format = .cuda_ptx,
341 .entry_name = "dup_first",
342 .argument_count = 1,
343 .payload = .{ .text = "// first" },
344 },
345 .{
346 .target = "accy.kernel.test.dup",
347 .version = 1,
348 .format = .cuda_ptx,
349 .entry_name = "dup_second",
350 .argument_count = 1,
351 .payload = .{ .text = "// second" },
352 },
353 };
354 try testing.expectError(error.DuplicateKernelCallArtifact, buildKernelCallRegistryIndex(testing.allocator, entries[0..]));
355 }
356
357 pub fn duplicateKernelCallShapeProfile(
358 allocator: std.mem.Allocator,
359 profile: KernelCallShapeProfile,
360 ) !KernelCallShapeProfile {
361 const name = try allocator.dupe(u8, profile.name);
362 errdefer allocator.free(name);
363 const dimensions = try allocator.alloc(KernelCallShapeProfileDimension, profile.dimensions.len);
364 var copied_count: usize = 0;
365 errdefer {
366 for (dimensions[0..copied_count]) |dimension| allocator.free(dimension.name);
367 allocator.free(dimensions);
368 }
369 for (profile.dimensions, dimensions) |source, *destination| {
370 destination.* = source;
371 destination.name = try allocator.dupe(u8, source.name);
372 copied_count += 1;
373 }
374 return .{
375 .name = name,
376 .fingerprint = profile.fingerprint,
377 .dimensions = dimensions,
378 };
379 }
380
381 pub fn deinitKernelCallShapeProfile(
382 allocator: std.mem.Allocator,
383 profile: KernelCallShapeProfile,
384 ) void {
385 allocator.free(profile.name);
386 for (profile.dimensions) |dimension| allocator.free(dimension.name);
387 allocator.free(profile.dimensions);
388 }
389
390 pub fn kernelCallShapeProfileEql(lhs: ?KernelCallShapeProfile, rhs: ?KernelCallShapeProfile) bool {
391 if (lhs == null and rhs == null) return true;
392 if (lhs == null or rhs == null) return false;
393 return kernelCallShapeProfileValueEql(lhs.?, rhs.?);
394 }
395
396 fn kernelCallShapeProfileValueEql(lhs: KernelCallShapeProfile, rhs: KernelCallShapeProfile) bool {
397 if (!std.mem.eql(u8, lhs.name, rhs.name)) return false;
398 if (lhs.fingerprint != rhs.fingerprint) return false;
399 if (lhs.dimensions.len != rhs.dimensions.len) return false;
400 for (lhs.dimensions, rhs.dimensions) |left, right| {
401 if (!std.mem.eql(u8, left.name, right.name)) return false;
402 if (left.runtime_scalar_argument_index != right.runtime_scalar_argument_index) return false;
403 if (!std.meta.eql(left.bounds, right.bounds)) return false;
404 }
405 return true;
406 }
407
408 fn runtimeU32CeilDiv(
409 runtime_scalar_arguments: []const choir_abi.ScalarArgument,
410 axis: KernelCallDerivedLaunchAxis.RuntimeU32CeilDiv,
411 ) gpu.BackendError!u32 {
412 if (axis.divisor == 0) return error.InvalidArtifact;
413 const argument_index: usize = @intCast(axis.argument_index);
414 if (argument_index >= runtime_scalar_arguments.len) return error.LaunchArgumentMismatch;
415 const value = switch (runtime_scalar_arguments[argument_index]) {
416 .u32 => |payload| payload,
417 else => return error.LaunchArgumentMismatch,
418 };
419 const biased = std.math.add(u32, value, axis.divisor - 1) catch return error.LaunchArgumentMismatch;
420 const extent = biased / axis.divisor;
421 if (extent == 0) return error.LaunchArgumentMismatch;
422 return extent;
423 }
424
425 pub fn validateThreadgroup(threadgroup: [3]u32) gpu.BackendError![3]u32 {
426 for (threadgroup) |extent| {
427 if (extent == 0) return error.InvalidArtifact;
428 }
429 return threadgroup;
430 }
431
432 pub fn defaultArtifactFormat(kind: gpu.BackendKind) ?gpu.ArtifactFormat {
433 return switch (kind) {
434 .cuda => .cuda_ptx,
435 .vulkan => .vulkan_spirv,
436 .metal => .metal_msl,
437 .webgpu => .webgpu_wgsl,
438 .cpu => .cpu_object,
439 .wasm => .webassembly_module,
440 else => null,
441 };
442 }
443
444 test "default artifact format includes host execution formats" {
445 try testing.expectEqual(gpu.ArtifactFormat.cuda_ptx, defaultArtifactFormat(.cuda).?);
446 try testing.expectEqual(gpu.ArtifactFormat.vulkan_spirv, defaultArtifactFormat(.vulkan).?);
447 try testing.expectEqual(gpu.ArtifactFormat.metal_msl, defaultArtifactFormat(.metal).?);
448 try testing.expectEqual(gpu.ArtifactFormat.webgpu_wgsl, defaultArtifactFormat(.webgpu).?);
449 try testing.expectEqual(gpu.ArtifactFormat.cpu_object, defaultArtifactFormat(.cpu).?);
450 try testing.expectEqual(gpu.ArtifactFormat.webassembly_module, defaultArtifactFormat(.wasm).?);
451 try testing.expect(defaultArtifactFormat(.external) == null);
452 }
453
454 test "kernel call derived launch computes geometry from runtime scalars" {
455 const launch = KernelCallDerivedLaunch{
456 .grid = .{
457 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 8 } },
458 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = 4 } },
459 .{ .fixed = 1 },
460 },
461 .threadgroup = .{ 8, 4, 1 },
462 .dynamic_shared_memory_bytes = 512,
463 };
464
465 const scalars = [_]choir_abi.ScalarArgument{
466 .{ .u32 = 17 },
467 .{ .u32 = 9 },
468 };
469 const geometry = try launch.geometry(scalars[0..]);
470
471 try std.testing.expectEqual(@as(u32, 3), geometry.grid[0]);
472 try std.testing.expectEqual(@as(u32, 3), geometry.grid[1]);
473 try std.testing.expectEqual(@as(u32, 1), geometry.grid[2]);
474 try std.testing.expectEqual(@as(u32, 8), geometry.threadgroup[0]);
475 try std.testing.expectEqual(@as(u32, 4), geometry.threadgroup[1]);
476 try std.testing.expectEqual(@as(u32, 1), geometry.threadgroup[2]);
477 try std.testing.expectEqual(@as(u32, 512), geometry.dynamic_shared_memory_bytes);
478 try std.testing.expectError(error.LaunchArgumentMismatch, launch.geometry(&.{.{ .u32 = 17 }}));
479 try std.testing.expectError(error.LaunchArgumentMismatch, launch.geometry(&.{ .{ .u32 = 17 }, .{ .i32 = 9 } }));
480 }