lib/accy/src/kernel/model/program/graph/model.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3
4 const core = @import("../../core/root.zig");
5 const plan = @import("../../plan/root.zig");
6
7 const builder = core.builder;
8 const schedule_mod = core.schedule;
9
10 const ir = choir.ir;
11
12 pub const product_name = "accy.kernel_program";
13
14 pub const Limits = core.Limits;
15
16 pub const Capacity = struct {
17 context_ownership: ContextOwnership,
18 body: builder.Builder.Capacity,
19 schedule: schedule_mod.Schedule.Capacity,
20 snapshot: schedule_mod.Snapshot.Capacity,
21 construction_bytes: usize,
22 maximum_overlap_bytes: usize,
23
24 pub const ContextOwnership = enum { owned, borrowed };
25
26 pub fn derive(limits: Limits) error{CapacityOverflow}!Capacity {
27 return deriveWithContextOwnership(limits, .owned);
28 }
29
30 pub fn deriveBorrowing(limits: Limits) error{CapacityOverflow}!Capacity {
31 return deriveWithContextOwnership(limits, .borrowed);
32 }
33
34 fn deriveWithContextOwnership(
35 limits: Limits,
36 context_ownership: ContextOwnership,
37 ) error{CapacityOverflow}!Capacity {
38 const body = try builder.Builder.Capacity.derive(limits.raw());
39 const schedule = try schedule_mod.Schedule.Capacity.derive(limits.schedule);
40 const snapshot = try schedule_mod.Snapshot.Capacity.derive(limits.snapshot);
41 const body_acquired_bytes = switch (context_ownership) {
42 .owned => body.owned_acquired_bytes,
43 .borrowed => body.borrowed_acquired_bytes,
44 };
45 const construction_bytes = std.math.add(
46 usize,
47 body_acquired_bytes,
48 schedule.storage_bytes,
49 ) catch return error.CapacityOverflow;
50 var maximum_overlap_bytes = std.math.add(
51 usize,
52 construction_bytes,
53 snapshot.storage_bytes,
54 ) catch return error.CapacityOverflow;
55 maximum_overlap_bytes = std.math.add(
56 usize,
57 maximum_overlap_bytes,
58 schedule.storage_bytes,
59 ) catch return error.CapacityOverflow;
60 return .{
61 .context_ownership = context_ownership,
62 .body = body,
63 .schedule = schedule,
64 .snapshot = snapshot,
65 .construction_bytes = construction_bytes,
66 .maximum_overlap_bytes = maximum_overlap_bytes,
67 };
68 }
69 };
70
71 const Storage = struct {
72 kernel: builder.Kernel,
73 schedule: schedule_mod.Schedule,
74 };
75
76 pub const Program = struct {
77 capacity: Capacity,
78 storage: Storage,
79
80 const Self = @This();
81
82 pub fn deinit(self: *Self) void {
83 self.storage.schedule.deinit(self.storage.kernel.allocator);
84 self.storage.kernel.deinit();
85 self.* = undefined;
86 }
87
88 pub fn verify(self: *Self) !void {
89 try self.storage.kernel.verify();
90 }
91
92 pub fn verifyWithDiagnostic(self: *Self, diagnostic: *builder.VerificationDiagnostic) !void {
93 try self.storage.kernel.verifyWithDiagnostic(diagnostic);
94 }
95
96 pub fn launch(self: *const Self) schedule_mod.ScheduleError!schedule_mod.Launch {
97 return self.storage.schedule.launch();
98 }
99
100 pub fn scheduleSnapshot(self: *const Self, allocator: std.mem.Allocator) schedule_mod.ScheduleError!schedule_mod.Snapshot {
101 return schedule_mod.createSnapshot(
102 allocator,
103 self.capacity.snapshot.asLimits(),
104 &self.storage.schedule,
105 );
106 }
107
108 pub fn createPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) plan.Error!plan.Plan {
109 return plan.create(
110 allocator,
111 &self.storage.kernel,
112 &self.storage.schedule,
113 self.capacity.snapshot.asLimits(),
114 options,
115 );
116 }
117
118 pub fn kernelModule(self: *const Self) *ir.Operation {
119 return self.storage.kernel.module();
120 }
121
122 pub fn params(self: *const Self) []const builder.Param {
123 return self.storage.kernel.params();
124 }
125
126 pub fn createCheckedPlan(self: *Self, allocator: std.mem.Allocator, options: plan.Options) !plan.Plan {
127 var out = try self.createPlan(allocator, options);
128 errdefer out.deinit();
129 try self.checkPlan(allocator, &out);
130 return out;
131 }
132
133 pub fn checkPlan(self: *Self, allocator: std.mem.Allocator, out: *const plan.Plan) !void {
134 if (out.version != plan.plan_version) return error.PlanVersionMismatch;
135 if (out.schedule_version != schedule_mod.snapshot_version) return error.ScheduleVersionMismatch;
136 if (out.body_fingerprint != try self.bodyFingerprint(allocator)) return error.BodyFingerprintMismatch;
137 if (out.schedule_fingerprint != self.scheduleFingerprint()) return error.ScheduleFingerprintMismatch;
138 if (out.schedule.fingerprint() != out.schedule_fingerprint) return error.ScheduleFingerprintMismatch;
139 if (!builder.paramsEql(out.params, self.storage.kernel.params())) return error.ParameterSchemaMismatch;
140 const argument_count = std.math.cast(u32, self.storage.kernel.params().len) orelse return error.ArgumentCountOverflow;
141 if (out.argument_count != argument_count) return error.ParameterSchemaMismatch;
142 const entry_name = self.storage.kernel.func().getName() orelse return error.MissingKernelName;
143 if (!std.mem.eql(u8, out.entry_name, entry_name)) return error.EntryNameMismatch;
144
145 var replayed = try schedule_mod.replaySnapshot(
146 allocator,
147 self.capacity.schedule.asLimits(),
148 &out.schedule,
149 );
150 defer replayed.deinit(allocator);
151 if (replayed.fingerprint() != out.schedule_fingerprint) return error.ScheduleReplayMismatch;
152
153 const launch_value = try self.launch();
154 if (!launchEqual(out.launch, launch_value)) return error.LaunchMismatch;
155 if (!launchEqual(try out.schedule.launch(), launch_value)) return error.LaunchMismatch;
156 }
157
158 pub fn bodyFingerprint(self: *const Self, allocator: std.mem.Allocator) !u64 {
159 return self.storage.kernel.fingerprint(allocator);
160 }
161
162 pub fn scheduleFingerprint(self: *const Self) u64 {
163 return self.storage.schedule.fingerprint();
164 }
165
166 pub fn fingerprint(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.Fingerprint {
167 const entry_name = self.storage.kernel.func().getName() orelse return error.MissingKernelName;
168 var fingerprint_builder = choir.product.incremental.FingerprintBuilder{};
169 fingerprint_builder.updateBytes(product_name);
170 fingerprint_builder.updateU32(plan.plan_version);
171 fingerprint_builder.updateU32(schedule_mod.snapshot_version);
172 fingerprint_builder.updateU64(try self.bodyFingerprint(allocator));
173 fingerprint_builder.updateU64(self.scheduleFingerprint());
174 fingerprint_builder.updateBytes(entry_name);
175 updateParamsFingerprint(&fingerprint_builder, self.storage.kernel.params());
176 return fingerprint_builder.finish();
177 }
178
179 pub fn productStamp(self: *const Self, allocator: std.mem.Allocator) !choir.product.incremental.ProductStamp {
180 return choir.product.incremental.productStamp(product_name, try self.fingerprint(allocator));
181 }
182
183 pub fn init(
184 kernel: builder.Kernel,
185 schedule: schedule_mod.Schedule,
186 capacity: Capacity,
187 ) Self {
188 return .{
189 .capacity = capacity,
190 .storage = .{
191 .kernel = kernel,
192 .schedule = schedule,
193 },
194 };
195 }
196 };
197
198 fn launchEqual(lhs: schedule_mod.Launch, rhs: schedule_mod.Launch) bool {
199 return std.mem.eql(u32, &lhs.grid, &rhs.grid) and
200 std.mem.eql(u32, &lhs.block, &rhs.block);
201 }
202
203 fn updateParamsFingerprint(fingerprint_builder: *choir.product.incremental.FingerprintBuilder, params: []const builder.Param) void {
204 fingerprint_builder.updateUsize(params.len);
205 for (params) |param| updateParamFingerprint(fingerprint_builder, param);
206 }
207
208 fn updateParamFingerprint(fingerprint_builder: *choir.product.incremental.FingerprintBuilder, param: builder.Param) void {
209 fingerprint_builder.updateEnumTag(std.meta.activeTag(param));
210 switch (param) {
211 .scalar => |dtype| fingerprint_builder.updateEnumTag(dtype),
212 .buffer => |buffer| {
213 fingerprint_builder.updateEnumTag(buffer.dtype);
214 fingerprint_builder.updateOptionalU64(buffer.size);
215 fingerprint_builder.updateEnumTag(buffer.space);
216 },
217 }
218 }
219
220 fn buildCopyProgram(
221 allocator: std.mem.Allocator,
222 name: []const u8,
223 extent: u64,
224 ) !Program {
225 const limits = Limits.testing;
226 const capacity = try Capacity.derive(limits);
227 var schedule = try schedule_mod.Schedule.init(allocator, limits.schedule);
228 errdefer schedule.deinit(allocator);
229
230 const i = try schedule.addAxis("i", extent);
231 try schedule.bind(i, .thread_x);
232
233 var b = try builder.Builder.init(allocator, limits.raw(), name, &.{
234 builder.dynamicBuffer(.i32),
235 builder.dynamicBuffer(.i32),
236 }, .{});
237 errdefer b.deinit();
238
239 const src = b.argument(0);
240 const dst = b.argument(1);
241 const index = try b.globalId(.x);
242 const value = try b.load(src, index);
243 try b.store(value, dst, index);
244 try b.return_();
245
246 const kernel = try b.finish();
247 schedule.activate();
248 return Program.init(kernel, schedule, capacity);
249 }
250
251 test "Program capacity distinguishes owned and borrowed Context acquisition" {
252 var limits = Limits.standard;
253 limits.parameters = 3;
254 limits.temporary_values = 5;
255 limits.temporary_types = 7;
256 limits.schedule = .{ .axes = 11, .steps = 13, .name_bytes = 17, .axis_ids = 19 };
257 limits.snapshot = .{ .axes = 23, .steps = 29, .name_bytes = 31 };
258
259 const owned = try Capacity.derive(limits);
260 const borrowed = try Capacity.deriveBorrowing(limits);
261 try std.testing.expectEqual(Capacity.ContextOwnership.owned, owned.context_ownership);
262 try std.testing.expectEqual(Capacity.ContextOwnership.borrowed, borrowed.context_ownership);
263 try std.testing.expectEqual(
264 owned.body.owned_acquired_bytes + owned.schedule.storage_bytes,
265 owned.construction_bytes,
266 );
267 try std.testing.expectEqual(
268 borrowed.body.borrowed_acquired_bytes + borrowed.schedule.storage_bytes,
269 borrowed.construction_bytes,
270 );
271 try std.testing.expectEqual(
272 owned.body.context.storage_bytes,
273 owned.construction_bytes - borrowed.construction_bytes,
274 );
275 try std.testing.expectEqual(
276 owned.construction_bytes + owned.snapshot.storage_bytes + owned.schedule.storage_bytes,
277 owned.maximum_overlap_bytes,
278 );
279 }
280
281 test "kernel Program exposes stable product identity" {
282 const allocator = std.testing.allocator;
283
284 var first = try buildCopyProgram(allocator, "kernel_program_product_identity", 4);
285 defer first.deinit();
286 var same = try buildCopyProgram(allocator, "kernel_program_product_identity", 4);
287 defer same.deinit();
288 var changed = try buildCopyProgram(allocator, "kernel_program_product_identity_changed", 8);
289 defer changed.deinit();
290
291 const first_stamp = try first.productStamp(allocator);
292 try std.testing.expectEqualStrings(product_name, first_stamp.name);
293 try std.testing.expectEqual(try first.fingerprint(allocator), first_stamp.fingerprint);
294 try std.testing.expectEqual(try first.fingerprint(allocator), try same.fingerprint(allocator));
295 try std.testing.expect(try first.fingerprint(allocator) != try changed.fingerprint(allocator));
296 }