src/backends/cpu/maxpool2d.cpp - with_cpu - vuron/adept

src/backends/cpu/maxpool2d.cpp
131 строка5 KB
kolkir
revert compiler compatibility to gcc12
03 мар 2025, 23:52
03 мар 2025, 23:5226a9274
#include <adept/backends/cpu/maxpool2d.hpp>
#include <adept/backends/cpu/pool_utils.hpp>
#include <adept/irange.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>

#include <cmath>

namespace adept::cpu {
Tensor max_pool2d_fwd(const Tensor& input,
                      Tensor& indices,
                      const ParamArray<2>& kernel,
                      const ParamArray<2>& stride,
                      const ParamArray<2>& padding,
                      const ParamArray<2>& dilation) {
  // TODO: consider to add check sizes
  if (input.shape().rang() != 3 && input.shape().rang() != 4) {
    THROW_ERROR("Invalid input shape ", input.shape(), " for MaxPool2d");
  }
  auto output_shape = pool2d_out_size(input.shape(), kernel, stride, padding);
  auto output =
      Tensor::empty({.shape = output_shape, .device = input.device(), .dtype = input.dtype()});
  if (!indices.defined()) {
    indices =
        Tensor::empty({.shape = output_shape, .device = input.device(), .dtype = dtype_t::Int32});
  }

  DISPATCH_TYPE(input.dtype(), [&]() {
    auto input_data = input.const_data_ptr<scalar_t>();
    auto output_data = output.mutable_data_ptr<scalar_t>();
    auto indices_data = indices.mutable_data_ptr<int32_t>();

    auto ndim = input.shape().rang();
    // treat batch size and channels as one dimension
    index_t channels =
        ndim == 3 ? input.shape().dim(0) : input.shape().dim(0) * input.shape().dim(1);
    int64_t input_height = input.shape().dim(-2);
    int64_t input_width = input.shape().dim(-1);
    auto output_height = output.shape().dim(-2);
    auto output_width = output.shape().dim(-1);

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

        for (auto oh : irange(output_height)) {
          int64_t ih0 = oh * stride[0] - padding[0];
          int64_t ih1 =
              std::min(static_cast<int64_t>(ih0 + (kernel[0] - 1) * dilation[0] + 1), input_height);
          while (ih0 < 0) {
            ih0 += dilation[0];
          }

          for (auto ow : irange(output_width)) {
            int64_t iw0 = ow * stride[1] - padding[1];
            int64_t iw1 = std::min(static_cast<int64_t>(iw0 + (kernel[1] - 1) * dilation[1] + 1),
                                   input_width);
            while (iw0 < 0) {
              iw0 += dilation[1];
            }

            // compute local max
            int64_t maxindex = ih0 * input_width + iw0;
            auto maxval = std::numeric_limits<scalar_t>::min();

            for (int64_t ih = ih0; ih < ih1; ih += dilation[0]) {
              for (int64_t iw = iw0; iw < iw1; iw += dilation[1]) {
                int64_t index = ih * input_width + iw;
                auto val = input_ptr[index];
                if ((val > maxval) || std::isnan(static_cast<double>(val))) {
                  maxval = val;
                  maxindex = index;
                }
              }
            }

            // set output to local max and store location of max
            auto i = oh * output_width + ow;
            output_ptr[i] = maxval;
            indices_ptr[i] = static_cast<int32_t>(maxindex);
          }
        }
      }
    });
  });
  return output;
}

Tensor max_pool2d_bwd(const Tensor& out_grad, const Tensor& indices, const Shape& input_shape) {
  auto input_grad =
      Tensor::zero({.shape = input_shape, .device = out_grad.device(), .dtype = out_grad.dtype()});
  DISPATCH_TYPE(out_grad.dtype(), [&]() {
    auto grad_output_data = out_grad.const_data_ptr<scalar_t>();
    auto grad_input_data = input_grad.mutable_data_ptr<scalar_t>();
    auto indices_data = indices.const_data_ptr<int32_t>();

    auto ndim = out_grad.shape().rang();
    // treat batch size and channels as one dimension
    auto channels =
        ndim == 3 ? out_grad.shape().dim(0) : out_grad.shape().dim(0) * out_grad.shape().dim(1);
    int64_t input_height = input_grad.shape().dim(-2);
    int64_t input_width = input_grad.shape().dim(-1);
    auto output_height = out_grad.shape().dim(-2);
    auto output_width = out_grad.shape().dim(-1);

    // parallel on dim of N, C
    parallel_for<scalar_t>(0, channels, output_height * output_width, [&](auto begin, auto end) {
      for (const auto c : irange(begin, end)) {
        scalar_t* in_grad_ptr = grad_input_data + c * input_height * input_width;
        const scalar_t* out_grad_ptr = grad_output_data + c * output_height * output_width;
        const int32_t* indices_ptr = indices_data + c * output_height * output_width;

        for (auto oh : irange(output_height)) {
          for (auto ow : irange(output_width)) {
            auto index = oh * output_width + ow;
            auto maxindex = indices_ptr[index];
            if (maxindex != -1) {
              in_grad_ptr[maxindex] += out_grad_ptr[index];
            }
          }
        }
      }
    });
  });
  return input_grad;
}

}  // namespace adept::cpu