src/backends/cpu/batchnorm2d-inl.hpp - master - vuron/adept

src/backends/cpu/batchnorm2d-inl.hpp
84 строки3 KB
kolkir
Fix compilation
28 фев 2025, 11:56
28 фев 2025, 11:569cc4d0f
#if defined(SIMD_ACTIVATIONS_INL_H_) == defined(HWY_TARGET_TOGGLE)
#ifdef SIMD_ACTIVATIONS_INL_H_
#undef SIMD_ACTIVATIONS_INL_H_
#else
#define SIMD_ACTIVATIONS_INL_H_
#endif

#include <hwy/highway.h>
#include "vector_ops-inl.hpp"

HWY_BEFORE_NAMESPACE();
namespace adept {
namespace HWY_NAMESPACE {
namespace hn = hwy::HWY_NAMESPACE;

template <typename DataType>
void simd_batchnorm_fwd(DataType* HWY_RESTRICT output,
                        const DataType* HWY_RESTRICT input,
                        const DataType alpha,
                        const DataType beta,
                        const size_t size,
                        bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto a = hn::Set(d, alpha);
  auto b = hn::Set(d, beta);
  auto op = [&](const hn::Vec<D>& in) { return in * a + b; };
  map(std::move(op), output, input, size, aligned);
}

template <typename DataType>
void simd_bn_sum(DataType& out,
                 const DataType* HWY_RESTRICT input,
                 const size_t size,
                 bool aligned) {
  using D = hn::ScalableTag<DataType>;
  auto op = [&](const hn::Vec<D>& a, const hn::Vec<D>& b) { return a + b; };
  out += reduce_all(std::move(op), input, size, aligned);
}

template <typename DataType>
void simd_bn_dotp(DataType& dotp,
                  DataType mean,
                  const DataType* HWY_RESTRICT x_ptr,
                  const DataType* HWY_RESTRICT dy_ptr,
                  const size_t size,
                  bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto mean_vec = hn::Set(d, mean);
  auto map_op = [&](const hn::Vec<D>& x, const hn::Vec<D>& dy) { return (x - mean_vec) * dy; };
  auto red_op = [&](const hn::Vec<D>& x, const hn::Vec<D>& y) { return x + y; };
  dotp += map2_reduce_all(std::move(map_op), std::move(red_op), x_ptr, dy_ptr, size, aligned);
}

template <typename DataType>
void simd_bn_in_grad(DataType grad_mean,
                     DataType mean,
                     DataType k,
                     DataType var,
                     DataType w,
                     DataType* HWY_RESTRICT dx_ptr,
                     const DataType* HWY_RESTRICT x_ptr,
                     const DataType* HWY_RESTRICT dy_ptr,
                     const size_t size,
                     bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto mean_vec = hn::Set(d, mean);
  auto grad_mean_vec = hn::Set(d, grad_mean);
  auto k_vec = hn::Set(d, k);
  auto var_vec = hn::Set(d, var);
  auto w_vec = hn::Set(d, w);
  auto op = [&](const hn::Vec<D>& x, const hn::Vec<D>& dy) {
    auto dx = (x - mean_vec) * k_vec;
    return (dy - grad_mean_vec - dx) * var_vec * w_vec;
  };
  map2(std::move(op), dx_ptr, x_ptr, dy_ptr, size, aligned);
}
}  // namespace HWY_NAMESPACE
}  // namespace adept
HWY_AFTER_NAMESPACE();

#endif  // SIMD_ACTIVATIONS_INL_H_