lib/accy/src/preparation/kernelization/lowering/work.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const choir_abi = @import("choir_abi");
 3 const choir = @import("choir");
 4 const lowering = @import("root.zig");
 5 const model = @import("../model/root.zig");
 6 const work = choir.passes.pass.work;
 7 const ir = choir.ir;
 8 const kernel = lowering.common.kernel_root;
 9 
10 /// The kernel stage calls this function before the pass runs to charge kernel generation against
11 /// the caller's limits. The function computes the work bound from counts read off the source
12 /// module, before any allocation and before any analysis is requested. The work bound is the costs
13 /// a pass declares before it runs, so compilation can refuse a pass that would exceed the caller's
14 /// limits. The call returns `error.MissingWorkContract` when the pass would run on more than one
15 /// thread or with a worker allocator. The analyses it depends on declare their own bounds and their
16 /// own retained storage.
17 pub fn analysis(input: work.Input) !work.Bounds {
18     if (input.options.max_threads != 1 or input.options.worker_allocator != null) {
19         return error.MissingWorkContract;
20     }
21     var source = try work.Census.inspect(input.operation);
22     source.values = try work.add(source.values, source.operands);
23     const limits = input.operation.context.capacity.asLimits();
24     const capacity = ir.Context.Capacity.derive(limits) catch return error.WorkOverflow;
25     const attempts = @max(
26         2,
27         lowering.row_pipeline.max_schedule_candidates + 2,
28         lowering.scan_lowering.max_schedule_candidates + 2,
29     );
30     const per_kernel = try work.multiply(attempts, try attemptStorage(source, limits));
31     const base = try model.KernelizationAnalysis.baseStorageBound(limits, source.operations);
32     const storage = try work.add(base, try work.multiply(source.operations, per_kernel));
33     const visits = try traversal(source, limits, attempts);
34     return .{
35         .work = .{
36             .input_bytes = source.input_bytes,
37             .output_bytes = capacity.storage_bytes,
38             .structural_visits = visits,
39             .analysis_computations = 1,
40             .allocation_capacity = storage,
41         },
42         .workspace = storage,
43         .retained_storage = storage,
44     };
45 }
46 
47 pub fn pass(_: work.Input) !work.Bounds {
48     return .{ .work = .{ .structural_visits = 1 } };
49 }
50 
51 /// `analysis` calls this function to size the storage of one generation attempt, one try at
52 /// emitting a kernel under one candidate schedule. One attempt owns one generated program, its name
53 /// and one parameter list. The scratch of every kernel family is added up so the bound holds
54 /// whichever family generation picks, including the vector fallback, the second body
55 /// element-by-element generation emits when the per-element form does not apply.
56 fn attemptStorage(source: work.Census, limits: ir.Context.Limits) !u64 {
57     var bytes = try lowering.builder.storageBound(limits);
58     bytes = try work.add(bytes, lowering.name.maximum_storage_bytes);
59     bytes = try work.add(bytes, try lowering.abi.storageBound(source.values));
60     bytes = try work.add(bytes, try lowering.elementwise.scratchStorageBound(source));
61     bytes = try work.add(bytes, try lowering.row_pipeline.scratchStorageBound(source.values));
62     const arrays = try work.multiply(source.values, @sizeOf(choir_abi.DType) +
63         @sizeOf(lowering.common.bufferization.BufferSlot) + @sizeOf(kernel.Value));
64     const alignment = @alignOf(choir_abi.DType) +
65         @alignOf(lowering.common.bufferization.BufferSlot) + @alignOf(kernel.Value);
66     bytes = try work.add(bytes, try work.add(arrays, alignment));
67     const memo = try work.hashMapGrowth(*ir.Value, kernel.Value, source.values);
68     const varying = try work.hashMapGrowth(*ir.Value, void, source.values);
69     bytes = try work.add(bytes, try work.add(try work.multiply(6, memo), varying));
70     if (bytes > std.math.maxInt(usize)) return error.WorkOverflow;
71     return bytes;
72 }
73 
74 /// `analysis` calls this function to count the visits kernel generation makes. The count covers
75 /// scans of the source and the plans, memo lookups, and building, checking and hashing the
76 /// generated code. The shared compiler context, the object that owns the operations and values of
77 /// the generated kernels, caps the generated records, and that cap holds even across failed
78 /// attempts.
79 fn traversal(source: work.Census, limits: ir.Context.Limits, attempts: u64) !u64 {
80     const source_units = try work.add(try work.add(source.atoms, source.input_bytes), 1);
81     const generated_bytes = try work.add(limits.operations.storage_bytes, limits.operations.nested_bytes);
82     const generated_values = generated_bytes / @sizeOf(ir.Value);
83     const generated_ops = limits.operations.storage_bytes / @sizeOf(ir.Operation);
84     const records = try work.add(try work.add(generated_values, generated_ops), 1);
85     const source_walks = try work.multiply(source_units, try work.add(source.operations, 1));
86     const lookups = try work.add(try work.hashMapCapacity(source.values), source_units);
87     const generated_walks = try work.multiply(records, try work.add(lookups, 256));
88     const total = try work.add(source_walks, generated_walks);
89     return work.multiply(128, try work.multiply(
90         try work.multiply(try work.add(source.operations, 1), attempts),
91         total,
92     ));
93 }