/ concept-collection / matmul-bench
Sign in
concept-collection / matmul-bench
matmul-bench / src / methods / wasmMatmul.ts
68 lines · 2.5 KBCodeBlameHistory
666db68matmul-bench: browser GEMM benchmark (JS, WebGPU, custom C WASM, libFLAME/BLIS WASM)Jeremy Magland 1// In a Vite bundle the .wasm cannot be located next to the glue at runtime, so
2// we resolve its URL via Vite's `?url` import and feed it to Emscripten's
3// `locateFile` (same pattern qhull-wasm-demo uses for qhull-wasm).
4import createMatmulModule, { type MatmulExports } from '../../wasm/dist/matmul.js'
5import wasmUrl from '../../wasm/dist/matmul.wasm?url'
6import type { MatmulMethod } from './types'
8let modulePromise: Promise<MatmulExports> | null = null
10function getModule(): Promise<MatmulExports> {
11 if (!modulePromise) {
12 modulePromise = createMatmulModule({ locateFile: () => wasmUrl })
13 }
14 return modulePromise
17async function run(kernel: '_matmul_naive' | '_matmul_blocked', n: number, a: Float64Array, b: Float64Array) {
18 const mod = await getModule()
19 const bytes = n * n * 8
20 const aPtr = mod._malloc(bytes)
21 const bPtr = mod._malloc(bytes)
22 const cPtr = mod._malloc(bytes)
23 try {
24 new Float64Array(mod.HEAPF64.buffer, aPtr, n * n).set(a)
25 new Float64Array(mod.HEAPF64.buffer, bPtr, n * n).set(b)
26 const t0 = performance.now()
27 mod[kernel](aPtr, bPtr, cPtr, n)
28 const ms = performance.now() - t0
29 // Memory may have grown (and the buffer been replaced) during the call,
30 // so re-view HEAPF64 rather than reuse a view captured before the call.
31 const sample = new Float64Array(mod.HEAPF64.buffer, cPtr, n * n)[0]
32 return { ms, sample }
33 } finally {
34 mod._free(aPtr)
35 mod._free(bPtr)
36 mod._free(cPtr)
37 }
40export const wasmNaiveMatmul: MatmulMethod = {
41 id: 'wasm-naive',
42 label: 'WASM (C, naive)',
43 precision: 'f64',
44 worker: true,
45 available: () => typeof WebAssembly !== 'undefined',
46 run: (n, a, b) => run('_matmul_naive', n, a, b),
49export const wasmBlockedMatmul: MatmulMethod = {
50 id: 'wasm-blocked',
51 label: 'WASM (C, blocked+SIMD)',
52 precision: 'f64',
53 worker: true,
54 available: () => typeof WebAssembly !== 'undefined',
55 run: (n, a, b) => run('_matmul_blocked', n, a, b),
58// The same blocked+SIMD kernel, parallelized over rows with WASM pthreads.
59// Runs through threadedClient's classic worker (see blisMatmul for the shared
60// mechanism); needs crossOriginIsolated for SharedArrayBuffer.
61export const wasmBlockedMtMatmul: MatmulMethod = {
62 id: 'wasm-blocked-mt',
63 label: 'WASM (C, blocked+SIMD, threaded)',
64 precision: 'f64',
65 threadedKind: 'matmul-mt',
66 note: 'WASM threads — needs cross-origin isolation (SharedArrayBuffer)',
67 available: () => typeof WebAssembly !== 'undefined' && globalThis.crossOriginIsolated === true,
moveopenescclose