lib/gpu/src/webgpu.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const builtin = @import("builtin");
3
4 const backend = @import("root.zig");
5 const host = @import("host.zig");
6
7 const Allocator = std.mem.Allocator;
8 const browserish = builtin.target.cpu.arch.isWasm() and builtin.target.os.tag == .freestanding;
9
10 pub const State = struct {
11 allocator: Allocator,
12
13 pub fn init(allocator: Allocator) State {
14 return .{
15 .allocator = allocator,
16 };
17 }
18
19 pub fn deinit(_: *State) void {}
20
21 pub fn handle(self: *State) backend.BackendHandle {
22 return .{
23 .ptr = self,
24 .vtable = &vtable,
25 .kind = .webgpu,
26 };
27 }
28 };
29
30 pub 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 }
63
64 fn queryCapabilities(_: *anyopaque) backend.BackendError!backend.BackendCapabilities {
65 return staticCapabilities();
66 }
67
68 fn 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 }
78
79 fn 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;
86
87 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 }
99
100 fn 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;
109
110 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 }
124
125 fn 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 }
136
137 fn 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 }
141
142 fn 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 }
146
147 fn 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;
152
153 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);
156
157 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);
160
161 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 }
176
177 fn synchronize(_: *anyopaque, _: backend.SyncRequest) backend.BackendError!void {}
178
179 fn destroyObject(_: *anyopaque, id: backend.BackendObjectId) void {
180 if (comptime browserish) Host.objectDestroy(id);
181 }
182
183 const 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;
191
192 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;
196
197 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;
211
212 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 }
215
216 fn bufferAlloc(byte_size: usize, alignment: u32) u64 {
217 return accy_webgpu_buffer_alloc(byte_size, alignment);
218 }
219
220 fn bufferWrite(id: u64, bytes_ptr: usize, byte_count: usize) i32 {
221 return accy_webgpu_buffer_write(id, bytes_ptr, byte_count);
222 }
223
224 fn bufferRead(id: u64, bytes_ptr: usize, byte_count: usize) i32 {
225 return accy_webgpu_buffer_read(id, bytes_ptr, byte_count);
226 }
227
228 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 }
257
258 fn objectDestroy(id: u64) void {
259 accy_webgpu_object_destroy(id);
260 }
261 };
262
263 const Host = if (browserish) BrowserHost else struct {};
264
265 const 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 };
276
277 test "webgpu backend reports wgsl capabilities" {
278 const testing = std.testing;
279
280 var state = State.init(testing.allocator);
281 defer state.deinit();
282 const caps = try state.handle().queryCapabilities();
283
284 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 }
295
296 test "webgpu backend creates owned wgsl artifacts" {
297 const testing = std.testing;
298
299 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();
309
310 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 }
316
317 test "webgpu backend rejects non-wgsl payloads" {
318 const testing = std.testing;
319
320 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 }
333
334 test "webgpu backend runtime entry points are unavailable off the browser host" {
335 const testing = std.testing;
336
337 var state = State.init(testing.allocator);
338 defer state.deinit();
339 const handle = state.handle();
340
341 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();
348
349 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 }