lib/accy/src/kernel/model/core/schedule/model.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_phase = @import("alloc_phase");
3 const choir = @import("choir");
4 const testing = std.testing;
5
6 const Allocator = std.mem.Allocator;
7
8 pub const ScheduleError = error{
9 UnknownAxis,
10 InvalidFactor,
11 ExtentNotDivisible,
12 BindTargetAlreadyUsed,
13 MissingBindingExtent,
14 LaunchDimensionOverflow,
15 UnsupportedVersion,
16 InvalidReplay,
17 CapacityOverflow,
18 AxisCapacityExceeded,
19 StepCapacityExceeded,
20 NameCapacityExceeded,
21 AxisIdCapacityExceeded,
22 ScheduleSealed,
23 ScheduleNotEmpty,
24 SnapshotNotReady,
25 } || Allocator.Error;
26
27 pub const snapshot_version: u32 = 1;
28
29 pub const Launch = struct {
30 grid: [3]u32 = .{ 1, 1, 1 },
31 block: [3]u32 = .{ 1, 1, 1 },
32
33 pub fn oneDim(extent: usize) Launch {
34 return .{ .grid = .{ @intCast(extent), 1, 1 } };
35 }
36 };
37
38 pub const AxisId = enum(u32) {
39 _,
40
41 pub fn index(self: AxisId) u32 {
42 return @backingInt(self);
43 }
44
45 fn fromIndex(raw_index: usize) AxisId {
46 return @fromBackingInt(@intCast(raw_index));
47 }
48 };
49
50 pub const BindTarget = enum {
51 block_x,
52 block_y,
53 block_z,
54 thread_x,
55 thread_y,
56 thread_z,
57
58 fn dim(self: BindTarget) usize {
59 return switch (self) {
60 .block_x, .thread_x => 0,
61 .block_y, .thread_y => 1,
62 .block_z, .thread_z => 2,
63 };
64 }
65
66 fn isGrid(self: BindTarget) bool {
67 return switch (self) {
68 .block_x, .block_y, .block_z => true,
69 .thread_x, .thread_y, .thread_z => false,
70 };
71 }
72 };
73
74 pub const Axis = struct {
75 id: AxisId,
76 name: []const u8,
77 extent: ?u64,
78 bind: ?BindTarget = null,
79 vector_width: ?u32 = null,
80 unroll_factor: ?u32 = null,
81 };
82
83 pub const Split = struct {
84 outer: AxisId,
85 inner: AxisId,
86 };
87
88 pub const AxisStep = struct {
89 id: AxisId,
90 name: []const u8,
91 extent: ?u64,
92 };
93
94 pub const Step = union(enum) {
95 axis: AxisStep,
96 split: struct {
97 source: AxisId,
98 outer: AxisId,
99 inner: AxisId,
100 factor: u64,
101 },
102 tile: struct {
103 source: AxisId,
104 outer: AxisId,
105 inner: AxisId,
106 factor: u64,
107 },
108 bind: struct {
109 axis: AxisId,
110 target: BindTarget,
111 },
112 vectorize: struct {
113 axis: AxisId,
114 width: u32,
115 },
116 unroll: struct {
117 axis: AxisId,
118 factor: u32,
119 },
120 };
121
122 /// A caller holds this record to save a kernel's schedule, the loop structure of one kernel, as
123 /// axes and the steps applied to them, or to replay it onto another kernel. The record holds a
124 /// format version, the schedule's axes, its steps in order, and the limits that bound a replay. The
125 /// record borrows its slices, either from the schedule snapshot that produced it or from a record
126 /// decoded on its own. Replaying it onto a schedule fails with `error.ScheduleNotEmpty` when that
127 /// schedule already has steps, `error.UnsupportedVersion` for another format version, and
128 /// `error.InvalidReplay` when an axis or step mismatches. Saved kernel programs embed this record.
129 pub const Record = struct {
130 version: u32,
131 axes: []const Axis,
132 steps: []const Step,
133 replay_limits: Schedule.Limits,
134 };
135
136 pub const Snapshot = struct {
137 phase: alloc_phase.capacity.Phase,
138 capacity: Capacity,
139 storage: [*]u8,
140 version: u32,
141 axes_storage: []Axis,
142 axes_len: usize,
143 steps_storage: []Step,
144 steps_len: usize,
145 names_storage: []u8,
146 names_len: usize,
147 replay_limits: Schedule.Limits,
148 captured: bool,
149
150 pub const Limits = struct {
151 axes: usize,
152 steps: usize,
153 name_bytes: usize,
154
155 pub const standard: Limits = .{
156 .axes = 64,
157 .steps = 256,
158 .name_bytes = 16 * 1024,
159 };
160
161 pub const testing: Limits = .{
162 .axes = 256,
163 .steps = 1024,
164 .name_bytes = 64 * 1024,
165 };
166 };
167
168 pub const Capacity = struct {
169 axes: usize,
170 steps: usize,
171 name_bytes: usize,
172 axes_offset: usize,
173 steps_offset: usize,
174 names_offset: usize,
175 storage_bytes: usize,
176 storage_alignment: std.mem.Alignment,
177
178 pub fn derive(limits: Limits) error{CapacityOverflow}!Capacity {
179 var cursor: usize = 0;
180 var alignment: usize = 1;
181 const axes_offset = try placeSlice(Axis, limits.axes, &cursor, &alignment);
182 const steps_offset = try placeSlice(Step, limits.steps, &cursor, &alignment);
183 const names_offset = try placeSlice(u8, limits.name_bytes, &cursor, &alignment);
184 if (cursor == 0) cursor = 1;
185 return .{
186 .axes = limits.axes,
187 .steps = limits.steps,
188 .name_bytes = limits.name_bytes,
189 .axes_offset = axes_offset,
190 .steps_offset = steps_offset,
191 .names_offset = names_offset,
192 .storage_bytes = cursor,
193 .storage_alignment = .fromByteUnits(alignment),
194 };
195 }
196
197 pub fn asLimits(self: Capacity) Limits {
198 return .{
199 .axes = self.axes,
200 .steps = self.steps,
201 .name_bytes = self.name_bytes,
202 };
203 }
204 };
205
206 pub const Exhaustion = error{
207 AxisCapacityExceeded,
208 StepCapacityExceeded,
209 NameCapacityExceeded,
210 };
211
212 pub const Usage = struct {
213 axes: usize,
214 steps: usize,
215 name_bytes: usize,
216 };
217
218 pub const claim: alloc_phase.capacity.Declaration = .{
219 .source = .{
220 .id = "accy.kernel_schedule_snapshot",
221 .kind = .phase_static,
222 .limit_source = .caller,
223 .storage = .{
224 .covered = &.{
225 .{
226 .id = "replayable_axis_records_and_owned_axis_names",
227 .lifetime = .steady,
228 .detail = "replayable axis records and owned axis names",
229 },
230 .{
231 .id = "replayable_schedule_steps_and_owned_axis_step_names",
232 .lifetime = .steady,
233 .detail = "replayable schedule steps and owned axis-step names",
234 },
235 },
236 .excluded = &.{
237 "source Schedule storage borrowed only during capture",
238 "caller allocator implementation state and replay Schedule storage",
239 },
240 },
241 .capacity = .{
242 .inputs = &.{
243 alloc_phase.capacity.bindInput(Limits, "axes", "axes"),
244 alloc_phase.capacity.bindInput(Limits, "steps", "steps"),
245 alloc_phase.capacity.bindInput(Limits, "name_bytes", "name_bytes"),
246 },
247 .type_selectors = &.{
248 alloc_phase.capacity.bindType(Axis, "axis"),
249 alloc_phase.capacity.bindType(Step, "step"),
250 },
251 .nodes = &.{
252 .{ .input = 0 },
253 .{ .scale = .{ .node = 0, .coefficient = .{ .size_of_concrete_type = 0 } } },
254 .{ .input = 1 },
255 .{ .scale = .{ .node = 2, .coefficient = .{ .size_of_concrete_type = 1 } } },
256 .{ .input = 2 },
257 .{ .constant = 0 },
258 .{ .alignment = .{ .node = 5, .alignment = .{ .literal = 16 } } },
259 .{ .add = .{ .left = 6, .right = 1 } },
260 .{ .alignment = .{ .node = 7, .alignment = .{ .literal = 16 } } },
261 .{ .add = .{ .left = 8, .right = 3 } },
262 .{ .alignment = .{ .node = 9, .alignment = .{ .literal = 1 } } },
263 .{ .add = .{ .left = 10, .right = 4 } },
264 .{ .alignment = .{ .node = 11, .alignment = .{ .literal = 16 } } },
265 },
266 .assertions = &.{.{
267 .scope = .closure_total,
268 .measure = .retained,
269 .relation = .exact,
270 .expression = 12,
271 }},
272 },
273 .overload = .{
274 .kind = .reject_before_mutation,
275 .detail = "capture preflights all record and name dimensions before copying any source state",
276 },
277 .risks = .{
278 .transitive = .{
279 .status = .witnessed,
280 .detail = "capture and immutable queries use only the single preacquired snapshot region",
281 },
282 .foreign = .{
283 .status = .excluded,
284 .detail = "snapshot capture and queries cross no foreign or operating-system boundary",
285 },
286 },
287 .obligations = &.{
288 .{ .key = "accy_kernel_schedule_snapshot_capacity", .role = .capacity_model },
289 .{ .key = "accy_kernel_schedule_snapshot_boundary", .role = .overload },
290 .{ .key = "accy_kernel_schedule_snapshot_lifetime_transitive_risk", .role = .transitive_risk },
291 .{ .key = "accy_kernel_schedule_snapshot_lifetime_foreign_risk", .role = .foreign_risk },
292 },
293 },
294 .bindings = .{
295 .owner = @This(),
296 .seal = .{
297 .family = alloc_phase.capacity.selector(@This().activate),
298 .premise = .{
299 .class = .checked_semantic_fact,
300 .authority = .checker,
301 },
302 },
303 .teardown = .{
304 .family = alloc_phase.capacity.selector(@This().deinit),
305 .premise = .{
306 .class = .checked_semantic_fact,
307 .authority = .checker,
308 },
309 },
310 },
311 };
312
313 pub fn init(allocator: Allocator, limits: Limits) ScheduleError!Snapshot {
314 const capacity = try Capacity.derive(limits);
315 const storage = allocator.rawAlloc(
316 capacity.storage_bytes,
317 capacity.storage_alignment,
318 @returnAddress(),
319 ) orelse return error.OutOfMemory;
320 return .{
321 .phase = .initialization,
322 .capacity = capacity,
323 .storage = storage,
324 .version = snapshot_version,
325 .axes_storage = typedSlice(Axis, storage, capacity.axes_offset, capacity.axes),
326 .axes_len = 0,
327 .steps_storage = typedSlice(Step, storage, capacity.steps_offset, capacity.steps),
328 .steps_len = 0,
329 .names_storage = typedSlice(u8, storage, capacity.names_offset, capacity.name_bytes),
330 .names_len = 0,
331 .replay_limits = .{
332 .axes = 0,
333 .steps = 0,
334 .name_bytes = 0,
335 .axis_ids = 0,
336 },
337 .captured = false,
338 };
339 }
340
341 pub fn capture(self: *Snapshot, schedule: *const Schedule) ScheduleError!void {
342 if (self.phase != .initialization or self.captured) return error.SnapshotNotReady;
343 const required = try schedule.snapshotLimits();
344 if (required.axes > self.capacity.axes) return error.AxisCapacityExceeded;
345 if (required.steps > self.capacity.steps) return error.StepCapacityExceeded;
346 if (required.name_bytes > self.capacity.name_bytes) return error.NameCapacityExceeded;
347
348 for (schedule.allAxes(), 0..) |source, index| {
349 self.axes_storage[index] = source;
350 self.axes_storage[index].name = self.writeName(source.name);
351 }
352 for (schedule.allSteps(), 0..) |source, index| {
353 self.steps_storage[index] = source;
354 if (source == .axis) {
355 self.steps_storage[index].axis.name = self.writeName(source.axis.name);
356 }
357 }
358 self.axes_len = required.axes;
359 self.steps_len = required.steps;
360 self.replay_limits = schedule.requiredReplayLimits();
361 self.captured = true;
362 }
363
364 pub fn activate(self: *Snapshot) void {
365 std.debug.assert(self.phase == .initialization);
366 std.debug.assert(self.captured);
367 self.phase = .steady;
368 }
369
370 pub fn deinit(self: *Snapshot, allocator: Allocator) void {
371 std.debug.assert(self.phase != .teardown);
372 const storage = self.storage;
373 const capacity = self.capacity;
374 self.phase = .teardown;
375 self.* = undefined;
376 allocator.rawFree(
377 storage[0..capacity.storage_bytes],
378 capacity.storage_alignment,
379 @returnAddress(),
380 );
381 }
382
383 pub fn capacityUsage(self: *const Snapshot) Usage {
384 return .{
385 .axes = self.axes_len,
386 .steps = self.steps_len,
387 .name_bytes = self.names_len,
388 };
389 }
390
391 pub fn allAxes(self: *const Snapshot) []const Axis {
392 return self.axes_storage[0..self.axes_len];
393 }
394
395 pub fn allSteps(self: *const Snapshot) []const Step {
396 return self.steps_storage[0..self.steps_len];
397 }
398
399 pub fn launch(self: *const Snapshot) ScheduleError!Launch {
400 if (!self.captured) return error.SnapshotNotReady;
401 return launchFromAxes(self.allAxes());
402 }
403
404 pub fn fingerprint(self: *const Snapshot) u64 {
405 return fingerprintSchedule(self.version, self.allAxes(), self.allSteps());
406 }
407
408 pub fn requiredReplayLimits(self: *const Snapshot) ScheduleError!Schedule.Limits {
409 if (!self.captured) return error.SnapshotNotReady;
410 if (self.version != snapshot_version) return error.UnsupportedVersion;
411 return self.replay_limits;
412 }
413
414 pub fn requiredReplayCapacity(self: *const Snapshot) ScheduleError!Schedule.Capacity {
415 return Schedule.Capacity.derive(try self.requiredReplayLimits());
416 }
417
418 pub fn record(self: *const Snapshot) ScheduleError!Record {
419 return .{
420 .version = self.version,
421 .axes = self.allAxes(),
422 .steps = self.allSteps(),
423 .replay_limits = try self.requiredReplayLimits(),
424 };
425 }
426
427 fn writeName(self: *Snapshot, name: []const u8) []const u8 {
428 const start = self.names_len;
429 const end = start + name.len;
430 @memcpy(self.names_storage[start..end], name);
431 self.names_len = end;
432 return self.names_storage[start..end];
433 }
434 };
435
436 pub const Schedule = struct {
437 phase: alloc_phase.capacity.Phase,
438 capacity: Capacity,
439 storage: [*]u8,
440 axes_storage: []Axis,
441 axes_len: usize,
442 steps_storage: []Step,
443 steps_len: usize,
444 names_storage: []u8,
445 names_len: usize,
446 next_axis_id: usize,
447
448 pub const Limits = struct {
449 axes: usize,
450 steps: usize,
451 name_bytes: usize,
452 axis_ids: usize,
453
454 pub const standard: Limits = .{
455 .axes = 64,
456 .steps = 256,
457 .name_bytes = 16 * 1024,
458 .axis_ids = 512,
459 };
460
461 pub const testing: Limits = .{
462 .axes = 256,
463 .steps = 1024,
464 .name_bytes = 64 * 1024,
465 .axis_ids = 2048,
466 };
467 };
468
469 pub const Capacity = struct {
470 axes: usize,
471 steps: usize,
472 name_bytes: usize,
473 axis_ids: usize,
474 axes_offset: usize,
475 steps_offset: usize,
476 names_offset: usize,
477 storage_bytes: usize,
478 storage_alignment: std.mem.Alignment,
479
480 pub fn derive(limits: Limits) error{CapacityOverflow}!Capacity {
481 if (limits.axis_ids > std.math.maxInt(u32)) return error.CapacityOverflow;
482 var cursor: usize = 0;
483 var alignment: usize = 1;
484 const axes_offset = try placeSlice(Axis, limits.axes, &cursor, &alignment);
485 const steps_offset = try placeSlice(Step, limits.steps, &cursor, &alignment);
486 const names_offset = try placeSlice(u8, limits.name_bytes, &cursor, &alignment);
487 if (cursor == 0) cursor = 1;
488 return .{
489 .axes = limits.axes,
490 .steps = limits.steps,
491 .name_bytes = limits.name_bytes,
492 .axis_ids = limits.axis_ids,
493 .axes_offset = axes_offset,
494 .steps_offset = steps_offset,
495 .names_offset = names_offset,
496 .storage_bytes = cursor,
497 .storage_alignment = .fromByteUnits(alignment),
498 };
499 }
500
501 pub fn asLimits(self: Capacity) Limits {
502 return .{
503 .axes = self.axes,
504 .steps = self.steps,
505 .name_bytes = self.name_bytes,
506 .axis_ids = self.axis_ids,
507 };
508 }
509 };
510
511 pub const Exhaustion = error{
512 AxisCapacityExceeded,
513 StepCapacityExceeded,
514 NameCapacityExceeded,
515 AxisIdCapacityExceeded,
516 };
517
518 pub const Usage = struct {
519 axes: usize,
520 steps: usize,
521 name_bytes: usize,
522 axis_ids: usize,
523 };
524
525 pub const claim: alloc_phase.capacity.Declaration = .{
526 .source = .{
527 .id = "accy.kernel_schedule",
528 .kind = .phase_static,
529 .limit_source = .caller,
530 .storage = .{
531 .covered = &.{
532 .{
533 .id = "bounded_mutable_axis_records",
534 .lifetime = .steady,
535 .detail = "bounded mutable axis records",
536 },
537 .{
538 .id = "bounded_replayable_schedule_steps",
539 .lifetime = .steady,
540 .detail = "bounded replayable schedule steps",
541 },
542 .{
543 .id = "live_duplicated_and_generated_axis_names",
544 .lifetime = .steady,
545 .detail = "live duplicated and generated axis names",
546 },
547 },
548 .excluded = &.{
549 "caller allocator implementation state",
550 "independently acquired snapshots and replay result Schedules",
551 },
552 },
553 .capacity = .{
554 .inputs = &.{
555 alloc_phase.capacity.bindInput(Limits, "axes", "axes"),
556 alloc_phase.capacity.bindInput(Limits, "steps", "steps"),
557 alloc_phase.capacity.bindInput(Limits, "axis_ids", "axis_ids"),
558 alloc_phase.capacity.bindInput(Limits, "name_bytes", "name_bytes"),
559 },
560 .type_selectors = &.{
561 alloc_phase.capacity.bindType(Axis, "axis"),
562 alloc_phase.capacity.bindType(Step, "step"),
563 alloc_phase.capacity.bindType(AxisId, "axisid"),
564 },
565 .nodes = &.{
566 .{ .input = 0 },
567 .{ .scale = .{ .node = 0, .coefficient = .{ .size_of_concrete_type = 0 } } },
568 .{ .input = 1 },
569 .{ .scale = .{ .node = 2, .coefficient = .{ .size_of_concrete_type = 1 } } },
570 .{ .input = 2 },
571 .{ .scale = .{ .node = 4, .coefficient = .{ .size_of_concrete_type = 2 } } },
572 .{ .input = 3 },
573 .{ .constant = 0 },
574 .{ .alignment = .{ .node = 7, .alignment = .{ .literal = 16 } } },
575 .{ .add = .{ .left = 8, .right = 1 } },
576 .{ .alignment = .{ .node = 9, .alignment = .{ .literal = 16 } } },
577 .{ .add = .{ .left = 10, .right = 3 } },
578 .{ .alignment = .{ .node = 11, .alignment = .{ .literal = 16 } } },
579 .{ .add = .{ .left = 12, .right = 5 } },
580 .{ .alignment = .{ .node = 13, .alignment = .{ .literal = 1 } } },
581 .{ .add = .{ .left = 14, .right = 6 } },
582 .{ .alignment = .{ .node = 15, .alignment = .{ .literal = 16 } } },
583 },
584 .assertions = &.{.{
585 .scope = .closure_total,
586 .measure = .retained,
587 .relation = .exact,
588 .expression = 16,
589 }},
590 },
591 .overload = .{
592 .kind = .reject_before_mutation,
593 .detail = "every mutation preflights axes, steps, name bytes, and issued IDs before publishing state",
594 },
595 .risks = .{
596 .transitive = .{
597 .status = .witnessed,
598 .detail = "all construction methods write only into the single preacquired Schedule region",
599 },
600 .foreign = .{
601 .status = .excluded,
602 .detail = "Schedule construction and queries cross no foreign or operating-system boundary",
603 },
604 },
605 .obligations = &.{
606 .{ .key = "accy_kernel_schedule_capacity", .role = .capacity_model },
607 .{ .key = "accy_kernel_schedule_single_acquisition", .role = .transitive_risk },
608 .{ .key = "accy_kernel_schedule_boundary_overload", .role = .overload },
609 .{ .key = "accy_kernel_schedule_boundary_foreign_risk", .role = .foreign_risk },
610 .{ .key = "accy_kernel_schedule_replay_boundary", .role = .overload },
611 },
612 },
613 .bindings = .{
614 .owner = @This(),
615 .seal = .{
616 .family = alloc_phase.capacity.selector(@This().activate),
617 .premise = .{
618 .class = .checked_semantic_fact,
619 .authority = .checker,
620 },
621 },
622 .teardown = .{
623 .family = alloc_phase.capacity.selector(@This().deinit),
624 .premise = .{
625 .class = .checked_semantic_fact,
626 .authority = .checker,
627 },
628 },
629 },
630 };
631
632 pub fn init(allocator: Allocator, limits: Limits) ScheduleError!Schedule {
633 const capacity = try Capacity.derive(limits);
634 const storage = allocator.rawAlloc(
635 capacity.storage_bytes,
636 capacity.storage_alignment,
637 @returnAddress(),
638 ) orelse return error.OutOfMemory;
639 return .{
640 .phase = .initialization,
641 .capacity = capacity,
642 .storage = storage,
643 .axes_storage = typedSlice(Axis, storage, capacity.axes_offset, capacity.axes),
644 .axes_len = 0,
645 .steps_storage = typedSlice(Step, storage, capacity.steps_offset, capacity.steps),
646 .steps_len = 0,
647 .names_storage = typedSlice(u8, storage, capacity.names_offset, capacity.name_bytes),
648 .names_len = 0,
649 .next_axis_id = 0,
650 };
651 }
652
653 pub fn activate(self: *Schedule) void {
654 std.debug.assert(self.phase == .initialization);
655 self.phase = .steady;
656 }
657
658 pub fn deinit(self: *Schedule, allocator: Allocator) void {
659 std.debug.assert(self.phase != .teardown);
660 const storage = self.storage;
661 const capacity = self.capacity;
662 self.phase = .teardown;
663 self.* = undefined;
664 allocator.rawFree(
665 storage[0..capacity.storage_bytes],
666 capacity.storage_alignment,
667 @returnAddress(),
668 );
669 }
670
671 pub fn capacityUsage(self: *const Schedule) Usage {
672 return .{
673 .axes = self.axes_len,
674 .steps = self.steps_len,
675 .name_bytes = self.names_len,
676 .axis_ids = self.next_axis_id,
677 };
678 }
679
680 pub fn addAxis(self: *Schedule, name: []const u8, extent: ?u64) ScheduleError!AxisId {
681 try self.requireBuilding();
682 const name_bytes = std.math.mul(usize, name.len, 2) catch return error.CapacityOverflow;
683 try self.requireGrowth(1, 1, name_bytes, 1);
684 const owned_name = self.writeName(name);
685 const step_name = self.writeName(name);
686 const id = self.allocAxisId();
687 self.axes_storage[self.axes_len] = .{
688 .id = id,
689 .name = owned_name,
690 .extent = extent,
691 };
692 self.axes_len += 1;
693 self.steps_storage[self.steps_len] = .{ .axis = .{
694 .id = id,
695 .name = step_name,
696 .extent = extent,
697 } };
698 self.steps_len += 1;
699 return id;
700 }
701
702 pub fn split(self: *Schedule, axis_id: AxisId, factor: u64) ScheduleError!Split {
703 if (factor == 0) return error.InvalidFactor;
704 try self.requireBuilding();
705 const index = try self.axisIndex(axis_id);
706 const source = self.axes_storage[index];
707 const extent = source.extent orelse return error.MissingBindingExtent;
708 if (extent % factor != 0) return error.ExtentNotDivisible;
709 return self.replaceAxis(index, source, factor, .split, "_outer", "_inner", extent / factor);
710 }
711
712 pub fn tile(self: *Schedule, axis_id: AxisId, factor: u64) ScheduleError!Split {
713 if (factor == 0) return error.InvalidFactor;
714 try self.requireBuilding();
715 const index = try self.axisIndex(axis_id);
716 const source = self.axes_storage[index];
717 const extent = source.extent orelse return error.MissingBindingExtent;
718 const outer_extent = (extent / factor) + @intFromBool(extent % factor != 0);
719 return self.replaceAxis(index, source, factor, .tile, "_tile", "_lane", outer_extent);
720 }
721
722 pub fn bind(self: *Schedule, axis_id: AxisId, target: BindTarget) ScheduleError!void {
723 try self.requireBuilding();
724 const index = try self.axisIndex(axis_id);
725 for (self.allAxes(), 0..) |existing, existing_index| {
726 if (existing_index == index) continue;
727 if (existing.bind) |bound| if (bound == target) return error.BindTargetAlreadyUsed;
728 }
729 try self.requireGrowth(0, 1, 0, 0);
730 self.axes_storage[index].bind = target;
731 self.steps_storage[self.steps_len] = .{ .bind = .{
732 .axis = axis_id,
733 .target = target,
734 } };
735 self.steps_len += 1;
736 }
737
738 pub fn vectorize(self: *Schedule, axis_id: AxisId, width: u32) ScheduleError!void {
739 if (width == 0) return error.InvalidFactor;
740 try self.requireBuilding();
741 const index = try self.axisIndex(axis_id);
742 if (self.axes_storage[index].extent) |extent| {
743 if (extent % width != 0) return error.ExtentNotDivisible;
744 }
745 try self.requireGrowth(0, 1, 0, 0);
746 self.axes_storage[index].vector_width = width;
747 self.steps_storage[self.steps_len] = .{ .vectorize = .{
748 .axis = axis_id,
749 .width = width,
750 } };
751 self.steps_len += 1;
752 }
753
754 pub fn unroll(self: *Schedule, axis_id: AxisId, factor: u32) ScheduleError!void {
755 if (factor == 0) return error.InvalidFactor;
756 try self.requireBuilding();
757 const index = try self.axisIndex(axis_id);
758 try self.requireGrowth(0, 1, 0, 0);
759 self.axes_storage[index].unroll_factor = factor;
760 self.steps_storage[self.steps_len] = .{ .unroll = .{
761 .axis = axis_id,
762 .factor = factor,
763 } };
764 self.steps_len += 1;
765 }
766
767 pub fn replay(self: *Schedule, source: Record) ScheduleError!void {
768 try self.requireBuilding();
769 if (self.axes_len != 0 or self.steps_len != 0 or self.names_len != 0 or self.next_axis_id != 0) {
770 return error.ScheduleNotEmpty;
771 }
772 if (source.version != snapshot_version) return error.UnsupportedVersion;
773 try self.requireLimits(source.replay_limits);
774 errdefer self.reset();
775
776 for (source.steps) |step| {
777 switch (step) {
778 .axis => |axis| {
779 const id = try self.addAxis(axis.name, axis.extent);
780 if (id != axis.id) return error.InvalidReplay;
781 },
782 .split => |split_value| {
783 const result = try self.split(split_value.source, split_value.factor);
784 if (result.outer != split_value.outer or result.inner != split_value.inner) {
785 return error.InvalidReplay;
786 }
787 },
788 .tile => |tile_value| {
789 const result = try self.tile(tile_value.source, tile_value.factor);
790 if (result.outer != tile_value.outer or result.inner != tile_value.inner) {
791 return error.InvalidReplay;
792 }
793 },
794 .bind => |bind_value| try self.bind(bind_value.axis, bind_value.target),
795 .vectorize => |vectorize_value| try self.vectorize(vectorize_value.axis, vectorize_value.width),
796 .unroll => |unroll_value| try self.unroll(unroll_value.axis, unroll_value.factor),
797 }
798 }
799 if (!axesEqual(self.allAxes(), source.axes)) return error.InvalidReplay;
800 if (!stepsEqual(self.allSteps(), source.steps)) return error.InvalidReplay;
801 if (!std.meta.eql(self.requiredReplayLimits(), source.replay_limits)) return error.InvalidReplay;
802 }
803
804 pub fn getAxis(self: *const Schedule, id: AxisId) ScheduleError!Axis {
805 return self.axes_storage[try self.axisIndex(id)];
806 }
807
808 pub fn allAxes(self: *const Schedule) []const Axis {
809 return self.axes_storage[0..self.axes_len];
810 }
811
812 pub fn allSteps(self: *const Schedule) []const Step {
813 return self.steps_storage[0..self.steps_len];
814 }
815
816 pub fn launch(self: *const Schedule) ScheduleError!Launch {
817 return launchFromAxes(self.allAxes());
818 }
819
820 pub fn fingerprint(self: *const Schedule) u64 {
821 return fingerprintSchedule(snapshot_version, self.allAxes(), self.allSteps());
822 }
823
824 pub fn snapshotLimits(self: *const Schedule) ScheduleError!Snapshot.Limits {
825 var name_bytes: usize = 0;
826 for (self.allAxes()) |axis_entry| {
827 name_bytes = std.math.add(usize, name_bytes, axis_entry.name.len) catch return error.CapacityOverflow;
828 }
829 for (self.allSteps()) |step| {
830 if (step == .axis) {
831 name_bytes = std.math.add(usize, name_bytes, step.axis.name.len) catch return error.CapacityOverflow;
832 }
833 }
834 return .{
835 .axes = self.axes_len,
836 .steps = self.steps_len,
837 .name_bytes = name_bytes,
838 };
839 }
840
841 pub fn requiredReplayLimits(self: *const Schedule) Limits {
842 return .{
843 .axes = self.axes_len,
844 .steps = self.steps_len,
845 .name_bytes = self.names_len,
846 .axis_ids = self.next_axis_id,
847 };
848 }
849
850 fn replaceAxis(
851 self: *Schedule,
852 index: usize,
853 source: Axis,
854 factor: u64,
855 step_tag: enum { split, tile },
856 outer_suffix: []const u8,
857 inner_suffix: []const u8,
858 outer_extent: u64,
859 ) ScheduleError!Split {
860 const outer_bytes = std.math.add(usize, source.name.len, outer_suffix.len) catch return error.CapacityOverflow;
861 const inner_bytes = std.math.add(usize, source.name.len, inner_suffix.len) catch return error.CapacityOverflow;
862 const replacement_bytes = std.math.add(usize, outer_bytes, inner_bytes) catch return error.CapacityOverflow;
863 const name_growth = std.math.sub(usize, replacement_bytes, source.name.len) catch unreachable;
864 try self.requireGrowth(1, 1, name_growth, 2);
865
866 const generated_names = self.replaceName(source, outer_suffix, inner_suffix, name_growth);
867 const outer_id = self.allocAxisId();
868 const inner_id = self.allocAxisId();
869
870 if (index + 1 < self.axes_len) {
871 std.mem.copyBackwards(
872 Axis,
873 self.axes_storage[index + 2 .. self.axes_len + 1],
874 self.axes_storage[index + 1 .. self.axes_len],
875 );
876 }
877 self.axes_storage[index] = .{
878 .id = outer_id,
879 .name = generated_names.outer,
880 .extent = outer_extent,
881 };
882 self.axes_storage[index + 1] = .{
883 .id = inner_id,
884 .name = generated_names.inner,
885 .extent = factor,
886 };
887 self.axes_len += 1;
888
889 self.steps_storage[self.steps_len] = switch (step_tag) {
890 .split => .{ .split = .{
891 .source = source.id,
892 .outer = outer_id,
893 .inner = inner_id,
894 .factor = factor,
895 } },
896 .tile => .{ .tile = .{
897 .source = source.id,
898 .outer = outer_id,
899 .inner = inner_id,
900 .factor = factor,
901 } },
902 };
903 self.steps_len += 1;
904 return .{ .outer = outer_id, .inner = inner_id };
905 }
906
907 fn replaceName(
908 self: *Schedule,
909 source: Axis,
910 outer_suffix: []const u8,
911 inner_suffix: []const u8,
912 growth: usize,
913 ) struct { outer: []const u8, inner: []const u8 } {
914 const names_start = @intFromPtr(self.names_storage.ptr);
915 const source_start = @intFromPtr(source.name.ptr) - names_start;
916 const source_end = source_start + source.name.len;
917 const old_names_len = self.names_len;
918 const new_names_len = old_names_len + growth;
919
920 std.mem.copyBackwards(
921 u8,
922 self.names_storage[source_end + growth .. new_names_len],
923 self.names_storage[source_end..old_names_len],
924 );
925 self.shiftNamesAfter(source.id, source_end, growth);
926
927 const outer_end = source_end;
928 const outer_suffix_end = outer_end + outer_suffix.len;
929 @memcpy(self.names_storage[outer_end..outer_suffix_end], outer_suffix);
930 const inner_end = outer_suffix_end + source.name.len;
931 @memcpy(self.names_storage[outer_suffix_end..inner_end], self.names_storage[source_start..outer_end]);
932 const inner_suffix_end = inner_end + inner_suffix.len;
933 @memcpy(self.names_storage[inner_end..inner_suffix_end], inner_suffix);
934 std.debug.assert(inner_suffix_end == source_end + growth);
935 self.names_len = new_names_len;
936 return .{
937 .outer = self.names_storage[source_start..outer_suffix_end],
938 .inner = self.names_storage[outer_suffix_end..inner_suffix_end],
939 };
940 }
941
942 fn shiftNamesAfter(self: *Schedule, replaced_axis: AxisId, source_end: usize, growth: usize) void {
943 const names_start = @intFromPtr(self.names_storage.ptr);
944 for (self.axes_storage[0..self.axes_len]) |*axis_entry| {
945 if (axis_entry.id == replaced_axis) continue;
946 self.shiftName(&axis_entry.name, names_start, source_end, growth);
947 }
948 for (self.steps_storage[0..self.steps_len]) |*step| {
949 if (step.* == .axis) self.shiftName(&step.axis.name, names_start, source_end, growth);
950 }
951 }
952
953 fn shiftName(
954 self: *Schedule,
955 name: *[]const u8,
956 names_start: usize,
957 source_end: usize,
958 growth: usize,
959 ) void {
960 const offset = @intFromPtr(name.ptr) - names_start;
961 if (offset < source_end) return;
962 name.* = self.names_storage[offset + growth .. offset + growth + name.len];
963 }
964
965 fn requireBuilding(self: *const Schedule) ScheduleError!void {
966 if (self.phase != .initialization) return error.ScheduleSealed;
967 }
968
969 fn requireLimits(self: *const Schedule, limits: Limits) ScheduleError!void {
970 if (limits.axes > self.capacity.axes) return error.AxisCapacityExceeded;
971 if (limits.steps > self.capacity.steps) return error.StepCapacityExceeded;
972 if (limits.name_bytes > self.capacity.name_bytes) return error.NameCapacityExceeded;
973 if (limits.axis_ids > self.capacity.axis_ids) return error.AxisIdCapacityExceeded;
974 }
975
976 fn requireGrowth(
977 self: *const Schedule,
978 axes: usize,
979 steps: usize,
980 name_bytes: usize,
981 axis_ids: usize,
982 ) ScheduleError!void {
983 const axes_after = std.math.add(usize, self.axes_len, axes) catch return error.CapacityOverflow;
984 const steps_after = std.math.add(usize, self.steps_len, steps) catch return error.CapacityOverflow;
985 const names_after = std.math.add(usize, self.names_len, name_bytes) catch return error.CapacityOverflow;
986 const ids_after = std.math.add(usize, self.next_axis_id, axis_ids) catch return error.CapacityOverflow;
987 if (axes_after > self.capacity.axes) return error.AxisCapacityExceeded;
988 if (steps_after > self.capacity.steps) return error.StepCapacityExceeded;
989 if (names_after > self.capacity.name_bytes) return error.NameCapacityExceeded;
990 if (ids_after > self.capacity.axis_ids) return error.AxisIdCapacityExceeded;
991 }
992
993 fn writeName(self: *Schedule, name: []const u8) []const u8 {
994 const start = self.names_len;
995 const end = start + name.len;
996 @memcpy(self.names_storage[start..end], name);
997 self.names_len = end;
998 return self.names_storage[start..end];
999 }
1000
1001 fn allocAxisId(self: *Schedule) AxisId {
1002 const id = AxisId.fromIndex(self.next_axis_id);
1003 self.next_axis_id += 1;
1004 return id;
1005 }
1006
1007 fn axisIndex(self: *const Schedule, id: AxisId) ScheduleError!usize {
1008 for (self.allAxes(), 0..) |axis_entry, index| {
1009 if (axis_entry.id == id) return index;
1010 }
1011 return error.UnknownAxis;
1012 }
1013
1014 fn reset(self: *Schedule) void {
1015 self.axes_len = 0;
1016 self.steps_len = 0;
1017 self.names_len = 0;
1018 self.next_axis_id = 0;
1019 }
1020 };
1021
1022 comptime {
1023 alloc_phase.capacity.requireAllocatorRejectingOwnerShape(Schedule);
1024 }
1025
1026 comptime {
1027 alloc_phase.capacity.requireAllocatorRejectingOwnerShape(Snapshot);
1028 }
1029
1030 pub fn createSnapshot(
1031 allocator: Allocator,
1032 limits: Snapshot.Limits,
1033 schedule: *const Schedule,
1034 ) ScheduleError!Snapshot {
1035 var snapshot = try Snapshot.init(allocator, limits);
1036 errdefer snapshot.deinit(allocator);
1037 try snapshot.capture(schedule);
1038 snapshot.activate();
1039 return snapshot;
1040 }
1041
1042 pub fn replaySnapshot(
1043 allocator: Allocator,
1044 limits: Schedule.Limits,
1045 snapshot: *const Snapshot,
1046 ) ScheduleError!Schedule {
1047 const required = try snapshot.requiredReplayLimits();
1048 try requireAvailableReplayLimits(limits, required);
1049 var schedule = try Schedule.init(allocator, limits);
1050 errdefer schedule.deinit(allocator);
1051 try schedule.replay(try snapshot.record());
1052 schedule.activate();
1053 return schedule;
1054 }
1055
1056 fn requireAvailableReplayLimits(
1057 available: Schedule.Limits,
1058 required: Schedule.Limits,
1059 ) ScheduleError!void {
1060 if (required.axes > available.axes) return error.AxisCapacityExceeded;
1061 if (required.steps > available.steps) return error.StepCapacityExceeded;
1062 if (required.name_bytes > available.name_bytes) return error.NameCapacityExceeded;
1063 if (required.axis_ids > available.axis_ids) return error.AxisIdCapacityExceeded;
1064 }
1065
1066 fn launchFromAxes(axes: []const Axis) ScheduleError!Launch {
1067 var grid: [3]u32 = .{ 1, 1, 1 };
1068 var block: [3]u32 = .{ 1, 1, 1 };
1069 for (axes) |axis_entry| {
1070 const target = axis_entry.bind orelse continue;
1071 const extent = try launchAxisExtent(axis_entry);
1072 const converted = std.math.cast(u32, extent) orelse return error.LaunchDimensionOverflow;
1073 if (target.isGrid()) {
1074 grid[target.dim()] = converted;
1075 } else {
1076 block[target.dim()] = converted;
1077 }
1078 }
1079 return .{ .grid = grid, .block = block };
1080 }
1081
1082 fn launchAxisExtent(axis_entry: Axis) ScheduleError!u64 {
1083 const extent = axis_entry.extent orelse return error.MissingBindingExtent;
1084 const width = axis_entry.vector_width orelse return extent;
1085 if (extent % width != 0) return error.ExtentNotDivisible;
1086 return extent / width;
1087 }
1088
1089 fn fingerprintSchedule(version: u32, axes: []const Axis, steps: []const Step) u64 {
1090 var builder = choir.product.incremental.FingerprintBuilder{};
1091 builder.updateBytes("accy.kernel.schedule");
1092 builder.updateU32(version);
1093 builder.updateUsize(axes.len);
1094 for (axes) |axis_entry| fingerprintAxis(&builder, axis_entry);
1095 builder.updateUsize(steps.len);
1096 for (steps) |step| fingerprintStep(&builder, step);
1097 return builder.finish();
1098 }
1099
1100 fn fingerprintAxis(builder: *choir.product.incremental.FingerprintBuilder, axis_entry: Axis) void {
1101 fingerprintAxisId(builder, axis_entry.id);
1102 builder.updateBytes(axis_entry.name);
1103 builder.updateOptionalU64(axis_entry.extent);
1104 builder.updateOptionalEnumTag(axis_entry.bind);
1105 builder.updateOptionalU32(axis_entry.vector_width);
1106 builder.updateOptionalU32(axis_entry.unroll_factor);
1107 }
1108
1109 fn fingerprintStep(builder: *choir.product.incremental.FingerprintBuilder, step: Step) void {
1110 builder.updateEnumTag(std.meta.activeTag(step));
1111 switch (step) {
1112 .axis => |axis| {
1113 fingerprintAxisId(builder, axis.id);
1114 builder.updateBytes(axis.name);
1115 builder.updateOptionalU64(axis.extent);
1116 },
1117 .split => |split_value| {
1118 fingerprintAxisId(builder, split_value.source);
1119 fingerprintAxisId(builder, split_value.outer);
1120 fingerprintAxisId(builder, split_value.inner);
1121 builder.updateU64(split_value.factor);
1122 },
1123 .tile => |tile_value| {
1124 fingerprintAxisId(builder, tile_value.source);
1125 fingerprintAxisId(builder, tile_value.outer);
1126 fingerprintAxisId(builder, tile_value.inner);
1127 builder.updateU64(tile_value.factor);
1128 },
1129 .bind => |bind_value| {
1130 fingerprintAxisId(builder, bind_value.axis);
1131 builder.updateEnumTag(bind_value.target);
1132 },
1133 .vectorize => |vectorize_value| {
1134 fingerprintAxisId(builder, vectorize_value.axis);
1135 builder.updateU32(vectorize_value.width);
1136 },
1137 .unroll => |unroll_value| {
1138 fingerprintAxisId(builder, unroll_value.axis);
1139 builder.updateU32(unroll_value.factor);
1140 },
1141 }
1142 }
1143
1144 fn fingerprintAxisId(builder: *choir.product.incremental.FingerprintBuilder, id: AxisId) void {
1145 builder.updateU32(id.index());
1146 }
1147
1148 fn axesEqual(lhs: []const Axis, rhs: []const Axis) bool {
1149 if (lhs.len != rhs.len) return false;
1150 for (lhs, rhs) |left, right| {
1151 if (!axisEqual(left, right)) return false;
1152 }
1153 return true;
1154 }
1155
1156 fn axisEqual(lhs: Axis, rhs: Axis) bool {
1157 return lhs.id == rhs.id and
1158 std.mem.eql(u8, lhs.name, rhs.name) and
1159 lhs.extent == rhs.extent and
1160 lhs.bind == rhs.bind and
1161 lhs.vector_width == rhs.vector_width and
1162 lhs.unroll_factor == rhs.unroll_factor;
1163 }
1164
1165 fn stepsEqual(lhs: []const Step, rhs: []const Step) bool {
1166 if (lhs.len != rhs.len) return false;
1167 for (lhs, rhs) |left, right| {
1168 if (!stepEqual(left, right)) return false;
1169 }
1170 return true;
1171 }
1172
1173 fn stepEqual(lhs: Step, rhs: Step) bool {
1174 if (std.meta.activeTag(lhs) != std.meta.activeTag(rhs)) return false;
1175 return switch (lhs) {
1176 .axis => |left| axisStepEqual(left, rhs.axis),
1177 .split => |left| std.meta.eql(left, rhs.split),
1178 .tile => |left| std.meta.eql(left, rhs.tile),
1179 .bind => |left| std.meta.eql(left, rhs.bind),
1180 .vectorize => |left| std.meta.eql(left, rhs.vectorize),
1181 .unroll => |left| std.meta.eql(left, rhs.unroll),
1182 };
1183 }
1184
1185 fn axisStepEqual(lhs: AxisStep, rhs: AxisStep) bool {
1186 return lhs.id == rhs.id and
1187 std.mem.eql(u8, lhs.name, rhs.name) and
1188 lhs.extent == rhs.extent;
1189 }
1190
1191 fn placeSlice(
1192 comptime T: type,
1193 count: usize,
1194 cursor: *usize,
1195 allocation_alignment: *usize,
1196 ) error{CapacityOverflow}!usize {
1197 const bytes = std.math.mul(usize, @sizeOf(T), count) catch return error.CapacityOverflow;
1198 return placeBytes(bytes, @alignOf(T), cursor, allocation_alignment);
1199 }
1200
1201 fn placeBytes(
1202 bytes: usize,
1203 alignment: usize,
1204 cursor: *usize,
1205 allocation_alignment: *usize,
1206 ) error{CapacityOverflow}!usize {
1207 const mask = std.math.sub(usize, alignment, 1) catch unreachable;
1208 const padded = std.math.add(usize, cursor.*, mask) catch return error.CapacityOverflow;
1209 const offset = padded & ~mask;
1210 cursor.* = std.math.add(usize, offset, bytes) catch return error.CapacityOverflow;
1211 allocation_alignment.* = @max(allocation_alignment.*, alignment);
1212 return offset;
1213 }
1214
1215 fn typedSlice(
1216 comptime T: type,
1217 storage: [*]u8,
1218 offset: usize,
1219 count: usize,
1220 ) []T {
1221 const pointer: [*]T = @ptrCast(@alignCast(storage + offset));
1222 return pointer[0..count];
1223 }
1224
1225 fn initTestingSchedule() !Schedule {
1226 return Schedule.init(testing.allocator, Schedule.Limits.testing);
1227 }
1228
1229 test "Schedule capacity follows an independent aligned byte model" {
1230 comptime {
1231 @stardustClaim(
1232 @import("alloc_phase").capacity.witness(Schedule, "accy_kernel_schedule_capacity"),
1233 null,
1234 null,
1235 null,
1236 null,
1237 null,
1238 null,
1239 );
1240 }
1241
1242 const limits = Schedule.Limits{
1243 .axes = 3,
1244 .steps = 5,
1245 .name_bytes = 7,
1246 .axis_ids = 9,
1247 };
1248 const capacity = try Schedule.Capacity.derive(limits);
1249 var expected: usize = 0;
1250 expected = std.mem.alignForward(usize, expected, @alignOf(Axis)) + 3 * @sizeOf(Axis);
1251 expected = std.mem.alignForward(usize, expected, @alignOf(Step)) + 5 * @sizeOf(Step);
1252 expected = std.mem.alignForward(usize, expected, @alignOf(u8)) + 7;
1253 try testing.expectEqual(expected, capacity.storage_bytes);
1254 try testing.expectEqual(@as(usize, 9), capacity.axis_ids);
1255
1256 var overflow = limits;
1257 overflow.axes = std.math.maxInt(usize);
1258 try testing.expectError(error.CapacityOverflow, Schedule.Capacity.derive(overflow));
1259 }
1260
1261 test "Schedule initialization retries after its one outer acquisition fails" {
1262 comptime {
1263 @stardustClaim(
1264 @import("alloc_phase").capacity.witness(Schedule, "accy_kernel_schedule_single_acquisition"),
1265 null,
1266 null,
1267 null,
1268 null,
1269 null,
1270 null,
1271 );
1272 }
1273
1274 var failing = testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
1275 try testing.expectError(
1276 error.OutOfMemory,
1277 Schedule.init(failing.allocator(), Schedule.Limits.standard),
1278 );
1279 failing.fail_index = std.math.maxInt(usize);
1280 var schedule = try Schedule.init(failing.allocator(), Schedule.Limits.standard);
1281 schedule.deinit(failing.allocator());
1282 }
1283
1284 test "Schedule records split bind and derives launch geometry" {
1285 var schedule = try initTestingSchedule();
1286 defer schedule.deinit(testing.allocator);
1287
1288 const i = try schedule.addAxis("i", 1024);
1289 const split_i = try schedule.split(i, 256);
1290 try schedule.bind(split_i.outer, .block_x);
1291 try schedule.bind(split_i.inner, .thread_x);
1292
1293 const launch = try schedule.launch();
1294 try testing.expectEqual(@as(u32, 4), launch.grid[0]);
1295 try testing.expectEqual(@as(u32, 256), launch.block[0]);
1296 try testing.expectEqual(@as(usize, 4), schedule.allSteps().len);
1297 }
1298
1299 test "Schedule vectorization derives packet launch geometry" {
1300 var schedule = try initTestingSchedule();
1301 defer schedule.deinit(testing.allocator);
1302
1303 const i = try schedule.addAxis("i", 32);
1304 const tiled_i = try schedule.tile(i, 16);
1305 try schedule.vectorize(tiled_i.inner, 4);
1306 try schedule.bind(tiled_i.outer, .block_x);
1307 try schedule.bind(tiled_i.inner, .thread_x);
1308
1309 const launch = try schedule.launch();
1310 try testing.expectEqual(@as(u32, 2), launch.grid[0]);
1311 try testing.expectEqual(@as(u32, 4), launch.block[0]);
1312 try testing.expectEqual(@as(?u32, 4), schedule.allAxes()[1].vector_width);
1313
1314 var snapshot = try createSnapshot(testing.allocator, Snapshot.Limits.testing, &schedule);
1315 defer snapshot.deinit(testing.allocator);
1316 try testing.expectEqual(@as(?u32, 4), snapshot.allAxes()[1].vector_width);
1317 const snapshot_launch = try snapshot.launch();
1318 try testing.expectEqual(@as(u32, 2), snapshot_launch.grid[0]);
1319 try testing.expectEqual(@as(u32, 4), snapshot_launch.block[0]);
1320 }
1321
1322 test "Schedule rejects duplicate bindings and non-divisible splits" {
1323 var schedule = try initTestingSchedule();
1324 defer schedule.deinit(testing.allocator);
1325
1326 const i = try schedule.addAxis("i", 10);
1327 try testing.expectError(error.ExtentNotDivisible, schedule.split(i, 4));
1328
1329 const j = try schedule.addAxis("j", 5);
1330 try schedule.bind(i, .block_x);
1331 try testing.expectError(error.BindTargetAlreadyUsed, schedule.bind(j, .block_x));
1332 }
1333
1334 test "Schedule rejects invalid vectorization factors" {
1335 var schedule = try initTestingSchedule();
1336 defer schedule.deinit(testing.allocator);
1337
1338 const i = try schedule.addAxis("i", 10);
1339 try testing.expectError(error.InvalidFactor, schedule.vectorize(i, 0));
1340 try testing.expectError(error.ExtentNotDivisible, schedule.vectorize(i, 4));
1341 }
1342
1343 test "Schedule tiles non-divisible axes with ceil outer extent" {
1344 var schedule = try initTestingSchedule();
1345 defer schedule.deinit(testing.allocator);
1346
1347 const i = try schedule.addAxis("i", 10);
1348 const tiled_i = try schedule.tile(i, 4);
1349 try schedule.bind(tiled_i.outer, .block_x);
1350 try schedule.bind(tiled_i.inner, .thread_x);
1351
1352 const launch = try schedule.launch();
1353 try testing.expectEqual(@as(u32, 3), launch.grid[0]);
1354 try testing.expectEqual(@as(u32, 4), launch.block[0]);
1355 try testing.expectEqualStrings("i_tile", schedule.allAxes()[0].name);
1356 try testing.expectEqualStrings("i_lane", schedule.allAxes()[1].name);
1357 }
1358
1359 test "Schedule split reclaims the replaced name and preserves shifted names" {
1360 const exact_limits = Schedule.Limits{
1361 .axes = 3,
1362 .steps = 3,
1363 .name_bytes = 17,
1364 .axis_ids = 4,
1365 };
1366 var schedule = try Schedule.init(testing.allocator, exact_limits);
1367 defer schedule.deinit(testing.allocator);
1368 const i = try schedule.addAxis("i", 8);
1369 const j = try schedule.addAxis("j", 4);
1370 _ = try schedule.split(i, 2);
1371 try testing.expectEqual(@as(usize, 17), schedule.capacityUsage().name_bytes);
1372 try testing.expectEqualStrings("j", (try schedule.getAxis(j)).name);
1373 try testing.expectEqualStrings("i", schedule.allSteps()[0].axis.name);
1374 try testing.expectEqualStrings("j", schedule.allSteps()[1].axis.name);
1375
1376 var short_limits = exact_limits;
1377 short_limits.name_bytes -= 1;
1378 var rejected = try Schedule.init(testing.allocator, short_limits);
1379 defer rejected.deinit(testing.allocator);
1380 const rejected_i = try rejected.addAxis("i", 8);
1381 _ = try rejected.addAxis("j", 4);
1382 const before = rejected.fingerprint();
1383 try testing.expectError(error.NameCapacityExceeded, rejected.split(rejected_i, 2));
1384 try testing.expectEqual(before, rejected.fingerprint());
1385 }
1386
1387 test "Schedule snapshot owns replayable schedule data" {
1388 comptime {
1389 @stardustClaim(
1390 @import("alloc_phase").capacity.witness(Snapshot, "accy_kernel_schedule_snapshot_lifetime_transitive_risk"),
1391 null,
1392 null,
1393 null,
1394 null,
1395 null,
1396 null,
1397 );
1398 }
1399 comptime {
1400 @stardustClaim(
1401 @import("alloc_phase").capacity.witness(Snapshot, "accy_kernel_schedule_snapshot_lifetime_foreign_risk"),
1402 null,
1403 null,
1404 null,
1405 null,
1406 null,
1407 null,
1408 );
1409 }
1410
1411 var snapshot: Snapshot = undefined;
1412 var schedule_fingerprint: u64 = 0;
1413 {
1414 var schedule = try initTestingSchedule();
1415 defer schedule.deinit(testing.allocator);
1416
1417 const i = try schedule.addAxis("i", 32);
1418 const split_i = try schedule.split(i, 8);
1419 try schedule.bind(split_i.outer, .block_x);
1420 try schedule.bind(split_i.inner, .thread_x);
1421
1422 snapshot = try createSnapshot(testing.allocator, Snapshot.Limits.testing, &schedule);
1423 schedule_fingerprint = schedule.fingerprint();
1424 }
1425 defer snapshot.deinit(testing.allocator);
1426
1427 try testing.expectEqual(snapshot_version, snapshot.version);
1428 try testing.expectEqual(@as(usize, 2), snapshot.allAxes().len);
1429 try testing.expectEqualStrings("i_outer", snapshot.allAxes()[0].name);
1430 try testing.expectEqualStrings("i_inner", snapshot.allAxes()[1].name);
1431 try testing.expectEqual(@as(usize, 4), snapshot.allSteps().len);
1432 try testing.expectEqualStrings("i", snapshot.allSteps()[0].axis.name);
1433
1434 const launch = try snapshot.launch();
1435 try testing.expectEqual(@as(u32, 4), launch.grid[0]);
1436 try testing.expectEqual(@as(u32, 8), launch.block[0]);
1437 try testing.expectEqual(schedule_fingerprint, snapshot.fingerprint());
1438
1439 var replayed = try replaySnapshot(
1440 testing.allocator,
1441 try snapshot.requiredReplayLimits(),
1442 &snapshot,
1443 );
1444 defer replayed.deinit(testing.allocator);
1445 try testing.expectEqual(snapshot.fingerprint(), replayed.fingerprint());
1446 try testing.expectEqual(@as(usize, 2), replayed.allAxes().len);
1447 try testing.expectEqualStrings("i_outer", replayed.allAxes()[0].name);
1448 try testing.expectEqualStrings("i_inner", replayed.allAxes()[1].name);
1449 try testing.expectEqual(@as(usize, 4), replayed.allSteps().len);
1450 }
1451
1452 test "Schedule fingerprint is stable and sensitive to replayable steps" {
1453 var first = try initTestingSchedule();
1454 defer first.deinit(testing.allocator);
1455 const first_i = try first.addAxis("i", 32);
1456 const first_split = try first.split(first_i, 8);
1457 try first.bind(first_split.outer, .block_x);
1458 try first.bind(first_split.inner, .thread_x);
1459
1460 var second = try initTestingSchedule();
1461 defer second.deinit(testing.allocator);
1462 const second_i = try second.addAxis("i", 32);
1463 const second_split = try second.split(second_i, 8);
1464 try second.bind(second_split.outer, .block_x);
1465 try second.bind(second_split.inner, .thread_x);
1466
1467 var changed = try initTestingSchedule();
1468 defer changed.deinit(testing.allocator);
1469 const changed_i = try changed.addAxis("i", 32);
1470 const changed_split = try changed.split(changed_i, 4);
1471 try changed.bind(changed_split.outer, .block_x);
1472 try changed.bind(changed_split.inner, .thread_x);
1473
1474 try testing.expectEqual(first.fingerprint(), second.fingerprint());
1475 try testing.expect(first.fingerprint() != changed.fingerprint());
1476 }
1477
1478 test "Schedule replay rejects unsupported snapshot versions" {
1479 var schedule = try initTestingSchedule();
1480 defer schedule.deinit(testing.allocator);
1481
1482 const i = try schedule.addAxis("i", 32);
1483 const split_i = try schedule.split(i, 8);
1484 try schedule.bind(split_i.outer, .block_x);
1485 try schedule.bind(split_i.inner, .thread_x);
1486
1487 var snapshot = try createSnapshot(testing.allocator, Snapshot.Limits.testing, &schedule);
1488 defer snapshot.deinit(testing.allocator);
1489 snapshot.version += 1;
1490
1491 try testing.expectError(error.UnsupportedVersion, snapshot.requiredReplayLimits());
1492 try testing.expectError(error.UnsupportedVersion, replaySnapshot(
1493 testing.allocator,
1494 Schedule.Limits.testing,
1495 &snapshot,
1496 ));
1497 }
1498
1499 test "Schedule admits independent maxima and preserves its fingerprint at max plus one" {
1500 comptime {
1501 @stardustClaim(
1502 @import("alloc_phase").capacity.witness(Schedule, "accy_kernel_schedule_boundary_overload"),
1503 null,
1504 null,
1505 null,
1506 null,
1507 null,
1508 null,
1509 );
1510 }
1511 comptime {
1512 @stardustClaim(
1513 @import("alloc_phase").capacity.witness(Schedule, "accy_kernel_schedule_boundary_foreign_risk"),
1514 null,
1515 null,
1516 null,
1517 null,
1518 null,
1519 null,
1520 );
1521 }
1522
1523 const cases = [_]struct {
1524 limits: Schedule.Limits,
1525 expected: ScheduleError,
1526 }{
1527 .{ .limits = .{ .axes = 1, .steps = 2, .name_bytes = 4, .axis_ids = 2 }, .expected = error.AxisCapacityExceeded },
1528 .{ .limits = .{ .axes = 2, .steps = 1, .name_bytes = 4, .axis_ids = 2 }, .expected = error.StepCapacityExceeded },
1529 .{ .limits = .{ .axes = 2, .steps = 2, .name_bytes = 2, .axis_ids = 2 }, .expected = error.NameCapacityExceeded },
1530 .{ .limits = .{ .axes = 2, .steps = 2, .name_bytes = 4, .axis_ids = 1 }, .expected = error.AxisIdCapacityExceeded },
1531 };
1532
1533 for (cases) |case| {
1534 var schedule = try Schedule.init(testing.allocator, case.limits);
1535 defer schedule.deinit(testing.allocator);
1536 _ = try schedule.addAxis("i", 8);
1537 const before = schedule.fingerprint();
1538 try testing.expectError(case.expected, schedule.addAxis("j", 8));
1539 try testing.expectEqual(before, schedule.fingerprint());
1540 }
1541 }
1542
1543 test "Snapshot capacity follows an independent aligned byte model" {
1544 comptime {
1545 @stardustClaim(
1546 @import("alloc_phase").capacity.witness(Snapshot, "accy_kernel_schedule_snapshot_capacity"),
1547 null,
1548 null,
1549 null,
1550 null,
1551 null,
1552 null,
1553 );
1554 }
1555
1556 const limits = Snapshot.Limits{ .axes = 3, .steps = 5, .name_bytes = 7 };
1557 const capacity = try Snapshot.Capacity.derive(limits);
1558 var expected: usize = 0;
1559 expected = std.mem.alignForward(usize, expected, @alignOf(Axis)) + 3 * @sizeOf(Axis);
1560 expected = std.mem.alignForward(usize, expected, @alignOf(Step)) + 5 * @sizeOf(Step);
1561 expected = std.mem.alignForward(usize, expected, @alignOf(u8)) + 7;
1562 try testing.expectEqual(expected, capacity.storage_bytes);
1563 }
1564
1565 test "Snapshot admits exact capacity and rejects max plus one before mutation" {
1566 comptime {
1567 @stardustClaim(
1568 @import("alloc_phase").capacity.witness(Snapshot, "accy_kernel_schedule_snapshot_boundary"),
1569 null,
1570 null,
1571 null,
1572 null,
1573 null,
1574 null,
1575 );
1576 }
1577
1578 var schedule = try initTestingSchedule();
1579 defer schedule.deinit(testing.allocator);
1580 const axis = try schedule.addAxis("i", 32);
1581 const split_axis = try schedule.split(axis, 8);
1582 try schedule.bind(split_axis.outer, .block_x);
1583 try schedule.bind(split_axis.inner, .thread_x);
1584
1585 const exact = try schedule.snapshotLimits();
1586 var accepted = try createSnapshot(testing.allocator, exact, &schedule);
1587 accepted.deinit(testing.allocator);
1588
1589 const cases = [_]struct {
1590 limits: Snapshot.Limits,
1591 expected: ScheduleError,
1592 }{
1593 .{ .limits = .{ .axes = exact.axes - 1, .steps = exact.steps, .name_bytes = exact.name_bytes }, .expected = error.AxisCapacityExceeded },
1594 .{ .limits = .{ .axes = exact.axes, .steps = exact.steps - 1, .name_bytes = exact.name_bytes }, .expected = error.StepCapacityExceeded },
1595 .{ .limits = .{ .axes = exact.axes, .steps = exact.steps, .name_bytes = exact.name_bytes - 1 }, .expected = error.NameCapacityExceeded },
1596 };
1597 for (cases) |case| {
1598 var rejected = try Snapshot.init(testing.allocator, case.limits);
1599 defer rejected.deinit(testing.allocator);
1600 const before_usage = rejected.capacityUsage();
1601 const before_fingerprint = rejected.fingerprint();
1602 try testing.expectError(case.expected, rejected.capture(&schedule));
1603 try testing.expectEqualDeep(before_usage, rejected.capacityUsage());
1604 try testing.expectEqual(before_fingerprint, rejected.fingerprint());
1605 try testing.expect(!rejected.captured);
1606 }
1607 }
1608
1609 test "Schedule replay admits exact capacity and rolls back max plus one" {
1610 comptime {
1611 @stardustClaim(
1612 @import("alloc_phase").capacity.witness(Schedule, "accy_kernel_schedule_replay_boundary"),
1613 null,
1614 null,
1615 null,
1616 null,
1617 null,
1618 null,
1619 );
1620 }
1621
1622 var source = try initTestingSchedule();
1623 defer source.deinit(testing.allocator);
1624 const axis = try source.addAxis("i", 32);
1625 const split_axis = try source.split(axis, 8);
1626 try source.bind(split_axis.outer, .block_x);
1627 try source.bind(split_axis.inner, .thread_x);
1628
1629 var snapshot = try createSnapshot(testing.allocator, Snapshot.Limits.testing, &source);
1630 defer snapshot.deinit(testing.allocator);
1631 const exact = try snapshot.requiredReplayLimits();
1632 const exact_capacity = try snapshot.requiredReplayCapacity();
1633 try testing.expectEqualDeep(exact, exact_capacity.asLimits());
1634 var accepted = try replaySnapshot(testing.allocator, exact, &snapshot);
1635 defer accepted.deinit(testing.allocator);
1636 try testing.expectEqual(snapshot.fingerprint(), accepted.fingerprint());
1637
1638 const cases = [_]struct {
1639 limits: Schedule.Limits,
1640 expected: ScheduleError,
1641 }{
1642 .{ .limits = .{ .axes = exact.axes - 1, .steps = exact.steps, .name_bytes = exact.name_bytes, .axis_ids = exact.axis_ids }, .expected = error.AxisCapacityExceeded },
1643 .{ .limits = .{ .axes = exact.axes, .steps = exact.steps - 1, .name_bytes = exact.name_bytes, .axis_ids = exact.axis_ids }, .expected = error.StepCapacityExceeded },
1644 .{ .limits = .{ .axes = exact.axes, .steps = exact.steps, .name_bytes = exact.name_bytes - 1, .axis_ids = exact.axis_ids }, .expected = error.NameCapacityExceeded },
1645 .{ .limits = .{ .axes = exact.axes, .steps = exact.steps, .name_bytes = exact.name_bytes, .axis_ids = exact.axis_ids - 1 }, .expected = error.AxisIdCapacityExceeded },
1646 };
1647 for (cases) |case| {
1648 var rejected = try Schedule.init(testing.allocator, case.limits);
1649 defer rejected.deinit(testing.allocator);
1650 const before_usage = rejected.capacityUsage();
1651 const before_fingerprint = rejected.fingerprint();
1652 try testing.expectError(case.expected, rejected.replay(try snapshot.record()));
1653 try testing.expectEqualDeep(before_usage, rejected.capacityUsage());
1654 try testing.expectEqual(before_fingerprint, rejected.fingerprint());
1655
1656 var failing = testing.FailingAllocator.init(testing.allocator, .{ .fail_index = 0 });
1657 try testing.expectError(
1658 case.expected,
1659 replaySnapshot(failing.allocator(), case.limits, &snapshot),
1660 );
1661 try testing.expectEqual(@as(usize, 0), failing.alloc_index);
1662 }
1663 }