@simulatte/doppler 0.1.6 → 0.1.7

package/src/gpu/kernels/silu.js

@@ -1,13 +1,26 @@
 
 
 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';
 import { getPipelineFast, createUniformBufferWithView } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
 
+function destroyAfterSubmit(device, buffer) {
+  if (!buffer) {
+    return;
+  }
+  device.queue.onSubmittedWorkDone()
+    .then(() => {
+      buffer.destroy();
+    })
+    .catch(() => {
+      buffer.destroy();
+    });
+}
+
 function canUseF16(input) {
   return input.dtype === 'f16';
 }
@@ -47,6 +60,12 @@ function createSiLUBindGroupEntries(uniformBuffer, input, output, gate) {
   ];
 }
 
+function cleanupRunResources(uniformBuffer, ownedOutput) {
+  if (ownedOutput) {
+    releaseBuffer(ownedOutput);
+  }
+}
+
 function planSiLUDispatch(device, size, useVec4) {
   const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
     ? device.limits.maxComputeWorkgroupsPerDimension
@@ -97,6 +116,7 @@ export async function runSiLU(
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
   const dispatchPlan = planSiLUDispatch(device, inferredSize, useVec4);
 
   // Create uniform buffer
@@ -116,17 +136,21 @@ export async function runSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -148,6 +172,7 @@ export async function runSwiGLURowsplitBias(
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'swiglu_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -164,23 +189,27 @@ export async function runSwiGLURowsplitBias(
   );
 
   // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'swiglu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: bias.buffer } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'swiglu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: bias.buffer } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+    dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
+    return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -202,6 +231,7 @@ export async function runSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -218,24 +248,28 @@ export async function runSiLURowSplit(
   );
 
   // Bind group: provide a dummy gate buffer to satisfy the fixed layout
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
-  dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 
@@ -258,6 +292,7 @@ export async function recordSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -272,22 +307,28 @@ export async function recordSiLURowSplit(
     recorder
   );
 
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 
@@ -328,6 +369,7 @@ export async function recordSiLU(
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
   const dispatchPlan = planSiLUDispatch(device, inferredSize, false);
 
   // Uniform buffer
@@ -346,13 +388,19 @@ export async function recordSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/softmax.js

@@ -1,6 +1,6 @@
 
 import { 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 { unifiedKernelWrapper } from './utils.js';
 import { createPipeline, createUniformBufferWithView, createBindGroupWithValidation } from './utils.js';
@@ -20,23 +20,34 @@ function selectSoftmaxVariant(innerSize) {
 
 async function _softmax(target, input, axis, options = {}) {
   const { batchSize = 1, size, seqLen, temperature = 1.0, outputBuffer = null } = options;
+  if (input.dtype !== 'f32') {
+    throw new Error(`Softmax requires f32 input, got ${input.dtype}.`);
+  }
 
-  const bytesPerElement = input.dtype === 'f16' ? 2 : 4;
+  const bytesPerElement = 4;
   const inferredSize = size || seqLen || (input.buffer.size / (batchSize * bytesPerElement));
   const variant = selectSoftmaxVariant(inferredSize);
   trace.kernels(`Softmax: size=${inferredSize}, variant=${variant}`);
 
   const outputSize = batchSize * inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'softmax_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'softmax', target, variant,
+      [input, output],
+      { inner_size: inferredSize, outer_size: batchSize, temperature },
+      batchSize
+    );
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
-  await unifiedKernelWrapper(
-    'softmax', target, variant,
-    [input, output],
-    { inner_size: inferredSize, outer_size: batchSize, temperature },
-    batchSize
-  );
-
-  return createTensor(output, input.dtype, [batchSize, inferredSize], 'softmax_output');
+  return createTensor(output, 'f32', [batchSize, inferredSize], 'softmax_output');
 }
 
 export async function runSoftmax(input, axis, options = {}) {
@@ -76,6 +87,7 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
 
   const indices = acquireBuffer(indicesSize, undefined, 'softmax_topk_indices');
   const weights = acquireBuffer(weightsSize, undefined, 'softmax_topk_weights');
+  let completed = false;
 
   const uniformBuffer = createUniformBufferWithView(
     'softmax_topk_uniforms', 16,
@@ -88,19 +100,26 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
     null, device
   );
 
-  const bindGroup = await createBindGroupWithValidation(device, {
-    label: 'softmax_topk_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: logits } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-    ],
-  }, `topk:${variant}`);
-
-  dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
-  uniformBuffer.destroy();
-
-  return { indices, weights };
+  try {
+    const bindGroup = await createBindGroupWithValidation(device, {
+      label: 'softmax_topk_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: logits } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+      ],
+    }, `topk:${variant}`);
+
+    dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
+    completed = true;
+    return { indices, weights };
+  } finally {
+    uniformBuffer.destroy();
+    if (!completed) {
+      releaseBuffer(indices);
+      releaseBuffer(weights);
+    }
+  }
 }

package/src/gpu/kernels/split_qkv.js

@@ -1,5 +1,5 @@
 
-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 { unifiedKernelWrapper } from './utils.js';
@@ -7,6 +7,9 @@ import { selectRuleValue } from './rule-registry.js';
 
 async function _splitQKV(target, qkvTensor, options) {
   const { numTokens, qSize, kSize, vSize, qTensor = null, kTensor = null, vTensor = null } = options;
+  const ownsQ = qTensor == null;
+  const ownsK = kTensor == null;
+  const ownsV = vTensor == null;
 
   const outputDtype = qkvTensor.dtype;
   const pipelineVariant = selectRuleValue('splitQkv', 'variant', { outputDtype });
@@ -18,18 +21,25 @@ async function _splitQKV(target, qkvTensor, options) {
 
   const totalElements = numTokens * (qSize + kSize + vSize);
 
-  await unifiedKernelWrapper(
-    'split_qkv', target, pipelineVariant,
-    [qkvTensor, qBuffer, kBuffer, vBuffer],
-    { num_tokens: numTokens, q_size: qSize, k_size: kSize, v_size: vSize },
-    Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
-  const K = kTensor || createTensor(kBuffer, outputDtype, [numTokens, kSize], 'K');
-  const V = vTensor || createTensor(vBuffer, outputDtype, [numTokens, vSize], 'V');
-
-  return { Q, K, V };
+  try {
+    await unifiedKernelWrapper(
+      'split_qkv', target, pipelineVariant,
+      [qkvTensor, qBuffer, kBuffer, vBuffer],
+      { num_tokens: numTokens, q_size: qSize, k_size: kSize, v_size: vSize },
+      Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT)
+    );
+
+    const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
+    const K = kTensor || createTensor(kBuffer, outputDtype, [numTokens, kSize], 'K');
+    const V = vTensor || createTensor(vBuffer, outputDtype, [numTokens, vSize], 'V');
+
+    return { Q, K, V };
+  } catch (error) {
+    if (ownsQ) releaseBuffer(qBuffer);
+    if (ownsK) releaseBuffer(kBuffer);
+    if (ownsV) releaseBuffer(vBuffer);
+    throw error;
+  }
 }
 
 export async function runSplitQKV(qkvTensor, options) {

package/src/gpu/kernels/transpose.js

@@ -1,4 +1,4 @@
-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 { unifiedKernelWrapper } from './utils.js';
@@ -20,18 +20,26 @@ async function _transpose(target, input, rows, cols, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = rows * cols * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'transpose_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
   const dispatchPlan = planTransposeDispatch(target, cols);
 
-  await unifiedKernelWrapper(
-    'transpose',
-    target,
-    'default',
-    [input, outputBuf],
-    { rows, cols, _pad0: dispatchPlan.dispatchStride, _pad1: 0 },
-    [dispatchPlan.workgroups[0], rows, 1]
-  );
+  try {
+    await unifiedKernelWrapper(
+      'transpose',
+      target,
+      'default',
+      [input, outputBuf],
+      { rows, cols, _pad0: dispatchPlan.dispatchStride, _pad1: 0 },
+      [dispatchPlan.workgroups[0], rows, 1]
+    );
 
-  return createTensor(outputBuf, input.dtype, [cols, rows], 'transpose_output');
+    return createTensor(outputBuf, input.dtype, [cols, rows], 'transpose_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runTranspose(input, rows, cols, options = {}) {

package/src/gpu/kernels/transpose.wgsl

@@ -20,11 +20,13 @@ struct Uniforms {
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     let dispatch_stride = max(u._pad0, 1u);
-    let row = gid.y;
-    let col = gid.x + row * dispatch_stride;
-    if (row >= u.rows || col >= u.cols) {
+    let linear_idx = gid.y * dispatch_stride + gid.x;
+    let total = u.rows * u.cols;
+    if (linear_idx >= total) {
         return;
     }
+    let row = linear_idx / u.cols;
+    let col = linear_idx % u.cols;
     let idx = row * u.cols + col;
     let out_idx = col * u.rows + row;
     output[out_idx] = input[idx];

package/src/gpu/kernels/upsample2d.js

@@ -1,4 +1,4 @@
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+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';
@@ -35,19 +35,27 @@ async function _upsample2d(target, input, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = channels * outHeight * outWidth * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'upsample2d_output');
+  const ownedOutput = outputBuffer ? null : output;
 
-  await unifiedKernelWrapper(
-    'upsample2d', target, selectUpsample2DVariant(input.dtype === 'f16'),
-    [input, output],
-    {
-      channels, in_height: resolvedHeight, in_width: resolvedWidth,
-      out_height: outHeight, out_width: outWidth, scale,
-      _pad0: 0, _pad1: 0,
-    },
-    [Math.ceil(outSpatial / WORKGROUP_SIZES.DEFAULT), channels, 1]
-  );
-
-  return createTensor(output, input.dtype, [channels, outHeight, outWidth], 'upsample2d_output');
+  try {
+    await unifiedKernelWrapper(
+      'upsample2d', target, selectUpsample2DVariant(input.dtype === 'f16'),
+      [input, output],
+      {
+        channels, in_height: resolvedHeight, in_width: resolvedWidth,
+        out_height: outHeight, out_width: outWidth, scale,
+        _pad0: 0, _pad1: 0,
+      },
+      [Math.ceil(outSpatial / WORKGROUP_SIZES.DEFAULT), channels, 1]
+    );
+
+    return createTensor(output, input.dtype, [channels, outHeight, outWidth], 'upsample2d_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runUpsample2D(input, options = {}) {

package/src/gpu/kernels/utils.js

@@ -137,21 +137,28 @@ export async function unifiedKernelWrapper(opName, target, variant, bindings, un
     });
   }
 
-  const bindGroup = device.createBindGroup({
-    label: `${opName}_bind_group`,
-    layout: pipeline.getBindGroupLayout(0),
-    entries: bindGroupEntries,
-  });
-
-  if (workgroups && typeof workgroups === 'object' && workgroups.indirectBuffer) {
-    const indirectOffset = workgroups.indirectOffset ?? 0;
-    if (recorder) {
-      recordDispatchIndirect(recorder, pipeline, bindGroup, workgroups.indirectBuffer, indirectOffset, opName);
+  try {
+    const bindGroup = device.createBindGroup({
+      label: `${opName}_bind_group`,
+      layout: pipeline.getBindGroupLayout(0),
+      entries: bindGroupEntries,
+    });
+
+    if (workgroups && typeof workgroups === 'object' && workgroups.indirectBuffer) {
+      const indirectOffset = workgroups.indirectOffset ?? 0;
+      if (recorder) {
+        recordDispatchIndirect(recorder, pipeline, bindGroup, workgroups.indirectBuffer, indirectOffset, opName);
+      } else {
+        dispatchIndirect(device, pipeline, bindGroup, workgroups.indirectBuffer, indirectOffset, opName);
+      }
     } else {
-      dispatchIndirect(device, pipeline, bindGroup, workgroups.indirectBuffer, indirectOffset, opName);
+      dispatchKernel(target, pipeline, bindGroup, workgroups, opName);
+    }
+  } catch (error) {
+    if (!recorder) {
+      uniformBuffer.destroy();
     }
-  } else {
-    dispatchKernel(target, pipeline, bindGroup, workgroups, opName);
+    throw error;
   }
 
   if (!recorder) {

package/src/gpu/partitioned-buffer-pool.js

@@ -11,10 +11,13 @@ export class PartitionedBufferPool {
   
   #expertPools;
 
+  #bufferOwners;
+
   
   constructor(partitions, schemaConfig = getRuntimeConfig().shared.bufferPool) {
     this.#sharedPool = new BufferPool(false, schemaConfig);
     this.#expertPools = new Map();
+    this.#bufferOwners = new WeakMap();
     for (const partition of partitions) {
       this.#expertPools.set(partition.id, new BufferPool(false, schemaConfig));
     }
@@ -28,12 +31,17 @@ export class PartitionedBufferPool {
     label
   ) {
     const pool = this.#expertPools.get(partitionId) || this.#sharedPool;
-    return pool.acquire(size, usage, label);
+    const buffer = pool.acquire(size, usage, label);
+    this.#bufferOwners.set(buffer, pool);
+    return buffer;
   }
 
   
   release(partitionId, buffer) {
-    const pool = this.#expertPools.get(partitionId) || this.#sharedPool;
+    const pool = this.#bufferOwners.get(buffer)
+      || this.#expertPools.get(partitionId)
+      || this.#sharedPool;
+    this.#bufferOwners.delete(buffer);
     pool.release(buffer);
   }
 

package/src/gpu/perf-guards.js

@@ -1,15 +1,8 @@
 
-
 import { log, trace } from '../debug/index.js';
+import { DEFAULT_PERF_GUARDS_CONFIG } from '../config/schema/debug.schema.js';
 
-// Initial config uses inline defaults; caller should configure via configurePerfGuards()
-let config = {
-  allowGPUReadback: true,
-  trackSubmitCount: false,
-  trackAllocations: false,
-  logExpensiveOps: false,
-  strictMode: false,
-};
+let config = { ...DEFAULT_PERF_GUARDS_CONFIG };
 
 
 let counters = {