tiny.gpu.wasm

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

1const std = @import("std");2const builtin = @import("builtin");34const backend = @import("root.zig");5const host = @import("host.zig");67const Allocator = std.mem.Allocator;8const browserish = builtin.target.cpu.arch.isWasm() and builtin.target.os.tag == .freestanding;910pub const State = struct {11    allocator: Allocator,1213    pub fn init(allocator: Allocator) State {14        return .{ .allocator = allocator };15    }1617    pub fn deinit(_: *State) void {}1819    pub fn handle(self: *State) backend.BackendHandle {20        return .{21            .ptr = self,22            .vtable = &vtable,23            .kind = .wasm,24        };25    }26};2728pub fn staticCapabilities() backend.BackendCapabilities {29    return .{30        .identity = .{31            .backend = .wasm,32            .family = .webassembly,33            .name = "webassembly",34        },35        .memory = .{36            .min_buffer_alignment = 1,37            .host_visible_device_memory = true,38        },39        .threadgroup = .{40            .max_threads = 1024,41            .max_blocks = .{ std.math.maxInt(u32), std.math.maxInt(u32), std.math.maxInt(u32) },42            .max_threads_per_dim = .{ 1024, 1024, 64 },43            .max_grid_per_dim = .{ std.math.maxInt(u32), std.math.maxInt(u32), std.math.maxInt(u32) },44        },45        .dtypes = backend.DTypeSet.init(&.{ .i1, .i32, .u32, .i64, .u64, .f32, .f64, .key }),46        .layouts = .{47            .row_major = true,48            .compact_strides = true,49            .broadcast_strides = true,50        },51        .runtime = .{52            .driver_loaded = true,53            .device_context = true,54        },55        .artifact_formats = backend.ArtifactFormatSet.init(&.{.webassembly_module}),56    };57}5859fn queryCapabilities(_: *anyopaque) backend.BackendError!backend.BackendCapabilities {60    return staticCapabilities();61}6263fn createArtifact(ptr: *anyopaque, request: backend.CompileRequest) backend.BackendError!backend.KernelArtifact {64    if (request.requested_format != .webassembly_module) return error.UnsupportedArtifactFormat;65    const state: *State = @ptrCast(@alignCast(ptr));66    const bytes = switch (request.payload) {67        .bytes => |payload| payload,68        else => return error.InvalidArtifact,69    };70    if (!isWasmModule(bytes)) return error.InvalidArtifact;7172    var artifact = backend.KernelArtifact.init(state.allocator, .{73        .backend = .wasm,74        .format = .webassembly_module,75        .entry_name = request.kernel_name,76        .argument_count = request.argument_count,77        .scalar_argument_count = request.scalar_argument_count,78        .diagnostic_id = request.diagnostic_id,79    }) catch return error.OutOfMemory;80    errdefer artifact.deinit();81    try artifact.setOwnedBytes(bytes);82    return artifact;83}8485fn loadArtifact(_: *anyopaque, artifact: *const backend.KernelArtifact) backend.BackendError!backend.LoadedArtifact {86    if (artifact.backend != .wasm or artifact.format != .webassembly_module) return error.InvalidArtifact;87    if (artifact.entry_name.len == 0) return error.InvalidArtifact;88    const bytes = switch (artifact.payload) {89        .bytes => |payload| payload,90        else => return error.InvalidArtifact,91    };92    if (!isWasmModule(bytes)) return error.InvalidArtifact;93    if (comptime !browserish) return error.RuntimeUnavailable;9495    const id = Host.artifactLoad(96        host.pointer(bytes),97        bytes.len,98        host.pointer(artifact.entry_name),99        artifact.entry_name.len,100        artifact.argument_count,101    );102    if (id == 0) return error.RuntimeUnavailable;103    return .{104        .id = id,105        .backend = .wasm,106        .format = .webassembly_module,107    };108}109110fn allocateBuffer(_: *anyopaque, request: backend.BufferAllocation) backend.BackendError!backend.BufferHandle {111    if (comptime !browserish) return error.RuntimeUnavailable;112    const id = Host.bufferAlloc(request.byte_size, request.alignment);113    if (id == 0) return error.OutOfMemory;114    return .{115        .id = id,116        .backend = .wasm,117        .byte_size = request.byte_size,118        .ownership = .backend,119    };120}121122fn writeBuffer(_: *anyopaque, request: backend.BufferWriteRequest) backend.BackendError!void {123    if (comptime !browserish) return error.RuntimeUnavailable;124    try host.status(Host.bufferWrite(request.handle.id, host.pointer(request.bytes), request.bytes.len));125}126127fn readBuffer(_: *anyopaque, request: backend.BufferReadRequest) backend.BackendError!void {128    if (comptime !browserish) return error.RuntimeUnavailable;129    try host.status(Host.bufferRead(request.handle.id, host.pointer(request.bytes), request.bytes.len));130}131132fn launch(ptr: *anyopaque, request: backend.LaunchRequest) backend.BackendError!void {133    if (comptime !browserish) return error.RuntimeUnavailable;134    const state: *State = @ptrCast(@alignCast(ptr));135    const loaded = request.loaded_artifact orelse return error.InvalidArtifact;136    if (loaded.backend != .wasm or loaded.format != .webassembly_module) return error.InvalidArtifact;137138    const bindings = state.allocator.alloc(host.Binding, request.buffers.len) catch return error.OutOfMemory;139    defer state.allocator.free(bindings);140    for (request.buffers, bindings) |source, *dest| dest.* = try host.binding(source);141142    const scalars = state.allocator.alloc(host.Scalar, request.scalar_arguments.len) catch return error.OutOfMemory;143    defer state.allocator.free(scalars);144    for (request.scalar_arguments, scalars) |source, *dest| dest.* = host.scalar(source);145146    try host.status(Host.launch(147        loaded.id,148        host.pointer(bindings),149        bindings.len,150        host.pointer(scalars),151        scalars.len,152        request.geometry.grid[0],153        request.geometry.grid[1],154        request.geometry.grid[2],155        request.geometry.threadgroup[0],156        request.geometry.threadgroup[1],157        request.geometry.threadgroup[2],158        request.geometry.dynamic_shared_memory_bytes,159    ));160}161162fn synchronize(_: *anyopaque, _: backend.SyncRequest) backend.BackendError!void {}163164fn destroyObject(_: *anyopaque, id: backend.BackendObjectId) void {165    if (comptime browserish) Host.objectDestroy(id);166}167168fn isWasmModule(bytes: []const u8) bool {169    return bytes.len >= 8 and170        std.mem.eql(u8, bytes[0..4], &.{ 0x00, 0x61, 0x73, 0x6d }) and171        std.mem.eql(u8, bytes[4..8], &.{ 0x01, 0x00, 0x00, 0x00 });172}173174const BrowserHost = struct {175    extern "accy" fn accy_wasm_artifact_load(176        bytes_ptr: usize,177        bytes_len: usize,178        entry_ptr: usize,179        entry_len: usize,180        argument_count: u32,181    ) callconv(.c) u64;182183    extern "accy" fn accy_wasm_buffer_alloc(byte_size: usize, alignment: u32) callconv(.c) u64;184    extern "accy" fn accy_wasm_buffer_write(id: u64, bytes_ptr: usize, byte_count: usize) callconv(.c) i32;185    extern "accy" fn accy_wasm_buffer_read(id: u64, bytes_ptr: usize, byte_count: usize) callconv(.c) i32;186    extern "accy" fn accy_wasm_object_destroy(id: u64) callconv(.c) void;187188    extern "accy" fn accy_wasm_launch(189        artifact_id: u64,190        bindings_ptr: usize,191        binding_count: usize,192        scalars_ptr: usize,193        scalar_count: usize,194        grid_x: u32,195        grid_y: u32,196        grid_z: u32,197        threadgroup_x: u32,198        threadgroup_y: u32,199        threadgroup_z: u32,200        dynamic_shared_memory_bytes: u32,201    ) callconv(.c) i32;202203    fn artifactLoad(bytes_ptr: usize, bytes_len: usize, entry_ptr: usize, entry_len: usize, argument_count: u32) u64 {204        return accy_wasm_artifact_load(bytes_ptr, bytes_len, entry_ptr, entry_len, argument_count);205    }206207    fn bufferAlloc(byte_size: usize, alignment: u32) u64 {208        return accy_wasm_buffer_alloc(byte_size, alignment);209    }210211    fn bufferWrite(id: u64, bytes_ptr: usize, byte_count: usize) i32 {212        return accy_wasm_buffer_write(id, bytes_ptr, byte_count);213    }214215    fn bufferRead(id: u64, bytes_ptr: usize, byte_count: usize) i32 {216        return accy_wasm_buffer_read(id, bytes_ptr, byte_count);217    }218219    fn launch(220        artifact_id: u64,221        bindings_ptr: usize,222        binding_count: usize,223        scalars_ptr: usize,224        scalar_count: usize,225        grid_x: u32,226        grid_y: u32,227        grid_z: u32,228        threadgroup_x: u32,229        threadgroup_y: u32,230        threadgroup_z: u32,231        dynamic_shared_memory_bytes: u32,232    ) i32 {233        return accy_wasm_launch(234            artifact_id,235            bindings_ptr,236            binding_count,237            scalars_ptr,238            scalar_count,239            grid_x,240            grid_y,241            grid_z,242            threadgroup_x,243            threadgroup_y,244            threadgroup_z,245            dynamic_shared_memory_bytes,246        );247    }248249    fn objectDestroy(id: u64) void {250        accy_wasm_object_destroy(id);251    }252};253254const Host = if (browserish) BrowserHost else struct {};255256const vtable = backend.BackendVTable{257    .query_capabilities = queryCapabilities,258    .create_artifact = createArtifact,259    .load_artifact = loadArtifact,260    .allocate_buffer = allocateBuffer,261    .write_buffer = writeBuffer,262    .read_buffer = readBuffer,263    .launch = launch,264    .synchronize = synchronize,265    .destroy_object = destroyObject,266};267268test "wasm backend reports module artifact capabilities" {269    var state = State.init(std.testing.allocator);270    defer state.deinit();271272    const caps = try state.handle().queryCapabilities();273    try std.testing.expectEqual(backend.BackendKind.wasm, caps.identity.backend);274    try std.testing.expectEqual(backend.DeviceFamily.webassembly, caps.identity.family);275    try std.testing.expect(caps.supportsDType(.f32));276    try std.testing.expect(caps.supportsDType(.u64));277    try std.testing.expect(caps.supportsDType(.key));278    try std.testing.expect(caps.supportsArtifactFormat(.webassembly_module));279    try std.testing.expect(!caps.supportsArtifactFormat(.webgpu_wgsl));280}281282test "wasm backend creates owned module artifacts" {283    var state = State.init(std.testing.allocator);284    defer state.deinit();285286    var artifact = try state.handle().createArtifact(.{287        .kernel_name = "copy",288        .requested_format = .webassembly_module,289        .argument_count = 2,290        .required_dtypes = backend.DTypeSet.init(&.{.f32}),291        .payload = .{ .bytes = &.{ 0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00 } },292    });293    defer artifact.deinit();294295    try std.testing.expectEqual(backend.BackendKind.wasm, artifact.backend);296    try std.testing.expectEqual(backend.ArtifactFormat.webassembly_module, artifact.format);297    try std.testing.expectEqual(backend.PayloadOwnership.owned, artifact.payload_ownership);298    try std.testing.expectEqual(@as(u32, 2), artifact.argument_count);299}300301test "wasm backend rejects non-module payloads" {302    var state = State.init(std.testing.allocator);303    defer state.deinit();304305    try std.testing.expectError(error.InvalidArtifact, state.handle().createArtifact(.{306        .kernel_name = "bad",307        .requested_format = .webassembly_module,308        .payload = .{ .bytes = &.{ 0x00, 0x61, 0x73, 0x6d } },309    }));310    try std.testing.expectError(error.UnsupportedArtifactFormat, state.handle().createArtifact(.{311        .kernel_name = "bad",312        .requested_format = .webgpu_wgsl,313        .payload = .{ .bytes = &.{ 0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00 } },314    }));315}