lib/accy/src/preparation/capture.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir = @import("choir");
3 const gpu = choir.backends.gpu;
4 const preparation = @import("root.zig");
5 const records = @import("../choir/root.zig").record;
6 const codec = records.codec;
7 const dispatch_record = records.dispatch;
8 const memory_record = records.memory;
9
10 pub fn dispatch(
11 allocator: std.mem.Allocator,
12 root: *choir.ir.Operation,
13 fusion: *const preparation.fusion.FusionPlanAnalysis,
14 schedule: *const preparation.schedule.SchedulePlanAnalysis,
15 entity_limit: u32,
16 ) ![]u8 {
17 comptime {
18 codec.coverage(preparation.fusion.FusionCluster, dispatch_record.Cluster, &.{});
19 codec.coverage(
20 preparation.fusion.FusionPlanAnalysis,
21 dispatch_record.Fusion,
22 &.{"allocator"},
23 );
24 codec.coverage(preparation.schedule.ScheduleWorkItem, dispatch_record.WorkItem, &.{});
25 codec.coverage(preparation.schedule.SchedulePlanAnalysis, dispatch_record.Schedule, &.{
26 "allocator", "root_to_item",
27 });
28 }
29 if (schedule.root_to_item.count() != schedule.work_items.items.len) {
30 return error.UnencodableProduct;
31 }
32 for (schedule.work_items.items, 0..) |item, index| {
33 if (schedule.root_to_item.get(item.root) != index) return error.UnencodableProduct;
34 }
35 var references = try records.reference.Index.init(allocator, root, entity_limit);
36 defer references.deinit();
37 return qualified(allocator, dispatch_record.Record, .dispatch, .{
38 .fusion = fusion,
39 .schedule = schedule,
40 }, &references);
41 }
42
43 pub fn memory(
44 allocator: std.mem.Allocator,
45 root: *choir.ir.Operation,
46 buffers: *const preparation.bufferization.BufferPlanAnalysis,
47 spaces: *const preparation.memory.MemorySpacePlanAnalysis,
48 layouts: *const preparation.layout.LayoutPlanAnalysis,
49 entity_limit: u32,
50 ) ![]u8 {
51 comptime memoryCoverage();
52 if (buffers.value_to_slot.count() < buffers.slots.items.len or
53 buffers.value_to_elision.count() != buffers.elisions.items.len or
54 spaces.slot_to_assignment.count() != spaces.assignments.items.len or
55 layouts.slot_to_assignment.count() != layouts.assignments.items.len)
56 {
57 return error.UnencodableProduct;
58 }
59 for (buffers.slots.items, 0..) |slot, index| {
60 if (buffers.value_to_slot.get(slot.value) != index) return error.UnencodableProduct;
61 }
62 for (buffers.elisions.items, 0..) |elision, index| {
63 if (buffers.value_to_elision.get(elision.value) != index) return error.UnencodableProduct;
64 }
65 for (spaces.assignments.items, 0..) |assignment, index| {
66 if (spaces.slot_to_assignment.get(assignment.slot_id) != index) {
67 return error.UnencodableProduct;
68 }
69 }
70 for (layouts.assignments.items, 0..) |assignment, index| {
71 if (layouts.slot_to_assignment.get(assignment.slot_id) != index) {
72 return error.UnencodableProduct;
73 }
74 }
75 var references = try records.reference.Index.init(allocator, root, entity_limit);
76 defer references.deinit();
77 const bindings = try memoryBindings(allocator, buffers, &references);
78 defer allocator.free(bindings);
79 const bytes = try qualified(allocator, memory_record.Record, .memory, .{
80 .buffers = buffers,
81 .bindings = bindings,
82 .spaces = spaces,
83 .layouts = layouts,
84 }, &references);
85 errdefer allocator.free(bytes);
86 try compareMemoryBindings(allocator, bytes, buffers, &references);
87 return bytes;
88 }
89
90 fn compareMemoryBindings(
91 allocator: std.mem.Allocator,
92 bytes: []const u8,
93 buffers: *const preparation.bufferization.BufferPlanAnalysis,
94 references: *const records.reference.Index,
95 ) !void {
96 var decoded = try codec.decode(allocator, memory_record.Record, .memory, bytes);
97 defer decoded.deinit();
98 if (decoded.value.bindings.len != buffers.value_to_slot.count()) return error.UnencodableProduct;
99 var entries = buffers.value_to_slot.iterator();
100 while (entries.next()) |entry| {
101 const expected = try references.value(entry.key_ptr.*);
102 const found = for (decoded.value.bindings) |binding| {
103 if (std.meta.eql(expected, binding.value)) break binding.slot_id;
104 } else return error.UnencodableProduct;
105 if (found != entry.value_ptr.*) return error.UnencodableProduct;
106 }
107 }
108
109 fn memoryBindings(
110 allocator: std.mem.Allocator,
111 buffers: *const preparation.bufferization.BufferPlanAnalysis,
112 references: *const records.reference.Index,
113 ) ![]memory_record.Binding {
114 const bindings = try allocator.alloc(memory_record.Binding, buffers.value_to_slot.count());
115 errdefer allocator.free(bindings);
116 var entries = buffers.value_to_slot.iterator();
117 var count: usize = 0;
118 while (entries.next()) |entry| {
119 if (entry.value_ptr.* >= buffers.slots.items.len) return error.UnencodableProduct;
120 bindings[count] = .{
121 .value = try references.value(entry.key_ptr.*),
122 .slot_id = entry.value_ptr.*,
123 };
124 count += 1;
125 }
126 std.debug.assert(count == bindings.len);
127 std.mem.sort(memory_record.Binding, bindings, {}, memory_record.Binding.lessThan);
128 return bindings;
129 }
130
131 fn memoryCoverage() void {
132 codec.coverage(preparation.bufferization.BufferSlot, memory_record.Slot, &.{});
133 codec.coverage(preparation.bufferization.FusionElision, memory_record.Elision, &.{});
134 codec.coverage(preparation.bufferization.BufferPlanAnalysis, memory_record.Buffers, &.{
135 "allocator", "value_to_slot", "value_to_elision",
136 });
137 codec.coverage(preparation.memory.MemorySpaceAssignment, memory_record.Assignment, &.{});
138 codec.coverage(preparation.memory.MemorySpacePlanAnalysis, memory_record.Spaces, &.{
139 "allocator", "slot_to_assignment",
140 });
141 codec.coverage(preparation.layout.LayoutAssignment, memory_record.Layout, &.{});
142 codec.coverage(preparation.layout.LayoutPlanAnalysis, memory_record.Layouts, &.{
143 "allocator", "slot_to_assignment",
144 });
145 }
146
147 pub fn kernel(
148 allocator: std.mem.Allocator,
149 root: *choir.ir.Operation,
150 outlines: *const preparation.kernelization.product.KernelOutlinePlanAnalysis,
151 generated: *const preparation.kernelization.KernelizationAnalysis,
152 comptime configuration: choir.product.operation.Configuration,
153 ) ![]u8 {
154 comptime kernelCoverage();
155 try requireWorkMap(outlines);
156 try requireWorkMap(generated);
157 var programs = try Generated.init(allocator, generated.kernels.items, configuration);
158 defer programs.deinit();
159 var references = try records.reference.Index.init(allocator, root, configuration.image.entities);
160 defer references.deinit();
161 return qualified(allocator, records.kernel.Record, .kernel, .{
162 .outlines = outlines,
163 .generated = .{ .kernels = programs.lowered },
164 }, &references);
165 }
166
167 const Generated = struct {
168 allocator: std.mem.Allocator,
169 programs: []codec.Decoded(records.program.Record),
170 lowered: []records.kernel.Lowered,
171
172 fn init(
173 allocator: std.mem.Allocator,
174 source: []const preparation.kernelization.LoweredKernel,
175 comptime configuration: choir.product.operation.Configuration,
176 ) !Generated {
177 comptime kernelCoverage();
178 if (source.len > configuration.codec.fields) return error.RecordLimit;
179 const programs = try allocator.alloc(
180 codec.Decoded(records.program.Record),
181 source.len,
182 );
183 errdefer allocator.free(programs);
184 var initialized: usize = 0;
185 errdefer for (programs[0..initialized]) |*program| program.deinit();
186 const lowered = try allocator.alloc(records.kernel.Lowered, programs.len);
187 errdefer allocator.free(lowered);
188 for (source, 0..) |*kernel_value, index| {
189 programs[index] = try records.program.capture(
190 allocator,
191 &kernel_value.program,
192 configuration,
193 );
194 initialized += 1;
195 inline for (@typeInfo(records.kernel.Lowered).@"struct".field_names) |name| {
196 @field(lowered[index], name) = if (comptime std.mem.eql(u8, name, "program"))
197 programs[index].value
198 else
199 @field(kernel_value, name);
200 }
201 }
202 return .{ .allocator = allocator, .programs = programs, .lowered = lowered };
203 }
204
205 fn deinit(self: *Generated) void {
206 for (self.programs) |*program| program.deinit();
207 self.allocator.free(self.programs);
208 self.allocator.free(self.lowered);
209 }
210 };
211
212 pub fn target(
213 allocator: std.mem.Allocator,
214 root: *choir.ir.Operation,
215 schedule: *const preparation.schedule.SchedulePlanAnalysis,
216 generated: []const preparation.kernelization.LoweredKernel,
217 comptime configuration: choir.product.operation.Configuration,
218 ) ![]u8 {
219 var programs = try Generated.init(allocator, generated, configuration);
220 defer programs.deinit();
221 const profile = preparation.readBackendTargetProfile(root);
222 const kernels = try allocator.alloc(records.target.Kernel, programs.lowered.len);
223 defer allocator.free(kernels);
224 var initialized: usize = 0;
225 defer for (kernels[0..initialized]) |kernel_value| {
226 if (kernel_value.abi) |abi| allocator.free(abi.static_arguments);
227 };
228 for (programs.lowered, generated, 0..) |lowered, *source, index| {
229 const work = for (schedule.work_items.items) |item| {
230 if (item.id == lowered.work_item_id) break item;
231 } else return error.UnencodableProduct;
232 kernels[index] = try targetKernel(allocator, lowered, source, work, profile);
233 initialized += 1;
234 }
235 var references = records.reference.Index{ .allocator = allocator, .limit = 0 };
236 defer references.deinit();
237 const row_schedules = preparation.target.readGeneratedRowPipelineSchedules(root);
238 return qualified(allocator, records.target.Record, .target, records.target.Record{
239 .profile = profile,
240 .math_tier = if (profile) |value| value.math_tier else .exact,
241 .generated_scan_schedules = preparation.target.readGeneratedScanSchedules(root),
242 .generated_row_pipeline_schedules = row_schedules,
243 .kernels = kernels,
244 }, &references);
245 }
246
247 fn targetKernel(
248 allocator: std.mem.Allocator,
249 lowered: records.kernel.Lowered,
250 source: *const preparation.kernelization.LoweredKernel,
251 work: preparation.schedule.ScheduleWorkItem,
252 profile: ?preparation.BackendTargetProfile,
253 ) !records.target.Kernel {
254 var dtypes = source.requiredDTypes();
255 dtypes.insert(work.dtype);
256 return .{
257 .lowered = lowered,
258 .work_dtype = work.dtype,
259 .element_count = work.element_count,
260 .required_dtype_bits = dtypes.bits,
261 .required_features = gpu.featureRequirementsForModule(source.program.kernelModule()),
262 .required_subgroup = gpu.subgroupRequirementsForModule(source.program.kernelModule()),
263 .runtime_scalar_argument_count = source.runtimeScalarArgumentCount(),
264 .abi = if (profile) |value| try targetAbi(allocator, source, work, value) else null,
265 };
266 }
267
268 fn targetAbi(
269 allocator: std.mem.Allocator,
270 source: *const preparation.kernelization.LoweredKernel,
271 work: preparation.schedule.ScheduleWorkItem,
272 profile: preparation.BackendTargetProfile,
273 ) !records.target.Abi {
274 const owner = @import("../target/root.zig");
275 const format = profile.artifact_format;
276 const count = try owner.abi.argumentCount(format, source.argument_count);
277 const launch = try owner.abi.launchGeometry(
278 format,
279 work.element_count,
280 source.launchGeometry(),
281 );
282 return .{
283 .argument_count = count,
284 .static_arguments = try owner.abi.staticArguments(
285 allocator,
286 format,
287 work.element_count,
288 launch,
289 ),
290 .launch = launch,
291 .compile_options = owner.compileOptionsForArtifactFormat(format, work.element_count),
292 };
293 }
294
295 fn requireWorkMap(analysis: anytype) !void {
296 if (analysis.work_to_kernel.count() != analysis.kernels.items.len) {
297 return error.UnencodableProduct;
298 }
299 for (analysis.kernels.items, 0..) |item, index| {
300 if (analysis.work_to_kernel.get(item.work_item_id) != index) {
301 return error.UnencodableProduct;
302 }
303 }
304 }
305
306 fn kernelCoverage() void {
307 const product = preparation.kernelization.product;
308 codec.coverage(product.KernelOutline, records.kernel.Outline, &.{});
309 codec.coverage(product.KernelOutlinePlanAnalysis, records.kernel.Outlines, &.{
310 "allocator", "work_to_kernel",
311 });
312 codec.coverage(product.LoweredKernel, records.kernel.Lowered, &.{});
313 codec.coverage(product.KernelizationAnalysis, records.kernel.Generated, &.{
314 "allocator", "context", "work_to_kernel",
315 });
316 }
317
318 fn qualified(
319 allocator: std.mem.Allocator,
320 comptime T: type,
321 stage: @import("../choir/root.zig").publication.Stage,
322 source: anytype,
323 references: *const records.reference.Index,
324 ) ![]u8 {
325 const bytes = try codec.encode(allocator, T, stage, source, references);
326 errdefer allocator.free(bytes);
327 var decoded = try codec.decode(allocator, T, stage, bytes);
328 defer decoded.deinit();
329 try codec.compare(decoded.value, source, references);
330 const compared = try codec.encode(allocator, T, stage, decoded.value, references);
331 defer allocator.free(compared);
332 if (!std.mem.eql(u8, bytes, compared)) return error.UnencodableProduct;
333 return bytes;
334 }