lib/accy/src/artifact/model/wire.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 //! A byte encoding for a table of prebuilt kernels and the chains of kernels
2 //! that call them, with a writer that produces the bytes and a reader that
3 //! checks them and builds the table again.
4 //!
5 //! A program compiled ahead of time has to carry its device code inside its own
6 //! binary and look up each kernel by name at startup, because the compiler is
7 //! absent. The bytes come from a separate build step, possibly under an older
8 //! version of this code, so the layout may differ from what the reader expects.
9 //! A count or a tag read from the input decides how much memory to take and
10 //! which case to parse, so a damaged input can point past its own end. The byte
11 //! layout and the meaning of a kernel change for different reasons and at
12 //! different times.
13 //!
14 //! The bytes carrying prebuilt kernels (*registry blob*) open with the eight
15 //! characters "ACCYKCRG" and a layout version, and every integer after them is
16 //! little-endian. The layout version counts changes to the encoding alone, and
17 //! each kernel entry carries its own version for what the kernel computes. The
18 //! decoder copies every field of variable length into one arena that the
19 //! decoded table owns. The decoder rejects a wrong magic, a version it does not
20 //! read, input that ends early, bytes left over at the end, two entries or two
21 //! chains with the same identity, a tag it does not know, and a push-constant
22 //! layout that lists too many members or whose offsets and sizes do not fit
23 //! together. Two counts, the static arguments of an entry and the dimensions of
24 //! a shape profile, size their allocation before any check against the bytes
25 //! that remain, so a damaged input can end in `error.OutOfMemory`. The same
26 //! little-endian writer and reader (`ByteWriter`, `ByteReader`) also serve the
27 //! schedule tuning files.
28 //!
29 //! - *registry entry*: one prebuilt kernel identified by its target name,
30 //! version and artifact format, with its entry point, argument count, code
31 //! payload and launch rule.
32 //! - *kernel pipeline*: an ordered chain of registry entries launched one after
33 //! another over shared scratch buffers.
34
35 const std = @import("std");
36 const gpu = @import("gpu");
37 const choir_abi = @import("choir_abi");
38
39 const pipeline_mod = @import("pipeline.zig");
40 const plan = @import("registry.zig");
41
42 /// The layout version written after the magic, and a registry blob that carries
43 /// any other number fails with `error.UnsupportedRegistryVersion`. A maintainer
44 /// reads this number before changing the encoding and raises it when the order
45 /// or width of an encoded field changes, when a tag changes meaning, or when
46 /// the number an enum value encodes as changes. A change to what a kernel
47 /// computes, with the bytes laid out as before, raises
48 /// `KernelCallArtifact.version` on that entry and leaves this number alone.
49 pub const format_version: u32 = 10;
50
51 const magic = "ACCYKCRG";
52 const minimum_entry_bytes: usize = 32;
53 const minimum_pipeline_bytes: usize = 20;
54
55 pub const Error = error{
56 InvalidRegistryBlob,
57 UnsupportedRegistryVersion,
58 TruncatedRegistryBlob,
59 InvalidRegistryTag,
60 } || std.mem.Allocator.Error;
61
62 /// The table `decode` builds: the entries, their lookup index, the kernel
63 /// chains and every payload byte, all allocated in one arena. A caller holds
64 /// this for as long as it looks kernels up in the decoded table. `registry`
65 /// gives a lookup view, and that view has no hash index when the table holds no
66 /// entries. Every slice stays valid until `deinit`, which frees the arena at
67 /// once.
68 pub const OwnedRegistry = struct {
69 arena: std.heap.ArenaAllocator,
70 entries: []plan.KernelCallArtifact = &.{},
71 index: plan.KernelCallRegistryIndex = .{},
72 pipelines: []pipeline_mod.KernelCallPipeline = &.{},
73
74 pub fn registry(self: *const OwnedRegistry) plan.KernelCallRegistry {
75 return .{
76 .entries = self.entries,
77 .index = if (self.index.slots.len == 0) null else self.index,
78 };
79 }
80
81 pub fn deinit(self: *OwnedRegistry) void {
82 self.arena.deinit();
83 self.* = undefined;
84 }
85 };
86
87 /// Writes the entries and the kernel chains into one new byte slice allocated
88 /// with `allocator`, and the caller frees that slice. A build step calls this
89 /// to produce the bytes it embeds in a program. Each list and each byte field
90 /// must hold fewer than 2^32 items, and the encoder does not check this limit.
91 pub fn encode(
92 allocator: std.mem.Allocator,
93 entries: []const plan.KernelCallArtifact,
94 pipelines: []const pipeline_mod.KernelCallPipeline,
95 ) ![]u8 {
96 var bytes: std.ArrayListUnmanaged(u8) = .empty;
97 errdefer bytes.deinit(allocator);
98
99 try bytes.appendSlice(allocator, magic);
100 try appendInt(&bytes, allocator, u32, format_version);
101 try appendInt(&bytes, allocator, u32, @intCast(entries.len));
102 for (entries) |entry| try encodeEntry(&bytes, allocator, entry);
103 try appendInt(&bytes, allocator, u32, @intCast(pipelines.len));
104 for (pipelines) |entry| try encodePipeline(&bytes, allocator, entry);
105 return bytes.toOwnedSlice(allocator);
106 }
107
108 /// Checks one whole registry blob and copies it into a new `OwnedRegistry`, so
109 /// the caller may free `bytes` as soon as the call returns. A program calls
110 /// this at startup on its embedded bytes to get a table it can search. The
111 /// errors are `error.InvalidRegistryBlob` for a wrong magic, leftover bytes, a
112 /// zero chain version, a repeated identity or a push-constant layout that does
113 /// not fit together, `error.UnsupportedRegistryVersion`,
114 /// `error.TruncatedRegistryBlob`, `error.InvalidRegistryTag`, and
115 /// `error.OutOfMemory`. A blob that claims more static arguments or shape
116 /// dimensions than it holds can fail with `error.OutOfMemory` before the
117 /// decoder sees that it is short.
118 pub fn decode(backing_allocator: std.mem.Allocator, bytes: []const u8) Error!OwnedRegistry {
119 var owned = OwnedRegistry{ .arena = std.heap.ArenaAllocator.init(backing_allocator) };
120 errdefer owned.arena.deinit();
121 const arena = owned.arena.allocator();
122
123 var cursor = ByteReader{ .bytes = bytes };
124 const magic_bytes = try take(&cursor, magic.len);
125 if (!std.mem.eql(u8, magic_bytes, magic)) return error.InvalidRegistryBlob;
126 const version = try readInt(&cursor, u32);
127 if (version != format_version) return error.UnsupportedRegistryVersion;
128
129 const count = try readInt(&cursor, u32);
130 if (count > cursor.remaining() / minimum_entry_bytes) return error.TruncatedRegistryBlob;
131 const entries = try arena.alloc(plan.KernelCallArtifact, count);
132 for (entries) |*entry| entry.* = try decodeEntry(&cursor, arena);
133 const pipeline_count = try readInt(&cursor, u32);
134 if (pipeline_count > cursor.remaining() / minimum_pipeline_bytes) return error.TruncatedRegistryBlob;
135 const pipelines = try arena.alloc(pipeline_mod.KernelCallPipeline, pipeline_count);
136 for (pipelines, 0..) |*entry, decoded| {
137 entry.* = try decodePipeline(&cursor, arena);
138 if (entry.version == 0) return error.InvalidRegistryBlob;
139 if (pipeline_mod.findPipeline(pipelines[0..decoded], entry.target, entry.version) != null) {
140 return error.InvalidRegistryBlob;
141 }
142 }
143 if (cursor.remaining() != 0) return error.InvalidRegistryBlob;
144 const index = plan.buildKernelCallRegistryIndex(arena, entries) catch |err| switch (err) {
145 error.DuplicateKernelCallArtifact => return error.InvalidRegistryBlob,
146 error.OutOfMemory => return error.OutOfMemory,
147 };
148
149 owned.entries = entries;
150 owned.index = index;
151 owned.pipelines = pipelines;
152 return owned;
153 }
154
155 fn encodeEntry(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, entry: plan.KernelCallArtifact) !void {
156 try appendSized(bytes, allocator, entry.target);
157 try appendInt(bytes, allocator, u32, entry.version);
158 try bytes.append(allocator, @backingInt(entry.format));
159 try appendSized(bytes, allocator, entry.entry_name);
160 try appendInt(bytes, allocator, u32, entry.argument_count);
161 try appendInt(bytes, allocator, u64, entry.required_dtypes.bits);
162 try appendInt(bytes, allocator, u16, packFeatures(entry.required_features));
163 try bytes.append(allocator, packSubgroupFlags(entry.required_subgroup));
164 try appendInt(bytes, allocator, u32, entry.required_subgroup.size_min);
165 try appendInt(bytes, allocator, u32, entry.required_subgroup.size_max);
166 try encodePayload(bytes, allocator, entry.payload);
167 try encodePushConstants(bytes, allocator, &entry.push_constants);
168 try encodeLaunch(bytes, allocator, entry.launch);
169 try bytes.append(allocator, @backingInt(entry.element_count_argument));
170 try appendInt(bytes, allocator, u32, entry.runtime_scalar_argument_count);
171 try appendInt(bytes, allocator, u32, @intCast(entry.static_arguments.len));
172 for (entry.static_arguments) |argument| try encodeScalar(bytes, allocator, argument);
173 try encodeOptionalU64(bytes, allocator, entry.shape_family_fingerprint);
174 try encodeShapeProfile(bytes, allocator, entry.shape_profile);
175 }
176
177 fn decodeEntry(cursor: *ByteReader, arena: std.mem.Allocator) Error!plan.KernelCallArtifact {
178 const target = try takeSizedCopy(cursor, arena);
179 const version = try readInt(cursor, u32);
180 const format = try readEnum(cursor, gpu.ArtifactFormat);
181 const entry_name = try takeSizedCopy(cursor, arena);
182 const argument_count = try readInt(cursor, u32);
183 const dtype_bits = try readInt(cursor, u64);
184 const features = unpackFeatures(try readInt(cursor, u16));
185 var subgroup = unpackSubgroupFlags(try readByte(cursor));
186 subgroup.size_min = try readInt(cursor, u32);
187 subgroup.size_max = try readInt(cursor, u32);
188 const payload = try decodePayload(cursor, arena);
189 const push_constants = try decodePushConstants(cursor);
190 const launch = try decodeLaunch(cursor);
191 const element_count_argument = try readEnum(cursor, plan.ElementCountArgument);
192 const runtime_scalar_argument_count = try readInt(cursor, u32);
193 const static_argument_count = try readInt(cursor, u32);
194 const static_arguments = try arena.alloc(choir_abi.ScalarArgument, static_argument_count);
195 for (static_arguments) |*argument| argument.* = try decodeScalar(cursor);
196 const shape_family_fingerprint = try decodeOptionalU64(cursor);
197 const shape_profile = try decodeShapeProfile(cursor, arena);
198
199 return .{
200 .target = target,
201 .version = version,
202 .format = format,
203 .entry_name = entry_name,
204 .argument_count = argument_count,
205 .shape_family_fingerprint = shape_family_fingerprint,
206 .shape_profile = shape_profile,
207 .required_dtypes = .{ .bits = dtype_bits },
208 .required_features = features,
209 .required_subgroup = subgroup,
210 .push_constants = push_constants,
211 .payload = payload,
212 .launch = launch,
213 .element_count_argument = element_count_argument,
214 .runtime_scalar_argument_count = runtime_scalar_argument_count,
215 .static_arguments = static_arguments,
216 };
217 }
218
219 fn encodePushConstants(
220 bytes: *std.ArrayListUnmanaged(u8),
221 allocator: std.mem.Allocator,
222 layout: *const choir_abi.PushConstants,
223 ) !void {
224 std.debug.assert(layout.valid());
225 try bytes.append(allocator, layout.count);
226 try bytes.append(allocator, layout.byte_size);
227 try bytes.appendSlice(allocator, layout.offsets[0..layout.count]);
228 try bytes.appendSlice(allocator, layout.sizes[0..layout.count]);
229 }
230
231 fn decodePushConstants(cursor: *ByteReader) Error!choir_abi.PushConstants {
232 var layout: choir_abi.PushConstants = .{};
233 layout.count = try readByte(cursor);
234 if (layout.count > choir_abi.max_push_constant_members) return error.InvalidRegistryBlob;
235 layout.byte_size = try readByte(cursor);
236 @memcpy(layout.offsets[0..layout.count], try take(cursor, layout.count));
237 @memcpy(layout.sizes[0..layout.count], try take(cursor, layout.count));
238 if (!layout.valid()) return error.InvalidRegistryBlob;
239 return layout;
240 }
241
242 fn encodePayload(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, payload: gpu.CompilePayload) !void {
243 switch (payload) {
244 .none => try bytes.append(allocator, 0),
245 .bytes => |data| {
246 try bytes.append(allocator, 1);
247 try appendInt(bytes, allocator, u64, data.len);
248 try bytes.appendSlice(allocator, data);
249 },
250 .words_u32 => |words| {
251 try bytes.append(allocator, 2);
252 try appendInt(bytes, allocator, u64, words.len);
253 for (words) |word| try appendInt(bytes, allocator, u32, word);
254 },
255 .text => |text| {
256 try bytes.append(allocator, 3);
257 try appendInt(bytes, allocator, u64, text.len);
258 try bytes.appendSlice(allocator, text);
259 },
260 }
261 }
262
263 fn decodePayload(cursor: *ByteReader, arena: std.mem.Allocator) Error!gpu.CompilePayload {
264 return switch (try readByte(cursor)) {
265 0 => .none,
266 1 => .{ .bytes = try takeCopy(cursor, arena, try readLength(cursor)) },
267 2 => blk: {
268 const count = try readLength(cursor);
269 if (count > cursor.remaining() / @sizeOf(u32)) return error.TruncatedRegistryBlob;
270 const words = try arena.alloc(u32, count);
271 for (words) |*word| word.* = try readInt(cursor, u32);
272 break :blk .{ .words_u32 = words };
273 },
274 3 => .{ .text = try takeCopy(cursor, arena, try readLength(cursor)) },
275 else => error.InvalidRegistryTag,
276 };
277 }
278
279 fn encodeLaunch(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, launch: plan.KernelCallLaunch) !void {
280 switch (launch) {
281 .derived => |derived| {
282 try bytes.append(allocator, 0);
283 try encodeDerivedLaunch(bytes, allocator, derived);
284 },
285 .fixed => |geometry| {
286 try bytes.append(allocator, 1);
287 for (geometry.grid) |extent| try appendInt(bytes, allocator, u32, extent);
288 for (geometry.threadgroup) |extent| try appendInt(bytes, allocator, u32, extent);
289 try appendInt(bytes, allocator, u32, geometry.dynamic_shared_memory_bytes);
290 },
291 }
292 }
293
294 fn decodeLaunch(cursor: *ByteReader) Error!plan.KernelCallLaunch {
295 return switch (try readByte(cursor)) {
296 0 => .{ .derived = try decodeDerivedLaunch(cursor) },
297 1 => blk: {
298 var geometry = choir_abi.LaunchGeometry{};
299 for (&geometry.grid) |*extent| extent.* = try readInt(cursor, u32);
300 for (&geometry.threadgroup) |*extent| extent.* = try readInt(cursor, u32);
301 geometry.dynamic_shared_memory_bytes = try readInt(cursor, u32);
302 break :blk .{ .fixed = geometry };
303 },
304 else => error.InvalidRegistryTag,
305 };
306 }
307
308 fn encodeDerivedLaunch(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, launch: plan.KernelCallDerivedLaunch) !void {
309 for (launch.grid) |axis| try encodeDerivedLaunchAxis(bytes, allocator, axis);
310 for (launch.threadgroup) |extent| try appendInt(bytes, allocator, u32, extent);
311 try appendInt(bytes, allocator, u32, launch.dynamic_shared_memory_bytes);
312 }
313
314 fn decodeDerivedLaunch(cursor: *ByteReader) Error!plan.KernelCallDerivedLaunch {
315 var launch = plan.KernelCallDerivedLaunch{};
316 for (&launch.grid) |*axis| axis.* = try decodeDerivedLaunchAxis(cursor);
317 for (&launch.threadgroup) |*extent| extent.* = try readInt(cursor, u32);
318 launch.dynamic_shared_memory_bytes = try readInt(cursor, u32);
319 return launch;
320 }
321
322 fn encodeDerivedLaunchAxis(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, axis: plan.KernelCallDerivedLaunchAxis) !void {
323 switch (axis) {
324 .fixed => |extent| {
325 try bytes.append(allocator, 0);
326 try appendInt(bytes, allocator, u32, extent);
327 },
328 .runtime_u32_ceil_div => |runtime| {
329 try bytes.append(allocator, 1);
330 try appendInt(bytes, allocator, u32, runtime.argument_index);
331 try appendInt(bytes, allocator, u32, runtime.divisor);
332 },
333 }
334 }
335
336 fn decodeDerivedLaunchAxis(cursor: *ByteReader) Error!plan.KernelCallDerivedLaunchAxis {
337 return switch (try readByte(cursor)) {
338 0 => .{ .fixed = try readInt(cursor, u32) },
339 1 => .{ .runtime_u32_ceil_div = .{
340 .argument_index = try readInt(cursor, u32),
341 .divisor = try readInt(cursor, u32),
342 } },
343 else => error.InvalidRegistryTag,
344 };
345 }
346
347 fn encodePipeline(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, entry: pipeline_mod.KernelCallPipeline) !void {
348 try appendSized(bytes, allocator, entry.target);
349 try appendInt(bytes, allocator, u32, entry.version);
350 try appendInt(bytes, allocator, u32, entry.operand_count);
351 try appendInt(bytes, allocator, u32, entry.result_count);
352 try appendInt(bytes, allocator, u32, entry.runtime_scalar_argument_count);
353 try appendInt(bytes, allocator, u32, @intCast(entry.runtime_scalar_bounds.len));
354 for (entry.runtime_scalar_bounds) |bound| {
355 try appendInt(bytes, allocator, u32, bound.argument_index);
356 try appendInt(bytes, allocator, u32, bound.max_u32);
357 }
358 try appendInt(bytes, allocator, u32, @intCast(entry.intermediates.len));
359 for (entry.intermediates) |intermediate| {
360 try bytes.append(allocator, @backingInt(intermediate.dtype));
361 try encodeScalarDerivation(bytes, allocator, intermediate.extent);
362 }
363 try appendInt(bytes, allocator, u32, @intCast(entry.stages.len));
364 for (entry.stages) |stage| {
365 try appendSized(bytes, allocator, stage.target);
366 try appendInt(bytes, allocator, u32, stage.version);
367 try appendInt(bytes, allocator, u32, @intCast(stage.buffers.len));
368 for (stage.buffers) |ref| try encodeValueRef(bytes, allocator, ref);
369 try appendInt(bytes, allocator, u32, @intCast(stage.scalars.len));
370 for (stage.scalars) |derivation| try encodeScalarDerivation(bytes, allocator, derivation);
371 }
372 }
373
374 fn decodePipeline(cursor: *ByteReader, arena: std.mem.Allocator) Error!pipeline_mod.KernelCallPipeline {
375 const target = try takeSizedCopy(cursor, arena);
376 const version = try readInt(cursor, u32);
377 const operand_count = try readInt(cursor, u32);
378 const result_count = try readInt(cursor, u32);
379 const runtime_scalar_argument_count = try readInt(cursor, u32);
380 const runtime_scalar_bound_count = try readInt(cursor, u32);
381 if (runtime_scalar_bound_count > cursor.remaining() / 8) return error.TruncatedRegistryBlob;
382 const runtime_scalar_bounds = try arena.alloc(pipeline_mod.PipelineRuntimeScalarBound, runtime_scalar_bound_count);
383 for (runtime_scalar_bounds) |*bound| {
384 bound.* = .{
385 .argument_index = try readInt(cursor, u32),
386 .max_u32 = try readInt(cursor, u32),
387 };
388 }
389 const intermediate_count = try readInt(cursor, u32);
390 if (intermediate_count > cursor.remaining()) return error.TruncatedRegistryBlob;
391 const intermediates = try arena.alloc(pipeline_mod.PipelineIntermediate, intermediate_count);
392 for (intermediates) |*intermediate| {
393 intermediate.* = .{
394 .dtype = try readEnum(cursor, choir_abi.DType),
395 .extent = try decodeScalarDerivation(cursor),
396 };
397 }
398 const stage_count = try readInt(cursor, u32);
399 if (stage_count > cursor.remaining()) return error.TruncatedRegistryBlob;
400 const stages = try arena.alloc(pipeline_mod.PipelineStage, stage_count);
401 for (stages) |*stage| {
402 const stage_target = try takeSizedCopy(cursor, arena);
403 const stage_version = try readInt(cursor, u32);
404 const buffer_count = try readInt(cursor, u32);
405 if (buffer_count > cursor.remaining()) return error.TruncatedRegistryBlob;
406 const buffers = try arena.alloc(pipeline_mod.PipelineValueRef, buffer_count);
407 for (buffers) |*ref| ref.* = try decodeValueRef(cursor);
408 const scalar_count = try readInt(cursor, u32);
409 if (scalar_count > cursor.remaining()) return error.TruncatedRegistryBlob;
410 const scalars = try arena.alloc(pipeline_mod.PipelineScalarDerivation, scalar_count);
411 for (scalars) |*derivation| derivation.* = try decodeScalarDerivation(cursor);
412 stage.* = .{
413 .target = stage_target,
414 .version = stage_version,
415 .buffers = buffers,
416 .scalars = scalars,
417 };
418 }
419 return .{
420 .target = target,
421 .version = version,
422 .operand_count = operand_count,
423 .result_count = result_count,
424 .runtime_scalar_argument_count = runtime_scalar_argument_count,
425 .runtime_scalar_bounds = runtime_scalar_bounds,
426 .intermediates = intermediates,
427 .stages = stages,
428 };
429 }
430
431 fn encodeValueRef(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, ref: pipeline_mod.PipelineValueRef) !void {
432 switch (ref) {
433 .operand => |index| {
434 try bytes.append(allocator, 0);
435 try appendInt(bytes, allocator, u32, index);
436 },
437 .result => |index| {
438 try bytes.append(allocator, 1);
439 try appendInt(bytes, allocator, u32, index);
440 },
441 .intermediate => |index| {
442 try bytes.append(allocator, 2);
443 try appendInt(bytes, allocator, u32, index);
444 },
445 }
446 }
447
448 fn decodeValueRef(cursor: *ByteReader) Error!pipeline_mod.PipelineValueRef {
449 return switch (try readByte(cursor)) {
450 0 => .{ .operand = try readInt(cursor, u32) },
451 1 => .{ .result = try readInt(cursor, u32) },
452 2 => .{ .intermediate = try readInt(cursor, u32) },
453 else => error.InvalidRegistryTag,
454 };
455 }
456
457 fn encodeScalarDerivation(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, derivation: pipeline_mod.PipelineScalarDerivation) !void {
458 switch (derivation) {
459 .forward => |index| {
460 try bytes.append(allocator, 0);
461 try appendInt(bytes, allocator, u32, index);
462 },
463 .ceil_div => |axis| {
464 try bytes.append(allocator, 1);
465 try appendInt(bytes, allocator, u32, axis.argument_index);
466 try appendInt(bytes, allocator, u32, axis.divisor);
467 },
468 .ceil_div_scaled => |scaled| {
469 try bytes.append(allocator, 2);
470 try appendInt(bytes, allocator, u32, scaled.argument_index);
471 try appendInt(bytes, allocator, u32, scaled.divisor);
472 try appendInt(bytes, allocator, u32, scaled.scale);
473 },
474 .ceil_div_scaled_by_arg => |scaled| {
475 try bytes.append(allocator, 3);
476 try appendInt(bytes, allocator, u32, scaled.argument_index);
477 try appendInt(bytes, allocator, u32, scaled.divisor);
478 try appendInt(bytes, allocator, u32, scaled.scale_argument_index);
479 },
480 }
481 }
482
483 fn decodeScalarDerivation(cursor: *ByteReader) Error!pipeline_mod.PipelineScalarDerivation {
484 return switch (try readByte(cursor)) {
485 0 => .{ .forward = try readInt(cursor, u32) },
486 1 => .{ .ceil_div = .{
487 .argument_index = try readInt(cursor, u32),
488 .divisor = try readInt(cursor, u32),
489 } },
490 2 => .{ .ceil_div_scaled = .{
491 .argument_index = try readInt(cursor, u32),
492 .divisor = try readInt(cursor, u32),
493 .scale = try readInt(cursor, u32),
494 } },
495 3 => .{ .ceil_div_scaled_by_arg = .{
496 .argument_index = try readInt(cursor, u32),
497 .divisor = try readInt(cursor, u32),
498 .scale_argument_index = try readInt(cursor, u32),
499 } },
500 else => error.InvalidRegistryTag,
501 };
502 }
503
504 fn encodeScalar(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, argument: choir_abi.ScalarArgument) !void {
505 try bytes.append(allocator, @backingInt(argument));
506 const value: u64 = switch (argument) {
507 .i32 => |v| @as(u32, @bitCast(v)),
508 .u32 => |v| v,
509 .i64 => |v| @as(u64, @bitCast(v)),
510 .u64 => |v| v,
511 .f32 => |v| @as(u32, @bitCast(v)),
512 .f64 => |v| @as(u64, @bitCast(v)),
513 };
514 try appendInt(bytes, allocator, u64, value);
515 }
516
517 fn encodeOptionalU64(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, value: ?u64) !void {
518 try bytes.append(allocator, @intFromBool(value != null));
519 if (value) |payload| try appendInt(bytes, allocator, u64, payload);
520 }
521
522 fn decodeOptionalU64(cursor: *ByteReader) Error!?u64 {
523 return switch (try readByte(cursor)) {
524 0 => null,
525 1 => try readInt(cursor, u64),
526 else => error.InvalidRegistryTag,
527 };
528 }
529
530 fn encodeShapeProfile(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, profile: ?plan.KernelCallShapeProfile) !void {
531 if (profile) |value| {
532 try bytes.append(allocator, 1);
533 try appendSized(bytes, allocator, value.name);
534 try appendInt(bytes, allocator, u64, value.fingerprint);
535 try appendInt(bytes, allocator, u32, @intCast(value.dimensions.len));
536 for (value.dimensions) |dimension| {
537 try appendSized(bytes, allocator, dimension.name);
538 try appendInt(bytes, allocator, u32, dimension.runtime_scalar_argument_index);
539 try encodeBounds(bytes, allocator, dimension.bounds);
540 }
541 } else {
542 try bytes.append(allocator, 0);
543 }
544 }
545
546 fn decodeShapeProfile(cursor: *ByteReader, arena: std.mem.Allocator) Error!?plan.KernelCallShapeProfile {
547 return switch (try readByte(cursor)) {
548 0 => null,
549 1 => blk: {
550 const name = try takeSizedCopy(cursor, arena);
551 const fingerprint = try readInt(cursor, u64);
552 const dimension_count = try readInt(cursor, u32);
553 const dimensions = try arena.alloc(plan.KernelCallShapeProfileDimension, dimension_count);
554 for (dimensions) |*dimension| {
555 dimension.* = .{
556 .name = try takeSizedCopy(cursor, arena),
557 .runtime_scalar_argument_index = try readInt(cursor, u32),
558 .bounds = try decodeBounds(cursor),
559 };
560 }
561 break :blk .{
562 .name = name,
563 .fingerprint = fingerprint,
564 .dimensions = dimensions,
565 };
566 },
567 else => error.InvalidRegistryTag,
568 };
569 }
570
571 fn encodeBounds(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, bounds: anytype) !void {
572 try encodeOptionalU64(bytes, allocator, bounds.min);
573 try encodeOptionalU64(bytes, allocator, bounds.opt);
574 try encodeOptionalU64(bytes, allocator, bounds.max);
575 }
576
577 fn decodeBounds(cursor: *ByteReader) Error!plan.KernelCallShapeProfileBounds {
578 return .{
579 .min = try decodeOptionalU64(cursor),
580 .opt = try decodeOptionalU64(cursor),
581 .max = try decodeOptionalU64(cursor),
582 };
583 }
584
585 fn decodeScalar(cursor: *ByteReader) Error!choir_abi.ScalarArgument {
586 const tag = try readEnum(cursor, std.meta.Tag(choir_abi.ScalarArgument));
587 const value = try readInt(cursor, u64);
588 return switch (tag) {
589 .i32 => .{ .i32 = @bitCast(@as(u32, @truncate(value))) },
590 .u32 => .{ .u32 = @truncate(value) },
591 .i64 => .{ .i64 = @bitCast(value) },
592 .u64 => .{ .u64 = value },
593 .f32 => .{ .f32 = @bitCast(@as(u32, @truncate(value))) },
594 .f64 => .{ .f64 = @bitCast(value) },
595 };
596 }
597
598 fn packFeatures(features: choir_abi.Features) u16 {
599 var packed_bits: u16 = 0;
600 if (features.atomic_i32) packed_bits |= 1 << 0;
601 if (features.atomic_u32) packed_bits |= 1 << 1;
602 if (features.atomic_index) packed_bits |= 1 << 2;
603 if (features.atomic_f32_add_device) packed_bits |= 1 << 3;
604 if (features.atomic_f32_add_shared) packed_bits |= 1 << 4;
605 if (features.unsupported_atomic) packed_bits |= 1 << 5;
606 if (features.async_copy) packed_bits |= 1 << 6;
607 if (features.tensor_cores) packed_bits |= 1 << 7;
608 if (features.cooperative_matrix) packed_bits |= 1 << 8;
609 if (features.dynamic_shared_memory) packed_bits |= 1 << 9;
610 if (features.indirect_launch) packed_bits |= 1 << 10;
611 return packed_bits;
612 }
613
614 fn unpackFeatures(packed_bits: u16) choir_abi.Features {
615 return .{
616 .atomic_i32 = packed_bits & (1 << 0) != 0,
617 .atomic_u32 = packed_bits & (1 << 1) != 0,
618 .atomic_index = packed_bits & (1 << 2) != 0,
619 .atomic_f32_add_device = packed_bits & (1 << 3) != 0,
620 .atomic_f32_add_shared = packed_bits & (1 << 4) != 0,
621 .unsupported_atomic = packed_bits & (1 << 5) != 0,
622 .async_copy = packed_bits & (1 << 6) != 0,
623 .tensor_cores = packed_bits & (1 << 7) != 0,
624 .cooperative_matrix = packed_bits & (1 << 8) != 0,
625 .dynamic_shared_memory = packed_bits & (1 << 9) != 0,
626 .indirect_launch = packed_bits & (1 << 10) != 0,
627 };
628 }
629
630 fn packSubgroupFlags(subgroup: choir_abi.SubgroupRequirements) u8 {
631 var packed_bits: u8 = 0;
632 if (subgroup.supported) packed_bits |= 1 << 0;
633 if (subgroup.shuffle) packed_bits |= 1 << 1;
634 if (subgroup.ballot) packed_bits |= 1 << 2;
635 if (subgroup.vote) packed_bits |= 1 << 3;
636 if (subgroup.arithmetic) packed_bits |= 1 << 4;
637 if (subgroup.scan) packed_bits |= 1 << 5;
638 return packed_bits;
639 }
640
641 fn unpackSubgroupFlags(packed_bits: u8) choir_abi.SubgroupRequirements {
642 return .{
643 .supported = packed_bits & (1 << 0) != 0,
644 .shuffle = packed_bits & (1 << 1) != 0,
645 .ballot = packed_bits & (1 << 2) != 0,
646 .vote = packed_bits & (1 << 3) != 0,
647 .arithmetic = packed_bits & (1 << 4) != 0,
648 .scan = packed_bits & (1 << 5) != 0,
649 };
650 }
651
652 fn appendInt(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, comptime T: type, value: T) !void {
653 var buffer: [@sizeOf(T)]u8 = undefined;
654 std.mem.writeInt(T, &buffer, value, .little);
655 try bytes.appendSlice(allocator, &buffer);
656 }
657
658 fn appendSized(bytes: *std.ArrayListUnmanaged(u8), allocator: std.mem.Allocator, data: []const u8) !void {
659 try appendInt(bytes, allocator, u32, @intCast(data.len));
660 try bytes.appendSlice(allocator, data);
661 }
662
663 fn readEnum(cursor: *ByteReader, comptime E: type) Error!E {
664 return std.enums.fromInt(E, try readByte(cursor)) orelse error.InvalidRegistryTag;
665 }
666
667 fn takeSizedCopy(cursor: *ByteReader, arena: std.mem.Allocator) Error![]u8 {
668 const length = try readInt(cursor, u32);
669 return takeCopy(cursor, arena, length);
670 }
671
672 fn takeCopy(cursor: *ByteReader, arena: std.mem.Allocator, length: usize) Error![]u8 {
673 return arena.dupe(u8, try take(cursor, length));
674 }
675
676 fn take(cursor: *ByteReader, count: usize) Error![]const u8 {
677 return cursor.readBytes(count) catch return error.TruncatedRegistryBlob;
678 }
679
680 fn readByte(cursor: *ByteReader) Error!u8 {
681 return cursor.readByte() catch return error.TruncatedRegistryBlob;
682 }
683
684 fn readInt(cursor: *ByteReader, comptime T: type) Error!T {
685 return cursor.readInt(T) catch return error.TruncatedRegistryBlob;
686 }
687
688 fn readLength(cursor: *ByteReader) Error!usize {
689 const length = try readInt(cursor, u64);
690 return std.math.cast(usize, length) orelse error.TruncatedRegistryBlob;
691 }
692
693 test "artifact wire round-trips a synthetic registry" {
694 const static_arguments = [_]choir_abi.ScalarArgument{
695 .{ .i32 = -7 },
696 .{ .u32 = 9 },
697 .{ .i64 = -1234567890123 },
698 .{ .u64 = 9876543210987 },
699 .{ .f32 = 1.5 },
700 .{ .f64 = -2.25 },
701 };
702 const words = [_]u32{ 0xdeadbeef, 1, 2, 3 };
703 const shape_profile_dimensions = [_]plan.KernelCallShapeProfileDimension{
704 .{
705 .name = "m",
706 .runtime_scalar_argument_index = 0,
707 .bounds = .{ .min = 1, .opt = 128, .max = 4096 },
708 },
709 .{
710 .name = "n",
711 .runtime_scalar_argument_index = 1,
712 .bounds = .{ .min = 1, .opt = 128, .max = 4096 },
713 },
714 };
715 var push_constants: choir_abi.PushConstants = .{};
716 _ = try push_constants.append(4, .{});
717 _ = try push_constants.append(8, .{});
718 const entries = [_]plan.KernelCallArtifact{
719 .{
720 .target = "accy.kernel.linalg.matmul5x7x3_4x2_f32",
721 .version = 1,
722 .format = .cuda_ptx,
723 .entry_name = "accy_kernel_linalg_matmul5x7x3_4x2_f32",
724 .argument_count = 3,
725 .required_dtypes = gpu.DTypeSet.init(&.{.f32}),
726 .payload = .{ .text = "// ptx text" },
727 .launch = .{ .fixed = .{
728 .grid = .{ 2, 3, 1 },
729 .threadgroup = .{ 4, 2, 1 },
730 } },
731 },
732 .{
733 .target = "accy.kernel.test.full",
734 .version = 7,
735 .format = .vulkan_spirv,
736 .entry_name = "accy_kernel_test_full",
737 .argument_count = 9,
738 .required_dtypes = gpu.DTypeSet.init(&.{ .f32, .i32 }),
739 .required_features = .{ .atomic_i32 = true, .atomic_f32_add_device = true, .dynamic_shared_memory = true },
740 .required_subgroup = .{
741 .supported = true,
742 .size_min = 16,
743 .size_max = 64,
744 .shuffle = true,
745 .scan = true,
746 },
747 .push_constants = push_constants,
748 .payload = .{ .words_u32 = words[0..] },
749 .launch = .{ .derived = .{
750 .grid = .{
751 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = 8 } },
752 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = 4 } },
753 .{ .fixed = 1 },
754 },
755 .threadgroup = .{ 8, 4, 1 },
756 .dynamic_shared_memory_bytes = 256,
757 } },
758 .element_count_argument = .scalar_u32,
759 .runtime_scalar_argument_count = 2,
760 .static_arguments = static_arguments[0..],
761 .shape_family_fingerprint = 0x1234_5678_90ab_cdef,
762 .shape_profile = .{
763 .name = "matrix_product",
764 .fingerprint = 0x1234_5678_90ab_cdef,
765 .dimensions = shape_profile_dimensions[0..],
766 },
767 },
768 .{
769 .target = "accy.kernel.test.bytes",
770 .version = 2,
771 .format = .metal_metallib,
772 .entry_name = "accy_kernel_test_bytes",
773 .argument_count = 1,
774 .payload = .{ .bytes = "\x00\x01\x02binary" },
775 },
776 };
777
778 const encoded = try encode(std.testing.allocator, entries[0..], &.{});
779 defer std.testing.allocator.free(encoded);
780
781 var decoded = try decode(std.testing.allocator, encoded);
782 defer decoded.deinit();
783
784 try std.testing.expectEqual(entries.len, decoded.entries.len);
785 try std.testing.expect(decoded.index.slots.len >= entries.len);
786 try std.testing.expect(decoded.registry().index != null);
787 for (entries[0..], decoded.entries) |expected, actual| {
788 try std.testing.expectEqualDeep(expected, actual);
789 }
790
791 const found = decoded.registry().find("accy.kernel.test.full", 7, .vulkan_spirv) orelse {
792 return error.TestExpectedRegistryEntry;
793 };
794 try std.testing.expectEqual(@as(u32, 9), found.argument_count);
795 try std.testing.expectEqual(plan.ElementCountArgument.scalar_u32, found.element_count_argument);
796 try std.testing.expectEqual(@as(u32, 2), found.runtime_scalar_argument_count);
797 try std.testing.expectEqual(@as(?u64, 0x1234_5678_90ab_cdef), found.shape_family_fingerprint);
798 const profile = found.shape_profile orelse return error.TestExpectedShapeProfile;
799 try std.testing.expectEqualStrings("matrix_product", profile.name);
800 try std.testing.expectEqual(@as(u64, 0x1234_5678_90ab_cdef), profile.fingerprint);
801 try std.testing.expectEqual(@as(usize, 2), profile.dimensions.len);
802 const n_dimension = profile.dimension("n") orelse return error.TestExpectedShapeProfile;
803 try std.testing.expectEqual(@as(u32, 1), n_dimension.runtime_scalar_argument_index);
804 try std.testing.expectEqual(@as(?u64, 4096), n_dimension.bounds.max);
805 switch (found.launch) {
806 .derived => |launch| {
807 try std.testing.expectEqual(@as(u32, 8), launch.threadgroup[0]);
808 try std.testing.expectEqual(@as(u32, 4), launch.threadgroup[1]);
809 try std.testing.expectEqual(@as(u32, 256), launch.dynamic_shared_memory_bytes);
810 },
811 .fixed => return error.TestExpectedDerivedLaunch,
812 }
813 }
814
815 test "artifact wire rejects malformed registry blobs" {
816 const entries = [_]plan.KernelCallArtifact{.{
817 .target = "accy.kernel.test.copy",
818 .version = 1,
819 .format = .cuda_ptx,
820 .entry_name = "accy_kernel_test_copy",
821 .argument_count = 2,
822 .payload = .{ .text = "// ptx" },
823 }};
824
825 const encoded = try encode(std.testing.allocator, entries[0..], &.{});
826 defer std.testing.allocator.free(encoded);
827
828 var bad_magic = try std.testing.allocator.dupe(u8, encoded);
829 defer std.testing.allocator.free(bad_magic);
830 bad_magic[0] = 'X';
831 try std.testing.expectError(error.InvalidRegistryBlob, decode(std.testing.allocator, bad_magic));
832
833 var bad_version = try std.testing.allocator.dupe(u8, encoded);
834 defer std.testing.allocator.free(bad_version);
835 std.mem.writeInt(u32, bad_version[magic.len..][0..4], format_version + 1, .little);
836 try std.testing.expectError(error.UnsupportedRegistryVersion, decode(std.testing.allocator, bad_version));
837
838 var prefix_length: usize = magic.len + @sizeOf(u32);
839 while (prefix_length < encoded.len) : (prefix_length += 7) {
840 try std.testing.expectError(
841 error.TruncatedRegistryBlob,
842 decode(std.testing.allocator, encoded[0..prefix_length]),
843 );
844 }
845
846 const trailing = try std.mem.concat(std.testing.allocator, u8, &.{ encoded, "junk" });
847 defer std.testing.allocator.free(trailing);
848 try std.testing.expectError(error.InvalidRegistryBlob, decode(std.testing.allocator, trailing));
849 }
850
851 test "artifact wire round-trips an empty registry" {
852 const encoded = try encode(std.testing.allocator, &.{}, &.{});
853 defer std.testing.allocator.free(encoded);
854
855 var decoded = try decode(std.testing.allocator, encoded);
856 defer decoded.deinit();
857 try std.testing.expectEqual(@as(usize, 0), decoded.entries.len);
858 try std.testing.expect(decoded.registry().find("missing", 1, .cuda_ptx) == null);
859 }
860
861 test "artifact wire rejects duplicate registry keys" {
862 const entries = [_]plan.KernelCallArtifact{
863 .{
864 .target = "accy.kernel.test.duplicate",
865 .version = 1,
866 .format = .cuda_ptx,
867 .entry_name = "duplicate_first",
868 .argument_count = 1,
869 .payload = .{ .text = "// first" },
870 },
871 .{
872 .target = "accy.kernel.test.duplicate",
873 .version = 1,
874 .format = .cuda_ptx,
875 .entry_name = "duplicate_second",
876 .argument_count = 1,
877 .payload = .{ .text = "// second" },
878 },
879 };
880
881 const encoded = try encode(std.testing.allocator, entries[0..], &.{});
882 defer std.testing.allocator.free(encoded);
883 try std.testing.expectError(error.InvalidRegistryBlob, decode(std.testing.allocator, encoded));
884 }
885
886 pub const ByteWriter = struct {
887 bytes: std.ArrayListUnmanaged(u8) = .empty,
888
889 pub fn deinit(self: *ByteWriter, allocator: std.mem.Allocator) void {
890 self.bytes.deinit(allocator);
891 self.* = undefined;
892 }
893
894 pub fn toOwnedSlice(
895 self: *ByteWriter,
896 allocator: std.mem.Allocator,
897 ) std.mem.Allocator.Error![]u8 {
898 return self.bytes.toOwnedSlice(allocator);
899 }
900
901 pub fn writeU32(
902 self: *ByteWriter,
903 allocator: std.mem.Allocator,
904 value: u32,
905 ) gpu.BackendError!void {
906 try self.writeInt(allocator, u32, value);
907 }
908
909 pub fn writeU64(
910 self: *ByteWriter,
911 allocator: std.mem.Allocator,
912 value: u64,
913 ) gpu.BackendError!void {
914 try self.writeInt(allocator, u64, value);
915 }
916
917 pub fn writeUsize(
918 self: *ByteWriter,
919 allocator: std.mem.Allocator,
920 value: usize,
921 ) gpu.BackendError!void {
922 try self.writeU64(allocator, @intCast(value));
923 }
924
925 pub fn writeBool(
926 self: *ByteWriter,
927 allocator: std.mem.Allocator,
928 value: bool,
929 ) gpu.BackendError!void {
930 self.bytes.append(allocator, @intFromBool(value)) catch return error.OutOfMemory;
931 }
932
933 pub fn writeLengthPrefixedBytes(
934 self: *ByteWriter,
935 allocator: std.mem.Allocator,
936 bytes: []const u8,
937 ) gpu.BackendError!void {
938 if (bytes.len > std.math.maxInt(u32)) return error.InvalidArtifact;
939 try self.writeU32(allocator, @intCast(bytes.len));
940 self.bytes.appendSlice(allocator, bytes) catch return error.OutOfMemory;
941 }
942
943 pub fn writeEnum(
944 self: *ByteWriter,
945 allocator: std.mem.Allocator,
946 comptime T: type,
947 value: T,
948 ) gpu.BackendError!void {
949 try self.writeU32(allocator, @intCast(@backingInt(value)));
950 }
951
952 fn writeInt(
953 self: *ByteWriter,
954 allocator: std.mem.Allocator,
955 comptime T: type,
956 value: T,
957 ) gpu.BackendError!void {
958 var buf: [@sizeOf(T)]u8 = undefined;
959 std.mem.writeInt(T, &buf, value, .little);
960 self.bytes.appendSlice(allocator, &buf) catch return error.OutOfMemory;
961 }
962 };
963
964 pub const ByteReader = struct {
965 bytes: []const u8,
966 offset: usize = 0,
967
968 pub fn remaining(self: *const ByteReader) usize {
969 if (self.offset > self.bytes.len) return 0;
970 return self.bytes.len - self.offset;
971 }
972
973 pub fn readByte(self: *ByteReader) gpu.BackendError!u8 {
974 const bytes = try self.readBytes(1);
975 return bytes[0];
976 }
977
978 pub fn readU32(self: *ByteReader) gpu.BackendError!u32 {
979 return try self.readInt(u32);
980 }
981
982 pub fn readU64(self: *ByteReader) gpu.BackendError!u64 {
983 return try self.readInt(u64);
984 }
985
986 pub fn readUsize(self: *ByteReader) gpu.BackendError!usize {
987 const value = try self.readU64();
988 if (value > std.math.maxInt(usize)) return error.InvalidArtifact;
989 return @intCast(value);
990 }
991
992 pub fn readLengthPrefixedBytes(self: *ByteReader) gpu.BackendError![]const u8 {
993 const length = try self.readU32();
994 return try self.readBytes(length);
995 }
996
997 pub fn readBool(self: *ByteReader) gpu.BackendError!bool {
998 const bytes = try self.readBytes(1);
999 return switch (bytes[0]) {
1000 0 => false,
1001 1 => true,
1002 else => error.InvalidArtifact,
1003 };
1004 }
1005
1006 pub fn readEnum(
1007 self: *ByteReader,
1008 comptime T: type,
1009 ) gpu.BackendError!T {
1010 return try byteEnumFromU32(T, try self.readU32());
1011 }
1012
1013 pub fn readInt(
1014 self: *ByteReader,
1015 comptime T: type,
1016 ) gpu.BackendError!T {
1017 const bytes = try self.readBytes(@sizeOf(T));
1018 return std.mem.readInt(T, bytes[0..@sizeOf(T)], .little);
1019 }
1020
1021 pub fn readBytes(
1022 self: *ByteReader,
1023 count: usize,
1024 ) gpu.BackendError![]const u8 {
1025 if (self.offset > self.bytes.len or count > self.bytes.len - self.offset) return error.InvalidArtifact;
1026 const start = self.offset;
1027 self.offset += count;
1028 return self.bytes[start..self.offset];
1029 }
1030
1031 pub fn expectDone(self: ByteReader) gpu.BackendError!void {
1032 if (self.offset != self.bytes.len) return error.InvalidArtifact;
1033 }
1034 };
1035
1036 fn byteEnumFromU32(comptime T: type, value: u32) gpu.BackendError!T {
1037 inline for (
1038 @typeInfo(T).@"enum".field_names,
1039 @typeInfo(T).@"enum".field_values,
1040 ) |field_name, field_name_value| {
1041 const field = .{ .name = field_name, .value = field_name_value };
1042 if (value == field.value) return @fromBackingInt(@intCast(field.value));
1043 }
1044 return error.InvalidArtifact;
1045 }
1046
1047 test "byte writer and reader round-trip primitive values" {
1048 const allocator = std.testing.allocator;
1049 var writer = ByteWriter{};
1050 errdefer writer.deinit(allocator);
1051
1052 try writer.writeU32(allocator, 7);
1053 try writer.writeU64(allocator, 1 << 40);
1054 try writer.writeBool(allocator, true);
1055 try writer.writeLengthPrefixedBytes(allocator, "target");
1056 try writer.writeEnum(allocator, gpu.BackendKind, .cuda);
1057
1058 const encoded = try writer.toOwnedSlice(allocator);
1059 defer allocator.free(encoded);
1060
1061 var reader = ByteReader{ .bytes = encoded };
1062 try std.testing.expectEqual(@as(u32, 7), try reader.readU32());
1063 try std.testing.expectEqual(@as(u64, 1 << 40), try reader.readU64());
1064 try std.testing.expectEqual(true, try reader.readBool());
1065 try std.testing.expectEqualStrings("target", try reader.readLengthPrefixedBytes());
1066 try std.testing.expectEqual(gpu.BackendKind.cuda, try reader.readEnum(gpu.BackendKind));
1067 try reader.expectDone();
1068
1069 var truncated = ByteReader{ .bytes = encoded[0 .. encoded.len - 1] };
1070 _ = try truncated.readU32();
1071 _ = try truncated.readU64();
1072 _ = try truncated.readBool();
1073 _ = try truncated.readLengthPrefixedBytes();
1074 try std.testing.expectError(error.InvalidArtifact, truncated.readEnum(gpu.BackendKind));
1075 }
1076
1077 test "artifact wire round-trips pipelines beside entries" {
1078 const entries = [_]plan.KernelCallArtifact{
1079 .{
1080 .target = "accy.kernel.test.block_scan",
1081 .version = 1,
1082 .format = .cuda_ptx,
1083 .entry_name = "block_scan",
1084 .argument_count = 4,
1085 .payload = .{ .text = "// block scan" },
1086 .runtime_scalar_argument_count = 1,
1087 },
1088 .{
1089 .target = "accy.kernel.test.add_base",
1090 .version = 1,
1091 .format = .cuda_ptx,
1092 .entry_name = "add_base",
1093 .argument_count = 3,
1094 .payload = .{ .text = "// add base" },
1095 .runtime_scalar_argument_count = 1,
1096 },
1097 };
1098 const pipelines = [_]pipeline_mod.KernelCallPipeline{.{
1099 .target = "accy.kernel.test.device_scan",
1100 .version = 3,
1101 .operand_count = 1,
1102 .result_count = 1,
1103 .runtime_scalar_argument_count = 1,
1104 .runtime_scalar_bounds = &.{.{ .argument_index = 0, .max_u32 = 6144 }},
1105 .intermediates = &.{
1106 .{ .dtype = .f32, .extent = .{ .ceil_div = .{ .argument_index = 0, .divisor = 64 } } },
1107 .{ .dtype = .f32, .extent = .{ .ceil_div_scaled = .{ .argument_index = 0, .divisor = 64, .scale = 16 } } },
1108 .{ .dtype = .f32, .extent = .{ .ceil_div_scaled_by_arg = .{ .argument_index = 0, .divisor = 64, .scale_argument_index = 0 } } },
1109 },
1110 .stages = &.{
1111 .{
1112 .target = "accy.kernel.test.block_scan",
1113 .version = 1,
1114 .buffers = &.{ .{ .result = 0 }, .{ .operand = 0 }, .{ .intermediate = 0 } },
1115 .scalars = &.{.{ .forward = 0 }},
1116 },
1117 .{
1118 .target = "accy.kernel.test.add_base",
1119 .version = 1,
1120 .buffers = &.{ .{ .result = 0 }, .{ .intermediate = 0 } },
1121 .scalars = &.{.{ .ceil_div = .{ .argument_index = 0, .divisor = 64 } }},
1122 },
1123 },
1124 }};
1125
1126 const encoded = try encode(std.testing.allocator, entries[0..], pipelines[0..]);
1127 defer std.testing.allocator.free(encoded);
1128
1129 var decoded = try decode(std.testing.allocator, encoded);
1130 defer decoded.deinit();
1131
1132 try std.testing.expectEqual(pipelines.len, decoded.pipelines.len);
1133 for (pipelines[0..], decoded.pipelines) |expected, actual| {
1134 try std.testing.expectEqualDeep(expected, actual);
1135 }
1136
1137 const found = pipeline_mod.findPipeline(decoded.pipelines, "accy.kernel.test.device_scan", 3) orelse {
1138 return error.TestExpectedPipeline;
1139 };
1140 try found.validate(decoded.registry(), .cuda_ptx);
1141
1142 var prefix_length: usize = magic.len + @sizeOf(u32);
1143 while (prefix_length < encoded.len) : (prefix_length += 7) {
1144 const result = decode(std.testing.allocator, encoded[0..prefix_length]);
1145 try std.testing.expectError(error.TruncatedRegistryBlob, result);
1146 }
1147 }
1148
1149 test "artifact wire rejects duplicate pipeline identities" {
1150 const duplicate = pipeline_mod.KernelCallPipeline{
1151 .target = "accy.kernel.test.device_scan",
1152 .version = 1,
1153 .stages = &.{.{ .target = "accy.kernel.test.stage", .version = 1 }},
1154 };
1155 const pipelines = [_]pipeline_mod.KernelCallPipeline{ duplicate, duplicate };
1156 const encoded = try encode(std.testing.allocator, &.{}, pipelines[0..]);
1157 defer std.testing.allocator.free(encoded);
1158 try std.testing.expectError(error.InvalidRegistryBlob, decode(std.testing.allocator, encoded));
1159 }
1160
1161 test "artifact wire rejects zero-version pipeline identities" {
1162 const zero_version = pipeline_mod.KernelCallPipeline{
1163 .target = "accy.kernel.test.device_scan",
1164 .version = 0,
1165 .stages = &.{.{ .target = "accy.kernel.test.stage", .version = 1 }},
1166 };
1167 const encoded = try encode(std.testing.allocator, &.{}, &.{zero_version});
1168 defer std.testing.allocator.free(encoded);
1169 try std.testing.expectError(error.InvalidRegistryBlob, decode(std.testing.allocator, encoded));
1170 }