Spaces:
Running
Running
cfahlgren1 HF Staff
convert talking pixel to a sveltekit app that runs pivo-1.4m tts in the browser
3ec6b9e | // WebAssembly SIMD GEMM for the onnx-lite interpreter: X[m, nin] @ W[nin, nout]. | |
| // Weights are copied into wasm memory once (keyed by the weight's data array). | |
| import { MATVEC_WASM_BASE64 } from './kernel-wasm.ts'; | |
| export interface KernelExports { | |
| matvec(w: number, x: number, y: number, nin: number, nout: number): void; | |
| gemm(w: number, x: number, y: number, m: number, nin: number, nout: number): void; | |
| } | |
| export function wasmSimdSupported(): boolean { | |
| const probe = new Uint8Array([0, 97, 115, 109, 1, 0, 0, 0, 1, 5, 1, 96, 0, 1, 123, 3, 2, 1, 0, 10, 10, 1, 8, 0, 65, 0, 253, 15, 253, 98, 11]); | |
| try { | |
| return typeof WebAssembly === 'object' && WebAssembly.validate(probe); | |
| } catch { | |
| return false; | |
| } | |
| } | |
| export function instantiateKernel(memory: WebAssembly.Memory): KernelExports { | |
| const bytes = Uint8Array.from(atob(MATVEC_WASM_BASE64), (c) => c.charCodeAt(0)); | |
| return new WebAssembly.Instance(new WebAssembly.Module(bytes), { env: { memory } }).exports as unknown as KernelExports; | |
| } | |
| export class WasmGemm { | |
| memory: WebAssembly.Memory; | |
| kernel: KernelExports; | |
| top: number; | |
| weights: WeakMap<Float32Array, number>; | |
| scratch: { off: number; bytes: number }; | |
| constructor() { | |
| this.memory = new WebAssembly.Memory({ initial: 64 }); | |
| this.kernel = instantiateKernel(this.memory); | |
| this.top = 0; | |
| this.weights = new WeakMap(); | |
| this.scratch = { off: 0, bytes: 0 }; | |
| } | |
| #reserve(bytes: number): number { | |
| const off = (this.top + 15) & ~15; | |
| const end = off + bytes; | |
| const have = this.memory.buffer.byteLength; | |
| if (end > have) this.memory.grow(Math.ceil((end - have) / 65536)); | |
| this.top = end; | |
| return off; | |
| } | |
| // Returns false when the shape is unsupported (caller falls back to JS). | |
| supports(nin: number, nout: number): boolean { | |
| return nin % 2 === 0 && nout % 4 === 0; | |
| } | |
| // Like matmul, but B is transient (copied each call instead of cached). | |
| matmulTransient(A: Float32Array, aOff: number, m: number, B: Float32Array, bOff: number, nin: number, nout: number, out: Float32Array, outOff: number): void { | |
| const need = (m * nin + nin * nout + m * nout) * 4; | |
| if (this.scratch.bytes < need) this.scratch = { off: this.#reserve(need * 2), bytes: need * 2 }; | |
| const x = this.scratch.off; | |
| const w = x + m * nin * 4; | |
| const y = w + nin * nout * 4; | |
| const buf = this.memory.buffer; | |
| new Float32Array(buf, x, m * nin).set(A.subarray(aOff, aOff + m * nin)); | |
| new Float32Array(buf, w, nin * nout).set(B.subarray(bOff, bOff + nin * nout)); | |
| this.kernel.gemm(w, x, y, m, nin, nout); | |
| out.set(new Float32Array(buf, y, m * nout), outOff); | |
| } | |
| matmul(A: Float32Array, m: number, B: Float32Array, nin: number, nout: number, out: Float32Array): void { | |
| let w = this.weights.get(B); | |
| if (w === undefined) { | |
| w = this.#reserve(B.length * 4); | |
| new Float32Array(this.memory.buffer, w, B.length).set(B); | |
| this.weights.set(B, w); | |
| this.scratch.bytes = 0; // weights were appended after scratch; re-reserve scratch | |
| } | |
| const need = (m * nin + m * nout) * 4; | |
| if (this.scratch.bytes < need) { | |
| this.scratch = { off: this.#reserve(need * 2), bytes: need * 2 }; | |
| } | |
| const x = this.scratch.off; | |
| const y = x + m * nin * 4; | |
| new Float32Array(this.memory.buffer, x, m * nin).set(A.subarray(0, m * nin)); | |
| this.kernel.gemm(w, x, y, m, nin, nout); | |
| out.set(new Float32Array(this.memory.buffer, y, m * nout)); | |
| } | |
| } | |