@simulatte/doppler 0.1.5 → 0.1.7

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

@@ -16,6 +16,7 @@ export interface EnergyUpdateOptions {
 export interface EnergyQuintelUpdateOptions {
   count?: number;
   size?: number;
+  flags?: number;
   stepSize?: number;
   gradientScale?: number;
   countDiff?: number;
@@ -26,48 +27,29 @@ export interface EnergyQuintelUpdateOptions {
   centerTarget?: number;
   clampMin?: number;
   clampMax?: number;
-  rules?: {
-    mirrorX?: boolean;
-    mirrorY?: boolean;
-    diagonal?: boolean;
-    count?: boolean;
-    center?: boolean;
-  };
 }
 
 export interface EnergyQuintelReduceOptions {
   count?: number;
   size?: number;
+  flags?: number;
   symmetryWeight?: number;
   centerWeight?: number;
   binarizeWeight?: number;
   centerTarget?: number;
-  rules?: {
-    mirrorX?: boolean;
-    mirrorY?: boolean;
-    diagonal?: boolean;
-    count?: boolean;
-    center?: boolean;
-  };
   outputBuffer?: GPUBuffer | null;
 }
 
 export interface EnergyQuintelGradOptions {
   count?: number;
   size?: number;
+  flags?: number;
   countDiff?: number;
   symmetryWeight?: number;
   countWeight?: number;
   centerWeight?: number;
   binarizeWeight?: number;
   centerTarget?: number;
-  rules?: {
-    mirrorX?: boolean;
-    mirrorY?: boolean;
-    diagonal?: boolean;
-    count?: boolean;
-    center?: boolean;
-  };
   outputBuffer?: GPUBuffer | null;
 }
 

package/src/gpu/kernels/energy.js

@@ -1,5 +1,5 @@
 import { getDevice } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { WORKGROUP_SIZES } from './constants.js';
 import { dispatch, recordDispatch } from './dispatch.js';
@@ -61,15 +61,14 @@ function resolveQuintelSize(state, sizeOverride) {
   return null;
 }
 
-function buildQuintelFlags(rules, binarizeWeight) {
-  let flags = 0;
-  if (rules?.mirrorX) flags |= 1;
-  if (rules?.mirrorY) flags |= 2;
-  if (rules?.diagonal) flags |= 4;
-  if (rules?.count) flags |= 8;
-  if (rules?.center) flags |= 16;
-  if (Number.isFinite(binarizeWeight) && binarizeWeight !== 0) flags |= 32;
-  return flags >>> 0;
+function resolveQuintelFlags(options, op) {
+  if (options.rules !== undefined) {
+    throw new Error(`${op}: quintel kernel flags must be resolved before dispatch.`);
+  }
+  if (!Number.isFinite(options.flags)) {
+    throw new Error(`${op}: flags is required for quintel kernels.`);
+  }
+  return options.flags >>> 0;
 }
 
 function resolveExecution(recorder) {
@@ -103,6 +102,12 @@ function releaseUniformBuffer(execution, uniformBuffer) {
   }
 }
 
+function releaseOwnedBuffer(ownedBuffer) {
+  if (ownedBuffer) {
+    releaseBuffer(ownedBuffer);
+  }
+}
+
 function writeQuintelUpdateUniform(view, params) {
   view.setUint32(0, params.elementCount, true);
   view.setUint32(4, params.boardSize, true);
@@ -149,6 +154,7 @@ async function executeEnergyEval(recorder, state, target, options = {}, op) {
 
   const outputSize = elementCount * 4;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'energy_eval_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   const variant = selectEnergyEvalVariant(state.dtype);
   const pipeline = await getPipelineFast('energy_eval', variant);
@@ -157,23 +163,27 @@ async function executeEnergyEval(recorder, state, target, options = {}, op) {
     view.setUint32(0, elementCount, true);
     view.setFloat32(4, scale, true);
   });
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'energy_eval_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: state.buffer } },
+        { binding: 2, resource: { buffer: target.buffer } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'energy_eval_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: state.buffer } },
-      { binding: 2, resource: { buffer: target.buffer } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
-  dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_eval');
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return createTensor(output, 'f32', [elementCount], 'energy_eval_output');
+    const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
+    dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_eval');
+    return createTensor(output, 'f32', [elementCount], 'energy_eval_output');
+  } catch (error) {
+    releaseOwnedBuffer(ownedOutput);
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 async function executeEnergyUpdate(recorder, state, target, options = {}, op) {
@@ -191,21 +201,23 @@ async function executeEnergyUpdate(recorder, state, target, options = {}, op) {
     view.setFloat32(8, gradientScale, true);
   });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'energy_update_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: state.buffer } },
-      { binding: 2, resource: { buffer: target.buffer } },
-    ],
-  });
-
-  const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
-  dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_update');
-  releaseUniformBuffer(execution, uniformBuffer);
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'energy_update_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: state.buffer } },
+        { binding: 2, resource: { buffer: target.buffer } },
+      ],
+    });
 
-  return state;
+    const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
+    dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_update');
+    return state;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 async function executeEnergyQuintelUpdate(recorder, state, options = {}, op) {
@@ -224,7 +236,6 @@ async function executeEnergyQuintelUpdate(recorder, state, options = {}, op) {
     centerTarget = 1.0,
     clampMin = 0.0,
     clampMax = 1.0,
-    rules = {},
   } = options;
   const elementCount = inferCount(state, count);
   const boardSize = resolveQuintelSize(state, size);
@@ -234,7 +245,7 @@ async function executeEnergyQuintelUpdate(recorder, state, options = {}, op) {
 
   const variant = selectEnergyQuintelUpdateVariant(state.dtype);
   const pipeline = await getPipelineFast('energy_quintel_update', variant);
-  const flags = buildQuintelFlags(rules, binarizeWeight);
+  const flags = resolveQuintelFlags(options, op);
 
   const uniformBuffer = createUniformBuffer(execution, 'energy_quintel_uniforms', 64, (view) => {
     writeQuintelUpdateUniform(view, {
@@ -254,20 +265,22 @@ async function executeEnergyQuintelUpdate(recorder, state, options = {}, op) {
     });
   });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'energy_quintel_update_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: state.buffer } },
-    ],
-  });
-
-  const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
-  dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_update');
-  releaseUniformBuffer(execution, uniformBuffer);
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'energy_quintel_update_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: state.buffer } },
+      ],
+    });
 
-  return state;
+    const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
+    dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_update');
+    return state;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 async function executeEnergyQuintelReduce(recorder, state, options = {}, op) {
@@ -280,7 +293,6 @@ async function executeEnergyQuintelReduce(recorder, state, options = {}, op) {
     centerWeight = 1.0,
     binarizeWeight = 0.0,
     centerTarget = 1.0,
-    rules = {},
     outputBuffer = null,
   } = options;
   const elementCount = inferCount(state, count);
@@ -291,7 +303,7 @@ async function executeEnergyQuintelReduce(recorder, state, options = {}, op) {
 
   const variant = selectEnergyQuintelReduceVariant(state.dtype);
   const pipeline = await getPipelineFast('energy_quintel_reduce', variant);
-  const flags = buildQuintelFlags(rules, binarizeWeight);
+  const flags = resolveQuintelFlags(options, op);
 
   const uniformBuffer = createUniformBuffer(execution, 'energy_quintel_reduce_uniforms', 48, (view) => {
     writeQuintelReduceUniform(view, {
@@ -308,21 +320,27 @@ async function executeEnergyQuintelReduce(recorder, state, options = {}, op) {
   const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
   const outputSize = workgroups * 16;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'energy_quintel_reduce_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'energy_quintel_reduce_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: state.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'energy_quintel_reduce_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: state.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_reduce');
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return createTensor(output, 'f32', [workgroups, 4], 'energy_quintel_reduce_output');
+    dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_reduce');
+    return createTensor(output, 'f32', [workgroups, 4], 'energy_quintel_reduce_output');
+  } catch (error) {
+    releaseOwnedBuffer(ownedOutput);
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 async function executeEnergyQuintelGrad(recorder, state, options = {}, op) {
@@ -337,7 +355,6 @@ async function executeEnergyQuintelGrad(recorder, state, options = {}, op) {
     centerWeight = 1.0,
     binarizeWeight = 0.0,
     centerTarget = 1.0,
-    rules = {},
     outputBuffer = null,
   } = options;
   const elementCount = inferCount(state, count);
@@ -348,7 +365,7 @@ async function executeEnergyQuintelGrad(recorder, state, options = {}, op) {
 
   const variant = selectEnergyQuintelGradVariant(state.dtype);
   const pipeline = await getPipelineFast('energy_quintel_grad', variant);
-  const flags = buildQuintelFlags(rules, binarizeWeight);
+  const flags = resolveQuintelFlags(options, op);
 
   const uniformBuffer = createUniformBuffer(execution, 'energy_quintel_grad_uniforms', 64, (view) => {
     writeQuintelGradUniform(view, {
@@ -366,22 +383,28 @@ async function executeEnergyQuintelGrad(recorder, state, options = {}, op) {
 
   const outputSize = elementCount * 4;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'energy_quintel_grad_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    const bindGroup = execution.device.createBindGroup({
+      label: 'energy_quintel_grad_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: state.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  const bindGroup = execution.device.createBindGroup({
-    label: 'energy_quintel_grad_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: state.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
-  dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_grad');
-  releaseUniformBuffer(execution, uniformBuffer);
-
-  return createTensor(output, 'f32', [elementCount], 'energy_quintel_grad_output');
+    const workgroups = Math.ceil(elementCount / WORKGROUP_SIZES.DEFAULT);
+    dispatchEnergy(execution, pipeline, bindGroup, workgroups, 'energy_quintel_grad');
+    return createTensor(output, 'f32', [elementCount], 'energy_quintel_grad_output');
+  } catch (error) {
+    releaseOwnedBuffer(ownedOutput);
+    throw error;
+  } finally {
+    releaseUniformBuffer(execution, uniformBuffer);
+  }
 }
 
 export async function runEnergyEval(state, target, options = {}) {

package/src/gpu/kernels/feature-check.js

@@ -16,7 +16,7 @@ export function hasRequiredFeatures(
   for (const feature of required) {
     if (feature === 'shader-f16' && !capabilities.hasF16) return false;
     if (feature === 'subgroups' && !capabilities.hasSubgroups) return false;
-    if (feature === 'subgroups-f16' && !capabilities.hasSubgroups) return false;
+    if (feature === 'subgroups-f16' && !capabilities.hasSubgroupsF16) return false;
   }
   return true;
 }

package/src/gpu/kernels/fused_ffn.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice, getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { KernelBase } from './kernel-base.js';
 import { createUniformBufferWithView } from './utils.js';
@@ -77,6 +77,17 @@ function resolveSwigluLimit(value, context) {
   return value;
 }
 
+function releaseRunResources(uniformBuffer, ownedBuffers) {
+  if (uniformBuffer) {
+    uniformBuffer.destroy();
+  }
+  for (const buffer of ownedBuffers) {
+    if (buffer) {
+      releaseBuffer(buffer);
+    }
+  }
+}
+
 
 export async function runFusedFFN(
   input,
@@ -132,7 +143,8 @@ export async function runFusedFFN(
   const outputBytesPerElement = isF16Native ? 2 : 4;
   const outputDtype = isF16Native ? 'f16' : 'f32';
   const outputSize = batchSize * intermediateSize * outputBytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'fused_ffn_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'fused_ffn_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer
   const uniformBuffer = createFFNUniformBuffer(device, null, {
@@ -145,41 +157,42 @@ export async function runFusedFFN(
     swigluLimit: activation === 'silu' ? swigluLimit : null,
   });
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'fused_ffn_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: getBuffer(W_gate) } },
-      { binding: 3, resource: { buffer: getBuffer(W_up) } },
-      { binding: 4, resource: { buffer: output } },
-    ],
-  });
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'fused_ffn_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: getBuffer(W_gate) } },
+        { binding: 3, resource: { buffer: getBuffer(W_up) } },
+        { binding: 4, resource: { buffer: output } },
+      ],
+    });
+
+    let workgroupsX;
+    let workgroupsY = 1;
+
+    if (variant === 'multi') {
+      const outputsPerWg = 4;
+      workgroupsX = Math.ceil(intermediateSize / outputsPerWg);
+    } else if (variant === 'q4k' || variant === 'q4k_batched') {
+      const colsPerWg = 32;
+      workgroupsX = Math.ceil(intermediateSize / colsPerWg);
+      workgroupsY = variant === 'q4k_batched' ? batchSize : 1;
+    } else if (variant === 'batched' || variant === 'f16_native_batched') {
+      workgroupsX = intermediateSize;
+      workgroupsY = batchSize;
+    } else {
+      workgroupsX = intermediateSize;
+    }
 
-  // Calculate workgroups
-
-  let workgroupsX;
-  let workgroupsY = 1;
-
-  if (variant === 'multi') {
-    const outputsPerWg = 4;
-    workgroupsX = Math.ceil(intermediateSize / outputsPerWg);
-  } else if (variant === 'q4k' || variant === 'q4k_batched') {
-    // Q4K uses multi-column: 32 columns per workgroup
-    const colsPerWg = 32;
-    workgroupsX = Math.ceil(intermediateSize / colsPerWg);
-    workgroupsY = variant === 'q4k_batched' ? batchSize : 1;
-  } else if (variant === 'batched' || variant === 'f16_native_batched') {
-    workgroupsX = intermediateSize;
-    workgroupsY = batchSize;
-  } else {
-    workgroupsX = intermediateSize;
+    kernel.dispatch(pipeline, bindGroup, workgroupsX, workgroupsY);
+  } catch (error) {
+    releaseRunResources(uniformBuffer, [ownedOutput]);
+    throw error;
   }
 
-  kernel.dispatch(pipeline, bindGroup, workgroupsX, workgroupsY);
-
   uniformBuffer.destroy();
 
   return createTensor(output, outputDtype, [batchSize, intermediateSize], 'fused_ffn_output');
@@ -240,7 +253,8 @@ export async function recordFusedFFN(
   const outputBytesPerElement = isF16Native ? 2 : 4;
   const outputDtype = isF16Native ? 'f16' : 'f32';
   const outputSize = batchSize * intermediateSize * outputBytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'fused_ffn_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'fused_ffn_output');
+  const output = outputBuffer || ownedOutput;
 
   const uniformBuffer = createFFNUniformBuffer(device, recorder, {
     M: batchSize,
@@ -252,39 +266,44 @@ export async function recordFusedFFN(
     swigluLimit: activation === 'silu' ? swigluLimit : null,
   });
 
-  const bindGroup = device.createBindGroup({
-    label: 'fused_ffn_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: getBuffer(W_gate) } },
-      { binding: 3, resource: { buffer: getBuffer(W_up) } },
-      { binding: 4, resource: { buffer: output } },
-    ],
-  });
-
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'fused_ffn_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: getBuffer(W_gate) } },
+        { binding: 3, resource: { buffer: getBuffer(W_up) } },
+        { binding: 4, resource: { buffer: output } },
+      ],
+    });
+
+    let workgroupsX;
+    let workgroupsY = 1;
+
+    if (variant === 'multi') {
+      const outputsPerWg = 4;
+      workgroupsX = Math.ceil(intermediateSize / outputsPerWg);
+    } else if (variant === 'q4k' || variant === 'q4k_batched') {
+      const colsPerWg = 32;
+      workgroupsX = Math.ceil(intermediateSize / colsPerWg);
+      workgroupsY = variant === 'q4k_batched' ? batchSize : 1;
+    } else if (variant === 'batched' || variant === 'f16_native_batched') {
+      workgroupsX = intermediateSize;
+      workgroupsY = batchSize;
+    } else {
+      workgroupsX = intermediateSize;
+    }
 
-  let workgroupsX;
-  let workgroupsY = 1;
-
-  if (variant === 'multi') {
-    const outputsPerWg = 4;
-    workgroupsX = Math.ceil(intermediateSize / outputsPerWg);
-  } else if (variant === 'q4k' || variant === 'q4k_batched') {
-    // Q4K uses multi-column: 32 columns per workgroup
-    const colsPerWg = 32;
-    workgroupsX = Math.ceil(intermediateSize / colsPerWg);
-    workgroupsY = variant === 'q4k_batched' ? batchSize : 1;
-  } else if (variant === 'batched' || variant === 'f16_native_batched') {
-    workgroupsX = intermediateSize;
-    workgroupsY = batchSize;
-  } else {
-    workgroupsX = intermediateSize;
+    kernel.record(recorder, pipeline, bindGroup, workgroupsX, workgroupsY);
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
   }
 
-  kernel.record(recorder, pipeline, bindGroup, workgroupsX, workgroupsY);
-
   return createTensor(output, outputDtype, [batchSize, intermediateSize], 'fused_ffn_output');
 }