cpu_adaptive_max_pool3d Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/AdaptiveMaxPoolKernel.cpp lines 483–556

template <typename scalar_t, typename accscalar_t>
void cpu_adaptive_max_pool3d(
    const Tensor& output_,
    const Tensor& indices_,
    const Tensor& input_,
    IntArrayRef output_size) {
  auto input = input_.contiguous();
  auto output = output_.contiguous();
  auto indices = indices_.contiguous();

  auto input_data = input.data_ptr<scalar_t>();
  auto output_data = output.data_ptr<scalar_t>();
  auto indices_data = indices.data_ptr<int64_t>();

  int64_t ndim = input.ndimension();
  // treat batch size and channels as one dimension
  int64_t channels = ndim == 4 ? input.size(0) : input.size(0) * input.size(1);
  int64_t input_depth = input.size(-3);
  int64_t input_height = input.size(-2);
  int64_t input_width = input.size(-1);
  int64_t output_depth = output_size[0];
  int64_t output_height = output_size[1];
  int64_t output_width = output_size[2];

  // parallel on dim of N, C
  at::parallel_for(0, channels, 0, [&](int64_t begin, int64_t end) {
    for (const auto c : c10::irange(begin, end)) {
      scalar_t* input_ptr = input_data + c * input_depth * input_height * input_width;
      scalar_t* output_ptr = output_data + c * output_depth * output_height * output_width;
      int64_t* indices_ptr = indices_data + c * output_depth * output_height * output_width;

      for (const auto od : c10::irange(output_depth)) {
        int64_t id0 = start_index(od, output_depth, input_depth);
        int64_t id1 = end_index(od, output_depth, input_depth);
          for (const auto oh : c10::irange(output_height)) {
            int64_t ih0 = start_index(oh, output_height, input_height);
            int64_t ih1 = end_index(oh, output_height, input_height);

          for (const auto ow : c10::irange(output_width)) {
            int64_t iw0 = start_index(ow, output_width, input_width);
            int64_t iw1 = end_index(ow, output_width, input_width);

            // compute local max
            int64_t maxindex = id0 * input_height * input_width + ih0 * input_width + iw0;
            accscalar_t maxval = -std::numeric_limits<accscalar_t>::infinity();
            for (int64_t id = id0; id < id1; id ++) {
              for (int64_t ih = ih0; ih < ih1; ih ++) {
                for (int64_t iw = iw0; iw < iw1; iw ++) {
                  int64_t index = id * input_height * input_width + ih * input_width + iw;
                  scalar_t val = input_ptr[index];
                  if ((val > maxval) || std::isnan(val)) {
                    maxval = val;
                    maxindex = index;
                  }
                }
              }
            }

            // set output to local max and store location of max
            output_ptr[od * output_height * output_width + oh * output_width + ow] = maxval;
            indices_ptr[od * output_height * output_width + oh * output_width + ow] = maxindex;
          }
        }
      }
    }
  });

  if (!output_.is_contiguous()) {
    output_.copy_(output);
  }
  if (!indices_.is_contiguous()) {
    indices_.copy_(indices);
  }
}