@simulatte/doppler 0.1.8 → 0.2.0

package/src/gpu/kernels/backward/utils.js

@@ -1,371 +0,0 @@
-import { getDevice } from '../../device.js';
-import { acquireBuffer } from '../../../memory/buffer-pool.js';
-import { createTensor, dtypeBytes } from '../../tensor.js';
-import { WORKGROUP_SIZES } from '../constants.js';
-import { dispatch, recordDispatch } from '../dispatch.js';
-import { createPipeline, createUniformBufferWithView } from '../utils.js';
-import { releaseUniformBuffer } from '../../uniform-cache.js';
-
-export async function runBackwardKernel(
-  opName,
-  input,
-  gradOutput,
-  uniformSize,
-  writeUniforms,
-  options = {}
-) {
-  const device = getDevice();
-  const { count, outputBuffer = null } = options;
-
-  const bytesPerElement = dtypeBytes(gradOutput.dtype);
-  const inferredCount = count ?? Math.floor(gradOutput.buffer.size / bytesPerElement);
-  const pipeline = await createPipeline(opName, 'default');
-
-  const outputSize = inferredCount * bytesPerElement;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, `${opName}_backward_output`);
-
-  const uniformBuffer = createUniformBufferWithView(
-    `${opName}_uniforms`,
-    uniformSize,
-    (view) => {
-      writeUniforms(view, inferredCount);
-    }
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: `${opName}_bind_group`,
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: gradOutput.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const workgroups = Math.ceil(inferredCount / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, opName);
-
-  releaseUniformBuffer(uniformBuffer);
-
-  return createTensor(outputBuf, gradOutput.dtype, [...gradOutput.shape], `${opName}_output`);
-}
-
-export async function recordBackwardKernel(
-  recorder,
-  opName,
-  input,
-  gradOutput,
-  uniformSize,
-  writeUniforms,
-  options = {}
-) {
-  const device = recorder.device;
-  const { count, outputBuffer = null } = options;
-
-  const bytesPerElement = dtypeBytes(gradOutput.dtype);
-  const inferredCount = count ?? Math.floor(gradOutput.buffer.size / bytesPerElement);
-  const pipeline = await createPipeline(opName, 'default');
-
-  const outputSize = inferredCount * bytesPerElement;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, `${opName}_backward_output`);
-
-  const uniformBuffer = createUniformBufferWithView(
-    `${opName}_uniforms`,
-    uniformSize,
-    (view) => {
-      writeUniforms(view, inferredCount);
-    },
-    recorder
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: `${opName}_bind_group`,
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: gradOutput.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const workgroups = Math.ceil(inferredCount / WORKGROUP_SIZES.DEFAULT);
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, opName);
-
-  return createTensor(outputBuf, gradOutput.dtype, [...gradOutput.shape], `${opName}_output`);
-}
-
-export function createBackwardKernel(opName, spec) {
-  const {
-    uniformSize,
-    writeUniforms,
-    calcWorkgroups,
-    outputBytes,
-    outputShape,
-    validate,
-    dtype: specDtype,
-    getDevice: useGetDevice,
-  } = spec;
-
-  async function run(...args) {
-    const opts = args[args.length - 1];
-    const inputs = args.slice(0, -1);
-
-    if (validate) validate(opts);
-
-    const { outputBuffer = null } = opts;
-    const device = useGetDevice ? getDevice() : inputs[0].buffer.device;
-    const outSize = outputBytes(opts);
-    const outputBuf = outputBuffer || acquireBuffer(outSize, undefined, `${opName}_output`);
-
-    const pipeline = await createPipeline(opName, 'default');
-
-    const uniformBuffer = createUniformBufferWithView(
-      `${opName}_uniforms`,
-      uniformSize,
-      (view) => writeUniforms(view, opts),
-      null,
-      device
-    );
-
-    const entries = [{ binding: 0, resource: { buffer: uniformBuffer } }];
-    for (let i = 0; i < inputs.length; i++) {
-      entries.push({ binding: i + 1, resource: { buffer: inputs[i].buffer } });
-    }
-    entries.push({ binding: inputs.length + 1, resource: { buffer: outputBuf } });
-
-    const bindGroup = device.createBindGroup({
-      label: `${opName}_bind_group`,
-      layout: pipeline.getBindGroupLayout(0),
-      entries,
-    });
-
-    const workgroups = calcWorkgroups(opts);
-    dispatch(device, pipeline, bindGroup, workgroups, opName);
-
-    uniformBuffer.destroy();
-
-    const dt = specDtype ? specDtype(opts, inputs) : (inputs[0] ? inputs[0].dtype : 'f32');
-    return createTensor(outputBuf, dt, outputShape(opts), `${opName}_output`);
-  }
-
-  async function record(recorder, ...args) {
-    const opts = args[args.length - 1];
-    const inputs = args.slice(0, -1);
-
-    if (validate) validate(opts);
-
-    const { outputBuffer = null } = opts;
-    const device = recorder.device;
-    const outSize = outputBytes(opts);
-    const outputBuf = outputBuffer || acquireBuffer(outSize, undefined, `${opName}_output`);
-
-    const pipeline = await createPipeline(opName, 'default');
-
-    const uniformBuffer = createUniformBufferWithView(
-      `${opName}_uniforms`,
-      uniformSize,
-      (view) => writeUniforms(view, opts),
-      recorder
-    );
-
-    const entries = [{ binding: 0, resource: { buffer: uniformBuffer } }];
-    for (let i = 0; i < inputs.length; i++) {
-      entries.push({ binding: i + 1, resource: { buffer: inputs[i].buffer } });
-    }
-    entries.push({ binding: inputs.length + 1, resource: { buffer: outputBuf } });
-
-    const bindGroup = device.createBindGroup({
-      label: `${opName}_bind_group`,
-      layout: pipeline.getBindGroupLayout(0),
-      entries,
-    });
-
-    const workgroups = calcWorkgroups(opts);
-    recordDispatch(recorder, pipeline, bindGroup, workgroups, opName);
-
-    const dt = specDtype ? specDtype(opts, inputs) : (inputs[0] ? inputs[0].dtype : 'f32');
-    return createTensor(outputBuf, dt, outputShape(opts), `${opName}_output`);
-  }
-
-  return { run, record };
-}
-
-export async function runMatmulBackwardDx(dY, W, M, K, N, options = {}) {
-  const { alpha = 1.0, transposeB = false, outputBuffer = null } = options;
-  const device = getDevice();
-  const outputSize = M * K * 4;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_backward_dx_output');
-
-  const pipeline = await createPipeline('matmul_backward', 'default');
-
-  const uniformBuffer = createUniformBufferWithView(
-    'matmul_backward_uniforms',
-    32,
-    (view) => {
-      view.setUint32(0, M, true);
-      view.setUint32(4, N, true);
-      view.setUint32(8, K, true);
-      view.setFloat32(12, alpha, true);
-      view.setUint32(16, transposeB ? 1 : 0, true);
-    },
-    null,
-    device
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_backward_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: dY.buffer } },
-      { binding: 2, resource: { buffer: W.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const TILE_SIZE = 16;
-  const workgroups = [
-    Math.ceil(M / TILE_SIZE),
-    Math.ceil(K / TILE_SIZE),
-    1,
-  ];
-
-  dispatch(device, pipeline, bindGroup, workgroups, 'matmul_backward');
-  releaseUniformBuffer(uniformBuffer);
-
-  return createTensor(outputBuf, 'f32', [M, K], 'matmul_backward_dx_output');
-}
-
-export async function recordMatmulBackwardDx(recorder, dY, W, M, K, N, options = {}) {
-  const { alpha = 1.0, transposeB = false, outputBuffer = null } = options;
-  const device = recorder.device;
-  const outputSize = M * K * 4;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_backward_dx_output');
-
-  const pipeline = await createPipeline('matmul_backward', 'default');
-
-  const uniformBuffer = createUniformBufferWithView(
-    'matmul_backward_uniforms',
-    32,
-    (view) => {
-      view.setUint32(0, M, true);
-      view.setUint32(4, N, true);
-      view.setUint32(8, K, true);
-      view.setFloat32(12, alpha, true);
-      view.setUint32(16, transposeB ? 1 : 0, true);
-    },
-    recorder
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_backward_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: dY.buffer } },
-      { binding: 2, resource: { buffer: W.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const TILE_SIZE = 16;
-  const workgroups = [
-    Math.ceil(M / TILE_SIZE),
-    Math.ceil(K / TILE_SIZE),
-    1,
-  ];
-
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'matmul_backward');
-
-  return createTensor(outputBuf, 'f32', [M, K], 'matmul_backward_dx_output');
-}
-
-export async function runMatmulTransposeA(A, B, M, N, K, options = {}) {
-  const { alpha = 1.0, outputBuffer = null } = options;
-  const device = getDevice();
-  const outputSize = M * N * 4;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_transpose_a_output');
-
-  const pipeline = await createPipeline('matmul_transpose_a', 'default');
-
-  const uniformBuffer = createUniformBufferWithView(
-    'matmul_transpose_a_uniforms',
-    32,
-    (view) => {
-      view.setUint32(0, M, true);
-      view.setUint32(4, N, true);
-      view.setUint32(8, K, true);
-      view.setFloat32(12, alpha, true);
-    },
-    null,
-    device
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_transpose_a_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: A.buffer } },
-      { binding: 2, resource: { buffer: B.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const TILE_SIZE = 16;
-  const workgroups = [
-    Math.ceil(M / TILE_SIZE),
-    Math.ceil(N / TILE_SIZE),
-    1,
-  ];
-
-  dispatch(device, pipeline, bindGroup, workgroups, 'matmul_transpose_a');
-  releaseUniformBuffer(uniformBuffer);
-
-  return createTensor(outputBuf, 'f32', [M, N], 'matmul_transpose_a_output');
-}
-
-export async function recordMatmulTransposeA(recorder, A, B, M, N, K, options = {}) {
-  const { alpha = 1.0, outputBuffer = null } = options;
-  const device = recorder.device;
-  const outputSize = M * N * 4;
-  const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_transpose_a_output');
-
-  const pipeline = await createPipeline('matmul_transpose_a', 'default');
-
-  const uniformBuffer = createUniformBufferWithView(
-    'matmul_transpose_a_uniforms',
-    32,
-    (view) => {
-      view.setUint32(0, M, true);
-      view.setUint32(4, N, true);
-      view.setUint32(8, K, true);
-      view.setFloat32(12, alpha, true);
-    },
-    recorder
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_transpose_a_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: A.buffer } },
-      { binding: 2, resource: { buffer: B.buffer } },
-      { binding: 3, resource: { buffer: outputBuf } },
-    ],
-  });
-
-  const TILE_SIZE = 16;
-  const workgroups = [
-    Math.ceil(M / TILE_SIZE),
-    Math.ceil(N / TILE_SIZE),
-    1,
-  ];
-
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'matmul_transpose_a');
-
-  return createTensor(outputBuf, 'f32', [M, N], 'matmul_transpose_a_output');
-}

package/src/gpu/kernels/bf16_to_f16.wgsl

@@ -1,54 +0,0 @@
-// BF16 to F16 Conversion Kernel
-//
-// Converts BF16 (bfloat16) data to F16.
-// BF16 is represented as the upper 16 bits of F32. Conversion is:
-//   bf16 -> f32 by shifting left 16 bits -> f16 cast.
-//
-// Supports 2D dispatch for large tensors.
-
-enable f16;
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    num_elements: u32,
-    input_offset: u32,   // Element offset for input (in BF16 elements)
-    output_offset: u32,  // Element offset for output (in F16 elements)
-    _pad: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> input: array<u32>;  // BF16 packed as u32 (2 per u32)
-@group(0) @binding(2) var<storage, read_write> output: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(
-    @builtin(global_invocation_id) global_id: vec3<u32>,
-    @builtin(num_workgroups) num_wg: vec3<u32>
-) {
-    // Support 2D dispatch for large tensors
-    let linear_idx = global_id.x + global_id.y * num_wg.x * WORKGROUP_SIZE;
-
-    // Each thread processes 2 BF16 values (one u32 contains 2 bf16)
-    let local_pair_idx = linear_idx;
-    let local_elem_idx = local_pair_idx * 2u;
-
-    if (local_elem_idx >= u.num_elements) {
-        return;
-    }
-
-    // Apply offsets for chunked processing
-    let input_pair_idx = (u.input_offset / 2u) + local_pair_idx;
-    let output_elem_idx = u.output_offset + local_elem_idx;
-
-    let packed = input[input_pair_idx];
-
-    let bf16_lo = packed & 0xFFFFu;
-    let bf16_hi = (packed >> 16u) & 0xFFFFu;
-
-    output[output_elem_idx] = f16(bitcast<f32>(bf16_lo << 16u));
-
-    if (local_elem_idx + 1u < u.num_elements) {
-        output[output_elem_idx + 1u] = f16(bitcast<f32>(bf16_hi << 16u));
-    }
-}

package/src/gpu/kernels/bf16_to_f32.wgsl

@@ -1,70 +0,0 @@
-// BF16 to F32 Conversion Kernel
-//
-// Converts BF16 (bfloat16) data to F32.
-// BF16 is just the upper 16 bits of F32, so conversion is a simple shift.
-//
-// Supports 2D dispatch for large tensors (>65535 workgroups).
-
-override WORKGROUP_SIZE: u32 = 256u;
-const MAX_WORKGROUPS_X: u32 = 65535u;
-
-struct Uniforms {
-    num_elements: u32,
-    input_offset: u32,   // Element offset for input (in BF16 elements)
-    output_offset: u32,  // Element offset for output (in F32 elements)
-    _pad: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> input: array<u32>;  // BF16 packed as u32 (2 per u32)
-@group(0) @binding(2) var<storage, read_write> output: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-    // Support 2D dispatch for large tensors (>65535*256 elements)
-    // When using 2D dispatch: linear_idx = x + y * MAX_WORKGROUPS_X * WORKGROUP_SIZE
-    let linear_idx = global_id.x + global_id.y * MAX_WORKGROUPS_X * WORKGROUP_SIZE;
-
-    // Each thread processes 2 BF16 values (one u32 contains 2 bf16)
-    let local_pair_idx = linear_idx;
-    let local_elem_idx = local_pair_idx * 2u;
-
-    if (local_elem_idx >= u.num_elements) {
-        return;
-    }
-
-    // Apply offsets for chunked processing
-    let input_pair_idx = (u.input_offset / 2u) + local_pair_idx;
-    let output_elem_idx = u.output_offset + local_elem_idx;
-
-    let packed = input[input_pair_idx];
-
-    // Extract two BF16 values and convert to F32
-    // BF16 is upper 16 bits of F32, so shift left by 16
-    let bf16_lo = packed & 0xFFFFu;
-    let bf16_hi = (packed >> 16u) & 0xFFFFu;
-
-    // Convert by shifting to F32 position
-    output[output_elem_idx] = bitcast<f32>(bf16_lo << 16u);
-
-    if (local_elem_idx + 1u < u.num_elements) {
-        output[output_elem_idx + 1u] = bitcast<f32>(bf16_hi << 16u);
-    }
-}
-
-// Single-element version for odd-sized tensors
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main_single(@builtin(global_invocation_id) global_id: vec3<u32>) {
-    let idx = global_id.x;
-
-    if (idx >= u.num_elements) {
-        return;
-    }
-
-    // Read as u16 pairs, extract the right one
-    let pair_idx = idx / 2u;
-    let packed = input[pair_idx];
-
-    let bf16 = select(packed & 0xFFFFu, (packed >> 16u) & 0xFFFFu, (idx & 1u) == 1u);
-    output[idx] = bitcast<f32>(bf16 << 16u);
-}

package/src/gpu/kernels/bias_add.wgsl

@@ -1,42 +0,0 @@
-// bias_add.wgsl
-
-/**
- * Bias Add Kernel
- *
- * Adds per-channel bias to a 2D tensor in-place.
- *
- * data layout: [numTokens, dim] flattened, with optional byte offset
- * bias layout: [N, dim] where we select slice at bias_offset
- */
-
-struct Uniforms {
-    num_tokens: u32,
-    dim: u32,
-    data_offset: u32,  // byte offset into data buffer (divide by 4 for F32)
-    bias_offset: u32,  // byte offset into bias buffer (divide by 4 for F32)
-    token_stride: u32,
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read_write> data: array<f32>;
-@group(0) @binding(2) var<storage, read> bias: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let d = gid.x;
-    let token = gid.z * max(u.token_stride, 1u) + gid.y;
-    if (token >= u.num_tokens || d >= u.dim) {
-        return;
-    }
-
-    // Convert byte offsets to F32 indices
-    let data_base = u.data_offset / 4u;
-    let bias_base = u.bias_offset / 4u;
-    let idx = token * u.dim + d;
-    data[data_base + idx] = data[data_base + idx] + bias[bias_base + d];
-}

package/src/gpu/kernels/bias_add_f16.wgsl

@@ -1,47 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/bias_add.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// bias_add_f16.wgsl
-
-/**
- * Bias Add Kernel (F16)
- *
- * Adds per-channel bias to a 2D tensor in-place.
- *
- * data layout: [numTokens, dim] flattened, with optional byte offset
- * bias layout: [N, dim] where we select slice at bias_offset
- */
-
-enable f16;
-
-struct Uniforms {
-    num_tokens: u32,
-    dim: u32,
-    data_offset: u32,  // byte offset into data buffer (divide by 2 for F16)
-    bias_offset: u32,  // byte offset into bias buffer (divide by 2 for F16)
-    token_stride: u32,
-    _pad0: u32,
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read_write> data: array<f16>;
-@group(0) @binding(2) var<storage, read> bias: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let d = gid.x;
-    let token = gid.z * max(u.token_stride, 1u) + gid.y;
-    if (token >= u.num_tokens || d >= u.dim) {
-        return;
-    }
-
-    // Convert byte offsets to F16 indices
-    let data_base = u.data_offset / 2u;
-    let bias_base = u.bias_offset / 2u;
-    let idx = token * u.dim + d;
-    let out = f32(data[data_base + idx]) + f32(bias[bias_base + d]);
-    data[data_base + idx] = f16(out);
-}

package/src/gpu/kernels/cast.d.ts

@@ -1,67 +0,0 @@
-/**
- * Type Casting Kernels
- *
- * Provides GPU-based type conversions:
- * - F32 to F16
- * - F16 to F32
- * - BF16 to F32
- */
-
-import type { Tensor } from '../tensor.js';
-import type { CommandRecorder } from '../command-recorder.js';
-import type { OutputBufferOptions } from './types.js';
-
-/** Cast kernel options */
-export interface CastOptions extends OutputBufferOptions {}
-
-/**
- * Cast F32 buffer to F16 on GPU
- */
-export declare function castF32ToF16(
-  input: Tensor,
-  options?: CastOptions
-): Promise<Tensor>;
-
-/**
- * Cast F16 buffer to F32 on GPU
- */
-export declare function castF16ToF32(
-  input: Tensor,
-  options?: CastOptions
-): Promise<Tensor>;
-
-/**
- * Record F32 to F16 cast (batched, no submit)
- */
-export declare function recordCastF32ToF16(
-  recorder: CommandRecorder,
-  input: Tensor,
-  options?: CastOptions
-): Promise<Tensor>;
-
-/**
- * Record F16 to F32 cast (batched, no submit)
- */
-export declare function recordCastF16ToF32(
-  recorder: CommandRecorder,
-  input: Tensor,
-  options?: CastOptions
-): Promise<Tensor>;
-
-/**
- * Convert BF16 buffer to F32 on GPU
- */
-export declare function runBF16ToF32(
-  input: GPUBuffer,
-  shape: readonly number[],
-  name?: string
-): Promise<Tensor>;
-
-/**
- * Convert BF16 buffer to F16 on GPU (no intermediate F32 buffer)
- */
-export declare function runBF16ToF16(
-  input: GPUBuffer,
-  shape: readonly number[],
-  name?: string
-): Promise<Tensor>;