lib/accy/src/preparation/recipe.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const stage_owner = @import("stage.zig");
4 const run = @import("run.zig");
5 const target = @import("target.zig");
6 const accy_choir = @import("../choir/root.zig");
7 const records = accy_choir.record;
8 const Stage = accy_choir.publication.Stage;
9 const Options = run.BackendPreparationRunOptions;
10 const library = @import("../kernel/library/root.zig");
11
12 pub const identity = choir.product.revision.record.Version{
13 .name = "accy-stage-recipe",
14 .version = 1,
15 };
16
17 pub const Policy = struct {
18 max_threads: usize = 1,
19 worker_allocator: ?void = null,
20 verification: choir.ir.VerifyOptions = choir.ir.verify.default_options,
21 };
22
23 pub const policy = Policy{};
24
25 pub fn runOptions() choir.passes.PassManagerRunOptions {
26 return .{ .max_threads = policy.max_threads, .worker_allocator = null };
27 }
28
29 pub fn pipelineName(comptime stage: Stage) []const u8 {
30 return @field(stage_owner, @tagName(stage) ++ "_pipeline_name");
31 }
32
33 /// The compile chain records this list with each stage so a change to the pass pipeline changes the
34 /// stage record, the sealed result of one compile stage. The function returns the ordered
35 /// identities of the passes in the stage's pipeline: each pass's declared name, version and cost
36 /// estimate, or its name at version 1 when it has none. The list is derived at compile time from
37 /// the same stage declarations that build the pipeline, so the recorded list and the pipeline that
38 /// runs cannot drift apart.
39 pub fn pipeline(comptime stage: Stage) []const choir.product.revision.record.Version {
40 const declarations = comptime stages(stage);
41 const identities = comptime blk: {
42 var result: [declarations.len]choir.product.revision.record.Version = undefined;
43 for (declarations, 0..) |entry, index| {
44 result[index] = if (entry.pass.work_contract) |contract|
45 contract.identity
46 else
47 .{ .name = entry.name, .version = 1 };
48 }
49 break :blk result;
50 };
51 return &identities;
52 }
53
54 fn stages(comptime stage: Stage) []const stage_owner.BackendPreparationStage {
55 return switch (stage) {
56 .semantic => &.{},
57 .contract => &stage_owner.contract_stages,
58 .tensor => &stage_owner.tensor_stages,
59 .dispatch => &stage_owner.dispatch_stages,
60 .memory => &stage_owner.memory_stages,
61 .kernel => &stage_owner.kernel_stages,
62 .target => &stage_owner.target_stages,
63 };
64 }
65
66 /// The compile chain calls this to fill a pass manager with one stage's pipeline. The call turns on
67 /// verification and adds the stage's passes to `manager`. The tensor stage's passes borrow
68 /// `options.tensor`, so `options` must outlive `manager`.
69 pub fn configure(
70 manager: *choir.passes.PassManager,
71 comptime stage: Stage,
72 options: *const Options,
73 ) !void {
74 manager.enableVerifierWithOptions(policy.verification);
75 switch (stage) {
76 .semantic => {},
77 .contract => try stage_owner.addContractPipeline(manager),
78 .tensor => try stage_owner.addTensorPipelineWithOptions(manager, &options.tensor),
79 .dispatch => try stage_owner.addDispatchPipeline(manager),
80 .memory => try stage_owner.addMemoryPipeline(manager),
81 .kernel => try stage_owner.addKernelPipeline(manager),
82 .target => try stage_owner.addTargetPipeline(manager),
83 }
84 }
85
86 /// The compile chain calls this to write a request's target settings, the device profile and
87 /// per-kernel schedule choices stored on the module, onto a module in a mutable compiler job. The
88 /// call writes the device profile when one is set, and the scan and row schedule choices when their
89 /// lists are non-empty. An unset profile or an empty list leaves the setting inherited from the
90 /// source exactly as it was.
91 pub fn applyTargetOptions(
92 allocator: std.mem.Allocator,
93 root: *choir.ir.Operation,
94 options: Options,
95 ) !void {
96 if (options.target_profile) |profile| {
97 try target.setBackendTargetProfile(root.context, root, profile);
98 }
99 if (options.generated_scan_schedules.len != 0) {
100 try target.setGeneratedScanSchedules(
101 allocator,
102 root.context,
103 root,
104 options.generated_scan_schedules,
105 );
106 }
107 if (options.generated_row_pipeline_schedules.len != 0) {
108 try target.setGeneratedRowPipelineSchedules(
109 allocator,
110 root.context,
111 root,
112 options.generated_row_pipeline_schedules,
113 );
114 }
115 }
116
117 const Einsum = struct {
118 strategy: accy_choir.einsum.Strategy,
119 exact_state_limit: usize,
120 beam_width: usize,
121 auto_beam_width: usize,
122 kernel_library: @import("library.zig").KernelLibraryLowering,
123 matrix_product_schedule: ?library.MatrixProductSchedule,
124 matrix_product_tuning: ?library.linalg.MatrixProductScheduleReader = null,
125 family_tuning: ?library.tuning.FamilyTuningReader = null,
126 };
127
128 const Indexing = struct {
129 kernel_library: @import("library.zig").KernelLibraryLowering,
130 gather_schedule: ?library.GatherSchedule,
131 scatter_schedule: ?library.ScatterSchedule,
132 scatter_add_schedule: ?library.ScatterAddSchedule,
133 family_tuning: ?library.tuning.FamilyTuningReader = null,
134 };
135
136 const Tensor = struct {
137 activation: stage_owner.ActivationLoweringOptions,
138 einsum: Einsum,
139 indexing: Indexing,
140 loss: stage_owner.LossLoweringOptions,
141 };
142
143 pub const TargetInputs = struct {
144 profile: ?target.BackendTargetProfile,
145 generated_scan_schedules: ?[]const u8,
146 generated_row_pipeline_schedules: ?[]const u8,
147 };
148
149 fn Lowering(comptime stage: Stage) type {
150 return switch (stage) {
151 .tensor => Tensor,
152 .kernel, .target => TargetInputs,
153 else => struct {},
154 };
155 }
156
157 pub fn Record(comptime stage: Stage) type {
158 return struct { execution: Policy, options: Lowering(stage) };
159 }
160
161 /// The compile chain calls this to encode the options a stage ran with, for its stage record. The
162 /// call reads the target settings already on `root` without changing the module, applies the
163 /// request's overrides, and returns bytes the caller owns. An empty override keeps the inherited
164 /// setting. A tuning table given in the options is recorded with its device facts and its records
165 /// in order, so a table that lacks a key is recorded as lacking it. Timing, the failure destination
166 /// and the clock are administrative and are left out, so they never change a record.
167 pub fn encode(
168 allocator: std.mem.Allocator,
169 comptime stage: Stage,
170 root: *choir.ir.Operation,
171 options: Options,
172 ) ![]u8 {
173 return encodeWithTarget(allocator, stage, .{
174 .profile = target.readBackendTargetProfile(root),
175 .generated_scan_schedules = target.readGeneratedScanSchedules(root),
176 .generated_row_pipeline_schedules = target.readGeneratedRowPipelineSchedules(root),
177 }, options);
178 }
179
180 /// The compile chain calls this to encode a stage's options from a retained source without a
181 /// mutable module at hand. The call starts from the target settings captured from the original
182 /// source, applies this request's overrides, and returns bytes the caller owns. The call decodes
183 /// its own output and compares it with the input before returning, so a value that does not survive
184 /// encoding is caught at once.
185 pub fn encodeWithTarget(
186 allocator: std.mem.Allocator,
187 comptime stage: Stage,
188 inherited: TargetInputs,
189 options: Options,
190 ) ![]u8 {
191 comptime classify();
192 const scan = if ((stage == .kernel or stage == .target) and
193 options.generated_scan_schedules.len != 0)
194 try target.encodeGeneratedScanSchedules(allocator, options.generated_scan_schedules)
195 else
196 null;
197 defer if (scan) |bytes| allocator.free(bytes);
198 const row = if ((stage == .kernel or stage == .target) and
199 options.generated_row_pipeline_schedules.len != 0)
200 try target.encodeGeneratedRowPipelineSchedules(
201 allocator,
202 options.generated_row_pipeline_schedules,
203 )
204 else
205 null;
206 defer if (row) |bytes| allocator.free(bytes);
207 const source = switch (stage) {
208 .tensor => options.tensor,
209 .kernel, .target => TargetInputs{
210 .profile = options.target_profile orelse inherited.profile,
211 .generated_scan_schedules = scan orelse inherited.generated_scan_schedules,
212 .generated_row_pipeline_schedules = row orelse
213 inherited.generated_row_pipeline_schedules,
214 },
215 else => Lowering(stage){},
216 };
217 var references = records.reference.Index{ .allocator = allocator, .limit = 0 };
218 defer references.deinit();
219 const effective = .{ .execution = policy, .options = source };
220 const bytes = try records.codec.encode(allocator, Record(stage), stage, effective, &references);
221 errdefer allocator.free(bytes);
222 var decoded = try records.codec.decode(allocator, Record(stage), stage, bytes);
223 defer decoded.deinit();
224 try records.codec.compare(decoded.value, effective, &references);
225 return bytes;
226 }
227
228 fn classify() void {
229 const fields = choir.product.revision.record.requireFields;
230 fields(Policy, &.{ "max_threads", "worker_allocator", "verification" });
231 fields(choir.passes.PassManagerRunOptions, &.{ "max_threads", "worker_allocator" });
232 fields(choir.ir.VerifyOptions, &.{
233 "check_terminators", "require_terminators", "recursive", "check_use_def",
234 "check_local_dominance", "check_cfg", "max_depth",
235 });
236 fields(Options, &.{
237 "timing", "failure", "target_profile", "generated_scan_schedules",
238 "generated_row_pipeline_schedules", "tensor", "now",
239 });
240 records.codec.coverage(stage_owner.TensorLoweringOptions, Tensor, &.{});
241 records.codec.coverage(stage_owner.EinsumLoweringOptions, Einsum, &.{});
242 records.codec.coverage(stage_owner.IndexingLoweringOptions, Indexing, &.{});
243 fields(stage_owner.ActivationLoweringOptions, &.{"kernel_library"});
244 fields(stage_owner.LossLoweringOptions, &.{
245 "kernel_library", "row_sparse_cross_entropy_schedule",
246 });
247 fields(target.BackendTargetProfile, &.{
248 "backend_kind", "artifact_format", "math_tier", "dtype_bits", "feature_bits",
249 });
250 }