q_avg_pool2d Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/AveragePool2d.cpp lines 174–249

template <typename scalar_t>
Tensor q_avg_pool2d(
    const Tensor& input,
    IntArrayRef kernel_size,
    IntArrayRef stride,
    IntArrayRef padding,
    bool ceil_mode,
    bool count_include_pad,
    std::optional<int64_t> divisor_override) {
  auto [kW, kH] = get_kernel(kernel_size);
  auto [dW, dH] = get_stride(stride, kW, kH);
  auto [padW, padH] = get_padding(padding);

  const int64_t nbatch = input.ndimension() == 4 ? input.size(-4) : 1;
  const int64_t nInputPlane = input.size(-3);
  const int64_t inputHeight = input.size(-2);
  const int64_t inputWidth = input.size(-1);

  TORCH_CHECK(
      !divisor_override.has_value() || divisor_override.value() != 0,
      "divisor must be not zero");

  auto output_shape =
      get_output_shape(input, kW, kH, dW, dH, padW, padH, ceil_mode);
  const int64_t outputHeight = output_shape[output_shape.size() - 2];
  const int64_t outputWidth = output_shape[output_shape.size() - 1];
  if (input.is_contiguous(c10::MemoryFormat::ChannelsLast)) {
    auto output = at::_empty_affine_quantized(
        output_shape,
        input.options().memory_format(input.suggest_memory_format()),
        input.q_scale(),
        input.q_zero_point(),
        std::nullopt);
    // fast path for channel last: qavg_pool_2d_nhwc_stub
    qavg_pool2d_nhwc_stub(
        input.device().type(),
        input,
        output,
        nbatch,
        nInputPlane,
        inputWidth,
        inputHeight,
        outputWidth,
        outputHeight,
        kW,
        kH,
        dW,
        dH,
        padW,
        padH,
        count_include_pad,
        divisor_override);
    return output;
  } else {
    auto output = at::_empty_affine_quantized(
        output_shape, input.options(), input.q_scale(), input.q_zero_point());
    avg_pool2d_out_frame<scalar_t>(
        input,
        output,
        // Contract batch and channels into one dimension
        nbatch * nInputPlane,
        inputWidth,
        inputHeight,
        outputWidth,
        outputHeight,
        kW,
        kH,
        dW,
        dH,
        padW,
        padH,
        count_include_pad,
        divisor_override);
    return output;
  }
}