/ concept-collection / matmul-bench
Sign in
concept-collection / matmul-bench
matmul-bench / public / matmul / worker.js
85 lines · 3.1 KBBlameHistoryRaw
1// Classic worker (NOT a module worker) that runs the pthread-capable WASM
2// matmul modules off the main thread. Threaded builds spawn their pthread
3// workers from here as nested workers — which is why this is a plain
4// importScripts worker served from public/, kept out of Vite's module graph.
5// (Mirrors the proven approach in libflame2wasm/web/bench_worker.js.)
6'use strict';
8// Modules and this worker all live in public/matmul/; resolve siblings
9// relative to it so each module and its pthread workers load from here.
10const DIR = self.location.href.replace(/[^/]*$/, '');
12// kind -> { script, factory export name, exported C function }. Every func has
13// signature (aPtr, bPtr, cPtr, n, nthreads) — single-threaded builds ignore
14// nthreads, so the call is uniform.
15const REGISTRY = {
16 'matmul-mt': { file: 'matmul_mt.js', name: 'createMatmulMT', fn: '_matmul_blocked_mt' },
17 'blis-st': { file: 'matmul_blis_st.js', name: 'createMatmulBlisST', fn: '_matmul_blis' },
18 'blis-mt': { file: 'matmul_blis_mt.js', name: 'createMatmulBlisMT', fn: '_matmul_blis' },
19};
21// Same deterministic PRNG as src/methods/random.ts, inlined so this worker is
22// self-contained and every method multiplies bit-identical inputs per size.
23function mulberry32(seed) {
24 let a = seed;
25 return () => {
26 a |= 0;
27 a = (a + 0x6d2b79f5) | 0;
28 let t = Math.imul(a ^ (a >>> 15), 1 | a);
29 t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
30 return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
31 };
34function generateMatrix(n, seed) {
35 const rand = mulberry32(seed);
36 const m = new Float64Array(n * n);
37 for (let i = 0; i < m.length; i++) m[i] = rand() - 0.5;
38 return m;
41const instances = {}; // kind -> { module, fn }
43async function getInstance(kind) {
44 if (!instances[kind]) {
45 const reg = REGISTRY[kind];
46 if (!reg) throw new Error(`Unknown kind: ${kind}`);
47 importScripts(DIR + reg.file);
48 const factory = self[reg.name];
49 const module = await factory({
50 locateFile: (path) => DIR + path,
51 mainScriptUrlOrBlob: DIR + reg.file,
52 });
53 instances[kind] = { module, fn: reg.fn };
54 }
55 return instances[kind];
58onmessage = async (e) => {
59 const { id, kind, n, seedA, seedB, threads } = e.data;
60 try {
61 const { module, fn } = await getInstance(kind);
62 const a = generateMatrix(n, seedA);
63 const b = generateMatrix(n, seedB);
64 const bytes = n * n * 8;
65 const aPtr = module._malloc(bytes);
66 const bPtr = module._malloc(bytes);
67 const cPtr = module._malloc(bytes);
68 try {
69 new Float64Array(module.HEAPF64.buffer, aPtr, n * n).set(a);
70 new Float64Array(module.HEAPF64.buffer, bPtr, n * n).set(b);
71 const t0 = performance.now();
72 module[fn](aPtr, bPtr, cPtr, n, threads);
73 const ms = performance.now() - t0;
74 // st builds allow memory growth (buffer may be replaced) — re-view.
75 const sample = new Float64Array(module.HEAPF64.buffer, cPtr, n * n)[0];
76 postMessage({ id, ms, sample });
77 } finally {
78 module._free(aPtr);
79 module._free(bPtr);
80 module._free(cPtr);
81 }
82 } catch (err) {
83 postMessage({ id, error: err && err.message ? err.message : String(err) });
84 }
85};
moveopenescclose