lib/accy/src/target/payload.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const gpu = @import("gpu");
3 const choir = @import("choir");
4 const choir_abi = @import("choir_abi");
5 const accy_root = @import("../root.zig");
6
7 const gpu_codegen = choir.backends.gpu;
8 const cpu = gpu_codegen.cpu;
9 const spirv = gpu_codegen.spirv;
10 const nvptx = gpu_codegen.nvptx;
11 const metal = gpu_codegen.metal;
12 const webgpu = gpu_codegen.webgpu;
13 const lowering = gpu_codegen.lowering;
14
15 const ir = choir.ir;
16 const passes = choir.passes;
17
18 pub const native_cpu_artifacts_supported = @hasDecl(choir.backends.x86_64, "backend");
19
20 pub const CompileOptions = accy_root.choir.record.target.CompileOptions;
21
22 pub fn compileOptionsForArtifactFormat(format: gpu.ArtifactFormat, element_count: u64) CompileOptions {
23 return .{ .cpu_vector_width = automaticCpuVectorWidth(format, element_count) };
24 }
25
26 fn automaticCpuVectorWidth(format: gpu.ArtifactFormat, element_count: u64) ?u32 {
27 const width: u32 = 4;
28 if (!gpu.artifactFormatIsNativeCpu(format)) return null;
29 if (element_count < width) return null;
30 return width;
31 }
32
33 /// One compiled kernel entry: the payload a backend loads, and the
34 /// push-constant layout its emitter produced (empty outside SPIR-V).
35 pub const Compilation = struct {
36 payload: gpu.CompilePayload,
37 push_constants: choir_abi.PushConstants = .{},
38 };
39
40 /// Vulkan emission leaves float modes unset because this call precedes device selection.
41 pub fn compileKernelForArtifactFormat(
42 allocator: std.mem.Allocator,
43 format: gpu.ArtifactFormat,
44 entry_name: []const u8,
45 module: *ir.Operation,
46 options: CompileOptions,
47 ) gpu.BackendError!Compilation {
48 const compile_module = switch (format) {
49 .cuda_ptx, .vulkan_spirv, .metal_msl, .webgpu_wgsl, .cpu_machine_code, .cpu_object, .webassembly_module => module.clone() catch |err| return compileModuleCloneError(err),
50 else => return error.UnsupportedArtifactFormat,
51 };
52 defer compile_module.erase();
53
54 const payload: gpu.CompilePayload = switch (format) {
55 .cuda_ptx => payload: {
56 try lowerKernelJobForArtifactFormat(allocator, compile_module, format);
57 break :payload .{ .text = try nvptx.ptx.emitPtx(allocator, entry_name, compile_module) };
58 },
59 .vulkan_spirv => {
60 try lowerKernelJobForArtifactFormat(allocator, compile_module, format);
61 const emission = try spirv.emitter.emitPlanWords(allocator, compile_module, entry_name, .{}, .{});
62 return .{
63 .payload = .{ .words_u32 = emission.words },
64 .push_constants = emission.push_constants,
65 };
66 },
67 .metal_msl => .{ .text = try metal.msl.emitMsl(allocator, entry_name, compile_module) },
68 .webgpu_wgsl => .{ .text = try webgpu.wgsl.emitWgsl(allocator, entry_name, compile_module) },
69 .cpu_machine_code => try compileKernelJobToCpuMachineCode(allocator, entry_name, compile_module, options),
70 .cpu_object => try compileKernelJobToCpuObject(allocator, entry_name, compile_module, options),
71 .webassembly_module => try compileKernelJobToWebAssemblyModule(allocator, entry_name, compile_module),
72 else => unreachable,
73 };
74 return .{ .payload = payload };
75 }
76
77 fn compileKernelJobToCpuMachineCode(
78 allocator: std.mem.Allocator,
79 entry_name: []const u8,
80 module: *ir.Operation,
81 options: CompileOptions,
82 ) gpu.BackendError!gpu.CompilePayload {
83 if (comptime !native_cpu_artifacts_supported) return error.UnsupportedOperation;
84
85 const host_module = try cpu.lowerKernelToHostLoop(allocator, module, .{
86 .entry_name = entry_name,
87 .vector_width = options.cpu_vector_width,
88 });
89 defer host_module.erase();
90
91 var x64_backend = choir.backends.x86_64.backend.Backend.init(
92 allocator,
93 host_module.getContext(),
94 .standard,
95 ) catch |err| return cpuCompileError(err);
96 defer x64_backend.deinit();
97
98 var machine_code = x64_backend.compileFunctionToMachineCodeWithRelocations(host_module, entry_name) catch |err| return cpuCompileError(err);
99 defer machine_code.deinit(allocator);
100 if (machine_code.relocations.len != 0) return error.UnsupportedOperation;
101 if (machine_code.data_symbols.len != 0) return error.UnsupportedOperation;
102 if (machine_code.data_relocations.len != 0) return error.UnsupportedOperation;
103 return .{ .bytes = allocator.dupe(u8, machine_code.code) catch return error.OutOfMemory };
104 }
105
106 fn compileKernelJobToCpuObject(
107 allocator: std.mem.Allocator,
108 entry_name: []const u8,
109 module: *ir.Operation,
110 options: CompileOptions,
111 ) gpu.BackendError!gpu.CompilePayload {
112 if (comptime !native_cpu_artifacts_supported) return error.UnsupportedOperation;
113
114 const host_module = try cpu.lowerKernelToHostLoop(allocator, module, .{
115 .entry_name = entry_name,
116 .vector_width = options.cpu_vector_width,
117 });
118 defer host_module.erase();
119
120 var x64_backend = choir.backends.x86_64.backend.Backend.init(
121 allocator,
122 host_module.getContext(),
123 .standard,
124 ) catch |err| return cpuCompileError(err);
125 defer x64_backend.deinit();
126
127 var object_artifact = x64_backend.compileModuleToObjectFile(host_module, entry_name) catch |err| return cpuCompileError(err);
128 defer object_artifact.deinit();
129
130 if (object_artifact.payload.buffers.items.len != 1) return error.InvalidArtifact;
131 if (object_artifact.payload.buffers.items[0].format != .object_file) return error.InvalidArtifact;
132 return .{ .bytes = allocator.dupe(u8, object_artifact.payload.buffers.items[0].bytes) catch return error.OutOfMemory };
133 }
134
135 fn compileKernelJobToWebAssemblyModule(
136 allocator: std.mem.Allocator,
137 entry_name: []const u8,
138 module: *ir.Operation,
139 ) gpu.BackendError!gpu.CompilePayload {
140 const host_module = try cpu.lowerKernelToHostLoop(allocator, module, .{ .entry_name = entry_name });
141 defer host_module.erase();
142
143 var wasm_backend = try choir.backends.wasm.Backend.init(allocator, host_module.getContext());
144 defer wasm_backend.deinit();
145
146 var wasm_artifact = wasm_backend.compileModuleToArtifact(host_module, .{ .entry = entry_name }) catch |err| return wasmCompileError(err);
147 defer wasm_artifact.deinit();
148
149 if (wasm_artifact.payload.buffers.items.len != 1) return error.InvalidArtifact;
150 if (wasm_artifact.payload.buffers.items[0].format != .webassembly_module) return error.InvalidArtifact;
151 return .{ .bytes = allocator.dupe(u8, wasm_artifact.payload.buffers.items[0].bytes) catch return error.OutOfMemory };
152 }
153
154 fn cpuCompileError(err: anyerror) gpu.BackendError {
155 return switch (err) {
156 error.OutOfMemory => error.OutOfMemory,
157 error.UnsupportedArchitecture => error.UnsupportedOperation,
158 else => error.CompilationFailed,
159 };
160 }
161
162 fn wasmCompileError(err: anyerror) gpu.BackendError {
163 return switch (err) {
164 error.OutOfMemory => error.OutOfMemory,
165 error.UnsupportedArchitecture => error.UnsupportedOperation,
166 else => error.CompilationFailed,
167 };
168 }
169
170 fn compileModuleCloneError(err: anyerror) gpu.BackendError {
171 return switch (err) {
172 error.OutOfMemory => error.OutOfMemory,
173 else => error.CompilationFailed,
174 };
175 }
176
177 fn lowerKernelJobForArtifactFormat(
178 allocator: std.mem.Allocator,
179 module: *ir.Operation,
180 format: gpu.ArtifactFormat,
181 ) gpu.BackendError!void {
182 switch (format) {
183 .cuda_ptx => try lowerKernelJobToTarget(allocator, module, .nvptx),
184 .vulkan_spirv => try lowerKernelJobToTarget(allocator, module, .spirv),
185 else => {},
186 }
187 }
188
189 fn lowerKernelJobToTarget(
190 allocator: std.mem.Allocator,
191 module: *ir.Operation,
192 target_lowering: lowering.TargetLowering,
193 ) gpu.BackendError!void {
194 const ctx = module.getContext();
195 if (!ctx.isFrozen()) {
196 gpu_codegen.registerTargetDialects(ctx) catch |err| return targetRegistrationError(err);
197 }
198
199 var pm = passes.PassManager.init(allocator);
200 defer pm.deinit();
201 pm.enableVerifier();
202 lowering.addTargetLoweringPipeline(&pm, target_lowering) catch |err| return targetRegistrationError(err);
203 if (pm.run(module, ctx) == .failure) {
204 emitPassFailure(module, &pm);
205 return error.CompilationFailed;
206 }
207 }
208
209 fn emitPassFailure(module: *ir.Operation, pm: *const passes.PassManager) void {
210 const failure = pm.getLastFailureReproducer() orelse return;
211 const metadata = [_]choir.diagnostics.Metadata{
212 .{ .name = "pipeline", .value = failure.pipeline },
213 .{ .name = "pass", .value = failure.pass_name orelse "dependency preparation" },
214 .{ .name = "target", .value = failure.target_op_name orelse "builtin.module" },
215 .{ .name = "ir", .value = failure.ir },
216 };
217 var diagnostic = module.getContext().emitDiagnostic(.{
218 .severity = .err,
219 .location = module.getLoc(),
220 .operation = module,
221 .message = "Accy target lowering failed",
222 .error_name = failure.verifier_error_name,
223 .metadata = &metadata,
224 });
225 defer diagnostic.deinit();
226 _ = diagnostic.emit() catch {};
227 }
228
229 fn targetRegistrationError(err: anyerror) gpu.BackendError {
230 return switch (err) {
231 error.OutOfMemory => error.OutOfMemory,
232 else => error.CompilationFailed,
233 };
234 }
235
236 fn deinitCompilePayload(allocator: std.mem.Allocator, payload: gpu.CompilePayload) void {
237 switch (payload) {
238 .bytes => |bytes| allocator.free(@constCast(bytes)),
239 .words_u32 => |words| allocator.free(@constCast(words)),
240 .text => |text| allocator.free(@constCast(text)),
241 .none => {},
242 }
243 }
244
245 test "target payload compilation rejects unsupported artifact formats" {
246 try std.testing.expectError(
247 error.UnsupportedArtifactFormat,
248 compileKernelForArtifactFormat(
249 std.testing.allocator,
250 .external,
251 "unsupported",
252 undefined,
253 .{},
254 ),
255 );
256 }
257
258 test "target payload compilation preserves source kernel module" {
259 const testing = std.testing;
260 const allocator = testing.allocator;
261 const kernel = accy_root.kernel;
262
263 var builder_state = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "target_payload_preserves_source_kernel_i32", &.{
264 kernel.dynamicBuffer(.i32),
265 kernel.dynamicBuffer(.i32),
266 });
267 errdefer builder_state.deinit();
268
269 const src = builder_state.argument(0);
270 const dst = builder_state.argument(1);
271 const index = try builder_state.globalId(.x);
272 const value = try builder_state.load(src, index);
273 try builder_state.store(value, dst, index);
274 try builder_state.return_();
275
276 var program = try builder_state.finish();
277 defer program.deinit();
278
279 const before = try program.bodyFingerprint(allocator);
280 const formats = [_]gpu.ArtifactFormat{
281 .cuda_ptx,
282 .vulkan_spirv,
283 .metal_msl,
284 .webgpu_wgsl,
285 .webassembly_module,
286 .cuda_ptx,
287 };
288 for (formats) |format| {
289 const compiled = try compileKernelForArtifactFormat(
290 allocator,
291 format,
292 "target_payload_preserves_source_kernel_i32",
293 program.kernelModule(),
294 .{},
295 );
296 defer deinitCompilePayload(allocator, compiled.payload);
297 try testing.expect(program.kernelModule().getContext().isFrozen());
298 try testing.expectEqual(before, try program.bodyFingerprint(allocator));
299 }
300 }
301
302 test "target payload compilation emits webassembly module for host loop kernel" {
303 const testing = std.testing;
304 const allocator = testing.allocator;
305 const kernel = accy_root.kernel;
306
307 var builder_state = try kernel.Builder.init(allocator, kernel.Builder.Limits.testing, "target_payload_wasm_copy_f32", &.{
308 kernel.dynamicBuffer(.f32),
309 kernel.dynamicBuffer(.f32),
310 });
311 errdefer builder_state.deinit();
312
313 const axis = try builder_state.axis("i", 4);
314 try builder_state.bind(axis, .thread_x);
315 const src = builder_state.argument(0);
316 const dst = builder_state.argument(1);
317 const index = try builder_state.globalId(.x);
318 const value = try builder_state.load(src, index);
319 try builder_state.store(value, dst, index);
320 try builder_state.return_();
321
322 var program = try builder_state.finish();
323 defer program.deinit();
324
325 const before = try program.bodyFingerprint(allocator);
326 const compiled = try compileKernelForArtifactFormat(
327 allocator,
328 .webassembly_module,
329 "target_payload_wasm_copy_f32",
330 program.kernelModule(),
331 .{},
332 );
333 defer deinitCompilePayload(allocator, compiled.payload);
334 try testing.expectEqual(before, try program.bodyFingerprint(allocator));
335
336 const bytes = switch (compiled.payload) {
337 .bytes => |payload| payload,
338 else => return error.ExpectedWebAssemblyModuleBytes,
339 };
340 try testing.expectEqualSlices(u8, &.{ 0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00 }, bytes[0..8]);
341 }