@simulatte/doppler 0.1.6 → 0.1.7

package/src/gpu/kernels/depthwise_conv2d.js

@@ -58,36 +58,46 @@ async function _depthwiseConv2D(target, input, weight, bias, options = {}) {
     device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
   }
 
-  await unifiedKernelWrapper(
-    'depthwise_conv2d',
-    target,
-    variant,
-    [input, weightBuffer, biasBuffer, output],
-    {
-      channels,
-      height,
-      width,
-      out_height: outHeight,
-      out_width: outWidth,
-      kernel_h: kernelH,
-      kernel_w: kernelW,
-      stride,
-      pad,
-      _pad0: 0,
-      _pad1: 0,
-    },
-    [Math.ceil(outSpatial / WORKGROUP_SIZES.DEFAULT), channels, 1]
-  );
+  try {
+    await unifiedKernelWrapper(
+      'depthwise_conv2d',
+      target,
+      variant,
+      [input, weightBuffer, biasBuffer, output],
+      {
+        channels,
+        height,
+        width,
+        out_height: outHeight,
+        out_width: outWidth,
+        kernel_h: kernelH,
+        kernel_w: kernelW,
+        stride,
+        pad,
+        _pad0: 0,
+        _pad1: 0,
+      },
+      [Math.ceil(outSpatial / WORKGROUP_SIZES.DEFAULT), channels, 1]
+    );
 
-  if (tempBias) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(tempBias);
-    } else {
+    if (tempBias) {
+      if (recorder) {
+        recorder.trackTemporaryBuffer(tempBias);
+      } else {
+        releaseBuffer(tempBias);
+      }
+    }
+
+    return createTensor(output, input.dtype, [channels, outHeight, outWidth], 'depthwise_conv2d_output');
+  } catch (error) {
+    if (tempBias) {
       releaseBuffer(tempBias);
     }
+    if (!outputBuffer) {
+      releaseBuffer(output);
+    }
+    throw error;
   }
-
-  return createTensor(output, input.dtype, [channels, outHeight, outWidth], 'depthwise_conv2d_output');
 }
 
 export async function runDepthwiseConv2D(input, weight, bias, options = {}) {

package/src/gpu/kernels/dequant.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 { GPU_LIMITS, TILE_SIZES, WORKGROUP_SIZES } from './constants.js';
 import { Q6K_BLOCK_BYTES, Q8_0_BLOCK_BYTES, Q8_0_BLOCK_SIZE } from '../../loader/quantization-constants.js';
@@ -69,6 +69,17 @@ export function createDequantBindGroupLayout() {
   ]);
 }
 
+function cleanupDequantResources(uniformBuffer, ownedBuffers) {
+  if (uniformBuffer) {
+    releaseUniformBuffer(uniformBuffer);
+  }
+  for (const buffer of ownedBuffers) {
+    if (buffer) {
+      releaseBuffer(buffer);
+    }
+  }
+}
+
 
 export async function dequantize(
   quantized,
@@ -76,12 +87,17 @@ export async function dequantize(
   options = {}
 ) {
   const device = getDevice();
+  const capabilities = getKernelCapabilities();
   const {
     outputOffset = 0,
     outputBuffer = null,
     outputDtype = 'f32',
   } = options;
 
+  if (outputDtype === 'f16' && capabilities?.hasF16 !== true) {
+    throw new Error('[dequantize] f16 output requires shader-f16 support.');
+  }
+
   // Select kernel
   const variant = selectDequantKernel({ ...options, outputDtype });
   const pipeline = await getPipelineFast('dequant', variant);
@@ -92,7 +108,8 @@ export async function dequantize(
   const outputSize = numBlocks * QK_K * bytesPerElem;
 
   // Create output buffer if not provided
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'dequant_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'dequant_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -108,21 +125,24 @@ export async function dequantize(
     device
   );
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'dequant_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: quantized } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = calculateDequantWorkgroups(variant, numBlocks);
-  dispatch(device, pipeline, bindGroup, workgroups, 'dequant');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'dequant_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: quantized } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = calculateDequantWorkgroups(variant, numBlocks);
+    dispatch(device, pipeline, bindGroup, workgroups, 'dequant');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
-  // Release uniform buffer back to cache (or destroy if not cached)
   releaseUniformBuffer(uniformBuffer);
 
   
@@ -140,7 +160,11 @@ export async function dequantizeRowwise(
   options = {}
 ) {
   const device = getDevice();
+  const capabilities = getKernelCapabilities();
   const { outputBuffer = null, outputDtype = 'f16' } = options;
+  if (outputDtype === 'f16' && capabilities?.hasF16 !== true) {
+    throw new Error('[dequantizeRowwise] f16 output requires shader-f16 support.');
+  }
   const finalOutputDtype = selectSharedRuleValue('shared', 'dtype', 'f16OrF32FromDtype', { dtype: outputDtype });
   const pipelineVariant = selectKernelRuleValue(
     'dequant',
@@ -157,7 +181,8 @@ export async function dequantizeRowwise(
   const bytesPerElem = finalOutputDtype === 'f16' ? 2 : 4;
   const outputSize = rows * K * bytesPerElem;
 
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'dequant_rowwise_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'dequant_rowwise_output');
+  const output = outputBuffer || ownedOutput;
 
   const uniformBuffer = createUniformBufferWithView(
     'dequant_rowwise_uniforms',
@@ -172,18 +197,23 @@ export async function dequantizeRowwise(
     device
   );
 
-  const bindGroup = device.createBindGroup({
-    label: 'dequant_rowwise_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: quantized } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = [numBlocks, 1, 1];
-  dispatch(device, pipeline, bindGroup, workgroups, 'dequant_rowwise');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'dequant_rowwise_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: quantized } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = [numBlocks, 1, 1];
+    dispatch(device, pipeline, bindGroup, workgroups, 'dequant_rowwise');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
   releaseUniformBuffer(uniformBuffer);
 
@@ -208,7 +238,8 @@ export async function dequantizeMXFP4(
   const pipeline = await getPipelineFast('dequant', 'mxfp4');
 
   const outputSize = totalElements * 4; // F32 output
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'mxfp4_dequant_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'mxfp4_dequant_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -224,26 +255,29 @@ export async function dequantizeMXFP4(
     device
   );
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'mxfp4_dequant_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: blocks } },
-      { binding: 2, resource: { buffer: scales } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT);
-  
-  const dispatchSize = [
-    Math.min(workgroups, GPU_LIMITS.MAX_WORKGROUPS),
-    Math.max(1, Math.ceil(workgroups / GPU_LIMITS.MAX_WORKGROUPS)),
-    1,
-  ];
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'mxfp4_dequant');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'mxfp4_dequant_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: blocks } },
+        { binding: 2, resource: { buffer: scales } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = [
+      Math.min(workgroups, GPU_LIMITS.MAX_WORKGROUPS),
+      Math.max(1, Math.ceil(workgroups / GPU_LIMITS.MAX_WORKGROUPS)),
+      1,
+    ];
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'mxfp4_dequant');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
   releaseUniformBuffer(uniformBuffer);
 
@@ -284,7 +318,8 @@ export async function dequantizeMXFP4Expert(
   const totalOutput = outDim * numGroups * 32;
   const bytesPerElement = outputDtype === 'f16' ? 2 : 4;
   const outputSize = totalOutput * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'mxfp4_expert_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'mxfp4_expert_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -301,26 +336,29 @@ export async function dequantizeMXFP4Expert(
     device
   );
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'mxfp4_expert_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: blocks } },
-      { binding: 2, resource: { buffer: scales } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil(totalOutput / WORKGROUP_SIZES.DEFAULT);
-  
-  const dispatchSize = [
-    Math.min(workgroups, GPU_LIMITS.MAX_WORKGROUPS),
-    Math.max(1, Math.ceil(workgroups / GPU_LIMITS.MAX_WORKGROUPS)),
-    1,
-  ];
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'mxfp4_expert');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'mxfp4_expert_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: blocks } },
+        { binding: 2, resource: { buffer: scales } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = Math.ceil(totalOutput / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = [
+      Math.min(workgroups, GPU_LIMITS.MAX_WORKGROUPS),
+      Math.max(1, Math.ceil(workgroups / GPU_LIMITS.MAX_WORKGROUPS)),
+      1,
+    ];
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'mxfp4_expert');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
   releaseUniformBuffer(uniformBuffer);
 
@@ -350,7 +388,8 @@ export async function dequantizeQ6K(
   const outputSize = numBlocks * QK_K * bytesPerElem;
 
   // Create output buffer if not provided
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'q6k_dequant_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'q6k_dequant_output');
+  const output = outputBuffer || ownedOutput;
 
   // Calculate workgroups for 2D dispatch
   const maxWorkgroups = GPU_LIMITS.MAX_WORKGROUPS;
@@ -370,26 +409,28 @@ export async function dequantizeQ6K(
     device
   );
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'q6k_dequant_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: quantized } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  // One workgroup per block, handle 2D dispatch for large counts
-  
-  const workgroups = [
-    workgroupsX,
-    numBlocks > maxWorkgroups ? Math.ceil(numBlocks / maxWorkgroups) : 1,
-    1
-  ];
-
-  dispatch(device, pipeline, bindGroup, workgroups, 'q6k_dequant');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'q6k_dequant_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: quantized } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = [
+      workgroupsX,
+      numBlocks > maxWorkgroups ? Math.ceil(numBlocks / maxWorkgroups) : 1,
+      1
+    ];
+
+    dispatch(device, pipeline, bindGroup, workgroups, 'q6k_dequant');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
   releaseUniformBuffer(uniformBuffer);
 
@@ -419,7 +460,8 @@ export async function dequantizeQ8_0(
   const outputSize = numBlocks * Q8_0_BLOCK_SIZE * bytesPerElem;
 
   // Create output buffer if not provided
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'q8_0_dequant_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'q8_0_dequant_output');
+  const output = outputBuffer || ownedOutput;
 
   // Calculate workgroups for 2D dispatch
   const maxWorkgroups = GPU_LIMITS.MAX_WORKGROUPS;
@@ -439,26 +481,28 @@ export async function dequantizeQ8_0(
     device
   );
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'q8_0_dequant_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: quantized } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  // One workgroup per block, handle 2D dispatch for large counts
-  
-  const workgroups = [
-    workgroupsX,
-    numBlocks > maxWorkgroups ? Math.ceil(numBlocks / maxWorkgroups) : 1,
-    1
-  ];
-
-  dispatch(device, pipeline, bindGroup, workgroups, 'q8_0_dequant');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'q8_0_dequant_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: quantized } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = [
+      workgroupsX,
+      numBlocks > maxWorkgroups ? Math.ceil(numBlocks / maxWorkgroups) : 1,
+      1
+    ];
+
+    dispatch(device, pipeline, bindGroup, workgroups, 'q8_0_dequant');
+  } catch (error) {
+    cleanupDequantResources(uniformBuffer, [ownedOutput]);
+    throw error;
+  }
 
   releaseUniformBuffer(uniformBuffer);
 
@@ -491,7 +535,8 @@ export async function recordDequantize(
   const outputSize = numBlocks * QK_K * bytesPerElem;
 
   // Output buffer
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'dequant_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'dequant_output');
+  const output = outputBuffer || ownedOutput;
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -505,18 +550,25 @@ export async function recordDequantize(
   );
 
   // Bind group
-  const bindGroup = device.createBindGroup({
-    label: 'dequant_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: quantized } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = calculateDequantWorkgroups(variant, numBlocks);
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'dequant');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'dequant_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: quantized } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = calculateDequantWorkgroups(variant, numBlocks);
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, 'dequant');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   
   const dtype = selectSharedRuleValue('shared', 'dtype', 'f16OrF32FromDtype', { dtype: outputDtype });

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;
 }