@mni-ml/framework 0.0.1

package/dist/gpu_kernels.js

@@ -0,0 +1,571 @@
+import * as operators from './operators.js';
+export const WORKGROUP_SIZE = 256;
+export const BLOCK_SIZE = 16; // 16x16 = 256 = WORKGROUP_SIZE, used for 2D tiled matmul
+const MAX_DIMS = 6;
+// ---- Operation registries ----
+export const UNARY_OPS = {
+    neg: 'return -x;',
+    id: 'return x;',
+    sigmoid: 'let s = select(1.0 / (1.0 + exp(-x)), exp(x) / (1.0 + exp(x)), x < 0.0); return s;',
+    relu: 'return max(x, 0.0);',
+    exp: 'return exp(x);',
+    log: 'return log(x);',
+    inv: 'return 1.0 / x;',
+};
+export const BINARY_OPS = {
+    add: 'return a + b;',
+    mul: 'return a * b;',
+    max: 'return max(a, b);',
+    lt: 'return select(0.0, 1.0, a < b);',
+    eq: 'return select(0.0, 1.0, abs(a - b) < 0.00001);',
+    isClose: 'return select(0.0, 1.0, abs(a - b) < 0.01);',
+};
+// Map TypeScript operator functions to their registry key.
+const unaryRegistry = new Map();
+unaryRegistry.set(operators.neg, 'neg');
+unaryRegistry.set(operators.id, 'id');
+unaryRegistry.set(operators.sigmoid, 'sigmoid');
+unaryRegistry.set(operators.relu, 'relu');
+unaryRegistry.set(operators.exp, 'exp');
+unaryRegistry.set(operators.log, 'log');
+unaryRegistry.set(operators.inv, 'inv');
+const binaryRegistry = new Map();
+binaryRegistry.set(operators.add, 'add');
+binaryRegistry.set(operators.mul, 'mul');
+binaryRegistry.set(operators.max, 'max');
+binaryRegistry.set(operators.lt, 'lt');
+binaryRegistry.set(operators.eq, 'eq');
+binaryRegistry.set(operators.isClose, 'isClose');
+export function resolveUnaryOp(fn) {
+    const name = unaryRegistry.get(fn);
+    if (name)
+        return name;
+    throw new Error(`Unknown GPU unary op: ${fn.name || fn.toString().slice(0, 60)}`);
+}
+export function resolveBinaryOp(fn) {
+    const name = binaryRegistry.get(fn);
+    if (name)
+        return name;
+    throw new Error(`Unknown GPU binary op: ${fn.name || fn.toString().slice(0, 60)}`);
+}
+export const REDUCE_IDENTITY = {
+    add: '0.0',
+    mul: '1.0',
+    max: '-1.0e+38',
+};
+// ---- WGSL utility: index helpers ported to WGSL ----
+// Uses storage buffers for shape/stride arrays to avoid uniform alignment restrictions.
+const WGSL_INDEX_HELPERS = `
+const MAX_DIMS: u32 = ${MAX_DIMS}u;
+
+fn toIndex(ordinal: u32, shape: array<u32, ${MAX_DIMS}>, dims: u32, out_idx: ptr<function, array<u32, ${MAX_DIMS}>>) {
+    var remaining = ordinal;
+    for (var i = i32(dims) - 1; i >= 0; i--) {
+        let d = shape[i];
+        (*out_idx)[i] = remaining % d;
+        remaining = remaining / d;
+    }
+}
+
+fn indexToPosition(idx: array<u32, ${MAX_DIMS}>, strd: array<u32, ${MAX_DIMS}>, dims: u32) -> u32 {
+    var pos: u32 = 0u;
+    for (var i: u32 = 0u; i < dims; i++) {
+        pos += idx[i] * strd[i];
+    }
+    return pos;
+}
+
+fn broadcastIndex(
+    bigIdx: array<u32, ${MAX_DIMS}>, bigDims: u32,
+    smallShape: array<u32, ${MAX_DIMS}>, smallDims: u32,
+    out_idx: ptr<function, array<u32, ${MAX_DIMS}>>
+) {
+    let off = bigDims - smallDims;
+    for (var i: u32 = 0u; i < smallDims; i++) {
+        let bigI = i + off;
+        if (smallShape[i] == 1u) {
+            (*out_idx)[i] = 0u;
+        } else {
+            (*out_idx)[i] = bigIdx[bigI];
+        }
+    }
+}
+`;
+// ---- Shader template builders ----
+/**
+ * Aligned map: shapes & strides match, simple 1:1 element mapping.
+ */
+export function buildAlignedMapShader(opBody) {
+    return `
+@group(0) @binding(0) var<storage, read> in_data: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out_data: array<f32>;
+
+struct Params { size: u32 }
+@group(0) @binding(2) var<uniform> params: Params;
+
+fn apply(x: f32) -> f32 { ${opBody} }
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+    let i = gid.x;
+    if (i >= params.size) { return; }
+    out_data[i] = apply(in_data[i]);
+}
+`;
+}
+/**
+ * Broadcast map: output and input may differ in shape.
+ * Uses storage buffer for params to avoid WGSL uniform array alignment rules.
+ */
+export function buildBroadcastMapShader(opBody) {
+    return `
+${WGSL_INDEX_HELPERS}
+
+@group(0) @binding(0) var<storage, read> in_data: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out_data: array<f32>;
+
+struct Params {
+    out_size: u32,
+    out_dims: u32,
+    in_dims: u32,
+    _pad: u32,
+    out_shape:   array<u32, ${MAX_DIMS}>,
+    out_strides: array<u32, ${MAX_DIMS}>,
+    in_shape:    array<u32, ${MAX_DIMS}>,
+    in_strides:  array<u32, ${MAX_DIMS}>,
+}
+@group(0) @binding(2) var<storage, read> params: Params;
+
+fn apply(x: f32) -> f32 { ${opBody} }
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+    let i = gid.x;
+    if (i >= params.out_size) { return; }
+
+    var outIdx: array<u32, ${MAX_DIMS}>;
+    toIndex(i, params.out_shape, params.out_dims, &outIdx);
+
+    var inIdx: array<u32, ${MAX_DIMS}>;
+    broadcastIndex(outIdx, params.out_dims, params.in_shape, params.in_dims, &inIdx);
+
+    let outPos = indexToPosition(outIdx, params.out_strides, params.out_dims);
+    let inPos  = indexToPosition(inIdx, params.in_strides, params.in_dims);
+    out_data[outPos] = apply(in_data[inPos]);
+}
+`;
+}
+/**
+ * Aligned zip: all three tensors share shape & strides.
+ */
+export function buildAlignedZipShader(opBody) {
+    return `
+@group(0) @binding(0) var<storage, read> a_data: array<f32>;
+@group(0) @binding(1) var<storage, read> b_data: array<f32>;
+@group(0) @binding(2) var<storage, read_write> out_data: array<f32>;
+
+struct Params { size: u32 }
+@group(0) @binding(3) var<uniform> params: Params;
+
+fn apply(a: f32, b: f32) -> f32 { ${opBody} }
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+    let i = gid.x;
+    if (i >= params.size) { return; }
+    out_data[i] = apply(a_data[i], b_data[i]);
+}
+`;
+}
+/**
+ * Broadcast zip: output, a, b may differ in shape.
+ * Uses storage buffer for params to avoid WGSL uniform array alignment rules.
+ */
+export function buildBroadcastZipShader(opBody) {
+    return `
+${WGSL_INDEX_HELPERS}
+
+@group(0) @binding(0) var<storage, read> a_data: array<f32>;
+@group(0) @binding(1) var<storage, read> b_data: array<f32>;
+@group(0) @binding(2) var<storage, read_write> out_data: array<f32>;
+
+struct Params {
+    out_size: u32,
+    out_dims: u32,
+    a_dims: u32,
+    b_dims: u32,
+    out_shape:   array<u32, ${MAX_DIMS}>,
+    out_strides: array<u32, ${MAX_DIMS}>,
+    a_shape:     array<u32, ${MAX_DIMS}>,
+    a_strides:   array<u32, ${MAX_DIMS}>,
+    b_shape:     array<u32, ${MAX_DIMS}>,
+    b_strides:   array<u32, ${MAX_DIMS}>,
+}
+@group(0) @binding(3) var<storage, read> params: Params;
+
+fn apply(a: f32, b: f32) -> f32 { ${opBody} }
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(@builtin(global_invocation_id) gid: vec3u) {
+    let i = gid.x;
+    if (i >= params.out_size) { return; }
+
+    var outIdx: array<u32, ${MAX_DIMS}>;
+    toIndex(i, params.out_shape, params.out_dims, &outIdx);
+
+    var aIdx: array<u32, ${MAX_DIMS}>;
+    broadcastIndex(outIdx, params.out_dims, params.a_shape, params.a_dims, &aIdx);
+
+    var bIdx: array<u32, ${MAX_DIMS}>;
+    broadcastIndex(outIdx, params.out_dims, params.b_shape, params.b_dims, &bIdx);
+
+    let outPos = indexToPosition(outIdx, params.out_strides, params.out_dims);
+    let aPos   = indexToPosition(aIdx, params.a_strides, params.a_dims);
+    let bPos   = indexToPosition(bIdx, params.b_strides, params.b_dims);
+    out_data[outPos] = apply(a_data[aPos], b_data[bPos]);
+}
+`;
+}
+/**
+ * Sum practice: block-level partial sums using shared memory.
+ * Input: array of length size. Output: array of length ceil(size / WORKGROUP_SIZE).
+ * Each workgroup sums WORKGROUP_SIZE contiguous elements into one output cell.
+ */
+export function buildSumPracticeShader() {
+    return `
+const BLOCK_DIM: u32 = ${WORKGROUP_SIZE}u;
+var<workgroup> sdata: array<f32, ${WORKGROUP_SIZE}>;
+
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read_write> result: array<f32>;
+
+struct Params { size: u32 }
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(
+    @builtin(local_invocation_index) tid: u32,
+    @builtin(workgroup_id) wid: vec3u,
+) {
+    let global_idx = wid.x * BLOCK_DIM + tid;
+
+    if (global_idx < params.size) {
+        sdata[tid] = a[global_idx];
+    } else {
+        sdata[tid] = 0.0;
+    }
+    workgroupBarrier();
+
+    for (var stride = BLOCK_DIM / 2u; stride > 0u; stride = stride >> 1u) {
+        if (tid < stride) {
+            sdata[tid] = sdata[tid] + sdata[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    if (tid == 0u) {
+        result[wid.x] = sdata[0];
+    }
+}
+`;
+}
+/**
+ * General reduce along one dimension.
+ * One workgroup per output element. Threads cooperatively reduce
+ * the reduction dimension using shared memory.
+ * Uses storage buffer for params to avoid WGSL uniform array alignment rules.
+ */
+export function buildReduceShader(opBody, identity) {
+    return `
+${WGSL_INDEX_HELPERS}
+
+const BLOCK_DIM: u32 = ${WORKGROUP_SIZE}u;
+var<workgroup> sdata: array<f32, ${WORKGROUP_SIZE}>;
+
+@group(0) @binding(0) var<storage, read> a_data: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out_data: array<f32>;
+
+struct Params {
+    out_size: u32,
+    out_dims: u32,
+    a_dims: u32,
+    reduce_dim: u32,
+    reduce_dim_size: u32,
+    reduce_stride: u32,
+    _pad0: u32,
+    _pad1: u32,
+    out_shape:   array<u32, ${MAX_DIMS}>,
+    out_strides: array<u32, ${MAX_DIMS}>,
+    a_shape:     array<u32, ${MAX_DIMS}>,
+    a_strides:   array<u32, ${MAX_DIMS}>,
+}
+@group(0) @binding(2) var<storage, read> params: Params;
+
+fn apply(a: f32, b: f32) -> f32 { ${opBody} }
+
+@compute @workgroup_size(${WORKGROUP_SIZE})
+fn main(
+    @builtin(local_invocation_index) tid: u32,
+    @builtin(workgroup_id) wid: vec3u,
+) {
+    let out_idx = wid.x;
+    if (out_idx >= params.out_size) { return; }
+
+    var outMI: array<u32, ${MAX_DIMS}>;
+    toIndex(out_idx, params.out_shape, params.out_dims, &outMI);
+
+    var aIdx: array<u32, ${MAX_DIMS}>;
+    for (var d: u32 = 0u; d < params.a_dims; d++) {
+        aIdx[d] = outMI[d];
+    }
+    aIdx[params.reduce_dim] = 0u;
+    let base_pos = indexToPosition(aIdx, params.a_strides, params.a_dims);
+
+    var local_acc: f32 = ${identity};
+    for (var j = tid; j < params.reduce_dim_size; j += BLOCK_DIM) {
+        local_acc = apply(local_acc, a_data[base_pos + j * params.reduce_stride]);
+    }
+    sdata[tid] = local_acc;
+    workgroupBarrier();
+
+    for (var s = BLOCK_DIM / 2u; s > 0u; s = s >> 1u) {
+        if (tid < s) {
+            sdata[tid] = apply(sdata[tid], sdata[tid + s]);
+        }
+        workgroupBarrier();
+    }
+
+    if (tid == 0u) {
+        let outPos = indexToPosition(outMI, params.out_strides, params.out_dims);
+        out_data[outPos] = sdata[0];
+    }
+}
+`;
+}
+/**
+ * Tiled matrix multiplication using workgroup shared memory.
+ * Dispatched as 3D: (ceil(N/BLOCK), ceil(M/BLOCK), batchSize).
+ * Each 16x16 workgroup computes one output tile, loading tiles of A and B
+ * into shared memory to satisfy:
+ *   - all data read from shared memory (not global) during accumulation
+ *   - each global cell of A and B read exactly once
+ *   - each thread writes to global memory exactly once
+ * Supports arbitrary broadcast batch dimensions via stride-based indexing.
+ */
+export function buildMatMulShader() {
+    const BLOCK = BLOCK_SIZE;
+    const SHARED = BLOCK * BLOCK;
+    return `
+${WGSL_INDEX_HELPERS}
+
+const BLOCK: u32 = ${BLOCK}u;
+var<workgroup> a_shared: array<f32, ${SHARED}>;
+var<workgroup> b_shared: array<f32, ${SHARED}>;
+
+@group(0) @binding(0) var<storage, read> a_data: array<f32>;
+@group(0) @binding(1) var<storage, read> b_data: array<f32>;
+@group(0) @binding(2) var<storage, read_write> out_data: array<f32>;
+
+struct Params {
+    batch_size: u32,
+    M: u32,
+    N: u32,
+    K: u32,
+    out_dims: u32,
+    a_dims: u32,
+    b_dims: u32,
+    _pad: u32,
+    out_shape:   array<u32, ${MAX_DIMS}>,
+    out_strides: array<u32, ${MAX_DIMS}>,
+    a_shape:     array<u32, ${MAX_DIMS}>,
+    a_strides:   array<u32, ${MAX_DIMS}>,
+    b_shape:     array<u32, ${MAX_DIMS}>,
+    b_strides:   array<u32, ${MAX_DIMS}>,
+}
+@group(0) @binding(3) var<storage, read> params: Params;
+
+@compute @workgroup_size(${BLOCK}, ${BLOCK}, 1)
+fn main(
+    @builtin(local_invocation_id) lid: vec3u,
+    @builtin(workgroup_id) wid: vec3u,
+) {
+    let tx = lid.x;
+    let ty = lid.y;
+    let row = wid.y * BLOCK + ty;
+    let col = wid.x * BLOCK + tx;
+    let batch = wid.z;
+
+    let batch_dims = params.out_dims - 2u;
+
+    // Decompose linear batch index into multi-dim output batch indices
+    var batch_idx: array<u32, ${MAX_DIMS}>;
+    if (batch_dims > 0u) {
+        var batch_shape: array<u32, ${MAX_DIMS}>;
+        for (var d: u32 = 0u; d < batch_dims; d++) {
+            batch_shape[d] = params.out_shape[d];
+        }
+        toIndex(batch, batch_shape, batch_dims, &batch_idx);
+    }
+
+    // Output base offset from batch indices
+    var out_base: u32 = 0u;
+    for (var d: u32 = 0u; d < batch_dims; d++) {
+        out_base += batch_idx[d] * params.out_strides[d];
+    }
+    let out_stride_row = params.out_strides[params.out_dims - 2u];
+    let out_stride_col = params.out_strides[params.out_dims - 1u];
+
+    // Broadcast batch indices into a's batch space and compute base offset
+    let a_batch_dims = params.a_dims - 2u;
+    var a_batch_idx: array<u32, ${MAX_DIMS}>;
+    if (a_batch_dims > 0u) {
+        var a_batch_shape: array<u32, ${MAX_DIMS}>;
+        for (var d: u32 = 0u; d < a_batch_dims; d++) {
+            a_batch_shape[d] = params.a_shape[d];
+        }
+        broadcastIndex(batch_idx, batch_dims, a_batch_shape, a_batch_dims, &a_batch_idx);
+    }
+    var a_base: u32 = 0u;
+    for (var d: u32 = 0u; d < a_batch_dims; d++) {
+        a_base += a_batch_idx[d] * params.a_strides[d];
+    }
+    let a_stride_row = params.a_strides[params.a_dims - 2u];
+    let a_stride_col = params.a_strides[params.a_dims - 1u];
+
+    // Broadcast batch indices into b's batch space and compute base offset
+    let b_batch_dims = params.b_dims - 2u;
+    var b_batch_idx: array<u32, ${MAX_DIMS}>;
+    if (b_batch_dims > 0u) {
+        var b_batch_shape: array<u32, ${MAX_DIMS}>;
+        for (var d: u32 = 0u; d < b_batch_dims; d++) {
+            b_batch_shape[d] = params.b_shape[d];
+        }
+        broadcastIndex(batch_idx, batch_dims, b_batch_shape, b_batch_dims, &b_batch_idx);
+    }
+    var b_base: u32 = 0u;
+    for (var d: u32 = 0u; d < b_batch_dims; d++) {
+        b_base += b_batch_idx[d] * params.b_strides[d];
+    }
+    let b_stride_row = params.b_strides[params.b_dims - 2u];
+    let b_stride_col = params.b_strides[params.b_dims - 1u];
+
+    // Tiled matmul: each tile loads BLOCK x BLOCK elements into shared memory
+    var acc: f32 = 0.0;
+    let num_tiles = (params.K + BLOCK - 1u) / BLOCK;
+
+    for (var t: u32 = 0u; t < num_tiles; t++) {
+        // Load A[row, t*BLOCK + tx] into a_shared[ty][tx]
+        let a_col = t * BLOCK + tx;
+        if (row < params.M && a_col < params.K) {
+            a_shared[ty * BLOCK + tx] = a_data[a_base + row * a_stride_row + a_col * a_stride_col];
+        } else {
+            a_shared[ty * BLOCK + tx] = 0.0;
+        }
+
+        // Load B[t*BLOCK + ty, col] into b_shared[ty][tx]
+        let b_row = t * BLOCK + ty;
+        if (b_row < params.K && col < params.N) {
+            b_shared[ty * BLOCK + tx] = b_data[b_base + b_row * b_stride_row + col * b_stride_col];
+        } else {
+            b_shared[ty * BLOCK + tx] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        // Accumulate partial dot products from shared memory
+        for (var k: u32 = 0u; k < BLOCK; k++) {
+            acc += a_shared[ty * BLOCK + k] * b_shared[k * BLOCK + tx];
+        }
+
+        workgroupBarrier();
+    }
+
+    // Single write to global memory per thread
+    if (row < params.M && col < params.N) {
+        out_data[out_base + row * out_stride_row + col * out_stride_col] = acc;
+    }
+}
+`;
+}
+// ---- Tensor core matmul via experimental subgroup matrix ----
+export const TC_TILE = 8; // Apple Metal simdgroup_matrix tile dimension
+/**
+ * Experimental tensor-core-accelerated matmul shader using the Dawn/Chrome
+ * `chromium_experimental_subgroup_matrix` extension.
+ *
+ * On Apple Silicon (Metal), this maps to `simdgroup_matrix` 8x8 f32 hardware
+ * instructions. On Vulkan with VK_KHR_cooperative_matrix it maps to the
+ * equivalent SPIR-V ops.
+ *
+ * @param workgroupX - must equal the device's maxSubgroupSize (e.g. 64 on
+ *   Apple M-series via Dawn, 32 or 64 on Vulkan). The WebGPU spec requires
+ *   the x-dimension of workgroup_size to be a multiple of maxSubgroupSize
+ *   when the shader uses subgroup matrices.
+ *
+ * Constraints (checked by the caller on the TS side):
+ *   - All tensors must be contiguous (natural row-major strides)
+ *   - M, N, K must all be multiples of TC_TILE (8)
+ *   - Batch dimensions of A and B must match exactly (no broadcasting)
+ *
+ * Dispatch: (N/TC_TILE, M/TC_TILE, batchSize).
+ * Each workgroup computes one 8x8 output tile by tiling over the K dimension
+ * with subgroupMatrixLoad / MultiplyAccumulate / Store.
+ */
+export function buildTensorCoreMatMulShader(workgroupX) {
+    return `
+enable chromium_experimental_subgroup_matrix;
+
+const TILE: u32 = ${TC_TILE}u;
+
+@group(0) @binding(0) var<storage, read> a_data: array<f32>;
+@group(0) @binding(1) var<storage, read> b_data: array<f32>;
+@group(0) @binding(2) var<storage, read_write> out_data: array<f32>;
+
+struct Params {
+    batch_size: u32,
+    M: u32,
+    N: u32,
+    K: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+@compute @workgroup_size(${workgroupX}, 1, 1)
+fn main(@builtin(workgroup_id) wid: vec3u) {
+    let tile_row = wid.y * TILE;
+    let tile_col = wid.x * TILE;
+    let batch = wid.z;
+
+    let a_base = batch * params.M * params.K;
+    let b_base = batch * params.K * params.N;
+    let out_base = batch * params.M * params.N;
+
+    // Zero-initialise 8x8 accumulator distributed across the subgroup
+    var acc = subgroup_matrix_result<f32, 8, 8>(0.0);
+
+    // Tile over the shared K dimension
+    let num_tiles = params.K / TILE;
+    for (var t: u32 = 0u; t < num_tiles; t++) {
+        let k = t * TILE;
+
+        // Load 8x8 tile of A (row-major, stride = K columns)
+        let a_tile = subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(
+            &a_data, a_base + tile_row * params.K + k, false, params.K
+        );
+
+        // Load 8x8 tile of B (row-major, stride = N columns)
+        let b_tile = subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(
+            &b_data, b_base + k * params.N + tile_col, false, params.N
+        );
+
+        // Hardware multiply-accumulate: acc += a_tile * b_tile
+        acc = subgroupMatrixMultiplyAccumulate(a_tile, b_tile, acc);
+    }
+
+    // Single store of the 8x8 result tile to global memory
+    subgroupMatrixStore(
+        &out_data, out_base + tile_row * params.N + tile_col, acc, false, params.N
+    );
+}
+`;
+}
+//# sourceMappingURL=gpu_kernels.js.map

package/dist/gpu_kernels.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"gpu_kernels.js","sourceRoot":"","sources":["../src/gpu_kernels.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,SAAS,MAAM,gBAAgB,CAAC;AAE5C,MAAM,CAAC,MAAM,cAAc,GAAG,GAAG,CAAC;AAClC,MAAM,CAAC,MAAM,UAAU,GAAG,EAAE,CAAC,CAAC,yDAAyD;AACvF,MAAM,QAAQ,GAAG,CAAC,CAAC;AAEnB,iCAAiC;AAEjC,MAAM,CAAC,MAAM,SAAS,GAA2B;IAC7C,GAAG,EAAM,YAAY;IACrB,EAAE,EAAO,WAAW;IACpB,OAAO,EAAE,oFAAoF;IAC7F,IAAI,EAAK,qBAAqB;IAC9B,GAAG,EAAM,gBAAgB;IACzB,GAAG,EAAM,gBAAgB;IACzB,GAAG,EAAM,iBAAiB;CAC7B,CAAC;AAEF,MAAM,CAAC,MAAM,UAAU,GAA2B;IAC9C,GAAG,EAAM,eAAe;IACxB,GAAG,EAAM,eAAe;IACxB,GAAG,EAAM,mBAAmB;IAC5B,EAAE,EAAO,iCAAiC;IAC1C,EAAE,EAAO,gDAAgD;IACzD,OAAO,EAAE,6CAA6C;CACzD,CAAC;AAEF,2DAA2D;AAC3D,MAAM,aAAa,GAAG,IAAI,GAAG,EAAoB,CAAC;AAClD,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC5C,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE,EAAO,IAAI,CAAC,CAAC;AAC3C,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;AAChD,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,IAAI,EAAK,MAAM,CAAC,CAAC;AAC7C,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC5C,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC5C,aAAa,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAE5C,MAAM,cAAc,GAAG,IAAI,GAAG,EAAoB,CAAC;AACnD,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC7C,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC7C,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,GAAG,EAAM,KAAK,CAAC,CAAC;AAC7C,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE,EAAO,IAAI,CAAC,CAAC;AAC5C,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,EAAE,EAAO,IAAI,CAAC,CAAC;AAC5C,cAAc,CAAC,GAAG,CAAC,SAAS,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;AAEjD,MAAM,UAAU,cAAc,CAAC,EAAY;IACvC,MAAM,IAAI,GAAG,aAAa,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IACnC,IAAI,IAAI;QAAE,OAAO,IAAI,CAAC;IACtB,MAAM,IAAI,KAAK,CAAC,yBAAyB,EAAE,CAAC,IAAI,IAAI,EAAE,CAAC,QAAQ,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC;AACtF,CAAC;AAED,MAAM,UAAU,eAAe,CAAC,EAAY;IACxC,MAAM,IAAI,GAAG,cAAc,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;IACpC,IAAI,IAAI;QAAE,OAAO,IAAI,CAAC;IACtB,MAAM,IAAI,KAAK,CAAC,0BAA0B,EAAE,CAAC,IAAI,IAAI,EAAE,CAAC,QAAQ,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC;AACvF,CAAC;AAED,MAAM,CAAC,MAAM,eAAe,GAA2B;IACnD,GAAG,EAAE,KAAK;IACV,GAAG,EAAE,KAAK;IACV,GAAG,EAAE,UAAU;CAClB,CAAC;AAEF,uDAAuD;AACvD,wFAAwF;AAExF,MAAM,kBAAkB,GAAG;wBACH,QAAQ;;6CAEa,QAAQ,mDAAmD,QAAQ;;;;;;;;;qCAS3E,QAAQ,uBAAuB,QAAQ;;;;;;;;;yBASnD,QAAQ;6BACJ,QAAQ;wCACG,QAAQ;;;;;;;;;;;;CAY/C,CAAC;AAEF,qCAAqC;AAErC;;GAEG;AACH,MAAM,UAAU,qBAAqB,CAAC,MAAc;IAChD,OAAO;;;;;;;4BAOiB,MAAM;;2BAEP,cAAc;;;;;;CAMxC,CAAC;AACF,CAAC;AAED;;;GAGG;AACH,MAAM,UAAU,uBAAuB,CAAC,MAAc;IAClD,OAAO;EACT,kBAAkB;;;;;;;;;;8BAUU,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;;;;4BAIV,MAAM;;2BAEP,cAAc;;;;;6BAKZ,QAAQ;;;4BAGT,QAAQ;;;;;;;CAOnC,CAAC;AACF,CAAC;AAED;;GAEG;AACH,MAAM,UAAU,qBAAqB,CAAC,MAAc;IAChD,OAAO;;;;;;;;oCAQyB,MAAM;;2BAEf,cAAc;;;;;;CAMxC,CAAC;AACF,CAAC;AAED;;;GAGG;AACH,MAAM,UAAU,uBAAuB,CAAC,MAAc;IAClD,OAAO;EACT,kBAAkB;;;;;;;;;;;8BAWU,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;;;;oCAIF,MAAM;;2BAEf,cAAc;;;;;6BAKZ,QAAQ;;;2BAGV,QAAQ;;;2BAGR,QAAQ;;;;;;;;CAQlC,CAAC;AACF,CAAC;AAED;;;;GAIG;AACH,MAAM,UAAU,sBAAsB;IAClC,OAAO;yBACc,cAAc;mCACJ,cAAc;;;;;;;;2BAQtB,cAAc;;;;;;;;;;;;;;;;;;;;;;;;;CAyBxC,CAAC;AACF,CAAC;AAED;;;;;GAKG;AACH,MAAM,UAAU,iBAAiB,CAAC,MAAc,EAAE,QAAgB;IAC9D,OAAO;EACT,kBAAkB;;yBAEK,cAAc;mCACJ,cAAc;;;;;;;;;;;;;;8BAcnB,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;;;;oCAIF,MAAM;;2BAEf,cAAc;;;;;;;;4BAQb,QAAQ;;;2BAGT,QAAQ;;;;;;;2BAOR,QAAQ;;;;;;;;;;;;;;;;;;;CAmBlC,CAAC;AACF,CAAC;AAED;;;;;;;;;GASG;AACH,MAAM,UAAU,iBAAiB;IAC7B,MAAM,KAAK,GAAG,UAAU,CAAC;IACzB,MAAM,MAAM,GAAG,KAAK,GAAG,KAAK,CAAC;IAC7B,OAAO;EACT,kBAAkB;;qBAEC,KAAK;sCACY,MAAM;sCACN,MAAM;;;;;;;;;;;;;;;8BAed,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;8BACR,QAAQ;;;;2BAIX,KAAK,KAAK,KAAK;;;;;;;;;;;;;;gCAcV,QAAQ;;sCAEF,QAAQ;;;;;;;;;;;;;;;;;kCAiBZ,QAAQ;;wCAEF,QAAQ;;;;;;;;;;;;;;;kCAed,QAAQ;;wCAEF,QAAQ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAiD/C,CAAC;AACF,CAAC;AAED,gEAAgE;AAEhE,MAAM,CAAC,MAAM,OAAO,GAAG,CAAC,CAAC,CAAC,8CAA8C;AAExE;;;;;;;;;;;;;;;;;;;;;GAqBG;AACH,MAAM,UAAU,2BAA2B,CAAC,UAAkB;IAC1D,OAAO;;;oBAGS,OAAO;;;;;;;;;;;;;;2BAcA,UAAU;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAqCpC,CAAC;AACF,CAAC"}

package/dist/gpu_ops.d.ts

@@ -0,0 +1,43 @@
+import type { Storage, Shape, Strides } from './tensor_data.js';
+/**
+ * Practice sum kernel: given array of length `size`, produce ceil(size / WORKGROUP_SIZE)
+ * partial sums. Each workgroup sums WORKGROUP_SIZE contiguous elements using shared memory.
+ */
+export declare function _sumPractice(out: Storage, a: Storage, size: number): Promise<void>;
+/**
+ * GPU higher-order tensor map.  fn must be a known operator from operators.ts.
+ */
+export declare function gpuTensorMap(fn: (x: number) => number): (outStorage: Storage, outShape: Shape, outStrides: Strides, inStorage: Storage, inShape: Shape, inStrides: Strides) => Promise<void>;
+/**
+ * GPU higher-order tensor zip (binary map). fn must be a known binary operator.
+ */
+export declare function gpuTensorZip(fn: (a: number, b: number) => number): (outStorage: Storage, outShape: Shape, outStrides: Strides, aStorage: Storage, aShape: Shape, aStrides: Strides, bStorage: Storage, bShape: Shape, bStrides: Strides) => Promise<void>;
+/**
+ * GPU higher-order tensor reduce. fn must be a known binary operator
+ * with an entry in REDUCE_IDENTITY.
+ * One workgroup per output element; threads cooperatively reduce
+ * the target dimension using shared memory tree reduction.
+ */
+export declare function gpuTensorReduce(fn: (acc: number, x: number) => number): (outStorage: Storage, outShape: Shape, outStrides: Strides, aStorage: Storage, aShape: Shape, aStrides: Strides, reduceDim: number) => Promise<void>;
+/**
+ * GPU tiled matrix multiplication with shared memory.
+ * Supports arbitrary broadcast batch dimensions as long as
+ * aShape[-1] === bShape[-2].
+ */
+export declare function gpuTensorMatrixMultiply(outStorage: Storage, outShape: Shape, outStrides: Strides, outSize: number, aStorage: Storage, aShape: Shape, aStrides: Strides, bStorage: Storage, bShape: Shape, bStrides: Strides): Promise<void>;
+/**
+ * Experimental tensor-core-accelerated matrix multiply.
+ *
+ * Uses the Dawn `chromium_experimental_subgroup_matrix` extension which maps
+ * to Apple simdgroup_matrix (8x8 f32) on Metal and VK_KHR_cooperative_matrix
+ * on Vulkan. Returns true if the fast path executed, false if a precondition
+ * was not met and the caller should fall back to `gpuTensorMatrixMultiply`.
+ *
+ * Preconditions (returns false when any is violated):
+ *   - Tensor core device available (hardware + experimental feature)
+ *   - All tensors contiguous (natural row-major strides)
+ *   - M, N, K all divisible by 8
+ *   - Batch dimensions of A and B identical (no broadcasting)
+ */
+export declare function gpuTensorMatrixMultiplyTensorCore(outStorage: Storage, outShape: Shape, outStrides: Strides, outSize: number, aStorage: Storage, aShape: Shape, aStrides: Strides, bStorage: Storage, bShape: Shape, bStrides: Strides): Promise<boolean>;
+//# sourceMappingURL=gpu_ops.d.ts.map

package/dist/gpu_ops.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"gpu_ops.d.ts","sourceRoot":"","sources":["../src/gpu_ops.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,OAAO,EAAE,KAAK,EAAE,OAAO,EAAE,MAAM,kBAAkB,CAAC;AAgFhE;;;GAGG;AACH,wBAAsB,YAAY,CAC9B,GAAG,EAAE,OAAO,EACZ,CAAC,EAAE,OAAO,EACV,IAAI,EAAE,MAAM,GACb,OAAO,CAAC,IAAI,CAAC,CAkCf;AAID;;GAEG;AACH,wBAAgB,YAAY,CACxB,EAAE,EAAE,CAAC,CAAC,EAAE,MAAM,KAAK,MAAM,GAC1B,CACC,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,KAAK,EACf,UAAU,EAAE,OAAO,EACnB,SAAS,EAAE,OAAO,EAClB,OAAO,EAAE,KAAK,EACd,SAAS,EAAE,OAAO,KACjB,OAAO,CAAC,IAAI,CAAC,CA4DjB;AAID;;GAEG;AACH,wBAAgB,YAAY,CACxB,EAAE,EAAE,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,KAAK,MAAM,GACrC,CACC,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,KAAK,EACf,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,EACjB,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,KAChB,OAAO,CAAC,IAAI,CAAC,CAsEjB;AAID;;;;;GAKG;AACH,wBAAgB,eAAe,CAC3B,EAAE,EAAE,CAAC,GAAG,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,KAAK,MAAM,GACvC,CACC,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,KAAK,EACf,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,EACjB,SAAS,EAAE,MAAM,KAChB,OAAO,CAAC,IAAI,CAAC,CA+DjB;AAID;;;;GAIG;AACH,wBAAsB,uBAAuB,CACzC,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,KAAK,EACf,UAAU,EAAE,OAAO,EACnB,OAAO,EAAE,MAAM,EACf,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,EACjB,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,GAClB,OAAO,CAAC,IAAI,CAAC,CA4Df;AAqBD;;;;;;;;;;;;;GAaG;AACH,wBAAsB,iCAAiC,CACnD,UAAU,EAAE,OAAO,EACnB,QAAQ,EAAE,KAAK,EACf,UAAU,EAAE,OAAO,EACnB,OAAO,EAAE,MAAM,EACf,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,EACjB,QAAQ,EAAE,OAAO,EACjB,MAAM,EAAE,KAAK,EACb,QAAQ,EAAE,OAAO,GAClB,OAAO,CAAC,OAAO,CAAC,CAkElB"}