lib/accy/src/artifact/input.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const fixed = @import("alloc_fixed");
3 const choir = @import("choir");
4 const preparation = @import("../preparation/root.zig");
5 const records = @import("../choir/root.zig").record;
6 const publication = @import("../choir/root.zig").publication;
7 const plan = @import("root.zig").plan;
8 const ir = choir.ir;
9 const Configuration = choir.product.operation.Configuration;
10 const Prepared = preparation.pipeline.BackendPreparedModule;
11
12 /// Short-lived compile state rebuilt from the dispatch, memory, kernel, and
13 /// target results of one finished preparation: the schedule plan, the buffer,
14 /// memory-space and layout plans, the kernel outlines, the generated kernel
15 /// programs, and the target settings. A compiler of device code uses this to
16 /// recover the plans of a finished compile inside memory the caller provides.
17 /// `create` allocates the job and all of its contents inside the caller's
18 /// `workspace`, and a workspace that runs out gives `error.WorkExhausted`. The
19 /// job requires each stage record to match: the target record must hold exactly
20 /// one program root (`error.InvalidStageRoots`), and the dispatch, memory, and
21 /// kernel records must decode to that same program
22 /// (`error.IncompatibleStageImages`). The job keeps no pointer to the prepared
23 /// module, no stage-record key, and no mutable semantic module. The job works
24 /// out each kernel's legality for the target again from the restored plans and
25 /// keeps that result with the rest. `destroy` releases the job's contents, and
26 /// the workspace itself stays with the caller.
27 pub const InputJob = struct {
28 fixed: fixed.Tracked,
29 context: ?*ir.Context = null,
30 decoded: ?choir.bytecode.DecodedModule = null,
31 target: ?records.codec.Decoded(records.target.Record) = null,
32 schedule: ?preparation.schedule.SchedulePlanAnalysis = null,
33 buffers: ?preparation.bufferization.BufferPlanAnalysis = null,
34 spaces: ?preparation.memory.MemorySpacePlanAnalysis = null,
35 layouts: ?preparation.layout.LayoutPlanAnalysis = null,
36 outlines: ?preparation.kernelization.product.KernelOutlinePlanAnalysis = null,
37 generated: ?preparation.kernelization.KernelizationAnalysis = null,
38 legal: ?preparation.backend.BackendLegalizationAnalysis = null,
39
40 pub fn create(
41 workspace: []u8,
42 prepared: *const Prepared,
43 comptime configuration: Configuration,
44 ) !*InputJob {
45 return initialize(workspace, prepared, configuration) catch |err| {
46 return if (err == error.OutOfMemory) error.WorkExhausted else err;
47 };
48 }
49
50 fn initialize(
51 workspace: []u8,
52 prepared: *const Prepared,
53 comptime configuration: Configuration,
54 ) !*InputJob {
55 var initial = fixed.Tracked.init(workspace);
56 const self = try initial.allocator().create(InputJob);
57 self.* = .{ .fixed = initial };
58 errdefer self.destroy();
59 const allocator = self.fixed.allocator();
60 self.context = try ir.Context.create(allocator, configuration.context);
61 const context = self.context.?;
62 (try choir.product.recipe.decode(prepared.stage(.target).inputs().policy)).restore(context);
63 try configuration.register(context);
64 const image = try open(prepared, allocator, .target, configuration);
65 defer image.destroy();
66 const root = image.root(0) orelse return error.InvalidStageRoots;
67 if (image.root(1) != null) return error.InvalidStageRoots;
68 self.decoded = try choir.bytecode.decodeModule(allocator, context, root.bytes);
69 self.target = try records.codec.decode(allocator, records.target.Record, .target, image.stage());
70 clearTargetAttributes(self.decoded.?.module);
71 inline for (.{ .dispatch, .memory, .kernel }) |stage| {
72 try self.requireSameProgram(prepared, stage, configuration);
73 }
74 var references = try References.init(allocator, self.decoded.?.module, configuration.image.entities);
75 defer references.deinit();
76 try self.restoreDispatch(prepared, &references, configuration);
77 try self.restoreMemory(prepared, &references, configuration);
78 try self.restoreKernel(prepared, &references, configuration);
79 try self.restoreTarget(configuration);
80 return self;
81 }
82
83 pub fn destroy(self: *InputJob) void {
84 const allocator = self.fixed.allocator();
85 if (self.legal) |*value| value.deinit();
86 if (self.generated) |*value| value.deinit();
87 if (self.outlines) |*value| value.deinit();
88 if (self.layouts) |*value| value.deinit();
89 if (self.spaces) |*value| value.deinit();
90 if (self.buffers) |*value| value.deinit();
91 if (self.schedule) |*value| value.deinit();
92 if (self.target) |*value| value.deinit();
93 if (self.decoded) |*value| value.deinit();
94 if (self.context) |context| {
95 context.deinit(allocator);
96 allocator.destroy(context);
97 }
98 }
99
100 /// Returns pointers into this job's own fields: the program root, the
101 /// target profile, and each restored plan. A caller passes the returned
102 /// plans to the artifact compiler. The pointers stay valid until `destroy`
103 /// runs or the caller writes over the workspace.
104 pub fn plans(self: *const InputJob) plan.PlanInputs {
105 return .{
106 .root = self.decoded.?.module,
107 .profile = self.target.?.value.profile,
108 .schedule = &self.schedule.?,
109 .buffers = &self.buffers.?,
110 .spaces = &self.spaces.?,
111 .layouts = &self.layouts.?,
112 .outlines = &self.outlines.?,
113 .generated = &self.generated.?,
114 .legal = &self.legal.?,
115 .target = &self.target.?.value,
116 };
117 }
118
119 fn requireSameProgram(
120 self: *InputJob,
121 prepared: *const Prepared,
122 comptime stage: publication.Stage,
123 comptime configuration: Configuration,
124 ) !void {
125 const allocator = self.fixed.allocator();
126 const image = try open(prepared, allocator, stage, configuration);
127 defer image.destroy();
128 if (image.root(1) != null) return error.InvalidStageRoots;
129 var context = try ir.Context.init(allocator, configuration.context);
130 defer context.deinit(allocator);
131 (try choir.product.recipe.decode(prepared.stage(stage).inputs().policy)).restore(&context);
132 try configuration.register(&context);
133 const root = image.root(0) orelse return error.InvalidStageRoots;
134 var decoded = try choir.bytecode.decodeModule(allocator, &context, root.bytes);
135 defer decoded.deinit();
136 clearTargetAttributes(decoded.module);
137 choir.bytecode.qualification.compare(allocator, decoded.module, self.decoded.?.module, decoded.resources, self.decoded.?.resources, configuration.codec) catch |err| {
138 return if (err == error.UnencodableProduct) error.IncompatibleStageImages else err;
139 };
140 }
141
142 fn restoreDispatch(self: *InputJob, prepared: *const Prepared, refs: *const References, comptime configuration: Configuration) !void {
143 const allocator = self.fixed.allocator();
144 const image = try open(prepared, allocator, .dispatch, configuration);
145 defer image.destroy();
146 var record = try records.codec.decode(allocator, records.dispatch.Record, .dispatch, image.stage());
147 defer record.deinit();
148 self.schedule = preparation.schedule.SchedulePlanAnalysis.init(allocator);
149 const output = &self.schedule.?;
150 copyCounters(output, record.value.schedule);
151 for (record.value.schedule.work_items, 0..) |item, index| {
152 const value = try project(preparation.schedule.ScheduleWorkItem, allocator, item, refs);
153 try output.work_items.append(allocator, value);
154 try output.root_to_item.putNoClobber(value.root, index);
155 }
156 }
157
158 fn restoreMemory(self: *InputJob, prepared: *const Prepared, refs: *const References, comptime configuration: Configuration) !void {
159 const allocator = self.fixed.allocator();
160 const image = try open(prepared, allocator, .memory, configuration);
161 defer image.destroy();
162 var record = try records.codec.decode(allocator, records.memory.Record, .memory, image.stage());
163 defer record.deinit();
164 self.buffers = preparation.bufferization.BufferPlanAnalysis.init(allocator);
165 self.spaces = preparation.memory.MemorySpacePlanAnalysis.init(allocator);
166 self.layouts = preparation.layout.LayoutPlanAnalysis.init(allocator);
167 const buffers = &self.buffers.?;
168 const spaces = &self.spaces.?;
169 const layouts = &self.layouts.?;
170 copyCounters(buffers, record.value.buffers);
171 copyCounters(spaces, record.value.spaces);
172 copyCounters(layouts, record.value.layouts);
173 for (record.value.buffers.slots) |item| {
174 const value = try project(preparation.bufferization.BufferSlot, allocator, item, refs);
175 try buffers.slots.append(allocator, value);
176 }
177 for (record.value.bindings) |binding| {
178 const value = try refs.value(binding.value);
179 try buffers.value_to_slot.putNoClobber(value, binding.slot_id);
180 }
181 for (record.value.buffers.elisions, 0..) |item, index| {
182 const value = try project(preparation.bufferization.FusionElision, allocator, item, refs);
183 try buffers.elisions.append(allocator, value);
184 try buffers.value_to_elision.putNoClobber(value.value, index);
185 }
186 for (record.value.spaces.assignments, 0..) |item, index| {
187 const value = try project(preparation.memory.MemorySpaceAssignment, allocator, item, refs);
188 try spaces.assignments.append(allocator, value);
189 try spaces.slot_to_assignment.putNoClobber(value.slot_id, index);
190 }
191 for (record.value.layouts.assignments, 0..) |item, index| {
192 const value = try project(preparation.layout.LayoutAssignment, allocator, item, refs);
193 try layouts.assignments.append(allocator, value);
194 try layouts.slot_to_assignment.putNoClobber(value.slot_id, index);
195 }
196 }
197
198 fn restoreKernel(self: *InputJob, prepared: *const Prepared, refs: *const References, comptime configuration: Configuration) !void {
199 const allocator = self.fixed.allocator();
200 const image = try open(prepared, allocator, .kernel, configuration);
201 defer image.destroy();
202 var record = try records.codec.decode(allocator, records.kernel.Record, .kernel, image.stage());
203 defer record.deinit();
204 self.outlines = preparation.kernelization.product.KernelOutlinePlanAnalysis.init(allocator);
205 const output = &self.outlines.?;
206 copyCounters(output, record.value.outlines);
207 for (record.value.outlines.kernels) |item| {
208 const value = try project(preparation.kernelization.product.KernelOutline, allocator, item, refs);
209 try output.kernels.append(allocator, value);
210 try output.work_to_kernel.putNoClobber(value.work_item_id, value.id);
211 }
212 }
213
214 fn restoreTarget(self: *InputJob, comptime configuration: Configuration) !void {
215 const allocator = self.fixed.allocator();
216 const source = self.target.?.value;
217 const context = self.context.?;
218 const root = self.decoded.?.module;
219 if (source.profile) |profile| try preparation.setBackendTargetProfile(context, root, profile);
220 if (source.generated_scan_schedules) |value| {
221 try root.setAttr(preparation.target.generated_scan_schedule_attr_name, try context.getStringAttr(value));
222 }
223 if (source.generated_row_pipeline_schedules) |value| {
224 try root.setAttr(preparation.target.generated_row_pipeline_schedule_attr_name, try context.getStringAttr(value));
225 }
226 self.generated = try preparation.kernelization.KernelizationAnalysis.init(allocator, configuration.context);
227 const output = &self.generated.?;
228 try output.reserveKernelCapacity(source.kernels.len);
229 for (source.kernels, 0..) |item, index| {
230 var value = try preparation.publication.restoreKernel(allocator, item.lowered, compilationConfiguration(configuration));
231 errdefer value.deinit(allocator);
232 try output.work_to_kernel.putNoClobber(value.work_item_id, index);
233 output.kernels.appendAssumeCapacity(value);
234 }
235 self.legal = preparation.backend.BackendLegalizationAnalysis.init(allocator, source.profile);
236 const legal = &self.legal.?;
237 for (self.outlines.?.kernels.items) |outline| {
238 const result = try preparation.backend.legalizeKernel(outline, &self.buffers.?, source.profile);
239 try legal.kernels.append(allocator, result);
240 if (result.isLegal()) {
241 legal.legal_kernel_count += 1;
242 legal.total_static_bytes = try std.math.add(u64, legal.total_static_bytes, result.static_bytes);
243 } else legal.illegal_kernel_count += 1;
244 }
245 }
246 };
247
248 fn open(prepared: *const Prepared, allocator: std.mem.Allocator, comptime stage: publication.Stage, comptime configuration: Configuration) !*choir.product.entity.Image {
249 return (choir.product.operation.Product{ .revision = prepared.stage(stage) }).open(allocator, configuration.image);
250 }
251
252 fn clearTargetAttributes(root: *ir.Operation) void {
253 inline for (.{ preparation.target.backend_kind_attr_name, preparation.target.artifact_format_attr_name, preparation.target.math_tier_attr_name, preparation.target.dtype_bits_attr_name, preparation.target.feature_bits_attr_name, preparation.target.generated_scan_schedule_attr_name, preparation.target.generated_row_pipeline_schedule_attr_name }) |name| {
254 _ = root.removeAttr(name);
255 }
256 }
257
258 fn copyCounters(output: anytype, source: anytype) void {
259 inline for (@typeInfo(@TypeOf(source)).@"struct".field_names) |name| {
260 const T = @TypeOf(@field(source, name));
261 if (@typeInfo(T) == .int) @field(output.*, name) = @field(source, name);
262 }
263 }
264
265 const References = struct {
266 allocator: std.mem.Allocator,
267 operations: []*ir.Operation,
268
269 fn init(allocator: std.mem.Allocator, root: *ir.Operation, limit: u32) !References {
270 var index = try records.reference.Index.init(allocator, root, limit);
271 defer index.deinit();
272 const operations = try allocator.alloc(*ir.Operation, index.operations.count());
273 var entries = index.operations.iterator();
274 while (entries.next()) |entry| operations[entry.value_ptr.ordinal] = entry.key_ptr.*;
275 return .{ .allocator = allocator, .operations = operations };
276 }
277
278 fn deinit(self: *References) void {
279 self.allocator.free(self.operations);
280 }
281
282 fn operation(self: *const References, ref: records.reference.Operation) !*ir.Operation {
283 if (ref.ordinal >= self.operations.len) return error.UnboundProductInput;
284 return self.operations[ref.ordinal];
285 }
286
287 fn value(self: *const References, ref: records.reference.Value) !*ir.Value {
288 switch (ref) {
289 .result => |item| {
290 const op = try self.operation(item.operation);
291 if (item.position >= op.results.items.len) return error.UnboundProductInput;
292 return &op.results.items[item.position];
293 },
294 .argument => |item| {
295 const op = try self.operation(item.operation);
296 if (item.region >= op.regions.items.len) return error.UnboundProductInput;
297 var blocks = op.regions.items[item.region].getBlocks();
298 var position: u32 = 0;
299 while (blocks.next()) |block| : (position += 1) {
300 if (position != item.block) continue;
301 if (item.position >= block.arguments.items.len) return error.UnboundProductInput;
302 return block.arguments.items[item.position];
303 }
304 return error.UnboundProductInput;
305 },
306 }
307 }
308 };
309
310 fn project(comptime T: type, allocator: std.mem.Allocator, source: anytype, refs: *const References) anyerror!T {
311 if (T == *ir.Operation) return refs.operation(source);
312 if (T == *ir.Value) return refs.value(source);
313 switch (@typeInfo(T)) {
314 .optional => |info| return if (source) |value| try project(info.child, allocator, value, refs) else null,
315 .pointer => |info| {
316 if (info.size != .slice) @compileError("uncaptured artifact pointer");
317 const values = try allocator.alloc(info.child, source.len);
318 for (values, source) |*out, value| out.* = try project(info.child, allocator, value, refs);
319 return values;
320 },
321 .@"struct" => |info| {
322 var value: T = undefined;
323 inline for (info.field_names, info.field_types) |name, Field| {
324 @field(value, name) = try project(Field, allocator, @field(source, name), refs);
325 }
326 return value;
327 },
328 else => return source,
329 }
330 }
331
332 fn compilationConfiguration(comptime original: Configuration) Configuration {
333 var result = original;
334 result.register = struct {
335 fn register(context: *ir.Context) !void {
336 try original.register(context);
337 try choir.backends.gpu.prepareCompilationDialects(context);
338 }
339 }.register;
340 return result;
341 }