lib/accy/src/choir/record/program.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const model = @import("../../kernel/model/root.zig");
4 const records = @import("root.zig");
5 const schedule = model.core.schedule;
6 const Configuration = choir.product.operation.Configuration;
7
8 /// The stored form of one generated kernel program: its own compiler bytecode,
9 /// the number of its function within that bytecode, its parameters, its loop
10 /// and thread structure, its launch shape and its arithmetic policy. A kernel
11 /// stage record carries one of these for each kernel it generated. The function
12 /// number counts operations in this program's own bytecode, which is separate
13 /// from the bytecode of the whole stage.
14 pub const Record = struct {
15 image: []const u8,
16 function: u32,
17 params: []const model.Param,
18 schedule: schedule.Record,
19 launch: schedule.Launch,
20 arithmetic: choir.product.recipe.ArithmeticPolicy,
21 };
22
23 /// Copies each generated program into its record for the kernel stage: verifies
24 /// the program, encodes its bytecode, and fails with `error.RecordLimit` when
25 /// that bytecode exceeds the configured image size. The call numbers the
26 /// program's operations, snapshots its schedule, encodes the whole record,
27 /// decodes it again and compares the result with the live program. The caller
28 /// receives the decoded copy, which owns its own memory apart from `source`.
29 pub fn capture(
30 allocator: std.mem.Allocator,
31 source: *const model.Program,
32 comptime configuration: Configuration,
33 ) !records.codec.Decoded(Record) {
34 comptime coverage();
35 const root = source.kernelModule();
36 var context = try choir.ir.Context.init(allocator, configuration.context);
37 defer context.deinit(allocator);
38 context.arithmetic_policy = root.context.arithmetic_policy;
39 try configuration.register(&context);
40 try choir.ir.verifyOperation(root, configuration.verify);
41 const image = try choir.bytecode.qualification.encode(
42 allocator,
43 root,
44 &.{},
45 &context,
46 configuration.codec,
47 );
48 defer allocator.free(image);
49 if (image.len > configuration.image.bytes) return error.RecordLimit;
50 const entity_limit = configuration.image.entities;
51 var references = try records.reference.Index.init(allocator, root, entity_limit);
52 defer references.deinit();
53 var snapshot = try source.scheduleSnapshot(allocator);
54 defer snapshot.deinit(allocator);
55 const projected = Record{
56 .image = image,
57 .function = (try references.operation(source.storage.kernel.func().op)).ordinal,
58 .params = source.params(),
59 .schedule = try snapshot.record(),
60 .launch = try source.launch(),
61 .arithmetic = root.context.arithmetic_policy,
62 };
63 const bytes = try records.codec.encode(allocator, Record, .kernel, projected, &references);
64 defer allocator.free(bytes);
65 var decoded = try records.codec.decode(allocator, Record, .kernel, bytes);
66 errdefer decoded.deinit();
67 try records.codec.compare(decoded.value, projected, &references);
68 return decoded;
69 }
70
71 /// Validates each stored program for the checker before a record is accepted:
72 /// decodes the bytecode into a scratch context, verifies it, and checks that
73 /// the function number is in range (`error.UnboundProductInput`). The function
74 /// must be a kernel with a body, named `entry_name`, whose arguments match the
75 /// stored parameters in count and type. The stored launch shape must match the
76 /// result of replaying the schedule, and any mismatch is
77 /// `error.InvalidStageRecord`.
78 pub fn validate(
79 allocator: std.mem.Allocator,
80 value: Record,
81 entry_name: []const u8,
82 comptime configuration: Configuration,
83 ) !void {
84 const image = try choir.bytecode.image.Index.create(
85 allocator,
86 value.image,
87 configuration.image,
88 );
89 defer image.destroy();
90 if (value.function >= image.view().operations.len) return error.UnboundProductInput;
91 var context = try choir.ir.Context.init(allocator, configuration.context);
92 defer context.deinit(allocator);
93 context.arithmetic_policy = value.arithmetic;
94 try configuration.register(&context);
95 var decoded = try choir.bytecode.decodeModule(allocator, &context, value.image);
96 defer decoded.deinit();
97 try choir.ir.verifyOperation(decoded.module, configuration.verify);
98 var references = try records.reference.Index.init(
99 allocator,
100 decoded.module,
101 configuration.image.entities,
102 );
103 defer references.deinit();
104 const function = try selectedFunction(&references, value.function);
105 const name = function.getName() orelse return error.InvalidStageRecord;
106 if (!std.mem.eql(u8, name, entry_name)) return error.InvalidStageRecord;
107 if (!function.isKernel() or !function.hasBody()) return error.InvalidStageRecord;
108 if (function.getNumArguments() != value.params.len) return error.InvalidStageRecord;
109 for (value.params, function.getArguments()) |param, argument| {
110 if (!argument.type.eql(try param.getType(&context))) return error.InvalidStageRecord;
111 }
112 var replayed = try schedule.Schedule.init(allocator, value.schedule.replay_limits);
113 defer replayed.deinit(allocator);
114 try replayed.replay(value.schedule);
115 if (!std.meta.eql(try replayed.launch(), value.launch)) return error.InvalidStageRecord;
116 }
117
118 /// Restores a working kernel program from a stored record for a later compile:
119 /// rebuilds the whole program in new storage sized from the record, then
120 /// verifies it, checks its name, and replays its schedule. Storage that runs
121 /// out gives `error.WorkExhausted`. The returned program copies the bytecode,
122 /// the parameters and the schedule, and keeps no pointer into `value`.
123 pub fn restore(
124 allocator: std.mem.Allocator,
125 value: Record,
126 entry_name: []const u8,
127 comptime configuration: Configuration,
128 ) !model.Program {
129 const limits = try restorationLimits(value, entry_name, configuration);
130 const capacity = model.program.Capacity.derive(limits) catch return error.WorkOverflow;
131 var storage = try model.core.builder.Storage.init(allocator, limits.raw());
132 var transferred = false;
133 defer if (!transferred) storage.deinit(allocator);
134 const program = restoreIn(
135 allocator,
136 value,
137 entry_name,
138 &storage,
139 capacity,
140 configuration,
141 ) catch |err| {
142 return if (storage.context.exhaustedSegment() != null) error.WorkExhausted else err;
143 };
144 transferred = true;
145 return program;
146 }
147
148 fn restorationLimits(
149 value: Record,
150 entry_name: []const u8,
151 comptime configuration: Configuration,
152 ) !model.program.Limits {
153 var names: usize = 0;
154 for (value.schedule.axes) |axis| {
155 names = std.math.add(usize, names, axis.name.len) catch return error.WorkOverflow;
156 }
157 for (value.schedule.steps) |step| {
158 if (step == .axis) {
159 names = std.math.add(usize, names, step.axis.name.len) catch return error.WorkOverflow;
160 }
161 }
162 return .{
163 .context = configuration.context,
164 .parameters = value.params.len,
165 .kernel_name_bytes = entry_name.len,
166 .temporary_values = 0,
167 .temporary_types = 0,
168 .schedule = value.schedule.replay_limits,
169 .snapshot = .{
170 .axes = value.schedule.axes.len,
171 .steps = value.schedule.steps.len,
172 .name_bytes = names,
173 },
174 };
175 }
176
177 fn restoreIn(
178 allocator: std.mem.Allocator,
179 value: Record,
180 entry_name: []const u8,
181 storage: *model.core.builder.Storage,
182 capacity: model.program.Capacity,
183 comptime configuration: Configuration,
184 ) !model.Program {
185 const image = try choir.bytecode.image.Index.create(
186 allocator,
187 value.image,
188 configuration.image,
189 );
190 defer image.destroy();
191 if (value.function >= image.view().operations.len) return error.UnboundProductInput;
192 storage.context.arithmetic_policy = value.arithmetic;
193 try configuration.register(storage.context);
194 var decoded = try choir.bytecode.decodeModule(
195 choir.ir.context.operationAllocator(storage.context),
196 storage.context,
197 value.image,
198 );
199 errdefer decoded.deinit();
200 errdefer decoded.module.erase();
201 try choir.ir.verifyOperation(decoded.module, configuration.verify);
202 var references = try records.reference.Index.init(
203 allocator,
204 decoded.module,
205 configuration.image.entities,
206 );
207 defer references.deinit();
208 const function = try selectedFunction(&references, value.function);
209 const name = function.getName() orelse return error.InvalidStageRecord;
210 if (!std.mem.eql(u8, name, entry_name)) return error.InvalidStageRecord;
211 var replayed = try schedule.Schedule.init(allocator, value.schedule.replay_limits);
212 errdefer replayed.deinit(allocator);
213 try replayed.replay(value.schedule);
214 if (!std.meta.eql(try replayed.launch(), value.launch)) return error.InvalidStageRecord;
215 const kernel = try model.Kernel.fromDecoded(
216 allocator,
217 storage.*,
218 decoded,
219 function,
220 value.params,
221 );
222 return model.Program.init(kernel, replayed, capacity);
223 }
224
225 fn selectedFunction(
226 references: *const records.reference.Index,
227 ordinal: u32,
228 ) !choir.dialects.FuncDialect.FuncOp {
229 var selected: ?*choir.ir.Operation = null;
230 var entries = references.operations.iterator();
231 for (0..references.operations.count()) |_| {
232 const entry = entries.next().?;
233 if (entry.value_ptr.ordinal == ordinal) selected = entry.key_ptr.*;
234 }
235 const operation = selected orelse return error.UnboundProductInput;
236 const FuncOp = choir.dialects.FuncDialect.FuncOp;
237 if (!std.mem.eql(u8, operation.name.name, FuncOp.operation_name)) {
238 return error.InvalidStageRecord;
239 }
240 return .{ .op = operation };
241 }
242
243 fn coverage() void {
244 const require = choir.product.revision.record.requireFields;
245 require(model.Program, &.{ "capacity", "storage" });
246 require(@FieldType(model.Program, "storage"), &.{ "kernel", "schedule" });
247 require(model.Kernel, &.{ "allocator", "capacity", "storage", "state" });
248 const State = @typeInfo(@FieldType(model.Kernel, "state")).pointer.child;
249 require(State, &.{ "ctx", "module", "func", "params", "decoded" });
250 require(schedule.Snapshot, &.{
251 "phase",
252 "capacity",
253 "storage",
254 "version",
255 "axes_storage",
256 "axes_len",
257 "steps_storage",
258 "steps_len",
259 "names_storage",
260 "names_len",
261 "replay_limits",
262 "captured",
263 });
264 }