lib/accy/src/preparation/execution.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 kernelization = @import("kernelization/root.zig");
7 const run_mod = @import("run.zig");
8 const stage_mod = @import("stage.zig");
9 const recipe = @import("recipe.zig");
10 const stage_fingerprint = @import("fingerprint.zig");
11
12 const contract = accy_choir.contract;
13 const dispatch = accy_choir.dispatch;
14 const kernel_product = accy_choir.gpu;
15 const memory_product = accy_choir.memory;
16 const semantic = accy_choir.semantic;
17 const tensor = accy_choir.tensor;
18
19 const ir = choir.ir;
20 const passes = choir.passes;
21
22 pub const PipelineError = error{
23 PassFailed,
24 };
25
26 const BackendPreparationRunOptions = run_mod.BackendPreparationRunOptions;
27 const BackendPreparationStats = run_mod.BackendPreparationStats;
28
29 pub fn capturePipelineFailure(
30 allocator: std.mem.Allocator,
31 failure: ?*run_mod.BackendPreparationFailure,
32 pipeline_name: []const u8,
33 pm: *const passes.PassManager,
34 ) void {
35 const out = failure orelse return;
36 if (pm.getLastFailureReproducer()) |reproducer| {
37 out.capture(allocator, pipeline_name, reproducer) catch {};
38 } else {
39 out.deinit(allocator);
40 out.pipeline_name = pipeline_name;
41 }
42 }
43
44 fn passFailed(
45 allocator: std.mem.Allocator,
46 failure: ?*run_mod.BackendPreparationFailure,
47 pipeline_name: []const u8,
48 pm: *const passes.PassManager,
49 ) PipelineError {
50 capturePipelineFailure(allocator, failure, pipeline_name, pm);
51 return PipelineError.PassFailed;
52 }
53
54 fn preparationPassFailed(
55 allocator: std.mem.Allocator,
56 options: BackendPreparationRunOptions,
57 pipeline_name: []const u8,
58 pm: *const passes.PassManager,
59 ) PipelineError {
60 return passFailed(allocator, options.failure, pipeline_name, pm);
61 }
62
63 pub const ContractPreparationResult = struct {
64 module: *contract.ContractJob,
65 elapsed_ns: u64,
66 stats: BackendPreparationStats,
67 initial_choir_ops: u64,
68 semantic_fingerprint: ?u64,
69 started_at: i128,
70 finished_at: i128,
71
72 pub fn deinit(self: *ContractPreparationResult) void {
73 self.module.deinit();
74 self.* = undefined;
75 }
76 };
77
78 pub const TensorPreparationResult = struct {
79 module: *tensor.TensorJob,
80 elapsed_ns: u64,
81 stats: BackendPreparationStats,
82
83 pub fn deinit(self: *TensorPreparationResult) void {
84 self.module.deinit();
85 self.* = undefined;
86 }
87 };
88
89 pub const DispatchPreparationResult = struct {
90 module: *dispatch.DispatchJob,
91 elapsed_ns: u64,
92 stats: BackendPreparationStats,
93 /// A caller reads this field to show whether the dispatch plans of two runs look alike. The
94 /// field holds a 64-bit summary of the plans made by the dispatch stage, for display in a run
95 /// stamp, the per-stage display record built for people reading a run. No product key or reuse
96 /// decision reads it.
97 plan_fingerprint: u64,
98
99 pub fn deinit(self: *DispatchPreparationResult) void {
100 self.module.deinit();
101 self.* = undefined;
102 }
103 };
104
105 pub const MemoryPreparationResult = struct {
106 module: *memory_product.MemoryJob,
107 elapsed_ns: u64,
108 stats: BackendPreparationStats,
109 /// A caller reads this field to show whether the memory plans of two runs look alike. The field
110 /// holds a 64-bit summary of the buffer, memory-space and layout plans made by the memory
111 /// stage, for display in run stamps. No product key or reuse decision reads it.
112 plan_fingerprint: u64,
113
114 pub fn deinit(self: *MemoryPreparationResult) void {
115 self.module.deinit();
116 self.* = undefined;
117 }
118 };
119
120 pub const KernelPreparationResult = struct {
121 module: *kernel_product.KernelJob,
122 elapsed_ns: u64,
123 stats: BackendPreparationStats,
124 /// A caller reads this field to show whether the kernel plans of two runs look alike. The field
125 /// holds a 64-bit summary of the kernel outline and generated kernels made by the kernel stage,
126 /// for display in run stamps. No product key or reuse decision reads it.
127 plan_fingerprint: u64,
128
129 pub fn deinit(self: *KernelPreparationResult) void {
130 self.module.deinit();
131 self.* = undefined;
132 }
133 };
134
135 pub const TargetPreparationResult = struct {
136 module: *target_product.TargetJob,
137 elapsed_ns: u64,
138 stats: BackendPreparationStats,
139 finished_at: i128,
140
141 pub fn deinit(self: *TargetPreparationResult) void {
142 self.module.deinit();
143 self.* = undefined;
144 }
145 };
146
147 pub const TargetProductPreparationResult = struct {
148 choir_module: *ir.Operation,
149 analysis_cache: passes.AnalysisCache,
150 kernelization_product: *const kernelization.KernelizationAnalysis,
151 elapsed_ns: u64,
152 stats: BackendPreparationStats,
153 finished_at: i128,
154
155 pub fn deinit(self: *TargetProductPreparationResult) void {
156 self.analysis_cache.deinit();
157 self.choir_module.erase();
158 self.* = undefined;
159 }
160 };
161
162 pub fn runTargetPipeline(
163 allocator: std.mem.Allocator,
164 choir_module: *ir.Operation,
165 ctx: *ir.Context,
166 ) !void {
167 return try runTargetPipelineWithOptions(allocator, choir_module, ctx, .{});
168 }
169
170 pub fn runTargetPipelineWithOptions(
171 allocator: std.mem.Allocator,
172 choir_module: *ir.Operation,
173 ctx: *ir.Context,
174 options: passes.PassManagerRunOptions,
175 ) !void {
176 return try runTargetPipelineWithDiagnostics(allocator, choir_module, ctx, options, null);
177 }
178
179 pub fn runTargetPipelineWithDiagnostics(
180 allocator: std.mem.Allocator,
181 choir_module: *ir.Operation,
182 ctx: *ir.Context,
183 options: passes.PassManagerRunOptions,
184 failure: ?*run_mod.BackendPreparationFailure,
185 ) !void {
186 var pm = passes.PassManager.init(allocator);
187 defer pm.deinit();
188 try recipe.configure(&pm, .target, &.{});
189 if (pm.runWithOptions(choir_module, ctx, options) == .failure) return passFailed(allocator, failure, stage_mod.target_pipeline_name, &pm);
190 }
191
192 pub fn prepareContractJobFromSemanticModule(
193 allocator: std.mem.Allocator,
194 module: *semantic.SemanticModule,
195 options: BackendPreparationRunOptions,
196 ) !*contract.ContractJob {
197 const result = try prepareContractJobFromSemanticModuleWithRun(
198 allocator,
199 module,
200 options,
201 );
202 return result.module;
203 }
204
205 pub fn prepareTensorJobFromContractJob(
206 allocator: std.mem.Allocator,
207 module: *contract.ContractJob,
208 options: BackendPreparationRunOptions,
209 ) !*tensor.TensorJob {
210 const result = try prepareTensorJobFromContractJobWithRun(
211 allocator,
212 module,
213 options,
214 );
215 return result.module;
216 }
217
218 pub fn prepareDispatchJobFromTensorJob(
219 allocator: std.mem.Allocator,
220 module: *tensor.TensorJob,
221 options: BackendPreparationRunOptions,
222 ) !*dispatch.DispatchJob {
223 const result = try prepareDispatchJobFromTensorJobWithRun(
224 allocator,
225 module,
226 options,
227 );
228 return result.module;
229 }
230
231 pub fn prepareMemoryJobFromDispatchJob(
232 allocator: std.mem.Allocator,
233 module: *dispatch.DispatchJob,
234 options: BackendPreparationRunOptions,
235 ) !*memory_product.MemoryJob {
236 const result = try prepareMemoryJobFromDispatchJobWithRun(
237 allocator,
238 module,
239 options,
240 );
241 return result.module;
242 }
243
244 pub fn prepareKernelJobFromMemoryJob(
245 allocator: std.mem.Allocator,
246 module: *memory_product.MemoryJob,
247 options: BackendPreparationRunOptions,
248 ) !*kernel_product.KernelJob {
249 const result = try prepareKernelJobFromMemoryJobWithRun(
250 allocator,
251 module,
252 options,
253 );
254 return result.module;
255 }
256
257 pub fn prepareTargetJobFromKernelJob(
258 allocator: std.mem.Allocator,
259 module: *kernel_product.KernelJob,
260 options: BackendPreparationRunOptions,
261 ) !*target_product.TargetJob {
262 const result = try prepareTargetJobFromKernelJobWithRun(
263 allocator,
264 module,
265 options,
266 );
267 return result.module;
268 }
269
270 pub fn prepareContractJobFromSemanticModuleWithRun(
271 allocator: std.mem.Allocator,
272 module: *semantic.SemanticModule,
273 options: BackendPreparationRunOptions,
274 ) !ContractPreparationResult {
275 var module_owned = true;
276 errdefer if (module_owned) module.deinit();
277
278 try module.verify();
279
280 const initial_choir_ops = countOperationTree(module.choir_module);
281 const semantic_fingerprint = try module.fingerprint(allocator);
282
283 var analysis_cache = passes.AnalysisCache.init(allocator, null);
284 defer analysis_cache.deinit();
285
286 var pm = passes.PassManager.init(allocator);
287 defer pm.deinit();
288 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
289 try recipe.configure(&pm, .contract, &options);
290
291 const contract_start = options.now();
292 if (pm.runWithAnalysisCache(
293 module.choir_module,
294 module.context(),
295 &analysis_cache,
296 recipe.runOptions(),
297 ) == .failure) {
298 return preparationPassFailed(allocator, options, stage_mod.contract_pipeline_name, &pm);
299 }
300 const contract_end = options.now();
301
302 const contract_module = try contract.ContractJob.init(allocator, module);
303 module_owned = false;
304 return .{
305 .module = contract_module,
306 .elapsed_ns = nsBetween(contract_start, contract_end),
307 .stats = run_mod.backendPreparationStats(pm.stats),
308 .initial_choir_ops = initial_choir_ops,
309 .semantic_fingerprint = semantic_fingerprint,
310 .started_at = contract_start,
311 .finished_at = contract_end,
312 };
313 }
314
315 pub fn prepareTensorJobFromContractJobWithRun(
316 allocator: std.mem.Allocator,
317 module: *contract.ContractJob,
318 options: BackendPreparationRunOptions,
319 ) !TensorPreparationResult {
320 var module_owned = true;
321 errdefer if (module_owned) module.deinit();
322
323 try module.verify();
324
325 var analysis_cache = passes.AnalysisCache.init(allocator, null);
326 var cache_owned = true;
327 errdefer if (cache_owned) analysis_cache.deinit();
328
329 var pm = passes.PassManager.init(allocator);
330 defer pm.deinit();
331 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
332 try recipe.configure(&pm, .tensor, &options);
333
334 const tensor_start = options.now();
335 if (pm.runWithAnalysisCache(
336 module.choir_module,
337 module.context(),
338 &analysis_cache,
339 recipe.runOptions(),
340 ) == .failure) {
341 return preparationPassFailed(allocator, options, stage_mod.tensor_pipeline_name, &pm);
342 }
343 const tensor_end = options.now();
344
345 const tensor_module = try tensor.TensorJob.init(allocator, module, analysis_cache);
346 module_owned = false;
347 cache_owned = false;
348 return .{
349 .module = tensor_module,
350 .elapsed_ns = nsBetween(tensor_start, tensor_end),
351 .stats = run_mod.backendPreparationStats(pm.stats),
352 };
353 }
354
355 pub fn prepareDispatchJobFromTensorJobWithRun(
356 allocator: std.mem.Allocator,
357 module: *tensor.TensorJob,
358 options: BackendPreparationRunOptions,
359 ) !DispatchPreparationResult {
360 var module_owned = true;
361 errdefer if (module_owned) module.deinit();
362
363 try module.verify();
364
365 var pm = passes.PassManager.init(allocator);
366 defer pm.deinit();
367 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
368 try recipe.configure(&pm, .dispatch, &options);
369
370 const dispatch_start = options.now();
371 if (pm.runWithAnalysisCache(
372 module.choir_module,
373 module.context(),
374 &module.analysis_cache,
375 recipe.runOptions(),
376 ) == .failure) {
377 return preparationPassFailed(allocator, options, stage_mod.dispatch_pipeline_name, &pm);
378 }
379 const dispatch_end = options.now();
380
381 const plan_fingerprint = try stage_fingerprint.dispatch(
382 allocator,
383 module,
384 stage_mod.dispatch_pipeline_name,
385 );
386 const dispatch_module = try dispatch.DispatchJob.init(allocator, module);
387 module_owned = false;
388 return .{
389 .module = dispatch_module,
390 .elapsed_ns = nsBetween(dispatch_start, dispatch_end),
391 .stats = run_mod.backendPreparationStats(pm.stats),
392 .plan_fingerprint = plan_fingerprint,
393 };
394 }
395
396 pub fn prepareMemoryJobFromDispatchJobWithRun(
397 allocator: std.mem.Allocator,
398 module: *dispatch.DispatchJob,
399 options: BackendPreparationRunOptions,
400 ) !MemoryPreparationResult {
401 var module_owned = true;
402 errdefer if (module_owned) module.deinit();
403
404 try module.verify();
405
406 var pm = passes.PassManager.init(allocator);
407 defer pm.deinit();
408 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
409 try recipe.configure(&pm, .memory, &options);
410
411 const memory_start = options.now();
412 if (pm.runWithAnalysisCache(
413 module.choir_module,
414 module.context(),
415 &module.tensor_module.analysis_cache,
416 recipe.runOptions(),
417 ) == .failure) {
418 return preparationPassFailed(allocator, options, stage_mod.memory_pipeline_name, &pm);
419 }
420 const memory_end = options.now();
421
422 const plan_fingerprint = try stage_fingerprint.memory(
423 allocator,
424 module,
425 stage_mod.memory_pipeline_name,
426 );
427 const memory_module = try memory_product.MemoryJob.init(allocator, module);
428 module_owned = false;
429 return .{
430 .module = memory_module,
431 .elapsed_ns = nsBetween(memory_start, memory_end),
432 .stats = run_mod.backendPreparationStats(pm.stats),
433 .plan_fingerprint = plan_fingerprint,
434 };
435 }
436
437 pub fn prepareKernelJobFromMemoryJobWithRun(
438 allocator: std.mem.Allocator,
439 module: *memory_product.MemoryJob,
440 options: BackendPreparationRunOptions,
441 ) !KernelPreparationResult {
442 var module_owned = true;
443 errdefer if (module_owned) module.deinit();
444
445 try module.verify();
446
447 var pm = passes.PassManager.init(allocator);
448 defer pm.deinit();
449 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
450 try recipe.configure(&pm, .kernel, &options);
451
452 try recipe.applyTargetOptions(allocator, module.choir_module, options);
453 const kernel_start = options.now();
454 if (pm.runWithAnalysisCache(
455 module.choir_module,
456 module.context(),
457 &module.dispatch_module.tensor_module.analysis_cache,
458 recipe.runOptions(),
459 ) == .failure) {
460 return preparationPassFailed(allocator, options, stage_mod.kernel_pipeline_name, &pm);
461 }
462 const kernel_end = options.now();
463
464 const plan_fingerprint = try stage_fingerprint.kernel(
465 allocator,
466 module,
467 stage_mod.kernel_pipeline_name,
468 );
469 const kernel_module = try kernel_product.KernelJob.init(allocator, module);
470 module_owned = false;
471 return .{
472 .module = kernel_module,
473 .elapsed_ns = nsBetween(kernel_start, kernel_end),
474 .stats = run_mod.backendPreparationStats(pm.stats),
475 .plan_fingerprint = plan_fingerprint,
476 };
477 }
478
479 pub fn prepareTargetJobFromKernelJobWithRun(
480 allocator: std.mem.Allocator,
481 module: *kernel_product.KernelJob,
482 options: BackendPreparationRunOptions,
483 ) !TargetPreparationResult {
484 var module_owned = true;
485 errdefer if (module_owned) module.deinit();
486
487 var target_result = try prepareTargetProductFromKernelJobWithRun(allocator, module, options);
488 var target_result_owned = true;
489 errdefer if (target_result_owned) target_result.deinit();
490 const target_module = try target_product.TargetJob.init(
491 allocator,
492 module,
493 target_result.choir_module,
494 target_result.analysis_cache,
495 target_result.kernelization_product,
496 );
497 target_result_owned = false;
498 module_owned = false;
499 return .{
500 .module = target_module,
501 .elapsed_ns = target_result.elapsed_ns,
502 .stats = target_result.stats,
503 .finished_at = target_result.finished_at,
504 };
505 }
506
507 pub fn prepareTargetProductFromKernelJobWithRun(
508 allocator: std.mem.Allocator,
509 module: *kernel_product.KernelJob,
510 options: BackendPreparationRunOptions,
511 ) !TargetProductPreparationResult {
512 try module.verify();
513
514 var pm = passes.PassManager.init(allocator);
515 defer pm.deinit();
516 if (options.timing) |timing| try run_mod.addTimingInstrumentation(&pm, timing);
517 try recipe.configure(&pm, .target, &options);
518
519 const target_start = options.now();
520 var kernel_pass_ctx = module.passContext();
521 defer kernel_pass_ctx.deinit();
522 const kernelization_product = try kernelization.getKernelizationAnalysis(&kernel_pass_ctx, module.choir_module);
523
524 const target_choir_module = try module.choir_module.clone();
525 var target_choir_module_owned = true;
526 errdefer if (target_choir_module_owned) target_choir_module.erase();
527
528 var target_analysis_cache = passes.AnalysisCache.init(allocator, null);
529 var target_cache_owned = true;
530 errdefer if (target_cache_owned) target_analysis_cache.deinit();
531
532 try recipe.applyTargetOptions(allocator, target_choir_module, options);
533 if (pm.runWithAnalysisCache(
534 target_choir_module,
535 module.context(),
536 &target_analysis_cache,
537 recipe.runOptions(),
538 ) == .failure) {
539 return preparationPassFailed(allocator, options, stage_mod.target_pipeline_name, &pm);
540 }
541 const target_end = options.now();
542
543 target_choir_module_owned = false;
544 target_cache_owned = false;
545 return .{
546 .choir_module = target_choir_module,
547 .analysis_cache = target_analysis_cache,
548 .kernelization_product = kernelization_product,
549 .elapsed_ns = nsBetween(target_start, target_end),
550 .stats = run_mod.backendPreparationStats(pm.stats),
551 .finished_at = target_end,
552 };
553 }
554
555 pub fn nsBetween(start: i128, end: i128) u64 {
556 if (end <= start) return 0;
557 const elapsed = end - start;
558 if (elapsed > std.math.maxInt(u64)) return std.math.maxInt(u64);
559 return @intCast(elapsed);
560 }
561
562 pub fn countOperationTree(op: *ir.Operation) u64 {
563 var count: u64 = 1;
564 for (op.regions.items) |*region| {
565 var block_iter = region.getBlocks();
566 while (block_iter.next()) |block| {
567 var op_node = block.operations.head;
568 while (op_node) |node| {
569 const child: *ir.Operation = @ptrCast(@alignCast(node));
570 count +|= countOperationTree(child);
571 op_node = child.next_op;
572 }
573 }
574 }
575 return count;
576 }