@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/cast.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { dispatch, recordDispatch } from './dispatch.js';
 import { createPipeline, createUniformBufferWithView } from './utils.js';
@@ -44,6 +44,7 @@ export async function castF32ToF16(
 ) {
   const device = getDevice();
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
   const numElements = input.shape.reduce((a, b) => a * b, 1);
 
   const pipeline = await createPipeline('cast', 'f32_to_f16');
@@ -51,35 +52,41 @@ export async function castF32ToF16(
   const outputSize = numElements * DTYPE_SIZES.f16;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f32_to_f16_output');
 
-  const uniformBuffer = createUniformBufferWithView(
-    'cast_f32_to_f16_uniforms',
-    16,
-    (view) => {
-      view.setUint32(0, numElements, true);
-    },
-    null,
-    device
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'cast_f32_to_f16_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  // Use 2D dispatch for large tensors (like embeddings with 300M+ elements)
-  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
-  const dispatchSize = calculate2DDispatch(workgroups);
+  let uniformBuffer = null;
+  try {
+    uniformBuffer = createUniformBufferWithView(
+      'cast_f32_to_f16_uniforms',
+      16,
+      (view) => {
+        view.setUint32(0, numElements, true);
+      },
+      null,
+      device
+    );
 
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+    const bindGroup = device.createBindGroup({
+      label: 'cast_f32_to_f16_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  uniformBuffer.destroy();
+    const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = calculate2DDispatch(workgroups);
 
-  return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+    return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  } finally {
+    uniformBuffer?.destroy();
+  }
 }
 
 
@@ -89,6 +96,7 @@ export async function castF16ToF32(
 ) {
   const device = getDevice();
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
   const numElements = input.shape.reduce((a, b) => a * b, 1);
 
   const pipeline = await createPipeline('cast', 'f16_to_f32');
@@ -96,34 +104,41 @@ export async function castF16ToF32(
   const outputSize = numElements * DTYPE_SIZES.f32;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f16_to_f32_output');
 
-  const uniformBuffer = createUniformBufferWithView(
-    'cast_f16_to_f32_uniforms',
-    16,
-    (view) => {
-      view.setUint32(0, numElements, true);
-    },
-    null,
-    device
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'cast_f16_to_f32_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
-  const dispatchSize = calculate2DDispatch(workgroups);
+  let uniformBuffer = null;
+  try {
+    uniformBuffer = createUniformBufferWithView(
+      'cast_f16_to_f32_uniforms',
+      16,
+      (view) => {
+        view.setUint32(0, numElements, true);
+      },
+      null,
+      device
+    );
 
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+    const bindGroup = device.createBindGroup({
+      label: 'cast_f16_to_f32_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  uniformBuffer.destroy();
+    const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = calculate2DDispatch(workgroups);
 
-  return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+    return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  } finally {
+    uniformBuffer?.destroy();
+  }
 }
 
 
@@ -134,6 +149,7 @@ export async function recordCastF32ToF16(
 ) {
   const device = recorder.device;
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
   const numElements = input.shape.reduce((a, b) => a * b, 1);
 
   const pipeline = await createPipeline('cast', 'f32_to_f16');
@@ -141,32 +157,37 @@ export async function recordCastF32ToF16(
   const outputSize = numElements * DTYPE_SIZES.f16;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f32_to_f16_output');
 
-  const uniformBuffer = createUniformBufferWithView(
-    'cast_f32_to_f16_uniforms',
-    16,
-    (view) => {
-      view.setUint32(0, numElements, true);
-    },
-    recorder
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'cast_f32_to_f16_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
+  try {
+    const uniformBuffer = createUniformBufferWithView(
+      'cast_f32_to_f16_uniforms',
+      16,
+      (view) => {
+        view.setUint32(0, numElements, true);
+      },
+      recorder
+    );
 
-  // Use 2D dispatch for large tensors
-  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
-  const dispatchSize = calculate2DDispatch(workgroups);
+    const bindGroup = device.createBindGroup({
+      label: 'cast_f32_to_f16_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+    const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = calculate2DDispatch(workgroups);
 
-  return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+    recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f32_to_f16');
+    return createTensor(output, 'f16', [...input.shape], input.label ? `${input.label}_f16` : 'cast_f32_to_f16_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  }
 }
 
 
@@ -177,6 +198,7 @@ export async function recordCastF16ToF32(
 ) {
   const device = recorder.device;
   const { outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
   const numElements = input.shape.reduce((a, b) => a * b, 1);
 
   const pipeline = await createPipeline('cast', 'f16_to_f32');
@@ -184,31 +206,37 @@ export async function recordCastF16ToF32(
   const outputSize = numElements * DTYPE_SIZES.f32;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'cast_f16_to_f32_output');
 
-  const uniformBuffer = createUniformBufferWithView(
-    'cast_f16_to_f32_uniforms',
-    16,
-    (view) => {
-      view.setUint32(0, numElements, true);
-    },
-    recorder
-  );
-
-  const bindGroup = device.createBindGroup({
-    label: 'cast_f16_to_f32_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-    ],
-  });
+  try {
+    const uniformBuffer = createUniformBufferWithView(
+      'cast_f16_to_f32_uniforms',
+      16,
+      (view) => {
+        view.setUint32(0, numElements, true);
+      },
+      recorder
+    );
 
-  const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
-  const dispatchSize = calculate2DDispatch(workgroups);
+    const bindGroup = device.createBindGroup({
+      label: 'cast_f16_to_f32_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+      ],
+    });
 
-  recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+    const workgroups = Math.ceil(numElements / WORKGROUP_SIZES.DEFAULT);
+    const dispatchSize = calculate2DDispatch(workgroups);
 
-  return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+    recordDispatch(recorder, pipeline, bindGroup, dispatchSize, 'cast_f16_to_f32');
+    return createTensor(output, 'f32', [...input.shape], input.label ? `${input.label}_f32` : 'cast_f16_to_f32_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  }
 }
 
 
@@ -276,11 +304,15 @@ export async function runBF16ToF32(
   const dispatchSize = calculate2DDispatch(workgroups);
 
   trace.kernels(`BF16ToF32: Dispatching ${dispatchSize[0]}x${dispatchSize[1]} workgroups for ${numPairs} pairs (${numElements} elements)`);
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f32');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, 'f32', [...shape], name);
+  try {
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f32');
+    return createTensor(output, 'f32', [...shape], name);
+  } catch (error) {
+    releaseBuffer(output);
+    throw error;
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 
@@ -337,11 +369,15 @@ export async function runBF16ToF16(
   const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
   const dispatchSize = calculate2DDispatch(workgroups);
 
-  dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f16');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, 'f16', [...shape], name);
+  try {
+    dispatch(device, pipeline, bindGroup, dispatchSize, 'bf16_to_f16');
+    return createTensor(output, 'f16', [...shape], name);
+  } catch (error) {
+    releaseBuffer(output);
+    throw error;
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 
@@ -375,48 +411,54 @@ async function runBF16ToF32Chunked(
 
   trace.kernels(`BF16ToF32: Chunking ${numElements} elements in ${numChunks} chunks`);
 
-  for (let chunkIdx = 0; chunkIdx < numChunks; chunkIdx++) {
-    const chunkStart = chunkIdx * maxElementsPerChunk;
-    const chunkEnd = Math.min((chunkIdx + 1) * maxElementsPerChunk, numElements);
-    const chunkSize = chunkEnd - chunkStart;
-
-    const uniformBuffer = createUniformBufferWithView(
-      `bf16_to_f32_chunk${chunkIdx}_uniforms`,
-      16,
-      (view) => {
-        view.setUint32(0, chunkSize, true);
-        view.setUint32(4, 0, true);
-        view.setUint32(8, 0, true);
-      },
-      null,
-      device
-    );
-
-    const inputOffsetBytes = chunkStart * DTYPE_SIZES.bf16;
-    const outputOffsetBytes = chunkStart * DTYPE_SIZES.f32;
-    const inputPairs = Math.ceil(chunkSize / 2);
-    const inputSizeBytes = inputPairs * DTYPE_SIZES.f32; // Pairs read as u32
-    const outputSizeBytes = chunkSize * DTYPE_SIZES.f32;
-
-    const bindGroup = device.createBindGroup({
-      label: `bf16_to_f32_chunk${chunkIdx}_bind_group`,
-      layout: pipeline.getBindGroupLayout(0),
-      entries: [
-        { binding: 0, resource: { buffer: uniformBuffer } },
-        { binding: 1, resource: { buffer: input, offset: inputOffsetBytes, size: inputSizeBytes } },
-        { binding: 2, resource: { buffer: output, offset: outputOffsetBytes, size: outputSizeBytes } },
-      ],
-    });
-
-    // Each thread processes 2 BF16 values
-    const numPairs = Math.ceil(chunkSize / 2);
-    const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
-    const dispatchSize = calculate2DDispatch(workgroups);
-
-    dispatch(device, pipeline, bindGroup, dispatchSize, `bf16_to_f32_chunk${chunkIdx}`);
+  try {
+    for (let chunkIdx = 0; chunkIdx < numChunks; chunkIdx++) {
+      const chunkStart = chunkIdx * maxElementsPerChunk;
+      const chunkEnd = Math.min((chunkIdx + 1) * maxElementsPerChunk, numElements);
+      const chunkSize = chunkEnd - chunkStart;
+
+      const uniformBuffer = createUniformBufferWithView(
+        `bf16_to_f32_chunk${chunkIdx}_uniforms`,
+        16,
+        (view) => {
+          view.setUint32(0, chunkSize, true);
+          view.setUint32(4, 0, true);
+          view.setUint32(8, 0, true);
+        },
+        null,
+        device
+      );
+
+      try {
+        const inputOffsetBytes = chunkStart * DTYPE_SIZES.bf16;
+        const outputOffsetBytes = chunkStart * DTYPE_SIZES.f32;
+        const inputPairs = Math.ceil(chunkSize / 2);
+        const inputSizeBytes = inputPairs * DTYPE_SIZES.f32;
+        const outputSizeBytes = chunkSize * DTYPE_SIZES.f32;
+
+        const bindGroup = device.createBindGroup({
+          label: `bf16_to_f32_chunk${chunkIdx}_bind_group`,
+          layout: pipeline.getBindGroupLayout(0),
+          entries: [
+            { binding: 0, resource: { buffer: uniformBuffer } },
+            { binding: 1, resource: { buffer: input, offset: inputOffsetBytes, size: inputSizeBytes } },
+            { binding: 2, resource: { buffer: output, offset: outputOffsetBytes, size: outputSizeBytes } },
+          ],
+        });
+
+        const numPairs = Math.ceil(chunkSize / 2);
+        const workgroups = Math.ceil(numPairs / WORKGROUP_SIZES.DEFAULT);
+        const dispatchSize = calculate2DDispatch(workgroups);
+
+        dispatch(device, pipeline, bindGroup, dispatchSize, `bf16_to_f32_chunk${chunkIdx}`);
+      } finally {
+        uniformBuffer.destroy();
+      }
+    }
 
-    uniformBuffer.destroy();
+    return createTensor(output, 'f32', [...shape], name);
+  } catch (error) {
+    releaseBuffer(output);
+    throw error;
   }
-
-  return createTensor(output, 'f32', [...shape], name);
 }

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

@@ -1,7 +1,7 @@
 
 
 import { getDevice, getDeviceEpoch } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, readBufferSlice } from '../../memory/buffer-pool.js';
 import { recordDispatch } from './dispatch.js';
 import { createUniformBufferFromData, getOrCreateBindGroupLayout, getOrCreatePipelineLayout } from './utils.js';
 import { allowReadback } from '../perf-guards.js';
@@ -133,49 +133,38 @@ export async function checkStop(params) {
     usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
   });
   const ownsStopBuffer = !params.shouldStopBuffer;
-  if (shouldStopBuffer.size < requiredBytes) {
-    throw new Error('[CheckStop] shouldStopBuffer too small for tokenIndex.');
-  }
 
-  const bindGroup = device.createBindGroup({
-    layout: getCheckStopBindGroupLayout(device),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: params.sampledTokenBuffer } },
-      { binding: 2, resource: { buffer: shouldStopBuffer } },
-    ],
-  });
-
-  const encoder = device.createCommandEncoder();
-  const pass = encoder.beginComputePass();
-  pass.setPipeline(pipeline);
-  pass.setBindGroup(0, bindGroup);
-  pass.dispatchWorkgroups(1, 1, 1);
-  pass.end();
-
-  // Readback result
-  const stagingBuffer = device.createBuffer({
-    size: U32_BYTES,
-    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
-  });
-  encoder.copyBufferToBuffer(
-    shouldStopBuffer,
-    tokenIndex * U32_BYTES,
-    stagingBuffer,
-    0,
-    U32_BYTES
-  );
-  device.queue.submit([encoder.finish()]);
-
-  await stagingBuffer.mapAsync(GPUMapMode.READ);
-  const result = new Uint32Array(stagingBuffer.getMappedRange())[0];
-  stagingBuffer.unmap();
-
-  uniformBuffer.destroy();
-  if (ownsStopBuffer) {
-    shouldStopBuffer.destroy();
+  try {
+    if (shouldStopBuffer.size < requiredBytes) {
+      throw new Error('[CheckStop] shouldStopBuffer too small for tokenIndex.');
+    }
+
+    const bindGroup = device.createBindGroup({
+      layout: getCheckStopBindGroupLayout(device),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: params.sampledTokenBuffer } },
+        { binding: 2, resource: { buffer: shouldStopBuffer } },
+      ],
+    });
+
+    const encoder = device.createCommandEncoder();
+    const pass = encoder.beginComputePass();
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, bindGroup);
+    pass.dispatchWorkgroups(1, 1, 1);
+    pass.end();
+
+    device.queue.submit([encoder.finish()]);
+
+    const result = new Uint32Array(
+      await readBufferSlice(shouldStopBuffer, tokenIndex * U32_BYTES, U32_BYTES)
+    )[0];
+    return result === 1;
+  } finally {
+    uniformBuffer.destroy();
+    if (ownsStopBuffer) {
+      shouldStopBuffer.destroy();
+    }
   }
-  stagingBuffer.destroy();
-
-  return result === 1;
 }

package/src/gpu/kernels/conv2d.js

@@ -49,27 +49,37 @@ async function _conv2d(target, input, weight, bias, options = {}) {
     device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
   }
 
-  await unifiedKernelWrapper(
-    'conv2d', target, variant,
-    [input, weightBuffer, biasBuffer, output],
-    {
-      in_channels: inChannels, out_channels: outChannels,
-      height, width, out_height: outHeight, out_width: outWidth,
-      kernel_h: kernelH, kernel_w: kernelW,
-      stride, pad, _pad0: 0, _pad1: 0,
-    },
-    Math.ceil((outChannels * outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT)
-  );
+  try {
+    await unifiedKernelWrapper(
+      'conv2d', target, variant,
+      [input, weightBuffer, biasBuffer, output],
+      {
+        in_channels: inChannels, out_channels: outChannels,
+        height, width, out_height: outHeight, out_width: outWidth,
+        kernel_h: kernelH, kernel_w: kernelW,
+        stride, pad, _pad0: 0, _pad1: 0,
+      },
+      [Math.ceil((outHeight * outWidth) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
 
-  if (tempBias) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(tempBias);
-    } else {
+    if (tempBias) {
+      if (recorder) {
+        recorder.trackTemporaryBuffer(tempBias);
+      } else {
+        releaseBuffer(tempBias);
+      }
+    }
+
+    return createTensor(output, input.dtype, [outChannels, outHeight, outWidth], 'conv2d_output');
+  } catch (error) {
+    if (tempBias) {
       releaseBuffer(tempBias);
     }
+    if (!outputBuffer) {
+      releaseBuffer(output);
+    }
+    throw error;
   }
-
-  return createTensor(output, input.dtype, [outChannels, outHeight, outWidth], 'conv2d_output');
 }
 
 export async function runConv2D(input, weight, bias, options = {}) {

package/src/gpu/kernels/conv2d.wgsl

@@ -27,19 +27,18 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let out_height = u.out_height;
     let out_width = u.out_width;
-    let out_size = u.out_channels * out_height * out_width;
-    if (idx >= out_size) {
+    let out_spatial = out_height * out_width;
+    let out_spatial_idx = gid.x;
+    let out_c = gid.y;
+    if (out_c >= u.out_channels || out_spatial_idx >= out_spatial) {
         return;
     }
 
-    let out_spatial = out_height * out_width;
-    let out_c = idx / out_spatial;
-    let rem = idx - out_c * out_spatial;
-    let out_y = rem / out_width;
-    let out_x = rem - out_y * out_width;
+    let out_y = out_spatial_idx / out_width;
+    let out_x = out_spatial_idx - out_y * out_width;
+    let idx = out_c * out_spatial + out_spatial_idx;
 
     var sum: f32 = bias[out_c];
 

package/src/gpu/kernels/conv2d_f16.wgsl

@@ -29,19 +29,18 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let out_height = u.out_height;
     let out_width = u.out_width;
-    let out_size = u.out_channels * out_height * out_width;
-    if (idx >= out_size) {
+    let out_spatial = out_height * out_width;
+    let out_spatial_idx = gid.x;
+    let out_c = gid.y;
+    if (out_c >= u.out_channels || out_spatial_idx >= out_spatial) {
         return;
     }
 
-    let out_spatial = out_height * out_width;
-    let out_c = idx / out_spatial;
-    let rem = idx - out_c * out_spatial;
-    let out_y = rem / out_width;
-    let out_x = rem - out_y * out_width;
+    let out_y = out_spatial_idx / out_width;
+    let out_x = out_spatial_idx - out_y * out_width;
+    let idx = out_c * out_spatial + out_spatial_idx;
 
     var sum: f32 = f32(bias[out_c]);