lib/gpu/src/wasm.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 .{ .allocator = allocator };
15 }
16
17 pub fn deinit(_: *State) void {}
18
19 pub fn handle(self: *State) backend.BackendHandle {
20 return .{
21 .ptr = self,
22 .vtable = &vtable,
23 .kind = .wasm,
24 };
25 }
26 };
27
28 pub 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 }
58
59 fn queryCapabilities(_: *anyopaque) backend.BackendError!backend.BackendCapabilities {
60 return staticCapabilities();
61 }
62
63 fn 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;
71
72 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 }
84
85 fn 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;
94
95 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 }
109
110 fn 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 }
121
122 fn 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 }
126
127 fn 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 }
131
132 fn 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;
137
138 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);
141
142 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);
145
146 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 }
161
162 fn synchronize(_: *anyopaque, _: backend.SyncRequest) backend.BackendError!void {}
163
164 fn destroyObject(_: *anyopaque, id: backend.BackendObjectId) void {
165 if (comptime browserish) Host.objectDestroy(id);
166 }
167
168 fn isWasmModule(bytes: []const u8) bool {
169 return bytes.len >= 8 and
170 std.mem.eql(u8, bytes[0..4], &.{ 0x00, 0x61, 0x73, 0x6d }) and
171 std.mem.eql(u8, bytes[4..8], &.{ 0x01, 0x00, 0x00, 0x00 });
172 }
173
174 const 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;
182
183 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;
187
188 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;
202
203 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 }
206
207 fn bufferAlloc(byte_size: usize, alignment: u32) u64 {
208 return accy_wasm_buffer_alloc(byte_size, alignment);
209 }
210
211 fn bufferWrite(id: u64, bytes_ptr: usize, byte_count: usize) i32 {
212 return accy_wasm_buffer_write(id, bytes_ptr, byte_count);
213 }
214
215 fn bufferRead(id: u64, bytes_ptr: usize, byte_count: usize) i32 {
216 return accy_wasm_buffer_read(id, bytes_ptr, byte_count);
217 }
218
219 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 }
248
249 fn objectDestroy(id: u64) void {
250 accy_wasm_object_destroy(id);
251 }
252 };
253
254 const Host = if (browserish) BrowserHost else struct {};
255
256 const 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 };
267
268 test "wasm backend reports module artifact capabilities" {
269 var state = State.init(std.testing.allocator);
270 defer state.deinit();
271
272 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 }
281
282 test "wasm backend creates owned module artifacts" {
283 var state = State.init(std.testing.allocator);
284 defer state.deinit();
285
286 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();
294
295 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 }
300
301 test "wasm backend rejects non-module payloads" {
302 var state = State.init(std.testing.allocator);
303 defer state.deinit();
304
305 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 }