Function conv2dWithIm2Col

tfjs-backend-webgpu/src/kernels/Conv2D_impl.ts:175–273

({
  x,
  filter,
  convInfo,
  backend,
  bias = null,
  preluActivationWeights = null,
  leakyreluAlpha = 0,
  activation = null
}: Conv2DConfig)

// Implements the im2col algorithm as outlined in "High Performance
// Convolutional Neural Networks for Document Processing" (Suvisoft, 2006)
function conv2dWithIm2Col({
  x,
  filter,
  convInfo,
  backend,
  bias = null,
  preluActivationWeights = null,
  leakyreluAlpha = 0,
  activation = null
}: Conv2DConfig) {
  // Rearranges conv2d input so each block to be convolved over forms the
  // row of a new matrix with shape [outHeight * outWidth,
  // filterWidth * filterHeight * inChannels]. The filter is also rearranged so
  // each output channel forms a col of a new matrix with shape [
  // filterWidth * filterHeight * inChannels, outChannels]. The convolution is
  // then computed by multiplying these matrices and reshaping the result.
  const {
    filterWidth,
    filterHeight,
    inChannels,
    strideWidth,
    strideHeight,
    padInfo,
    outWidth,
    outHeight,
    dilationWidth,
    dilationHeight,
    dataFormat
  } = convInfo;

  const isChannelsLast = dataFormat === 'channelsLast';

  const sharedDim = filterWidth * filterHeight * inChannels;
  const numCols = outHeight * outWidth;
  const x2ColShape = isChannelsLast ? [convInfo.batchSize, numCols, sharedDim] :
                                      [convInfo.batchSize, sharedDim, numCols];

  const im2ColProgram = new Im2ColProgram(x2ColShape, isChannelsLast);
  const dimensions = [
    {type: 'int32', data: [padInfo.top, padInfo.left]},      // Padding.
    {type: 'int32', data: [strideHeight, strideWidth]},      // Stride.
    {type: 'int32', data: [dilationHeight, dilationWidth]},  // Dilation.
    {type: 'int32', data: [outWidth]},
    {type: 'int32', data: [inChannels * filterWidth]},  // itemsPerBlockRow.
    {type: 'int32', data: [inChannels]}
  ];
  const x2Col =
      backend.runWebGPUProgram(im2ColProgram, [x], x.dtype, dimensions);

  const intermediates: TensorInfo[] = [];
  intermediates.push(x2Col);

  const filterReshaped = reshape(
      {inputs: {x: filter}, backend, attrs: {shape: [1, sharedDim, -1]}});
  intermediates.push(filterReshaped);

  if (preluActivationWeights != null) {
    const targetShape =