666db68matmul-bench: browser GEMM benchmark (JS, WebGPU, custom C WASM, libFLAME/BLIS WASM)Jeremy Magland 1import type { MatmulMethod } from './types'
3const TILE = 16
5// Tiled sgemm: each workgroup computes one TILE×TILE block of C, staging
6// TILE×TILE tiles of A and B through workgroup (shared) memory.
7const SHADER = /* wgsl */ `
8struct Dims { n: u32 }
9@group(0) @binding(0) var<uniform> dims: Dims;
10@group(0) @binding(1) var<storage, read> A: array<f32>;
11@group(0) @binding(2) var<storage, read> B: array<f32>;
12@group(0) @binding(3) var<storage, read_write> C: array<f32>;
14var<workgroup> tileA: array<array<f32, ${TILE}>, ${TILE}>;
15var<workgroup> tileB: array<array<f32, ${TILE}>, ${TILE}>;
17@compute @workgroup_size(${TILE}, ${TILE})
18fn main(
19 @builtin(global_invocation_id) gid: vec3<u32>,
20 @builtin(local_invocation_id) lid: vec3<u32>,
21) {
22 let n = dims.n;
23 let row = gid.y;
24 let col = gid.x;
25 var sum: f32 = 0.0;
26 let numTiles = (n + ${TILE}u - 1u) / ${TILE}u;
28 for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
29 let aCol = t * ${TILE}u + lid.x;
30 let bRow = t * ${TILE}u + lid.y;
31 tileA[lid.y][lid.x] = select(0.0, A[row * n + aCol], row < n && aCol < n);
32 tileB[lid.y][lid.x] = select(0.0, B[bRow * n + col], bRow < n && col < n);
33 workgroupBarrier();
34 for (var k: u32 = 0u; k < ${TILE}u; k = k + 1u) {
35 sum = sum + tileA[lid.y][k] * tileB[k][lid.x];
36 }
37 workgroupBarrier();
38 }
40 if (row < n && col < n) {
41 C[row * n + col] = sum;
42 }
43}
44`
46let devicePromise: Promise<GPUDevice> | null = null
48function getDevice(): Promise<GPUDevice> {
49 if (!devicePromise) {
50 devicePromise = (async () => {
51 if (!navigator.gpu) throw new Error('WebGPU not supported')
52 const adapter = await navigator.gpu.requestAdapter()
53 if (!adapter) throw new Error('No WebGPU adapter available')
54 return adapter.requestDevice()
55 })()
56 }
57 return devicePromise
58}
60let pipelinePromise: Promise<{ device: GPUDevice; pipeline: GPUComputePipeline }> | null = null
62function getPipeline() {
63 if (!pipelinePromise) {
64 pipelinePromise = getDevice().then((device) => {
65 const module = device.createShaderModule({ code: SHADER })
66 const pipeline = device.createComputePipeline({
67 layout: 'auto',
68 compute: { module, entryPoint: 'main' },
69 })
70 return { device, pipeline }
71 })
72 }
73 return pipelinePromise
74}
76export const webgpuMatmul: MatmulMethod = {
77 id: 'webgpu',
78 label: 'WebGPU',
79 precision: 'f32',
80 note: 'single precision only — WGSL has no f64',
81 available: () => typeof navigator !== 'undefined' && !!navigator.gpu,
82 async run(n, af64, bf64) {
83 const { device, pipeline } = await getPipeline()
85 const a = Float32Array.from(af64)
86 const b = Float32Array.from(bf64)
87 const bytes = n * n * 4
89 const dimsBuf = device.createBuffer({ size: 4, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST })
90 device.queue.writeBuffer(dimsBuf, 0, new Uint32Array([n]))
92 const aBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST })
93 const bBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST })
94 const cBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC })
95 const readBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ })
97 device.queue.writeBuffer(aBuf, 0, a)
98 device.queue.writeBuffer(bBuf, 0, b)
100 const bindGroup = device.createBindGroup({
101 layout: pipeline.getBindGroupLayout(0),
102 entries: [
103 { binding: 0, resource: { buffer: dimsBuf } },
104 { binding: 1, resource: { buffer: aBuf } },
105 { binding: 2, resource: { buffer: bBuf } },
106 { binding: 3, resource: { buffer: cBuf } },
107 ],
108 })
110 const t0 = performance.now()
111 const encoder = device.createCommandEncoder()
112 const pass = encoder.beginComputePass()
113 pass.setPipeline(pipeline)
114 pass.setBindGroup(0, bindGroup)
115 const groups = Math.ceil(n / TILE)
116 pass.dispatchWorkgroups(groups, groups)
117 pass.end()
118 encoder.copyBufferToBuffer(cBuf, 0, readBuf, 0, bytes)
119 device.queue.submit([encoder.finish()])
121 await readBuf.mapAsync(GPUMapMode.READ)
122 // Timing includes the GPU->CPU readback, since that's what a real caller pays.
123 const ms = performance.now() - t0
124 const sample = new Float32Array(readBuf.getMappedRange())[0]
125 readBuf.unmap()
127 dimsBuf.destroy()
128 aBuf.destroy()
129 bBuf.destroy()
130 cBuf.destroy()
131 readBuf.destroy()
133 return { ms, sample }
134 },
135}