Function conv2DImpl

tfjs-backend-webgpu/src/kernels/Conv2D_impl.ts:275–394

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

}

export function conv2DImpl({
  x,
  filter,
  convInfo,
  backend,
  bias = null,
  preluActivationWeights = null,
  leakyreluAlpha = 0,
  activation = null
}: Conv2DConfig) {
  const hasBias = bias != null;
  const hasPreluActivationWeights = preluActivationWeights != null;
  const isChannelsLast = convInfo.dataFormat === 'channelsLast';
  const sameSize = isChannelsLast &&
      convInfo.filterHeight === convInfo.inHeight &&
      convInfo.filterWidth === convInfo.inWidth &&
      convInfo.padInfo.type === 'VALID';
  const useNaiveConv2d = env().getBool('WEBGPU_USE_NAIVE_CONV2D_DEBUG');

  if (!useNaiveConv2d &&
      (sameSize ||
       (convInfo.filterHeight === 1 && convInfo.filterWidth === 1 &&
        convInfo.dilationHeight === 1 && convInfo.dilationWidth === 1 &&
        convInfo.strideHeight === 1 && convInfo.strideWidth === 1 &&
        (convInfo.padInfo.type === 'SAME' ||
         convInfo.padInfo.type === 'VALID')))) {
    return conv2dByMatMul({
      x,
      filter,
      convInfo,
      backend,
      bias,
      activation,
      preluActivationWeights,
      leakyreluAlpha
    });
  }

  const thresholdFlagValue =
      env().getNumber('WEBGPU_THRESHOLD_TO_INCREASE_WORKGROUPS_FOR_MATMUL');
  const thresholdToIncreaseWorkgroups = thresholdFlagValue > -1 ?
      thresholdFlagValue :
      backend.thresholdToIncreaseWorkgroups;
  const workgroupsBy32x32 = convInfo.batchSize *
      Math.ceil((convInfo.outHeight * convInfo.outWidth) / 32) *
      Math.ceil(convInfo.outChannels / 32);
  if (env().getBool('WEBGPU_CONV_SEPARATE_IM2COL_SHADER') ||
      workgroupsBy32x32 <= thresholdToIncreaseWorkgroups) {
    return conv2dWithIm2Col({
      x,
      filter,
      convInfo,
      backend,
      bias,
      preluActivationWeights,
      leakyreluAlpha,
      activation
    });