@simulatte/doppler 0.1.0

package/src/gpu/kernels/backward/attention_backward.js

@@ -0,0 +1,276 @@
+import { CommandRecorder } from '../../command-recorder.js';
+import { getDevice } from '../../device.js';
+import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { createTensor, dtypeBytes } from '../../tensor.js';
+import { castF16ToF32, recordCastF16ToF32 } from '../cast.js';
+import { runMatmul, recordMatmul } from '../matmul.js';
+import { runTranspose, recordTranspose } from '../transpose.js';
+import { runSoftmaxBackward, recordSoftmaxBackward } from './softmax_backward.js';
+import { runBackwardKernel, recordBackwardKernel } from './utils.js';
+
+async function ensureF32(tensor, recorder = null) {
+  if (tensor.dtype !== 'f16') {
+    return tensor;
+  }
+  if (!recorder) {
+    return castF16ToF32(tensor);
+  }
+  const casted = await recordCastF16ToF32(recorder, tensor);
+  recorder.trackTemporaryBuffer(casted.buffer);
+  return casted;
+}
+
+function createHeadSliceBuffers(recorder, headBytes, softmaxBytes) {
+  const qHeadBuf = acquireBuffer(headBytes, undefined, 'attn_q_head');
+  const kHeadBuf = acquireBuffer(headBytes, undefined, 'attn_k_head');
+  const vHeadBuf = acquireBuffer(headBytes, undefined, 'attn_v_head');
+  const sHeadBuf = acquireBuffer(softmaxBytes, undefined, 'attn_s_head');
+  const dHeadBuf = acquireBuffer(headBytes, undefined, 'attn_d_head');
+
+  recorder.trackTemporaryBuffer(qHeadBuf);
+  recorder.trackTemporaryBuffer(kHeadBuf);
+  recorder.trackTemporaryBuffer(vHeadBuf);
+  recorder.trackTemporaryBuffer(sHeadBuf);
+  recorder.trackTemporaryBuffer(dHeadBuf);
+
+  return { qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf };
+}
+
+function createHeadTensors(qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf, seqLen, headDim) {
+  const qHead = createTensor(qHeadBuf, 'f32', [seqLen, headDim], 'attn_q_head');
+  const kHead = createTensor(kHeadBuf, 'f32', [seqLen, headDim], 'attn_k_head');
+  const vHead = createTensor(vHeadBuf, 'f32', [seqLen, headDim], 'attn_v_head');
+  const sHead = createTensor(sHeadBuf, 'f32', [seqLen, seqLen], 'attn_s_head');
+  const dHead = createTensor(dHeadBuf, 'f32', [seqLen, headDim], 'attn_d_head');
+  return { qHead, kHead, vHead, sHead, dHead };
+}
+
+function trackTensorBuffer(recorder, tensor) {
+  recorder.trackTemporaryBuffer(tensor.buffer);
+}
+
+async function runAttentionBackwardCore(
+  q,
+  k,
+  v,
+  softmax,
+  gradOutput,
+  options = {},
+  recorder = null
+) {
+  const { seqLen, numHeads, headDim, scale = 1.0, causal = false } = options;
+  if (!seqLen || !numHeads || !headDim) {
+    throw new Error('attention backward requires seqLen, numHeads, and headDim');
+  }
+
+  const qTensor = await ensureF32(q, recorder);
+  const kTensor = await ensureF32(k, recorder);
+  const vTensor = await ensureF32(v, recorder);
+  const sTensor = await ensureF32(softmax, recorder);
+  const dTensor = await ensureF32(gradOutput, recorder);
+
+  const headElements = seqLen * headDim;
+  const headBytes = headElements * dtypeBytes(qTensor.dtype);
+  const softmaxBytes = seqLen * seqLen * dtypeBytes(sTensor.dtype);
+
+  const totalBytes = numHeads * headBytes;
+  const gradQBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_q');
+  const gradKBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_k');
+  const gradVBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_v');
+
+  if (!recorder) {
+    for (let h = 0; h < numHeads; h += 1) {
+      const qOffset = h * headBytes;
+      const kOffset = h * headBytes;
+      const vOffset = h * headBytes;
+      const dOffset = h * headBytes;
+      const sOffset = h * softmaxBytes;
+
+      const qHeadBuf = acquireBuffer(headBytes, undefined, 'attn_q_head');
+      const kHeadBuf = acquireBuffer(headBytes, undefined, 'attn_k_head');
+      const vHeadBuf = acquireBuffer(headBytes, undefined, 'attn_v_head');
+      const sHeadBuf = acquireBuffer(softmaxBytes, undefined, 'attn_s_head');
+      const dHeadBuf = acquireBuffer(headBytes, undefined, 'attn_d_head');
+
+      const sliceEncoder = getDevice().createCommandEncoder();
+      sliceEncoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
+      sliceEncoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
+      sliceEncoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
+      sliceEncoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
+      sliceEncoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
+      getDevice().queue.submit([sliceEncoder.finish()]);
+
+      const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
+        qHeadBuf,
+        kHeadBuf,
+        vHeadBuf,
+        sHeadBuf,
+        dHeadBuf,
+        seqLen,
+        headDim
+      );
+
+      const sTransposed = await runTranspose(sHead, seqLen, seqLen);
+      const dV = await runMatmul(sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        bDtype: 'f32',
+      });
+
+      const vTransposed = await runTranspose(vHead, seqLen, headDim);
+      const dS = await runMatmul(dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
+        transposeB: false,
+        bDtype: 'f32',
+      });
+      const dQK = causal
+        ? await runBackwardKernel(
+          'attention_backward',
+          sHead,
+          dS,
+          16,
+          (view) => {
+            view.setUint32(0, seqLen, true);
+            view.setUint32(4, seqLen, true);
+            view.setUint32(8, 1, true);
+          }
+        )
+        : await runSoftmaxBackward(sHead, dS, { rows: seqLen, cols: seqLen });
+
+      const dQ = await runMatmul(dQK, kHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        alpha: scale,
+        bDtype: 'f32',
+      });
+      const dQKTransposed = await runTranspose(dQK, seqLen, seqLen);
+      const dK = await runMatmul(dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        alpha: scale,
+        bDtype: 'f32',
+      });
+
+      const copyEncoder = getDevice().createCommandEncoder();
+      copyEncoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
+      copyEncoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
+      copyEncoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
+      getDevice().queue.submit([copyEncoder.finish()]);
+    }
+  } else {
+    const encoder = recorder.getEncoder();
+    for (let h = 0; h < numHeads; h += 1) {
+      const qOffset = h * headBytes;
+      const kOffset = h * headBytes;
+      const vOffset = h * headBytes;
+      const dOffset = h * headBytes;
+      const sOffset = h * softmaxBytes;
+
+      const { qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf } = createHeadSliceBuffers(
+        recorder,
+        headBytes,
+        softmaxBytes
+      );
+
+      encoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
+      encoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
+      encoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
+      encoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
+      encoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
+
+      const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
+        qHeadBuf,
+        kHeadBuf,
+        vHeadBuf,
+        sHeadBuf,
+        dHeadBuf,
+        seqLen,
+        headDim
+      );
+
+      const sTransposed = await recordTranspose(recorder, sHead, seqLen, seqLen);
+      const dV = await recordMatmul(recorder, sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        bDtype: 'f32',
+      });
+
+      const vTransposed = await recordTranspose(recorder, vHead, seqLen, headDim);
+      const dS = await recordMatmul(recorder, dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
+        transposeB: false,
+        bDtype: 'f32',
+      });
+      const dQK = causal
+        ? await recordBackwardKernel(
+          recorder,
+          'attention_backward',
+          sHead,
+          dS,
+          16,
+          (view) => {
+            view.setUint32(0, seqLen, true);
+            view.setUint32(4, seqLen, true);
+            view.setUint32(8, 1, true);
+          }
+        )
+        : await recordSoftmaxBackward(recorder, sHead, dS, { rows: seqLen, cols: seqLen });
+
+      const dQ = await recordMatmul(recorder, dQK, kHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        alpha: scale,
+        bDtype: 'f32',
+      });
+      const dQKTransposed = await recordTranspose(recorder, dQK, seqLen, seqLen);
+      const dK = await recordMatmul(recorder, dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
+        transposeB: false,
+        alpha: scale,
+        bDtype: 'f32',
+      });
+
+      encoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
+      encoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
+      encoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
+
+      trackTensorBuffer(recorder, sTransposed);
+      trackTensorBuffer(recorder, dV);
+      trackTensorBuffer(recorder, vTransposed);
+      trackTensorBuffer(recorder, dS);
+      trackTensorBuffer(recorder, dQK);
+      trackTensorBuffer(recorder, dQ);
+      trackTensorBuffer(recorder, dQKTransposed);
+      trackTensorBuffer(recorder, dK);
+    }
+  }
+
+  return {
+    gradQ: createTensor(gradQBuf, 'f32', [...q.shape], 'attn_grad_q'),
+    gradK: createTensor(gradKBuf, 'f32', [...k.shape], 'attn_grad_k'),
+    gradV: createTensor(gradVBuf, 'f32', [...v.shape], 'attn_grad_v'),
+  };
+}
+
+export async function runAttentionBackward(
+  q,
+  k,
+  v,
+  softmax,
+  gradOutput,
+  options = {}
+) {
+  const device = getDevice();
+  if (!device) {
+    throw new Error('runAttentionBackward requires a GPU device');
+  }
+
+  const recorder = new CommandRecorder(device, 'attention_backward');
+  const result = await runAttentionBackwardCore(q, k, v, softmax, gradOutput, options, recorder);
+  recorder.submit();
+  return result;
+}
+
+export async function recordAttentionBackward(
+  recorder,
+  q,
+  k,
+  v,
+  softmax,
+  gradOutput,
+  options = {}
+) {
+  return runAttentionBackwardCore(q, k, v, softmax, gradOutput, options, recorder);
+}

package/src/gpu/kernels/backward/attention_backward.wgsl

@@ -0,0 +1,49 @@
+// attention_backward.wgsl
+
+/**
+ * Attention Backward Kernel (GPU)
+ *
+ * Computes gradients for attention outputs.
+ * Placeholder implementation until full GPU backward is wired.
+ */
+override WORKGROUP_SIZE: u32 = 256u;
+
+struct Uniforms {
+    rows: u32,
+    cols: u32,
+    causal: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> softmax: array<f32>;
+@group(0) @binding(2) var<storage, read> grad_output: 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 idx = gid.x;
+    let total = u.rows * u.cols;
+    if (idx >= total) {
+        return;
+    }
+
+    let row = idx / u.cols;
+    let col = idx % u.cols;
+    if (u.causal != 0u && col > row) {
+        output[idx] = 0.0;
+        return;
+    }
+
+    var row_sum: f32 = 0.0;
+    for (var j: u32 = 0u; j < u.cols; j = j + 1u) {
+        if (u.causal != 0u && j > row) {
+            continue;
+        }
+        let s = softmax[row * u.cols + j];
+        row_sum = row_sum + s * grad_output[row * u.cols + j];
+    }
+
+    let s = softmax[idx];
+    output[idx] = s * (grad_output[idx] - row_sum);
+}

package/src/gpu/kernels/backward/bias_add_backward.d.ts

@@ -0,0 +1,17 @@
+import type { BackwardKernelOptions } from './utils.js';
+
+export interface BiasAddBackwardOptions extends BackwardKernelOptions {
+  numTokens: number;
+  dim: number;
+}
+
+export declare function runBiasAddBackward(
+  gradOutput: any,
+  options: BiasAddBackwardOptions
+): Promise<any>;
+
+export declare function recordBiasAddBackward(
+  recorder: any,
+  gradOutput: any,
+  options: BiasAddBackwardOptions
+): Promise<any>;

package/src/gpu/kernels/backward/bias_add_backward.js

@@ -0,0 +1,24 @@
+import { createBackwardKernel } from './utils.js';
+
+const { run, record } = createBackwardKernel('bias_add_backward', {
+  uniformSize: 16,
+  writeUniforms: (view, opts) => {
+    view.setUint32(0, opts.numTokens, true);
+    view.setUint32(4, opts.dim, true);
+  },
+  calcWorkgroups: (opts) => Math.ceil(opts.dim / 256),
+  outputBytes: (opts) => opts.dim * 4,
+  outputShape: (opts) => [opts.dim],
+  getDevice: true,
+  validate: (opts) => {
+    if (!opts.numTokens || !opts.dim) throw new Error('bias_add backward requires numTokens and dim');
+  },
+});
+
+export async function runBiasAddBackward(gradOutput, options = {}) {
+  return run(gradOutput, options);
+}
+
+export async function recordBiasAddBackward(recorder, gradOutput, options = {}) {
+  return record(recorder, gradOutput, options);
+}

package/src/gpu/kernels/backward/bias_add_backward.wgsl

@@ -0,0 +1,33 @@
+// bias_add_backward.wgsl
+//
+// Computes gradient for bias: dB = sum_over_tokens(dY)
+//
+// For dY [num_tokens, dim], result dB is [dim].
+
+struct Uniforms {
+    num_tokens: u32,
+    dim: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> grad_output: array<f32>;
+@group(0) @binding(2) var<storage, read_write> grad_bias: array<f32>;
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let d = gid.x;
+    if (d >= u.dim) {
+        return;
+    }
+
+    var sum: f32 = 0.0;
+    for (var t: u32 = 0u; t < u.num_tokens; t = t + 1u) {
+        sum = sum + grad_output[t * u.dim + d];
+    }
+
+    grad_bias[d] = sum;
+}

package/src/gpu/kernels/backward/conv2d_backward.d.ts

@@ -0,0 +1,31 @@
+import type { BackwardKernelOptions } from './utils.js';
+
+export interface Conv2DBackwardOptions extends BackwardKernelOptions {
+  inChannels: number;
+  outChannels: number;
+  height: number;
+  width: number;
+  outHeight: number;
+  outWidth: number;
+  kernelH: number;
+  kernelW: number;
+  stride: number;
+  pad: number;
+  computeGradInput?: boolean;
+  computeGradWeight?: boolean;
+}
+
+export declare function runConv2DBackward(
+  input: any,
+  weight: any,
+  gradOutput: any,
+  options: Conv2DBackwardOptions
+): Promise<{ gradInput: any; gradWeight: any }>;
+
+export declare function recordConv2DBackward(
+  recorder: any,
+  input: any,
+  weight: any,
+  gradOutput: any,
+  options: Conv2DBackwardOptions
+): Promise<{ gradInput: any; gradWeight: any }>;

package/src/gpu/kernels/backward/conv2d_backward.js

@@ -0,0 +1,135 @@
+import { createTensor } from '../../tensor.js';
+import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { createPipeline, createUniformBufferWithView } from '../utils.js';
+import { dispatch, recordDispatch } from '../dispatch.js';
+import { getDevice } from '../../device.js';
+
+export async function runConv2DBackward(input, weight, gradOutput, options = {}) {
+  const { inChannels, outChannels, height, width, outHeight, outWidth, kernelH, kernelW, stride, pad, computeGradInput = true, computeGradWeight = true } = options;
+
+  let gradInput = null;
+  let gradWeight = null;
+
+  const device = getDevice();
+
+  const uniformBuffer = createUniformBufferWithView(
+    'conv2d_backward_uniforms',
+    48,
+    (view) => {
+      view.setUint32(0, inChannels, true);
+      view.setUint32(4, outChannels, true);
+      view.setUint32(8, height, true);
+      view.setUint32(12, width, true);
+      view.setUint32(16, outHeight, true);
+      view.setUint32(20, outWidth, true);
+      view.setUint32(24, kernelH, true);
+      view.setUint32(28, kernelW, true);
+      view.setUint32(32, stride, true);
+      view.setUint32(36, pad, true);
+    },
+    null,
+    device
+  );
+
+  if (computeGradInput) {
+    const outputSize = inChannels * height * width * 4;
+    const outputBuf = acquireBuffer(outputSize, undefined, 'conv2d_backward_input_output');
+    const pipeline = await createPipeline('conv2d_backward_input', 'default');
+    const bindGroup = device.createBindGroup({
+      label: 'conv2d_backward_input_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: gradOutput.buffer } },
+        { binding: 2, resource: { buffer: weight.buffer } },
+        { binding: 3, resource: { buffer: outputBuf } },
+      ],
+    });
+    dispatch(device, pipeline, bindGroup, Math.ceil((inChannels * height * width) / 256), 'conv2d_backward_input');
+    gradInput = createTensor(outputBuf, 'f32', [inChannels, height, width], 'conv2d_grad_input');
+  }
+
+  if (computeGradWeight) {
+    const outputSize = outChannels * inChannels * kernelH * kernelW * 4;
+    const outputBuf = acquireBuffer(outputSize, undefined, 'conv2d_backward_weight_output');
+    const pipeline = await createPipeline('conv2d_backward_weight', 'default');
+    const bindGroup = device.createBindGroup({
+      label: 'conv2d_backward_weight_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: gradOutput.buffer } },
+        { binding: 2, resource: { buffer: input.buffer } },
+        { binding: 3, resource: { buffer: outputBuf } },
+      ],
+    });
+    dispatch(device, pipeline, bindGroup, Math.ceil((outChannels * inChannels * kernelH * kernelW) / 256), 'conv2d_backward_weight');
+    gradWeight = createTensor(outputBuf, 'f32', [outChannels, inChannels, kernelH, kernelW], 'conv2d_grad_weight');
+  }
+
+  uniformBuffer.destroy();
+  return { gradInput, gradWeight };
+}
+
+export async function recordConv2DBackward(recorder, input, weight, gradOutput, options = {}) {
+  const { inChannels, outChannels, height, width, outHeight, outWidth, kernelH, kernelW, stride, pad, computeGradInput = true, computeGradWeight = true } = options;
+
+  let gradInput = null;
+  let gradWeight = null;
+
+  const uniformBuffer = createUniformBufferWithView(
+    'conv2d_backward_uniforms',
+    48,
+    (view) => {
+      view.setUint32(0, inChannels, true);
+      view.setUint32(4, outChannels, true);
+      view.setUint32(8, height, true);
+      view.setUint32(12, width, true);
+      view.setUint32(16, outHeight, true);
+      view.setUint32(20, outWidth, true);
+      view.setUint32(24, kernelH, true);
+      view.setUint32(28, kernelW, true);
+      view.setUint32(32, stride, true);
+      view.setUint32(36, pad, true);
+    },
+    recorder
+  );
+
+  if (computeGradInput) {
+    const outputSize = inChannels * height * width * 4;
+    const outputBuf = acquireBuffer(outputSize, undefined, 'conv2d_backward_input_output');
+    const pipeline = await createPipeline('conv2d_backward_input', 'default');
+    const bindGroup = recorder.device.createBindGroup({
+      label: 'conv2d_backward_input_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: gradOutput.buffer } },
+        { binding: 2, resource: { buffer: weight.buffer } },
+        { binding: 3, resource: { buffer: outputBuf } },
+      ],
+    });
+    recordDispatch(recorder, pipeline, bindGroup, Math.ceil((inChannels * height * width) / 256), 'conv2d_backward_input');
+    gradInput = createTensor(outputBuf, 'f32', [inChannels, height, width], 'conv2d_grad_input');
+  }
+
+  if (computeGradWeight) {
+    const outputSize = outChannels * inChannels * kernelH * kernelW * 4;
+    const outputBuf = acquireBuffer(outputSize, undefined, 'conv2d_backward_weight_output');
+    const pipeline = await createPipeline('conv2d_backward_weight', 'default');
+    const bindGroup = recorder.device.createBindGroup({
+      label: 'conv2d_backward_weight_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: gradOutput.buffer } },
+        { binding: 2, resource: { buffer: input.buffer } },
+        { binding: 3, resource: { buffer: outputBuf } },
+      ],
+    });
+    recordDispatch(recorder, pipeline, bindGroup, Math.ceil((outChannels * inChannels * kernelH * kernelW) / 256), 'conv2d_backward_weight');
+    gradWeight = createTensor(outputBuf, 'f32', [outChannels, inChannels, kernelH, kernelW], 'conv2d_grad_weight');
+  }
+
+  return { gradInput, gradWeight };
+}

package/src/gpu/kernels/backward/conv2d_backward_input.wgsl

@@ -0,0 +1,83 @@
+// conv2d_backward_input.wgsl
+//
+// Computes gradient wrt input (dX).
+// This is essentially a transposed convolution (or deconvolution).
+// For each input pixel [c, y, x], it sums the contributions from all output pixels 
+// that it contributed to during the forward pass.
+
+struct Uniforms {
+    in_channels: u32,
+    out_channels: u32,
+    height: u32,
+    width: u32,
+    out_height: u32,
+    out_width: u32,
+    kernel_h: u32,
+    kernel_w: u32,
+    stride: u32,
+    pad: u32,
+    _pad0: u32,
+    _pad1: u32,
+}
+
+@group(0) @binding(0) var<uniform> u: Uniforms;
+@group(0) @binding(1) var<storage, read> grad_output: array<f32>; // dY [out_channels, out_height, out_width]
+@group(0) @binding(2) var<storage, read> weight: array<f32>;      // W [out_channels, in_channels, kernel_h, kernel_w]
+@group(0) @binding(3) var<storage, read_write> grad_input: array<f32>; // dX [in_channels, height, width]
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let in_idx = gid.x;
+    let total_in_elements = u.in_channels * u.height * u.width;
+    
+    if (in_idx >= total_in_elements) {
+        return;
+    }
+
+    let in_spatial = u.height * u.width;
+    let ic = in_idx / in_spatial;
+    let rem = in_idx % in_spatial;
+    let iy = rem / u.width;
+    let ix = rem % u.width;
+
+    var sum: f32 = 0.0;
+
+    // Weight is [OC, IC, KH, KW]
+    let weight_ic_stride = u.kernel_h * u.kernel_w;
+    let weight_oc_stride = u.in_channels * weight_ic_stride;
+
+    // Iterate over all output pixels that this input pixel could have contributed to.
+    // Forward mapping: oy = (iy + pad - ky) / stride, ox = (ix + pad - kx) / stride
+    // We iterate over the kernel [ky, kx] and see if the resulting [oy, ox] is valid.
+    
+    for (var oc: u32 = 0u; oc < u.out_channels; oc = oc + 1u) {
+        let out_c_offset = oc * u.out_height * u.out_width;
+        let weight_oc_offset = oc * weight_oc_stride + ic * weight_ic_stride;
+        
+        for (var ky: u32 = 0u; ky < u.kernel_h; ky = ky + 1u) {
+            let py = iy + u.pad;
+            if (py < ky) { continue; }
+            let y_after_pad = py - ky;
+            if (y_after_pad % u.stride != 0u) { continue; }
+            let oy = y_after_pad / u.stride;
+            if (oy >= u.out_height) { continue; }
+
+            for (var kx: u32 = 0u; kx < u.kernel_w; kx = kx + 1u) {
+                let px = ix + u.pad;
+                if (px < kx) { continue; }
+                let x_after_pad = px - kx;
+                if (x_after_pad % u.stride != 0u) { continue; }
+                let ox = x_after_pad / u.stride;
+                if (ox >= u.out_width) { continue; }
+
+                let dy = grad_output[out_c_offset + oy * u.out_width + ox];
+                let w = weight[weight_oc_offset + ky * u.kernel_w + kx];
+                sum = sum + dy * w;
+            }
+        }
+    }
+
+    grad_input[in_idx] = sum;
+}