cpu_adaptive_max_pool3d_backward_channels_last Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/AdaptiveMaxPoolKernel.cpp lines 881–932

template <typename scalar_t>
void cpu_adaptive_max_pool3d_backward_channels_last(
    const Tensor& grad_input_,
    const Tensor& grad_output_,
    const Tensor& indices_) {
  TORCH_CHECK(grad_output_.ndimension() == 5,
              "3d adaptive max pooling backward with channels last format supports tensors with 5 dims.");
  auto memory_format = at::MemoryFormat::ChannelsLast3d;
  auto grad_input = grad_input_.contiguous(memory_format);
  auto grad_output = grad_output_.contiguous(memory_format);
  auto indices = indices_.contiguous(memory_format);

  auto grad_input_data = grad_input.mutable_data_ptr<scalar_t>();
  auto grad_output_data = grad_output.data_ptr<scalar_t>();
  auto indices_data = indices.data_ptr<int64_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);

  // 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_depth * output_height * output_width * channels;
      int64_t* indices_ptr = indices_data + n * output_depth * 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)) {
            scalar_t* gout = grad_output_ptr + od * output_height * output_width * channels + oh * output_width * channels + ow * channels;
            int64_t* ind = indices_ptr + od * output_height * output_width * channels + oh * output_width * channels + ow * channels;
            // TODO: gcc vectorization
            for (const auto c : c10::irange(channels)) {
              int64_t maxindex = ind[c];
              grad_input_ptr[maxindex * channels + c] += gout[c];
            }
          }
        }
      }
    }
  });

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