tiny.gpu.webgpu

1pub const webgpu = @import("webgpu.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 .{15            .allocator = allocator,16        };17    }1819    pub fn deinit(_: *State) void {}2021    pub fn handle(self: *State) backend.BackendHandle {22        return .{23            .ptr = self,24            .vtable = &vtable,25            .kind = .webgpu,26        };27    }28};2930pub fn staticCapabilities() backend.BackendCapabilities {31    return .{32        .identity = .{33            .backend = .webgpu,34            .family = .webgpu,35            .name = "webgpu",36        },37        .memory = .{38            .shared_memory_per_threadgroup_bytes = 16 * 1024,39            .min_buffer_alignment = 4,40        },41        .threadgroup = .{42            .max_threads = 256,43            .max_blocks = .{ 65_535, 65_535, 65_535 },44            .max_threads_per_dim = .{ 256, 256, 64 },45            .max_grid_per_dim = .{ 65_535, 65_535, 65_535 },46            .shared_memory_bytes = 16 * 1024,47        },48        .dtypes = backend.DTypeSet.init(&.{ .i1, .i32, .u32, .f32, .key }),49        .layouts = .{50            .row_major = true,51            .compact_strides = true,52            .broadcast_strides = true,53            .tiled = true,54            .opaque_backend_layouts = true,55        },56        .runtime = .{57            .driver_loaded = browserish,58            .device_context = browserish,59        },60        .artifact_formats = backend.ArtifactFormatSet.init(&.{.webgpu_wgsl}),61    };62}6364fn queryCapabilities(_: *anyopaque) backend.BackendError!backend.BackendCapabilities {65    return staticCapabilities();66}6768fn createArtifact(ptr: *anyopaque, request: backend.CompileRequest) backend.BackendError!backend.KernelArtifact {69    if (request.requested_format != .webgpu_wgsl) return error.UnsupportedOperation;70    const state: *State = @ptrCast(@alignCast(ptr));71    return switch (request.payload) {72        .text => |source| createWgslArtifact(state, request, source),73        .bytes => |source| createWgslArtifact(state, request, source),74        .none => error.UnsupportedOperation,75        .words_u32 => error.UnsupportedArtifactFormat,76    };77}7879fn createWgslArtifact(80    state: *State,81    request: backend.CompileRequest,82    source: []const u8,83) backend.BackendError!backend.KernelArtifact {84    if (request.kernel_name.len == 0) return error.InvalidArtifact;85    if (source.len == 0) return error.InvalidArtifact;8687    var artifact = backend.KernelArtifact.init(state.allocator, .{88        .backend = .webgpu,89        .format = .webgpu_wgsl,90        .entry_name = request.kernel_name,91        .argument_count = request.argument_count,92        .scalar_argument_count = request.scalar_argument_count,93        .diagnostic_id = request.diagnostic_id,94    }) catch return error.OutOfMemory;95    errdefer artifact.deinit();96    try artifact.setOwnedText(source);97    return artifact;98}99100fn loadArtifact(_: *anyopaque, artifact: *const backend.KernelArtifact) backend.BackendError!backend.LoadedArtifact {101    if (artifact.backend != .webgpu or artifact.format != .webgpu_wgsl) return error.InvalidArtifact;102    if (artifact.entry_name.len == 0) return error.InvalidArtifact;103    const source = switch (artifact.payload) {104        .text => |payload| payload,105        else => return error.InvalidArtifact,106    };107    if (source.len == 0) return error.InvalidArtifact;108    if (comptime !browserish) return error.RuntimeUnavailable;109110    const id = Host.artifactLoad(111        host.pointer(source),112        source.len,113        host.pointer(artifact.entry_name),114        artifact.entry_name.len,115        artifact.argument_count,116    );117    if (id == 0) return error.RuntimeUnavailable;118    return .{119        .id = id,120        .backend = .webgpu,121        .format = .webgpu_wgsl,122    };123}124125fn allocateBuffer(_: *anyopaque, request: backend.BufferAllocation) backend.BackendError!backend.BufferHandle {126    if (comptime !browserish) return error.RuntimeUnavailable;127    const id = Host.bufferAlloc(request.byte_size, request.alignment);128    if (id == 0) return error.OutOfMemory;129    return .{130        .id = id,131        .backend = .webgpu,132        .byte_size = request.byte_size,133        .ownership = .backend,134    };135}136137fn writeBuffer(_: *anyopaque, request: backend.BufferWriteRequest) backend.BackendError!void {138    if (comptime !browserish) return error.RuntimeUnavailable;139    try host.status(Host.bufferWrite(request.handle.id, host.pointer(request.bytes), request.bytes.len));140}141142fn readBuffer(_: *anyopaque, request: backend.BufferReadRequest) backend.BackendError!void {143    if (comptime !browserish) return error.RuntimeUnavailable;144    try host.status(Host.bufferRead(request.handle.id, host.pointer(request.bytes), request.bytes.len));145}146147fn launch(ptr: *anyopaque, request: backend.LaunchRequest) backend.BackendError!void {148    if (comptime !browserish) return error.RuntimeUnavailable;149    const state: *State = @ptrCast(@alignCast(ptr));150    const loaded = request.loaded_artifact orelse return error.InvalidArtifact;151    if (loaded.backend != .webgpu or loaded.format != .webgpu_wgsl) return error.InvalidArtifact;152153    const bindings = state.allocator.alloc(host.Binding, request.buffers.len) catch return error.OutOfMemory;154    defer state.allocator.free(bindings);155    for (request.buffers, bindings) |source, *dest| dest.* = try host.binding(source);156157    const scalars = state.allocator.alloc(host.Scalar, request.scalar_arguments.len) catch return error.OutOfMemory;158    defer state.allocator.free(scalars);159    for (request.scalar_arguments, scalars) |source, *dest| dest.* = host.scalar(source);160161    try host.status(Host.launch(162        loaded.id,163        host.pointer(bindings),164        bindings.len,165        host.pointer(scalars),166        scalars.len,167        request.geometry.grid[0],168        request.geometry.grid[1],169        request.geometry.grid[2],170        request.geometry.threadgroup[0],171        request.geometry.threadgroup[1],172        request.geometry.threadgroup[2],173        request.geometry.dynamic_shared_memory_bytes,174    ));175}176177fn synchronize(_: *anyopaque, _: backend.SyncRequest) backend.BackendError!void {}178179fn destroyObject(_: *anyopaque, id: backend.BackendObjectId) void {180    if (comptime browserish) Host.objectDestroy(id);181}182183const BrowserHost = struct {184    extern "accy" fn accy_webgpu_artifact_load(185        source_ptr: usize,186        source_len: usize,187        entry_ptr: usize,188        entry_len: usize,189        argument_count: u32,190    ) callconv(.c) u64;191192    extern "accy" fn accy_webgpu_buffer_alloc(byte_size: usize, alignment: u32) callconv(.c) u64;193    extern "accy" fn accy_webgpu_buffer_write(id: u64, bytes_ptr: usize, byte_count: usize) callconv(.c) i32;194    extern "accy" fn accy_webgpu_buffer_read(id: u64, bytes_ptr: usize, byte_count: usize) callconv(.c) i32;195    extern "accy" fn accy_webgpu_object_destroy(id: u64) callconv(.c) void;196197    extern "accy" fn accy_webgpu_launch(198        artifact_id: u64,199        bindings_ptr: usize,200        binding_count: usize,201        scalars_ptr: usize,202        scalar_count: usize,203        grid_x: u32,204        grid_y: u32,205        grid_z: u32,206        threadgroup_x: u32,207        threadgroup_y: u32,208        threadgroup_z: u32,209        dynamic_shared_memory_bytes: u32,210    ) callconv(.c) i32;211212    fn artifactLoad(source_ptr: usize, source_len: usize, entry_ptr: usize, entry_len: usize, argument_count: u32) u64 {213        return accy_webgpu_artifact_load(source_ptr, source_len, entry_ptr, entry_len, argument_count);214    }215216    fn bufferAlloc(byte_size: usize, alignment: u32) u64 {217        return accy_webgpu_buffer_alloc(byte_size, alignment);218    }219220    fn bufferWrite(id: u64, bytes_ptr: usize, byte_count: usize) i32 {221        return accy_webgpu_buffer_write(id, bytes_ptr, byte_count);222    }223224    fn bufferRead(id: u64, bytes_ptr: usize, byte_count: usize) i32 {225        return accy_webgpu_buffer_read(id, bytes_ptr, byte_count);226    }227228    fn launch(229        artifact_id: u64,230        bindings_ptr: usize,231        binding_count: usize,232        scalars_ptr: usize,233        scalar_count: usize,234        grid_x: u32,235        grid_y: u32,236        grid_z: u32,237        threadgroup_x: u32,238        threadgroup_y: u32,239        threadgroup_z: u32,240        dynamic_shared_memory_bytes: u32,241    ) i32 {242        return accy_webgpu_launch(243            artifact_id,244            bindings_ptr,245            binding_count,246            scalars_ptr,247            scalar_count,248            grid_x,249            grid_y,250            grid_z,251            threadgroup_x,252            threadgroup_y,253            threadgroup_z,254            dynamic_shared_memory_bytes,255        );256    }257258    fn objectDestroy(id: u64) void {259        accy_webgpu_object_destroy(id);260    }261};262263const Host = if (browserish) BrowserHost else struct {};264265const vtable = backend.BackendVTable{266    .query_capabilities = queryCapabilities,267    .create_artifact = createArtifact,268    .load_artifact = loadArtifact,269    .allocate_buffer = allocateBuffer,270    .write_buffer = writeBuffer,271    .read_buffer = readBuffer,272    .launch = launch,273    .synchronize = synchronize,274    .destroy_object = destroyObject,275};276277test "webgpu backend reports wgsl capabilities" {278    const testing = std.testing;279280    var state = State.init(testing.allocator);281    defer state.deinit();282    const caps = try state.handle().queryCapabilities();283284    try testing.expectEqual(backend.BackendKind.webgpu, caps.identity.backend);285    try testing.expectEqual(backend.DeviceFamily.webgpu, caps.identity.family);286    try testing.expect(caps.supportsArtifactFormat(.webgpu_wgsl));287    try testing.expect(caps.supportsDType(.f32));288    try testing.expect(caps.supportsDType(.i32));289    try testing.expect(caps.supportsDType(.key));290    try testing.expect(!caps.supportsDType(.f16));291    try testing.expect(!caps.subgroup.supported);292    try testing.expectEqual(@as(u32, 256), caps.threadgroup.max_threads);293    try testing.expect(!caps.runtime.driver_loaded);294}295296test "webgpu backend creates owned wgsl artifacts" {297    const testing = std.testing;298299    var state = State.init(testing.allocator);300    defer state.deinit();301    var artifact = try state.handle().createArtifact(.{302        .kernel_name = "accy_test",303        .requested_format = .webgpu_wgsl,304        .argument_count = 2,305        .required_dtypes = backend.DTypeSet.init(&.{.f32}),306        .payload = .{ .text = "@compute fn accy_test() {}" },307    });308    defer artifact.deinit();309310    try testing.expectEqual(backend.BackendKind.webgpu, artifact.backend);311    try testing.expectEqual(backend.ArtifactFormat.webgpu_wgsl, artifact.format);312    try testing.expectEqual(@as(u32, 2), artifact.argument_count);313    try testing.expectEqual(backend.PayloadOwnership.owned, artifact.payload_ownership);314    try testing.expectEqualStrings("@compute fn accy_test() {}", artifact.payload.text);315}316317test "webgpu backend rejects non-wgsl payloads" {318    const testing = std.testing;319320    var state = State.init(testing.allocator);321    defer state.deinit();322    try testing.expectError(error.UnsupportedArtifactFormat, state.handle().createArtifact(.{323        .kernel_name = "accy_test",324        .requested_format = .webgpu_wgsl,325        .payload = .{ .words_u32 = &.{ 0x07230203, 0 } },326    }));327    try testing.expectError(error.UnsupportedArtifactFormat, state.handle().createArtifact(.{328        .kernel_name = "accy_test",329        .requested_format = .vulkan_spirv,330        .payload = .{ .text = "@compute fn accy_test() {}" },331    }));332}333334test "webgpu backend runtime entry points are unavailable off the browser host" {335    const testing = std.testing;336337    var state = State.init(testing.allocator);338    defer state.deinit();339    const handle = state.handle();340341    var artifact = try handle.createArtifact(.{342        .kernel_name = "accy_test",343        .requested_format = .webgpu_wgsl,344        .argument_count = 1,345        .payload = .{ .text = "@compute fn accy_test() {}" },346    });347    defer artifact.deinit();348349    try testing.expectError(error.RuntimeUnavailable, handle.loadArtifact(&artifact));350    try testing.expectError(error.RuntimeUnavailable, handle.allocateBuffer(.{351        .byte_size = 16,352        .alignment = 4,353        .dtype = .f32,354        .element_count = 4,355    }));356    var bytes = @as([16]u8, @splat(0));357    const buffer = backend.BufferHandle{ .id = 1, .backend = .webgpu, .byte_size = 16, .ownership = .backend };358    try testing.expectError(error.RuntimeUnavailable, handle.writeBuffer(.{ .handle = buffer, .bytes = bytes[0..] }));359    try testing.expectError(error.RuntimeUnavailable, handle.readBuffer(.{ .handle = buffer, .bytes = bytes[0..] }));360    const launch_bindings = [_]backend.BufferBinding{.{361        .handle = buffer,362        .access = .read_write,363        .ownership = .backend,364        .byte_size = buffer.byte_size,365    }};366    try testing.expectError(error.RuntimeUnavailable, handle.launch(.{367        .artifact = &artifact,368        .loaded_artifact = .{ .id = 1, .backend = .webgpu, .format = .webgpu_wgsl },369        .buffers = launch_bindings[0..],370        .scalar_arguments = &.{},371        .geometry = .{ .grid = .{ 1, 1, 1 }, .threadgroup = .{ 1, 1, 1 } },372    }));373}