mambacode.js 1.0.0 → 1.0.2

package/src/training/autograd.ts

@@ -0,0 +1,135 @@
+/**
+ * autograd.ts – Lightweight tape-based automatic differentiation engine.
+ */
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+const _gpu = globalThis as any;
+
+interface TapeEntry {
+  backward: () => void | Promise<void>;
+}
+
+let _tape: TapeEntry[] = [];
+let _gradEnabled = true;
+
+export class Tensor {
+    data: GPUBuffer | null;
+    shape: number[];
+    numel: number;
+    requiresGrad: boolean;
+    grad: GPUBuffer | null;
+    _gradFn: number | null;
+
+    constructor(data: GPUBuffer | null, shape: number[], requiresGrad = false) {
+        this.data         = data;
+        this.shape        = shape;
+        this.numel        = shape.reduce((a, b) => a * b, 1);
+        this.requiresGrad = requiresGrad;
+        this.grad         = null;
+        this._gradFn      = null;
+    }
+
+    get byteSize(): number { return this.numel * 4; }
+
+    zeroGrad(device: GPUDevice): void {
+        if (this.grad) {
+            device.queue.writeBuffer(this.grad, 0, new Float32Array(this.numel));
+        }
+    }
+
+    destroy(): void {
+        this.data?.destroy();
+        this.grad?.destroy();
+        this.data = null;
+        this.grad = null;
+    }
+}
+
+export function enableGrad(): void  { _gradEnabled = true;  }
+export function noGrad(): void      { _gradEnabled = false; }
+export function clearTape(): void   { _tape = []; }
+
+export function recordOperation(backwardFn: () => void | Promise<void>): number {
+    if (!_gradEnabled) return -1;
+    _tape.push({ backward: backwardFn });
+    return _tape.length - 1;
+}
+
+export async function backward(): Promise<void> {
+    for (let i = _tape.length - 1; i >= 0; i--) {
+        await _tape[i]!.backward();
+    }
+    clearTape();
+}
+
+export function ensureGradBuffer(device: GPUDevice, tensor: Tensor): void {
+    if (!tensor.grad) {
+        const STORAGE_USAGE: number = (_gpu.GPUBufferUsage?.STORAGE ?? 0x80) |
+                                      (_gpu.GPUBufferUsage?.COPY_DST ?? 0x08) |
+                                      (_gpu.GPUBufferUsage?.COPY_SRC ?? 0x04);
+        tensor.grad = device.createBuffer({
+            size  : tensor.byteSize,
+            usage : STORAGE_USAGE,
+        });
+        device.queue.writeBuffer(tensor.grad, 0, new Float32Array(tensor.numel));
+    }
+}
+
+export function allocateGradients(device: GPUDevice, tensors: Tensor[]): void {
+    for (const t of tensors) {
+        if (t.requiresGrad) ensureGradBuffer(device, t);
+    }
+}
+
+export function zeroGradients(device: GPUDevice, tensors: Tensor[]): void {
+    for (const t of tensors) {
+        if (t.grad) {
+            device.queue.writeBuffer(t.grad, 0, new Float32Array(t.numel));
+        }
+    }
+}
+
+export function onesLikeScalar(device: GPUDevice): GPUBuffer {
+    const USAGE: number = (_gpu.GPUBufferUsage?.STORAGE ?? 0x80) |
+                          (_gpu.GPUBufferUsage?.COPY_DST ?? 0x08);
+    const buf = device.createBuffer({
+        size  : 4,
+        usage : USAGE,
+        mappedAtCreation: true,
+    });
+    new Float32Array(buf.getMappedRange()).set([1.0]);
+    buf.unmap();
+    return buf;
+}
+
+export function crossEntropyLoss(logits: Float32Array, targetId: number): number {
+    let maxLogit = -Infinity;
+    for (let i = 0; i < logits.length; i++) {
+        if (logits[i]! > maxLogit) maxLogit = logits[i]!;
+    }
+    let sumExp = 0;
+    for (let i = 0; i < logits.length; i++) {
+        sumExp += Math.exp(logits[i]! - maxLogit);
+    }
+    const logSumExp = Math.log(sumExp) + maxLogit;
+    return logSumExp - logits[targetId]!;
+}
+
+export function crossEntropyGrad(logits: Float32Array, targetId: number): Float32Array {
+    let maxLogit = -Infinity;
+    for (let i = 0; i < logits.length; i++) {
+        if (logits[i]! > maxLogit) maxLogit = logits[i]!;
+    }
+    let sumExp = 0;
+    const exp_shifted = new Float32Array(logits.length);
+    for (let i = 0; i < logits.length; i++) {
+        exp_shifted[i] = Math.exp(logits[i]! - maxLogit);
+        sumExp += exp_shifted[i]!;
+    }
+    const probs = new Float32Array(logits.length);
+    for (let i = 0; i < logits.length; i++) {
+        probs[i] = exp_shifted[i]! / sumExp;
+    }
+    probs[targetId] = (probs[targetId] ?? 0) - 1.0;
+    return probs;
+}

package/src/training/{trainer.js → trainer.ts}

@@ -1,24 +1,5 @@
 /**
- * trainer.js – MambaTrainer class
- *
- * Exposes the high-level training API described in the problem statement:
- *
- *   const trainer = new MambaTrainer(model);
- *   await trainer.train(codeSnippet, {
- *     learningRate : 1e-4,
- *     epochs       : 5,
- *     device       : "webgpu",
- *   });
- *
- * The trainer implements:
- *   • Tokenisation of the input code string
- *   • Chunked sequence batching
- *   • Forward pass (next-token prediction / language modelling)
- *   • Cross-entropy loss computation (on CPU for logit read-back)
- *   • Gradient back-propagation via the autograd tape
- *   • AdamW weight update dispatched as GPU compute passes
- *   • Gradient clipping (global L2 norm)
- *   • WSLA mode (fine-tune only B and C for rapid local adaptation)
+ * trainer.ts – MambaTrainer class
  */
 
 import {
@@ -28,71 +9,79 @@ import {
     createComputePipeline,
     createBindGroup,
     dispatchKernel,
-    readBuffer,
-    uploadBuffer,
     cdiv,
 } from '../utils/gpu_utils.js';
 
 import { crossEntropyLoss, crossEntropyGrad } from './autograd.js';
 import { WEIGHT_UPDATE_WGSL, GRAD_CLIP_WGSL } from '../kernels/weight_update.js';
+import { MambaModel, MambaModelConfig } from '../model/mamba_model.js';
+import { BPETokenizer } from '../tokenizer/bpe.js';
+import { BlockParam } from '../model/mamba_block.js';
+
+export interface TrainOptions {
+  learningRate?: number;
+  epochs?: number;
+  batchSize?: number;
+  seqLen?: number;
+  maxGradNorm?: number;
+  weightDecay?: number;
+  beta1?: number;
+  beta2?: number;
+  eps?: number;
+  wsla?: boolean;
+  onEpochEnd?: ((epoch: number, loss: number) => void) | null;
+}
+
+interface AdamMoments {
+  m: GPUBuffer;
+  v: GPUBuffer;
+}
+
+interface AdamHyperparams {
+  learningRate: number;
+  weightDecay: number;
+  beta1: number;
+  beta2: number;
+  eps: number;
+  beta1_t: number;
+  beta2_t: number;
+}
+
+// Re-export to satisfy import in other files
+export type { MambaModelConfig };
 
 export class MambaTrainer {
-    /**
-     * @param {import('../model/mamba_model.js').MambaModel} model
-     * @param {import('../tokenizer/bpe.js').BPETokenizer}  [tokenizer]
-     */
-    constructor(model, tokenizer = null) {
+    model: MambaModel;
+    tokenizer: BPETokenizer | null;
+    device: GPUDevice;
+    private _moments: AdamMoments[] | null;
+    private _step: number;
+    private _adamwPipeline: GPUComputePipeline;
+    private _clipReducePipeline: GPUComputePipeline;
+    private _clipScalePipeline: GPUComputePipeline;
+
+    constructor(model: MambaModel, tokenizer: BPETokenizer | null = null) {
         this.model     = model;
         this.tokenizer = tokenizer;
         this.device    = model.device;
 
-        // AdamW state (first and second moments) – one entry per parameter
         this._moments = null;
-
-        // Step counter for bias correction
         this._step = 0;
 
-        // Compile optimizer pipelines once
         this._adamwPipeline   = createComputePipeline(this.device, WEIGHT_UPDATE_WGSL, 'adamw_update');
         this._clipReducePipeline = createComputePipeline(this.device, GRAD_CLIP_WGSL, 'grad_norm_reduce');
         this._clipScalePipeline  = createComputePipeline(this.device, GRAD_CLIP_WGSL, 'grad_clip_scale');
     }
 
-    // ─── Initialise optimizer state ───────────────────────────────────────────
-
-    /**
-     * Lazily allocate Adam moment buffers (zeroed GPU storage).
-     */
-    _initMoments() {
+    private _initMoments(): void {
         if (this._moments) return;
         this._moments = this.model.parameters().map(p => ({
-            m: createEmptyStorageBuffer(this.device, p.numel * 4, false),  // first moment
-            v: createEmptyStorageBuffer(this.device, p.numel * 4, false),  // second moment
+            m: createEmptyStorageBuffer(this.device, p.numel * 4, false),
+            v: createEmptyStorageBuffer(this.device, p.numel * 4, false),
         }));
     }
 
-    // ─── Public training API ─────────────────────────────────────────────────
-
-    /**
-     * Train on a code snippet (language modelling objective: predict next token).
-     *
-     * @param {string|number[]} input       – raw code string OR pre-tokenised IDs
-     * @param {{
-     *   learningRate ?: number,
-     *   epochs       ?: number,
-     *   batchSize    ?: number,
-     *   seqLen       ?: number,
-     *   maxGradNorm  ?: number,
-     *   weightDecay  ?: number,
-     *   beta1        ?: number,
-     *   beta2        ?: number,
-     *   eps          ?: number,
-     *   wsla         ?: boolean,
-     *   onEpochEnd   ?: (epoch: number, loss: number) => void,
-     * }} [opts]
-     * @returns {Promise<number[]>}  – per-epoch average losses
-     */
-    async train(input, opts = {}) {
+    async train(input: string | number[], opts: TrainOptions = {}): Promise<number[]> {
         const {
             learningRate = 1e-4,
             epochs       = 5,
@@ -107,11 +96,9 @@ export class MambaTrainer {
             onEpochEnd   = null,
         } = opts;
 
-        // Enable WSLA mode if requested (fine-tune only B/C matrices)
         if (wsla) this.model.setWSLAMode(true);
 
-        // Tokenize
-        let tokenIds;
+        let tokenIds: number[];
         if (typeof input === 'string') {
             if (!this.tokenizer) {
                 throw new Error(
@@ -128,7 +115,6 @@ export class MambaTrainer {
             throw new Error('Input must contain at least 2 tokens to form a training pair.');
         }
 
-        // Build (input, target) sequence chunks of length seqLen
         const chunks = buildChunks(tokenIds, seqLen);
         if (chunks.length === 0) {
             throw new Error('Input is too short to form any training chunk.');
@@ -136,7 +122,7 @@ export class MambaTrainer {
 
         this._initMoments();
 
-        const epochLosses = [];
+        const epochLosses: number[] = [];
 
         for (let epoch = 0; epoch < epochs; epoch++) {
             let epochLoss = 0;
@@ -161,94 +147,66 @@ export class MambaTrainer {
         return epochLosses;
     }
 
-    // ─── Single training step ─────────────────────────────────────────────────
-
-    /**
-     * @param {number[]} inputs   – token IDs (length seqLen)
-     * @param {number[]} targets  – target token IDs (length seqLen, inputs shifted by 1)
-     * @param {number}   batch
-     * @param {Object}   hyperparams
-     * @returns {Promise<number>}  – scalar loss
-     */
-    async _trainStep(inputs, targets, batch, hyperparams) {
+    private async _trainStep(
+        inputs: number[],
+        targets: number[],
+        batch: number,
+        hyperparams: TrainOptions & { learningRate: number; maxGradNorm: number; weightDecay: number; beta1: number; beta2: number; eps: number }
+    ): Promise<number> {
         const { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps } = hyperparams;
 
         this._step++;
         const seqLen    = inputs.length;
         const vocabSize = this.model.config.vocabSize;
 
-        // ── Forward pass ──────────────────────────────────────────────────────
         const { logits, gpuLogits } = await this.model.forward(
             new Uint32Array(inputs), batch, seqLen
         );
 
-        // ── Compute loss (CPU) ────────────────────────────────────────────────
         let totalLoss = 0;
         const dLogits = new Float32Array(batch * seqLen * vocabSize);
 
         for (let i = 0; i < seqLen; i++) {
             const offset = i * vocabSize;
             const logitSlice = logits.slice(offset, offset + vocabSize);
-            const target = targets[i];
+            const target = targets[i]!;
             totalLoss += crossEntropyLoss(logitSlice, target);
             const grad  = crossEntropyGrad(logitSlice, target);
-            // Average over sequence length
             for (let v = 0; v < vocabSize; v++) {
-                dLogits[offset + v] = grad[v] / seqLen;
+                dLogits[offset + v] = grad[v]! / seqLen;
             }
         }
         const loss = totalLoss / seqLen;
 
-        // ── Upload gradients to GPU ───────────────────────────────────────────
         const dLogitsBuf = createStorageBuffer(this.device, dLogits, false);
 
-        // ── Gradient clipping ─────────────────────────────────────────────────
-        // (Applied after backward pass, but for the LM-head grad we do it now)
         await this._clipGradients(dLogitsBuf, dLogits.length, maxGradNorm);
 
-        // ── Parameter update (AdamW) ──────────────────────────────────────────
         const params  = this.model.parameters();
         const beta1_t = Math.pow(beta1, this._step);
         const beta2_t = Math.pow(beta2, this._step);
 
-        // For each parameter we need its gradient buffer.
-        // In a full implementation we'd run a proper backward pass through all
-        // layers by replaying the autograd tape.  Here we use the upstream
-        // gradient signal (dLogits) and update the LM head embedding with it,
-        // then propagate a synthetic gradient into the block parameters.
-        //
-        // Full backprop through all Mamba blocks is wired through the autograd
-        // tape (see autograd.js + backward kernels in selective_scan.js).
-        // For conciseness here we demonstrate the optimizer step using the
-        // available gradient buffer.
-
         await this._adamwStep(
             params, [dLogitsBuf],
             { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t }
         );
 
-        // Cleanup
         dLogitsBuf.destroy();
         gpuLogits.destroy();
 
         return loss;
     }
 
-    // ─── AdamW update ─────────────────────────────────────────────────────────
-
-    /**
-     * Apply AdamW update to each parameter using its gradient buffer.
-     *
-     * @param {Array<{buf: GPUBuffer, numel: number}>} params
-     * @param {GPUBuffer[]}                            gradBufs   – one per param
-     * @param {Object}                                 hp         – hyperparameters
-     */
-    async _adamwStep(params, gradBufs, hp) {
+    private async _adamwStep(
+        params: BlockParam[],
+        gradBufs: GPUBuffer[],
+        hp: AdamHyperparams
+    ): Promise<void> {
         const { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t } = hp;
 
         for (let i = 0; i < params.length; i++) {
-            const p       = params[i];
-            const gradBuf = gradBufs[Math.min(i, gradBufs.length - 1)];
+            const p       = params[i]!;
+            const gradBuf = gradBufs[Math.min(i, gradBufs.length - 1)]!;
 
             if (!gradBuf || gradBuf.size < p.numel * 4) continue;
 
@@ -260,8 +218,8 @@ export class MambaTrainer {
                 paramsBuf,
                 p.buf,
                 gradBuf,
-                this._moments[i].m,
-                this._moments[i].v,
+                this._moments![i]!.m,
+                this._moments![i]!.v,
             ]);
 
             dispatchKernel(this.device, this._adamwPipeline, bg,
@@ -271,17 +229,7 @@ export class MambaTrainer {
         }
     }
 
-    // ─── Gradient clipping ────────────────────────────────────────────────────
-
-    /**
-     * Clip gradient buffer in-place to max_norm (global L2 norm).
-     *
-     * @param {GPUBuffer} gradBuf
-     * @param {number}    numel
-     * @param {number}    maxNorm
-     */
-    async _clipGradients(gradBuf, numel, maxNorm) {
-        // Allocate norm_sq accumulator (single float, zeroed)
+    private async _clipGradients(gradBuf: GPUBuffer, numel: number, maxNorm: number): Promise<void> {
         const normSqBuf = createEmptyStorageBuffer(this.device, 4, true);
         this.device.queue.writeBuffer(normSqBuf, 0, new Float32Array([0.0]));
 
@@ -290,13 +238,11 @@ export class MambaTrainer {
         new Float32Array(clipParams, 4, 1).set([maxNorm * maxNorm]);
         const pBuf = createUniformBuffer(this.device, clipParams);
 
-        // Pass 1: compute norm squared
         const bg1 = createBindGroup(this.device, this._clipReducePipeline,
             [pBuf, gradBuf, normSqBuf]);
         dispatchKernel(this.device, this._clipReducePipeline, bg1,
             [cdiv(numel, 256), 1, 1]);
 
-        // Pass 2: scale gradients
         const bg2 = createBindGroup(this.device, this._clipScalePipeline,
             [pBuf, gradBuf, normSqBuf]);
         dispatchKernel(this.device, this._clipScalePipeline, bg2,
@@ -306,14 +252,8 @@ export class MambaTrainer {
         normSqBuf.destroy();
     }
 
-    /**
-     * Evaluate perplexity on a held-out code string.
-     *
-     * @param {string|number[]} input
-     * @returns {Promise<number>}  – perplexity (exp(average_loss))
-     */
-    async evaluate(input) {
-        let tokenIds;
+    async evaluate(input: string | number[]): Promise<number> {
+        let tokenIds: number[];
         if (typeof input === 'string') {
             if (!this.tokenizer) throw new Error('Tokenizer required for string input.');
             tokenIds = this.tokenizer.encode(input);
@@ -333,7 +273,7 @@ export class MambaTrainer {
             const offset = i * vocabSize;
             totalLoss += crossEntropyLoss(
                 logits.slice(offset, offset + vocabSize),
-                tokenIds[i + 1]
+                tokenIds[i + 1]!
             );
         }
 
@@ -342,25 +282,14 @@ export class MambaTrainer {
     }
 }
 
-// ─── Helpers ──────────────────────────────────────────────────────────────────
-
-/**
- * Split a flat token ID array into overlapping (input, target) pairs.
- * Each chunk is seqLen long; target is input shifted by 1.
- *
- * @param {number[]} ids
- * @param {number}   seqLen
- * @returns {Array<{inputs: number[], targets: number[]}>}
- */
-function buildChunks(ids, seqLen) {
-    const chunks = [];
+function buildChunks(ids: number[], seqLen: number): Array<{inputs: number[], targets: number[]}> {
+    const chunks: Array<{inputs: number[], targets: number[]}> = [];
     for (let start = 0; start + seqLen < ids.length; start += seqLen) {
         chunks.push({
             inputs : ids.slice(start, start + seqLen),
             targets: ids.slice(start + 1, start + seqLen + 1),
         });
     }
-    // Final partial chunk
     const rem = ids.length % seqLen;
     if (rem > 1) {
         const start = ids.length - rem;
@@ -372,21 +301,10 @@ function buildChunks(ids, seqLen) {
     return chunks;
 }
 
-/**
- * Pack AdamW hyperparameters into an ArrayBuffer matching the WGSL uniform struct.
- * Layout (byte offsets):
- *   0  : u32  num_elements
- *   4  : f32  lr
- *   8  : f32  beta1
- *   12 : f32  beta2
- *   16 : f32  eps
- *   20 : f32  weight_decay
- *   24 : f32  beta1_t
- *   28 : f32  beta2_t
- *
- * @returns {ArrayBuffer}
- */
-function packAdamParams(numElements, lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t) {
+function packAdamParams(
+    numElements: number, lr: number, beta1: number, beta2: number,
+    eps: number, weightDecay: number, beta1_t: number, beta2_t: number
+): ArrayBuffer {
     const buf = new ArrayBuffer(32);
     new Uint32Array(buf, 0, 1).set([numElements]);
     new Float32Array(buf, 4, 7).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t]);

package/src/utils/gpu_utils.ts

@@ -0,0 +1,147 @@
+/**
+ * gpu_utils.ts – WebGPU device management and buffer helpers.
+ */
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+const _gpu = globalThis as any;
+const UNIFORM: number  = _gpu.GPUBufferUsage?.UNIFORM  ?? 0x40;
+const STORAGE: number  = _gpu.GPUBufferUsage?.STORAGE  ?? 0x80;
+const COPY_SRC: number = _gpu.GPUBufferUsage?.COPY_SRC ?? 0x04;
+const COPY_DST: number = _gpu.GPUBufferUsage?.COPY_DST ?? 0x08;
+const MAP_READ: number = _gpu.GPUBufferUsage?.MAP_READ ?? 0x01;
+
+export interface InitWebGPUOptions {
+  powerPreference?: 'high-performance' | 'low-power';
+}
+
+export interface InitWebGPUResult {
+  device: GPUDevice;
+  adapter: GPUAdapter;
+}
+
+export async function initWebGPU(opts: InitWebGPUOptions = {}): Promise<InitWebGPUResult> {
+    if (typeof navigator === 'undefined' || !navigator.gpu) {
+        throw new Error(
+            'WebGPU is not available in this environment. ' +
+            'Use Chrome 113+, Edge 113+, or Firefox Nightly with WebGPU enabled.'
+        );
+    }
+
+    const adapter = await navigator.gpu.requestAdapter({
+        powerPreference: opts.powerPreference ?? 'high-performance',
+    });
+
+    if (!adapter) {
+        throw new Error('Failed to acquire a GPUAdapter. Your GPU may not support WebGPU.');
+    }
+
+    const adapterLimits = adapter.limits;
+    const requested3GB  = 3 * 1024 * 1024 * 1024;
+    const device = await adapter.requestDevice({
+        requiredLimits: {
+            maxBufferSize: Math.min(
+                requested3GB,
+                adapterLimits.maxBufferSize
+            ),
+            maxStorageBufferBindingSize: Math.min(
+                requested3GB,
+                adapterLimits.maxStorageBufferBindingSize
+            ),
+            maxComputeInvocationsPerWorkgroup: Math.min(
+                256,
+                adapterLimits.maxComputeInvocationsPerWorkgroup
+            ),
+        },
+    });
+
+    device.lost.then((info) => {
+        console.error('WebGPU device lost:', info.message);
+    });
+
+    return { device, adapter };
+}
+
+export function createStorageBuffer(device: GPUDevice, data: Float32Array | Uint32Array | number[], readable = false): GPUBuffer {
+    const arr    = data instanceof Float32Array || data instanceof Uint32Array ? data : new Float32Array(data);
+    const usage  = STORAGE | COPY_DST | (readable ? COPY_SRC : 0);
+    const buffer = device.createBuffer({ size: arr.byteLength, usage, mappedAtCreation: true });
+    if (arr instanceof Uint32Array) {
+        new Uint32Array(buffer.getMappedRange()).set(arr);
+    } else {
+        new Float32Array(buffer.getMappedRange()).set(arr as Float32Array);
+    }
+    buffer.unmap();
+    return buffer;
+}
+
+export function createEmptyStorageBuffer(device: GPUDevice, byteSize: number, readable = false): GPUBuffer {
+    const usage = STORAGE | COPY_DST | (readable ? COPY_SRC : 0);
+    return device.createBuffer({ size: byteSize, usage });
+}
+
+export function createUniformBuffer(device: GPUDevice, data: ArrayBuffer | ArrayBufferView): GPUBuffer {
+    const bytes  = ArrayBuffer.isView(data) ? data.buffer : data;
+    const buffer = device.createBuffer({
+        size  : bytes.byteLength,
+        usage : UNIFORM | COPY_DST,
+        mappedAtCreation: true,
+    });
+    new Uint8Array(buffer.getMappedRange()).set(new Uint8Array(bytes));
+    buffer.unmap();
+    return buffer;
+}
+
+export async function readBuffer(device: GPUDevice, srcBuffer: GPUBuffer, byteSize: number): Promise<Float32Array> {
+    const MAP_READ_FLAG: number = _gpu.GPUMapMode?.READ ?? 0x01;
+    const stagingBuffer = device.createBuffer({
+        size  : byteSize,
+        usage : MAP_READ | COPY_DST,
+    });
+
+    const encoder = device.createCommandEncoder();
+    encoder.copyBufferToBuffer(srcBuffer, 0, stagingBuffer, 0, byteSize);
+    device.queue.submit([encoder.finish()]);
+
+    await stagingBuffer.mapAsync(MAP_READ_FLAG);
+    const result = new Float32Array(stagingBuffer.getMappedRange().slice(0));
+    stagingBuffer.unmap();
+    stagingBuffer.destroy();
+    return result;
+}
+
+export function uploadBuffer(device: GPUDevice, buffer: GPUBuffer, data: Float32Array, byteOffset = 0): void {
+    device.queue.writeBuffer(buffer, byteOffset, data.buffer, data.byteOffset, data.byteLength);
+}
+
+export function createComputePipeline(device: GPUDevice, wgslSource: string, entryPoint: string): GPUComputePipeline {
+    const shaderModule = device.createShaderModule({ code: wgslSource });
+    return device.createComputePipeline({
+        layout : 'auto',
+        compute: { module: shaderModule, entryPoint },
+    });
+}
+
+export function createBindGroup(device: GPUDevice, pipeline: GPUComputePipeline, buffers: GPUBuffer[], groupIndex = 0): GPUBindGroup {
+    const entries = buffers.map((buf, i) => ({
+        binding : i,
+        resource: { buffer: buf },
+    }));
+    return device.createBindGroup({
+        layout : pipeline.getBindGroupLayout(groupIndex),
+        entries,
+    });
+}
+
+export function dispatchKernel(device: GPUDevice, pipeline: GPUComputePipeline, bindGroup: GPUBindGroup, workgroups: [number, number, number]): void {
+    const encoder = device.createCommandEncoder();
+    const pass    = encoder.beginComputePass();
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, bindGroup);
+    pass.dispatchWorkgroups(...workgroups);
+    pass.end();
+    device.queue.submit([encoder.finish()]);
+}
+
+export function cdiv(a: number, b: number): number {
+    return Math.ceil(a / b);
+}