conv2dSimpleGPU function

GPUTensor<Matrix> conv2dSimpleGPU(GPUTensor<Matrix> input, GPUTensor<Matrix> kernel, CommandBuffer tape, { GPUTensor<Matrix>? outTensor, })

Computes a basic, single-channel 2D convolution between an input matrix and a kernel matrix.

Applies strictly 'valid' padding with a stride of 1. Records the forward operation to the tape and registers individual backward pass kernels for both the input and the kernel gradients.

Implementation

GPUTensor<Matrix> conv2dSimpleGPU(
    GPUTensor<Matrix> input, GPUTensor<Matrix> kernel, CommandBuffer tape,
    {GPUTensor<Matrix>? outTensor}) {
  int inH = input.shape[0];
  int inW = input.shape[1];
  int kH = kernel.shape[0];
  int kW = kernel.shape[1];

  int outH = inH - kH + 1;
  int outW = inW - kW + 1;

  GPUTensor<Matrix> out =
      outTensor ?? GPUTensor<Matrix>.empty(<int>[outH, outW]);

  tape.putInt(OP_CONV2D_FORWARD);
  tape.putString(input.id);
  tape.putString(kernel.id);
  tape.putString(out.id);

  out.creator = GPUNode(
    <GPUTensor>[input, kernel],
    (CommandBuffer bTape) {
      bTape.putInt(OP_CONV2D_BACKWARD_INPUT);
      bTape.putString(out.gradId);
      bTape.putString(kernel.id);
      bTape.putString(input.gradId);

      bTape.putInt(OP_CONV2D_BACKWARD_KERNEL);
      bTape.putString(input.id);
      bTape.putString(out.gradId);
      bTape.putString(kernel.gradId);
    },
    opName: 'conv2dSimpleGPU',
    cost: outH * outW * kH * kW,
  );

  return out;
}