tiny.accy.target.payload

1const std = @import("std");2const gpu = @import("gpu");3const choir = @import("choir");4const choir_abi = @import("choir_abi");5const accy_root = @import("../root.zig");67const gpu_codegen = choir.backends.gpu;8const cpu = gpu_codegen.cpu;9const spirv = gpu_codegen.spirv;10const nvptx = gpu_codegen.nvptx;11const metal = gpu_codegen.metal;12const webgpu = gpu_codegen.webgpu;13const lowering = gpu_codegen.lowering;1415const ir = choir.ir;16const passes = choir.passes;1718pub const native_cpu_artifacts_supported = @hasDecl(choir.backends.x86_64, "backend");1920pub const CompileOptions = accy_root.choir.record.target.CompileOptions;2122pub fn compileOptionsForArtifactFormat(format: gpu.ArtifactFormat, element_count: u64) CompileOptions {23    return .{ .cpu_vector_width = automaticCpuVectorWidth(format, element_count) };24}2526fn 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}3233/// One compiled kernel entry: the payload a backend loads, and the34/// push-constant layout its emitter produced (empty outside SPIR-V).35pub const Compilation = struct {36    payload: gpu.CompilePayload,37    push_constants: choir_abi.PushConstants = .{},38};3940/// Vulkan emission leaves float modes unset because this call precedes device selection.41pub 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();5354    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}7677fn 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;8485    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();9091    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();9798    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}105106fn 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;113114    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();119120    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();126127    var object_artifact = x64_backend.compileModuleToObjectFile(host_module, entry_name) catch |err| return cpuCompileError(err);128    defer object_artifact.deinit();129130    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}134135fn 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();142143    var wasm_backend = try choir.backends.wasm.Backend.init(allocator, host_module.getContext());144    defer wasm_backend.deinit();145146    var wasm_artifact = wasm_backend.compileModuleToArtifact(host_module, .{ .entry = entry_name }) catch |err| return wasmCompileError(err);147    defer wasm_artifact.deinit();148149    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}153154fn 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}161162fn 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}169170fn compileModuleCloneError(err: anyerror) gpu.BackendError {171    return switch (err) {172        error.OutOfMemory => error.OutOfMemory,173        else => error.CompilationFailed,174    };175}176177fn 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}188189fn 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    }198199    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}208209fn 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}228229fn targetRegistrationError(err: anyerror) gpu.BackendError {230    return switch (err) {231        error.OutOfMemory => error.OutOfMemory,232        else => error.CompilationFailed,233    };234}235236fn 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}244245test "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}257258test "target payload compilation preserves source kernel module" {259    const testing = std.testing;260    const allocator = testing.allocator;261    const kernel = accy_root.kernel;262263    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();268269    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_();275276    var program = try builder_state.finish();277    defer program.deinit();278279    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}301302test "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;306307    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();312313    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_();321322    var program = try builder_state.finish();323    defer program.deinit();324325    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));335336    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}

1pub const payload = @import("payload.zig");