@simulatte/doppler 0.1.9 → 0.2.0

package/src/gpu/kernels/repeat_channels.js

@@ -1,68 +0,0 @@
-import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
-import { createTensor, dtypeBytes } from '../tensor.js';
-import { unifiedKernelWrapper } from './utils.js';
-import { selectRuleValue } from './rule-registry.js';
-import { WORKGROUP_SIZES } from './constants.js';
-
-function selectRepeatChannelsVariant(dtype) {
-  return selectRuleValue('repeatChannels', 'variant', { dtype });
-}
-
-async function _repeatChannels(target, input, options = {}) {
-  const {
-    inChannels,
-    height,
-    width,
-    repeats,
-    outputBuffer = null,
-  } = options;
-
-  if (
-    !Number.isFinite(inChannels) ||
-    !Number.isFinite(height) ||
-    !Number.isFinite(width) ||
-    !Number.isFinite(repeats) ||
-    repeats < 1
-  ) {
-    throw new Error('RepeatChannels requires inChannels, height, width, and repeats.');
-  }
-
-  const outChannels = inChannels * repeats;
-  const variant = selectRepeatChannelsVariant(input.dtype);
-  const bytesPerElement = dtypeBytes(input.dtype);
-  const outputSize = outChannels * height * width * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'repeat_channels_output');
-  const ownedOutput = outputBuffer ? null : output;
-
-  try {
-    await unifiedKernelWrapper(
-      'repeat_channels',
-      target,
-      variant,
-      [input, output],
-      {
-        in_channels: inChannels,
-        height,
-        width,
-        repeats,
-        _pad0: 0,
-      },
-      [Math.ceil((height * width) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
-    );
-
-    return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
-  } catch (error) {
-    if (ownedOutput) {
-      releaseBuffer(ownedOutput);
-    }
-    throw error;
-  }
-}
-
-export async function runRepeatChannels(input, options = {}) {
-  return _repeatChannels(null, input, options);
-}
-
-export async function recordRepeatChannels(recorder, input, options = {}) {
-  return _repeatChannels(recorder, input, options);
-}

package/src/gpu/kernels/repeat_channels.wgsl

@@ -1,28 +0,0 @@
-override WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    in_channels: u32,
-    height: u32,
-    width: u32,
-    repeats: u32,
-    _pad0: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> input: array<f32>;
-@group(0) @binding(2) var<storage, read_write> output: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let spatial = u.height * u.width;
-    let out_channels = u.in_channels * u.repeats;
-    let spatial_idx = gid.x;
-    let out_channel = gid.y;
-    if (out_channel >= out_channels || spatial_idx >= spatial) {
-        return;
-    }
-
-    let channel = out_channel / u.repeats;
-    let idx = out_channel * spatial + spatial_idx;
-    output[idx] = input[channel * spatial + spatial_idx];
-}

package/src/gpu/kernels/repeat_channels_f16.wgsl

@@ -1,30 +0,0 @@
-enable f16;
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-struct Uniforms {
-    in_channels: u32,
-    height: u32,
-    width: u32,
-    repeats: u32,
-    _pad0: u32,
-}
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> input: array<f16>;
-@group(0) @binding(2) var<storage, read_write> output: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let spatial = u.height * u.width;
-    let out_channels = u.in_channels * u.repeats;
-    let spatial_idx = gid.x;
-    let out_channel = gid.y;
-    if (out_channel >= out_channels || spatial_idx >= spatial) {
-        return;
-    }
-
-    let channel = out_channel / u.repeats;
-    let idx = out_channel * spatial + spatial_idx;
-    output[idx] = input[channel * spatial + spatial_idx];
-}

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

@@ -1,74 +0,0 @@
-/**
- * Residual Connection Kernels
- *
- * Provides element-wise addition operations for:
- * - Residual connections (add two tensors)
- * - Bias addition
- */
-
-import type { Tensor } from '../tensor.js';
-import type { CommandRecorder } from '../command-recorder.js';
-import type { OutputBufferOptions } from './types.js';
-
-/** Residual kernel options */
-export interface ResidualOptions extends OutputBufferOptions {
-  useVec4?: boolean;
-  dataOffset?: number;
-  biasOffset?: number;
-}
-
-/** Context for residual variant selection rules. */
-export interface ResidualSelectionContext {
-  outputDtype: 'f16' | 'f32';
-  useVec4: boolean;
-}
-
-/** Context for bias add variant selection rules. */
-export interface BiasAddSelectionContext {
-  dataDtype: 'f16' | 'f32';
-  biasDtype: 'f16' | 'f32';
-}
-
-/**
- * Run residual add (element-wise addition)
- */
-export declare function runResidualAdd(
-  a: Tensor,
-  b: Tensor,
-  size: number,
-  options?: ResidualOptions
-): Promise<Tensor>;
-
-/**
- * Run bias add
- */
-export declare function runBiasAdd(
-  data: Tensor,
-  bias: Tensor,
-  numTokens: number,
-  dim: number,
-  options?: ResidualOptions
-): Promise<Tensor>;
-
-/**
- * Record residual add (batched, no submit)
- */
-export declare function recordResidualAdd(
-  recorder: CommandRecorder,
-  a: Tensor,
-  b: Tensor,
-  size: number,
-  options?: ResidualOptions
-): Promise<Tensor>;
-
-/**
- * Record bias add (batched, no submit)
- */
-export declare function recordBiasAdd(
-  recorder: CommandRecorder,
-  data: Tensor,
-  bias: Tensor,
-  numTokens: number,
-  dim: number,
-  options?: ResidualOptions
-): Promise<Tensor>;

package/src/gpu/kernels/residual.js

@@ -1,173 +0,0 @@
-
-import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
-import { createTensor, inferOutputDtype, dtypeBytes } from '../tensor.js';
-import { WORKGROUP_SIZES, VEC4_ELEMENTS_PER_WG } from './constants.js';
-import { unifiedKernelWrapper } from './utils.js';
-import { castF16ToF32, castF32ToF16, recordCastF16ToF32, recordCastF32ToF16 } from './cast.js';
-import { selectRuleValue } from './rule-registry.js';
-
-
-function selectResidualVariant(outputDtype, useVec4) {
-  return selectRuleValue(
-    'residual',
-    'residualVariant',
-    { outputDtype, useVec4 }
-  );
-}
-
-function selectBiasAddVariant(dataDtype, biasDtype) {
-  return selectRuleValue(
-    'residual',
-    'biasAddVariant',
-    { dataDtype, biasDtype }
-  );
-}
-
-async function alignResidualInputs(a, b, recorder) {
-  if (a.dtype === b.dtype) {
-    return { a, b, temps: [] };
-  }
-  if (a.dtype === 'f16' && b.dtype === 'f32') {
-    const casted = recorder ? await recordCastF16ToF32(recorder, a) : await castF16ToF32(a);
-    return { a: casted, b, temps: [casted.buffer] };
-  }
-  if (a.dtype === 'f32' && b.dtype === 'f16') {
-    const casted = recorder ? await recordCastF16ToF32(recorder, b) : await castF16ToF32(b);
-    return { a, b: casted, temps: [casted.buffer] };
-  }
-  return { a, b, temps: [] };
-}
-
-async function alignBiasTensor(data, bias, recorder) {
-  if (data.dtype === bias.dtype) {
-    return { bias, temps: [] };
-  }
-  if (data.dtype === 'f16' && bias.dtype === 'f32') {
-    const casted = recorder ? await recordCastF32ToF16(recorder, bias) : await castF32ToF16(bias);
-    return { bias: casted, temps: [casted.buffer] };
-  }
-  if (data.dtype === 'f32' && bias.dtype === 'f16') {
-    const casted = recorder ? await recordCastF16ToF32(recorder, bias) : await castF16ToF32(bias);
-    return { bias: casted, temps: [casted.buffer] };
-  }
-  return { bias, temps: [] };
-}
-
-function cleanupTemps(temps, recorder) {
-  for (const temp of temps) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(temp);
-    } else {
-      releaseBuffer(temp);
-    }
-  }
-}
-
-function planResidualDispatch(target, size, elementsPerWorkgroup) {
-  const device = target?.device;
-  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
-    ? device.limits.maxComputeWorkgroupsPerDimension
-    : 65535;
-  const dispatchStride = Math.min(size, maxPerDim * elementsPerWorkgroup);
-  return {
-    dispatchStride,
-    workgroups: [
-      Math.ceil(dispatchStride / elementsPerWorkgroup),
-      Math.ceil(size / dispatchStride),
-      1,
-    ],
-  };
-}
-
-async function _residualAdd(target, a, b, size, options = {}) {
-  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
-  const { useVec4 = true, outputBuffer = null } = options;
-  const ownsOutput = outputBuffer == null;
-
-  const { a: aAligned, b: bAligned, temps } = await alignResidualInputs(a, b, recorder);
-  const outputDtype = inferOutputDtype(aAligned, bAligned);
-  const bytesPerElement = dtypeBytes(outputDtype);
-
-  const variant = selectResidualVariant(outputDtype, useVec4);
-  const outputSize = size * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'residual_output');
-
-  const dispatchPlan = planResidualDispatch(
-    target,
-    size,
-    useVec4 ? VEC4_ELEMENTS_PER_WG : WORKGROUP_SIZES.DEFAULT
-  );
-
-  try {
-    await unifiedKernelWrapper(
-      'residual', target, variant,
-      [aAligned, bAligned, output],
-      { size, scale: 1, _pad1: dispatchPlan.dispatchStride, _pad2: 0 },
-      dispatchPlan.workgroups
-    );
-    return createTensor(output, outputDtype, [size], 'residual_output');
-  } catch (error) {
-    if (ownsOutput) {
-      releaseBuffer(output);
-    }
-    throw error;
-  } finally {
-    cleanupTemps(temps, recorder);
-  }
-}
-
-async function _biasAdd(target, data, bias, numTokens, dim, options = {}) {
-  const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
-  const { dataOffset = 0, biasOffset = 0 } = options;
-
-  const { bias: biasAligned, temps } = await alignBiasTensor(data, bias, recorder);
-  const variant = selectBiasAddVariant(data.dtype, biasAligned.dtype);
-  const device = target?.device;
-  const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
-    ? device.limits.maxComputeWorkgroupsPerDimension
-    : 65535;
-  const tokenStride = Math.min(numTokens, maxPerDim);
-
-  const workgroups = [
-    Math.ceil(dim / WORKGROUP_SIZES.DEFAULT),
-    tokenStride,
-    Math.ceil(numTokens / tokenStride),
-  ];
-
-  try {
-    await unifiedKernelWrapper(
-      'bias_add', target, variant,
-      [data, biasAligned],
-      {
-        num_tokens: numTokens,
-        dim,
-        data_offset: dataOffset,
-        bias_offset: biasOffset,
-        token_stride: tokenStride,
-        _pad0: 0,
-        _pad1: 0,
-        _pad2: 0,
-      },
-      workgroups
-    );
-    return createTensor(data.buffer, data.dtype, [numTokens, dim], 'bias_add_output');
-  } finally {
-    cleanupTemps(temps, recorder);
-  }
-}
-
-export async function runResidualAdd(a, b, size, options = {}) {
-  return _residualAdd(null, a, b, size, options);
-}
-
-export async function recordResidualAdd(recorder, a, b, size, options = {}) {
-  return _residualAdd(recorder, a, b, size, options);
-}
-
-export async function runBiasAdd(data, bias, numTokens, dim, options = {}) {
-  return _biasAdd(null, data, bias, numTokens, dim, options);
-}
-
-export async function recordBiasAdd(recorder, data, bias, numTokens, dim, options = {}) {
-  return _biasAdd(recorder, data, bias, numTokens, dim, options);
-}

package/src/gpu/kernels/residual.wgsl

@@ -1,56 +0,0 @@
-// residual.wgsl
-
-/**
- * Residual Add Kernel
- *
- * Performs element-wise addition for residual connections.
- * output = a + b
- */
-
-struct Uniforms {
-    size: u32,     // Total number of elements
-    scale: f32,    // Scale factor for add_scaled
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> a: array<f32>;
-@group(0) @binding(2) var<storage, read> b: array<f32>;
-@group(0) @binding(3) var<storage, read_write> output: array<f32>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 1u);
-    let idx = gid.y * dispatch_stride + gid.x;
-    if (idx >= u.size) {
-        return;
-    }
-    output[idx] = a[idx] + b[idx];
-}
-
-// In-place version: output = output + b
-// Note: Caller should copy 'a' to 'output' first, then call this kernel
-// This avoids requiring a different bind group layout with read_write on 'a'
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn add_inplace(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 1u);
-    let idx = gid.y * dispatch_stride + gid.x;
-    if (idx >= u.size) {
-        return;
-    }
-    output[idx] = output[idx] + b[idx];
-}
-
-// Fused residual + scale: output = a + scale * b
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn add_scaled(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 1u);
-    let idx = gid.y * dispatch_stride + gid.x;
-    if (idx >= u.size) {
-        return;
-    }
-    output[idx] = a[idx] + u.scale * b[idx];
-}

package/src/gpu/kernels/residual_f16.wgsl

@@ -1,36 +0,0 @@
-// AUTO-GENERATED from src/gpu/kernels/residual.wgsl.
-// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
-// residual_f16.wgsl
-
-/**
- * Residual Add Kernel (F16)
- *
- * Performs element-wise addition for residual connections.
- * output = a + b
- */
-
-enable f16;
-
-struct Uniforms {
-    size: u32,     // Total number of elements
-    scale: f32,    // Scale factor for add_scaled
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE: u32 = 256u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> a: array<f16>;
-@group(0) @binding(2) var<storage, read> b: array<f16>;
-@group(0) @binding(3) var<storage, read_write> output: array<f16>;
-
-@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 1u);
-    let idx = gid.y * dispatch_stride + gid.x;
-    if (idx >= u.size) {
-        return;
-    }
-    output[idx] = f16(f32(a[idx]) + f32(b[idx]));
-}

package/src/gpu/kernels/residual_f16_vec4.wgsl

@@ -1,48 +0,0 @@
-// residual_f16_vec4.wgsl
-
-/**
- * Residual Add Kernel (F16, vec4)
- *
- * Vectorized element-wise addition for residual connections in F16.
- */
-
-enable f16;
-
-struct Uniforms {
-    size: u32,     // Total number of elements
-    scale: f32,    // Scale factor (unused here, kept for layout parity)
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> a: array<f16>;
-@group(0) @binding(2) var<storage, read> b: array<f16>;
-@group(0) @binding(3) var<storage, read_write> output: array<f16>;
-
-// Vectorized version for better throughput
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn add_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 4u);
-    let idx = gid.y * dispatch_stride + gid.x * 4u;
-    let size = u.size;
-
-    if (idx >= size) {
-        return;
-    }
-
-    let remaining = min(4u, size - idx);
-
-    if (remaining >= 4u) {
-        output[idx] = f16(f32(a[idx]) + f32(b[idx]));
-        output[idx + 1u] = f16(f32(a[idx + 1u]) + f32(b[idx + 1u]));
-        output[idx + 2u] = f16(f32(a[idx + 2u]) + f32(b[idx + 2u]));
-        output[idx + 3u] = f16(f32(a[idx + 3u]) + f32(b[idx + 3u]));
-    } else {
-        for (var i = 0u; i < remaining; i = i + 1u) {
-            output[idx + i] = f16(f32(a[idx + i]) + f32(b[idx + i]));
-        }
-    }
-}

package/src/gpu/kernels/residual_vec4.wgsl

@@ -1,47 +0,0 @@
-// residual_vec4.wgsl
-
-/**
- * Residual Add Kernel (vec4)
- *
- * Vectorized element-wise addition for residual connections.
- */
-
-struct Uniforms {
-    size: u32,     // Total number of elements
-    scale: f32,    // Scale factor (unused here, kept for layout parity)
-    _pad1: u32,
-    _pad2: u32,
-}
-
-override WORKGROUP_SIZE_VEC4: u32 = 64u;
-
-@group(0) @binding(0) var<uniform> u: Uniforms;
-@group(0) @binding(1) var<storage, read> a: array<f32>;
-@group(0) @binding(2) var<storage, read> b: array<f32>;
-@group(0) @binding(3) var<storage, read_write> output: array<f32>;
-
-// Vectorized version for better throughput
-@compute @workgroup_size(WORKGROUP_SIZE_VEC4, 1, 1)
-fn add_vec4(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dispatch_stride = max(u._pad1, 4u);
-    let idx = gid.y * dispatch_stride + gid.x * 4u;
-    let size = u.size;
-
-    if (idx >= size) {
-        return;
-    }
-
-    // Handle up to 4 elements at a time
-    let remaining = min(4u, size - idx);
-
-    if (remaining >= 4u) {
-        output[idx] = a[idx] + b[idx];
-        output[idx + 1u] = a[idx + 1u] + b[idx + 1u];
-        output[idx + 2u] = a[idx + 2u] + b[idx + 2u];
-        output[idx + 3u] = a[idx + 3u] + b[idx + 3u];
-    } else {
-        for (var i = 0u; i < remaining; i = i + 1u) {
-            output[idx + i] = a[idx + i] + b[idx + i];
-        }
-    }
-}

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

@@ -1,53 +0,0 @@
-/**
- * RMSNorm Kernels
- *
- * Provides RMS normalization with optional residual connection.
- */
-
-import type { Tensor } from '../tensor.js';
-import type { CommandRecorder } from '../command-recorder.js';
-import type { OutputBufferOptions } from './types.js';
-import type { WeightBuffer, TensorLike } from '../weight-buffer.js';
-
-/** RMSNorm kernel options */
-export interface RMSNormOptions extends OutputBufferOptions {
-  batchSize?: number;
-  hiddenSize?: number | null;
-  residual?: Tensor | null;
-  /** Use (1+w)*x normalization for Gemma 2/3 */
-  rmsNormWeightOffset?: boolean;
-}
-
-/** Context for RMSNorm variant selection rules. */
-export interface RMSNormSelectionContext {
-  isF16: boolean;
-  residual: boolean;
-  hasSubgroups: boolean;
-  isSmall: boolean;
-}
-
-/**
- * Select RMSNorm kernel variant based on options, tensor dtypes, and GPU capabilities.
- */
-export declare function selectRMSNormKernel(options?: RMSNormOptions, isF16?: boolean): string;
-
-/**
- * Run RMSNorm
- */
-export declare function runRMSNorm(
-  input: Tensor,
-  weight: GPUBuffer | WeightBuffer | TensorLike,
-  eps: number,
-  options?: RMSNormOptions
-): Promise<Tensor>;
-
-/**
- * Record RMSNorm (batched, no submit)
- */
-export declare function recordRMSNorm(
-  recorder: CommandRecorder,
-  input: Tensor,
-  weight: GPUBuffer | WeightBuffer | TensorLike,
-  eps: number,
-  options?: RMSNormOptions
-): Promise<Tensor>;