lib/accy/src/preparation/prepared.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3
4 const accy_choir = @import("../choir/root.zig");
5 const target_product = @import("../target/root.zig");
6 const execution_mod = @import("execution.zig");
7 const kernelization = @import("kernelization/root.zig");
8 const product_mod = @import("product.zig");
9 const run_mod = @import("run.zig");
10
11 const contract = accy_choir.contract;
12 const dispatch = accy_choir.dispatch;
13 const semantic = accy_choir.semantic;
14 const tensor = accy_choir.tensor;
15
16 const ir = choir.ir;
17 const passes = choir.passes;
18
19 const BackendPreparationProductStamps = product_mod.BackendPreparationProductStamps;
20 const BackendPreparationRun = run_mod.BackendPreparationRun;
21 const BackendPreparationRunOptions = run_mod.BackendPreparationRunOptions;
22 const BackendPreparationStats = run_mod.BackendPreparationStats;
23 const ContractPreparationResult = execution_mod.ContractPreparationResult;
24 const DispatchPreparationResult = execution_mod.DispatchPreparationResult;
25 const KernelPreparationResult = execution_mod.KernelPreparationResult;
26 const MemoryPreparationResult = execution_mod.MemoryPreparationResult;
27 const TargetPreparationResult = execution_mod.TargetPreparationResult;
28 const TensorPreparationResult = execution_mod.TensorPreparationResult;
29
30 const countOperationTree = execution_mod.countOperationTree;
31 const nsBetween = execution_mod.nsBetween;
32 const prepareContractJobFromSemanticModuleWithRun =
33 execution_mod.prepareContractJobFromSemanticModuleWithRun;
34 const prepareDispatchJobFromTensorJobWithRun =
35 execution_mod.prepareDispatchJobFromTensorJobWithRun;
36 const prepareKernelJobFromMemoryJobWithRun =
37 execution_mod.prepareKernelJobFromMemoryJobWithRun;
38 const prepareMemoryJobFromDispatchJobWithRun =
39 execution_mod.prepareMemoryJobFromDispatchJobWithRun;
40 const prepareTargetJobFromKernelJobWithRun =
41 execution_mod.prepareTargetJobFromKernelJobWithRun;
42 const prepareTensorJobFromContractJobWithRun =
43 execution_mod.prepareTensorJobFromContractJobWithRun;
44
45 pub const BackendPreparedJob = struct {
46 allocator: std.mem.Allocator,
47 ctx: *ir.Context,
48 choir_module: *ir.Operation,
49 target_module: ?*target_product.TargetJob,
50 run: BackendPreparationRun,
51
52 pub fn deinit(self: *BackendPreparedJob) void {
53 if (self.target_module) |module| module.deinit();
54 self.* = undefined;
55 }
56
57 pub fn passContext(self: *BackendPreparedJob) passes.PassContext {
58 return self.target_module.?.passContext();
59 }
60
61 pub fn targetModule(self: *const BackendPreparedJob) !*target_product.TargetJob {
62 return self.target_module orelse error.MissingTargetProduct;
63 }
64
65 pub fn kernelizationProduct(self: *const BackendPreparedJob) !*const kernelization.KernelizationAnalysis {
66 return (try self.targetModule()).kernelizationProduct();
67 }
68
69 pub fn generatedKernelCount(self: *const BackendPreparedJob) !usize {
70 const product = try self.kernelizationProduct();
71 return product.kernelCount();
72 }
73
74 pub fn productStamps(self: *const BackendPreparedJob) !BackendPreparationProductStamps {
75 return product_mod.stampsFromRun(self.run);
76 }
77
78 pub fn generatedKernelSummary(
79 self: *const BackendPreparedJob,
80 kernel_index: usize,
81 ) !kernelization.GeneratedKernelSummary {
82 const product = try self.kernelizationProduct();
83 return product.kernelSummary(kernel_index);
84 }
85
86 pub fn generatedKernelProgram(
87 self: *const BackendPreparedJob,
88 kernel_index: usize,
89 ) !*const kernelization.GeneratedKernelProgram {
90 const product = try self.kernelizationProduct();
91 return product.kernelProgram(kernel_index);
92 }
93
94 pub fn generatedKernelSummaryForWork(
95 self: *const BackendPreparedJob,
96 work_item_id: usize,
97 ) !kernelization.GeneratedKernelSummary {
98 const product = try self.kernelizationProduct();
99 return product.kernelSummaryForWork(work_item_id);
100 }
101
102 pub fn generatedKernelProgramForWork(
103 self: *const BackendPreparedJob,
104 work_item_id: usize,
105 ) !*const kernelization.GeneratedKernelProgram {
106 const product = try self.kernelizationProduct();
107 return product.kernelProgramForWork(work_item_id);
108 }
109
110 pub fn copyGeneratedKernelSummaries(
111 self: *const BackendPreparedJob,
112 result_allocator: std.mem.Allocator,
113 ) !kernelization.GeneratedKernelSummaries {
114 const product = try self.kernelizationProduct();
115 return product.copyKernelSummaries(result_allocator);
116 }
117 };
118
119 pub fn runBackendPreparationPipelineFromSemanticModule(
120 allocator: std.mem.Allocator,
121 module: *semantic.SemanticModule,
122 options: BackendPreparationRunOptions,
123 ) !BackendPreparationRun {
124 var prepared = try prepareBackendJobFromSemanticModule(allocator, module, options);
125 defer prepared.deinit();
126 return prepared.run;
127 }
128
129 pub fn prepareBackendJobFromSemanticModule(
130 allocator: std.mem.Allocator,
131 module: *semantic.SemanticModule,
132 options: BackendPreparationRunOptions,
133 ) !BackendPreparedJob {
134 return try prepareOwnedSemanticModuleForBackend(
135 allocator,
136 module,
137 options,
138 options.now(),
139 );
140 }
141
142 pub fn prepareBackendJobFromContractJob(
143 allocator: std.mem.Allocator,
144 module: *contract.ContractJob,
145 options: BackendPreparationRunOptions,
146 ) !BackendPreparedJob {
147 const initial_choir_ops = countOperationTree(module.choir_module);
148 const total_start = options.now();
149 const tensor_result = try prepareTensorJobFromContractJobWithRun(
150 allocator,
151 module,
152 options,
153 );
154 return try prepareBackendJobFromTensorPreparationResult(
155 allocator,
156 tensor_result,
157 options,
158 0,
159 .{},
160 initial_choir_ops,
161 null,
162 total_start,
163 );
164 }
165
166 fn prepareOwnedSemanticModuleForBackend(
167 allocator: std.mem.Allocator,
168 module: *semantic.SemanticModule,
169 options: BackendPreparationRunOptions,
170 total_start: i128,
171 ) !BackendPreparedJob {
172 const contract_result = try prepareContractJobFromSemanticModuleWithRun(allocator, module, options);
173 return try prepareBackendJobFromContractPreparationResult(allocator, contract_result, options, total_start);
174 }
175
176 pub fn prepareBackendJobFromContractPreparationResult(
177 allocator: std.mem.Allocator,
178 result: ContractPreparationResult,
179 options: BackendPreparationRunOptions,
180 total_start: i128,
181 ) !BackendPreparedJob {
182 var contract_result = result;
183 var contract_owned = true;
184 errdefer if (contract_owned) contract_result.deinit();
185
186 const contract_ns = result.elapsed_ns;
187 const contract_stats = result.stats;
188 const initial_choir_ops = result.initial_choir_ops;
189 const semantic_fingerprint = result.semantic_fingerprint;
190 contract_owned = false;
191 const tensor_result = try prepareTensorJobFromContractJobWithRun(
192 allocator,
193 contract_result.module,
194 options,
195 );
196 return try prepareBackendJobFromTensorPreparationResult(
197 allocator,
198 tensor_result,
199 options,
200 contract_ns,
201 contract_stats,
202 initial_choir_ops,
203 semantic_fingerprint,
204 total_start,
205 );
206 }
207
208 pub fn prepareBackendJobFromTensorPreparationResult(
209 allocator: std.mem.Allocator,
210 result: TensorPreparationResult,
211 options: BackendPreparationRunOptions,
212 contract_ns: u64,
213 contract_stats: BackendPreparationStats,
214 initial_choir_ops: u64,
215 semantic_fingerprint: ?u64,
216 total_start: i128,
217 ) !BackendPreparedJob {
218 var tensor_result = result;
219 var tensor_owned = true;
220 errdefer if (tensor_owned) tensor_result.deinit();
221
222 const tensor_ns = result.elapsed_ns;
223 const tensor_stats = result.stats;
224 tensor_owned = false;
225 const dispatch_result = try prepareDispatchJobFromTensorJobWithRun(
226 allocator,
227 tensor_result.module,
228 options,
229 );
230 return try prepareBackendJobFromDispatchPreparationResult(
231 allocator,
232 dispatch_result,
233 options,
234 contract_ns,
235 contract_stats,
236 tensor_ns,
237 tensor_stats,
238 initial_choir_ops,
239 semantic_fingerprint,
240 total_start,
241 );
242 }
243
244 pub fn prepareBackendJobFromDispatchPreparationResult(
245 allocator: std.mem.Allocator,
246 result: DispatchPreparationResult,
247 options: BackendPreparationRunOptions,
248 contract_ns: u64,
249 contract_stats: BackendPreparationStats,
250 tensor_ns: u64,
251 tensor_stats: BackendPreparationStats,
252 initial_choir_ops: u64,
253 semantic_fingerprint: ?u64,
254 total_start: i128,
255 ) !BackendPreparedJob {
256 var dispatch_result = result;
257 var dispatch_owned = true;
258 errdefer if (dispatch_owned) dispatch_result.deinit();
259
260 const dispatch_ns = result.elapsed_ns;
261 const dispatch_stats = result.stats;
262 const dispatch_fingerprint = result.plan_fingerprint;
263 dispatch_owned = false;
264 const memory_result = try prepareMemoryJobFromDispatchJobWithRun(
265 allocator,
266 dispatch_result.module,
267 options,
268 );
269 return try prepareBackendJobFromMemoryPreparationResult(
270 allocator,
271 memory_result,
272 options,
273 contract_ns,
274 contract_stats,
275 tensor_ns,
276 tensor_stats,
277 dispatch_ns,
278 dispatch_stats,
279 dispatch_fingerprint,
280 initial_choir_ops,
281 semantic_fingerprint,
282 total_start,
283 );
284 }
285
286 pub fn prepareBackendJobFromMemoryPreparationResult(
287 allocator: std.mem.Allocator,
288 result: MemoryPreparationResult,
289 options: BackendPreparationRunOptions,
290 contract_ns: u64,
291 contract_stats: BackendPreparationStats,
292 tensor_ns: u64,
293 tensor_stats: BackendPreparationStats,
294 dispatch_ns: u64,
295 dispatch_stats: BackendPreparationStats,
296 dispatch_fingerprint: u64,
297 initial_choir_ops: u64,
298 semantic_fingerprint: ?u64,
299 total_start: i128,
300 ) !BackendPreparedJob {
301 var memory_result = result;
302 var memory_owned = true;
303 errdefer if (memory_owned) memory_result.deinit();
304
305 const memory_ns = result.elapsed_ns;
306 const memory_stats = result.stats;
307 const memory_fingerprint = result.plan_fingerprint;
308 memory_owned = false;
309 const kernel_result = try prepareKernelJobFromMemoryJobWithRun(
310 allocator,
311 memory_result.module,
312 options,
313 );
314 return try prepareBackendJobFromKernelPreparationResult(
315 allocator,
316 kernel_result,
317 options,
318 contract_ns,
319 contract_stats,
320 tensor_ns,
321 tensor_stats,
322 dispatch_ns,
323 dispatch_stats,
324 dispatch_fingerprint,
325 memory_ns,
326 memory_stats,
327 memory_fingerprint,
328 initial_choir_ops,
329 semantic_fingerprint,
330 total_start,
331 );
332 }
333
334 pub fn prepareBackendJobFromKernelPreparationResult(
335 allocator: std.mem.Allocator,
336 result: KernelPreparationResult,
337 options: BackendPreparationRunOptions,
338 contract_ns: u64,
339 contract_stats: BackendPreparationStats,
340 tensor_ns: u64,
341 tensor_stats: BackendPreparationStats,
342 dispatch_ns: u64,
343 dispatch_stats: BackendPreparationStats,
344 dispatch_fingerprint: u64,
345 memory_ns: u64,
346 memory_stats: BackendPreparationStats,
347 memory_fingerprint: u64,
348 initial_choir_ops: u64,
349 semantic_fingerprint: ?u64,
350 total_start: i128,
351 ) !BackendPreparedJob {
352 var kernel_result = result;
353 var kernel_owned = true;
354 errdefer if (kernel_owned) kernel_result.deinit();
355
356 const kernel_ns = result.elapsed_ns;
357 const kernel_stats = result.stats;
358 const kernel_fingerprint = result.plan_fingerprint;
359 kernel_owned = false;
360 const target_result = try prepareTargetJobFromKernelJobWithRun(
361 allocator,
362 kernel_result.module,
363 options,
364 );
365 return try prepareBackendJobFromTargetPreparationResult(
366 allocator,
367 target_result,
368 options,
369 contract_ns,
370 contract_stats,
371 tensor_ns,
372 tensor_stats,
373 dispatch_ns,
374 dispatch_stats,
375 dispatch_fingerprint,
376 memory_ns,
377 memory_stats,
378 memory_fingerprint,
379 kernel_ns,
380 kernel_stats,
381 kernel_fingerprint,
382 initial_choir_ops,
383 semantic_fingerprint,
384 total_start,
385 );
386 }
387
388 pub fn prepareBackendJobFromTargetPreparationResult(
389 allocator: std.mem.Allocator,
390 result: TargetPreparationResult,
391 options: BackendPreparationRunOptions,
392 contract_ns: u64,
393 contract_stats: BackendPreparationStats,
394 tensor_ns: u64,
395 tensor_stats: BackendPreparationStats,
396 dispatch_ns: u64,
397 dispatch_stats: BackendPreparationStats,
398 dispatch_fingerprint: u64,
399 memory_ns: u64,
400 memory_stats: BackendPreparationStats,
401 memory_fingerprint: u64,
402 kernel_ns: u64,
403 kernel_stats: BackendPreparationStats,
404 kernel_fingerprint: u64,
405 initial_choir_ops: u64,
406 semantic_fingerprint: ?u64,
407 total_start: i128,
408 ) !BackendPreparedJob {
409 var target_result = result;
410 var target_owned = true;
411 errdefer if (target_owned) target_result.deinit();
412
413 target_owned = false;
414 return try prepareOwnedTargetJob(
415 allocator,
416 target_result.module,
417 options,
418 contract_ns,
419 contract_stats,
420 tensor_ns,
421 tensor_stats,
422 dispatch_ns,
423 dispatch_stats,
424 dispatch_fingerprint,
425 memory_ns,
426 memory_stats,
427 memory_fingerprint,
428 kernel_ns,
429 kernel_stats,
430 kernel_fingerprint,
431 result.elapsed_ns,
432 result.stats,
433 initial_choir_ops,
434 semantic_fingerprint,
435 total_start,
436 result.finished_at,
437 );
438 }
439
440 fn prepareOwnedTargetJob(
441 allocator: std.mem.Allocator,
442 module: *target_product.TargetJob,
443 options: BackendPreparationRunOptions,
444 contract_ns: u64,
445 contract_stats: BackendPreparationStats,
446 tensor_ns: u64,
447 tensor_stats: BackendPreparationStats,
448 dispatch_ns: u64,
449 dispatch_stats: BackendPreparationStats,
450 dispatch_fingerprint: u64,
451 memory_ns: u64,
452 memory_stats: BackendPreparationStats,
453 memory_fingerprint: u64,
454 kernel_ns: u64,
455 kernel_stats: BackendPreparationStats,
456 kernel_fingerprint: u64,
457 target_ns: u64,
458 target_stats: BackendPreparationStats,
459 initial_choir_ops: u64,
460 semantic_fingerprint: ?u64,
461 total_start: i128,
462 total_end: i128,
463 ) !BackendPreparedJob {
464 var module_owned = true;
465 errdefer if (module_owned) module.deinit();
466
467 try module.verify();
468 const final_choir_ops = countOperationTree(module.choir_module);
469
470 var prepared = BackendPreparedJob{
471 .allocator = allocator,
472 .ctx = module.context(),
473 .choir_module = module.choir_module,
474 .target_module = module,
475 .run = undefined,
476 };
477
478 prepared.run = .{
479 .total_ns = nsBetween(total_start, total_end),
480 .contract_ns = contract_ns,
481 .tensor_ns = tensor_ns,
482 .dispatch_ns = dispatch_ns,
483 .memory_ns = memory_ns,
484 .kernel_ns = kernel_ns,
485 .target_ns = target_ns,
486 .initial_choir_ops = initial_choir_ops,
487 .final_choir_ops = final_choir_ops,
488 .semantic_fingerprint = semantic_fingerprint,
489 .contract_fingerprint = module.kernel_module.memory_module.dispatch_module.tensor_module.contract_module.fingerprint(),
490 .tensor_fingerprint = module.kernel_module.memory_module.dispatch_module.tensor_module.fingerprint(),
491 .dispatch_fingerprint = dispatch_fingerprint,
492 .memory_fingerprint = memory_fingerprint,
493 .kernel_fingerprint = kernel_fingerprint,
494 .target_fingerprint = module.fingerprint(),
495 .contract_stats = contract_stats,
496 .tensor_stats = tensor_stats,
497 .dispatch_stats = dispatch_stats,
498 .memory_stats = memory_stats,
499 .kernel_stats = kernel_stats,
500 .target_stats = target_stats,
501 .target_profile = options.target_profile,
502 };
503 module_owned = false;
504 return prepared;
505 }