cpu_avg_pool3d_backward_channels_last Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/AvgPoolKernel.cpp lines 985–1073

template <typename scalar_t>
void cpu_avg_pool3d_backward_channels_last(
    const Tensor& grad_input_,
    const Tensor& grad_output_,
    int kW, int kH, int kD,
    int dW, int dH, int dD,
    int padW, int padH, int padD,
    bool count_include_pad,
    std::optional<int64_t> divisor_override) {
  auto memory_format = at::MemoryFormat::ChannelsLast3d;
  auto grad_input = grad_input_.contiguous(memory_format);
  auto grad_output = grad_output_.contiguous(memory_format);

  auto grad_input_data = grad_input.mutable_data_ptr<scalar_t>();
  auto grad_output_data = grad_output.data_ptr<scalar_t>();

  int64_t nbatch = grad_input.size(0);
  int64_t channels = grad_input.size(1);
  int64_t input_depth = grad_input.size(2);
  int64_t input_height = grad_input.size(3);
  int64_t input_width = grad_input.size(4);
  int64_t output_depth = grad_output.size(2);
  int64_t output_height = grad_output.size(3);
  int64_t output_width = grad_output.size(4);

  using Vec = vec::Vectorized<scalar_t>;
  // parallel on dim N
  at::parallel_for(0, nbatch, 0, [&](int64_t begin, int64_t end) {
    for (const auto n : c10::irange(begin, end)) {
      scalar_t* grad_input_ptr = grad_input_data + n * input_depth * input_height * input_width * channels;
      scalar_t* grad_output_ptr = grad_output_data + n * output_height * output_width * channels;

      for (const auto od : c10::irange(output_depth)) {
        for (const auto oh : c10::irange(output_height)) {
          for (const auto ow : c10::irange(output_width)) {
            int64_t id0 = od * dD - padD;
            int64_t ih0 = oh * dH - padH;
            int64_t iw0 = ow * dW - padW;
            int64_t id1 = std::min(id0 + kD, input_depth + padD);
            int64_t ih1 = std::min(ih0 + kH, input_height + padH);
            int64_t iw1 = std::min(iw0 + kW, input_width + padW);
            int64_t pool_size = (id1 - id0) * (ih1 - ih0) * (iw1 - iw0);
            id0 = std::max(id0, (int64_t) 0);
            ih0 = std::max(ih0, (int64_t) 0);
            iw0 = std::max(iw0, (int64_t) 0);
            id1 = std::min(id1, input_depth);
            ih1 = std::min(ih1, input_height);
            iw1 = std::min(iw1, input_width);

            int64_t divide_factor = 0;
            if (divisor_override.has_value()) {
              divide_factor = divisor_override.value();
            } else {
              if(count_include_pad) {
                divide_factor = pool_size;
              } else {
                divide_factor = (id1 - id0) * (ih1 - ih0) * (iw1 - iw0);
              }
            }

            scalar_t* gout = grad_output_ptr + od * output_height * output_width * channels + oh * output_width * channels + ow * channels;
            int64_t size = channels;
            int64_t len = size - (size % Vec::size());
            for (const auto id : c10::irange(id0, id1)) {
              for (const auto ih : c10::irange(ih0, ih1)) {
                for (const auto iw : c10::irange(iw0, iw1)) {
                  scalar_t* gin = grad_input_ptr + id * input_height * input_width * channels + ih * input_width * channels + iw * channels;

                  int64_t d = 0;
                  for (; d < len; d += Vec::size()) {
                    Vec gin_vec = Vec::loadu(gin + d) + Vec::loadu(gout + d) / Vec(scalar_t(divide_factor));
                    gin_vec.store(gin + d);
                  }
                  for (; d < size; d++) {
                    gin[d] += gout[d] / divide_factor;
                  }
                }
              }
            }
          }
        }
      }
    }
  });

  if (!grad_input_.is_contiguous(memory_format)) {
    grad_input_.copy_(grad_input);
  }
}