src/backends/cpu/avgpool2d.cpp - master - vuron/adept

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

namespace adept::cpu {
Tensor avg_pool2d_fwd(const Tensor& input,
                      const ParamArray<2>& kernel,
                      const ParamArray<2>& stride,
                      const ParamArray<2>& padding) {
  // TODO: consider to add check sizes
  if (input.shape().rang() != 3 && input.shape().rang() != 4) {
    THROW_ERROR("Invalid input shape ", input.shape(), " for AvgPool2d");
  }
  auto output_shape = pool2d_out_size(input.shape(), kernel, stride, padding);
  auto output =
      Tensor::empty({.shape = output_shape, .device = input.device(), .dtype = input.dtype()});
  DISPATCH_TYPE(input.dtype(), [&]() {
    auto input_data = input.const_data_ptr<scalar_t>();
    auto output_data = output.mutable_data_ptr<scalar_t>();

    auto numel = output.shape().numel();
    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, H, W
    parallel_for<scalar_t>(0, numel, [&](auto begin, auto end) {
      index_t c = 0;
      index_t oh = 0;
      index_t ow = 0;
      data_index_init(begin, c, channels, oh, output_height, ow, output_width);

      for (const auto i : irange(begin, end)) {
        output_data[i] = static_cast<scalar_t>(0);
        const scalar_t* input_ptr = input_data + c * input_height * input_width;

        int64_t ih0 = oh * stride[0] - padding[0];
        int64_t iw0 = ow * stride[1] - padding[1];
        int64_t ih1 = std::min(ih0 + kernel[0], input_height + padding[0]);
        int64_t iw1 = std::min(iw0 + kernel[1], input_width + padding[1]);
        auto pool_size = (ih1 - ih0) * (iw1 - iw0);
        ih0 = std::max(ih0, (int64_t)0);
        iw0 = std::max(iw0, (int64_t)0);
        ih1 = std::min(ih1, input_height);
        iw1 = std::min(iw1, input_width);

        if (ih0 >= ih1 || iw0 >= iw1) {
          data_index_step(c, channels, oh, output_height, ow, output_width);
          continue;
        }

        scalar_t sum = 0;
        for (const auto ih : irange(ih0, ih1)) {
          for (const auto iw : irange(iw0, iw1)) {
            sum += input_ptr[ih * input_width + iw];
          }
        }
        output_data[i] += static_cast<scalar_t>(sum / pool_size);

        data_index_step(c, channels, oh, output_height, ow, output_width);
      }
    });
  });
  return output;
}

Tensor avg_pool2d_bwd(const Tensor& out_grad,
                      const Shape& input_shape,
                      const ParamArray<2>& kernel,
                      const ParamArray<2>& stride,
                      const ParamArray<2>& padding) {
  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 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;

        for (const auto oh : irange(output_height)) {
          for (const auto ow : irange(output_width)) {
            int64_t ih0 = oh * stride[0] - padding[0];
            int64_t iw0 = ow * stride[1] - padding[1];
            int64_t ih1 = std::min(ih0 + kernel[0], input_height + padding[0]);
            int64_t iw1 = std::min(iw0 + kernel[1], input_width + padding[1]);
            auto pool_size = (ih1 - ih0) * (iw1 - iw0);
            ih0 = std::max(ih0, (int64_t)0);
            iw0 = std::max(iw0, (int64_t)0);
            ih1 = std::min(ih1, input_height);
            iw1 = std::min(iw1, input_width);

            scalar_t grad_delta = out_grad_ptr[oh * output_width + ow] / pool_size;
            for (const auto ih : irange(ih0, ih1)) {
              for (const auto iw : irange(iw0, iw1)) {
                in_grad_ptr[ih * input_width + iw] += grad_delta;
              }
            }
          }
        }
      }
    });
  });
  return input_grad;
}

}  // namespace adept::cpu