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

src/backends/cpu/logic-inl.hpp
93 строки3 KB
kolkir
Refactor tensor logic simd ops with map
16 мар 2025, 23:44
16 мар 2025, 23:448996bad

#if defined(SIMD_LOGIC_INL_H_) == defined(HWY_TARGET_TOGGLE)
#ifdef SIMD_LOGIC_INL_H_
#undef SIMD_LOGIC_INL_H_
#else
#define SIMD_LOGIC_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 gt_vector_scalar(const DataType* HWY_RESTRICT x_array,
                      float scalar,
                      DataType* HWY_RESTRICT out_array,
                      const size_t size,
                      bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto y = hn::Set(d, hwy::ConvertScalarTo<DataType>(scalar));
  auto one = hn::Set(d, hwy::ConvertScalarTo<DataType>(1));
  auto zero = hn::Set(d, hwy::ConvertScalarTo<DataType>(0));
  auto op = [y, one, zero](const hn::Vec<D>& vec) {
    auto z = hn::Gt(vec, y);
    return hn::IfThenElse(z, one, zero);
  };
  map(std::move(op), out_array, x_array, size, aligned);
}

template <typename DataType>
void ge_vector_scalar(const DataType* HWY_RESTRICT x_array,
                      float scalar,
                      DataType* HWY_RESTRICT out_array,
                      const size_t size,
                      bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto y = hn::Set(d, hwy::ConvertScalarTo<DataType>(scalar));
  auto one = hn::Set(d, hwy::ConvertScalarTo<DataType>(1));
  auto zero = hn::Set(d, hwy::ConvertScalarTo<DataType>(0));
  auto op = [y, one, zero](const hn::Vec<D>& vec) {
    auto z = hn::Ge(vec, y);
    return hn::IfThenElse(z, one, zero);
  };
  map(std::move(op), out_array, x_array, size, aligned);
}

template <typename DataType>
void lt_vector_scalar(const DataType* HWY_RESTRICT x_array,
                      float scalar,
                      DataType* HWY_RESTRICT out_array,
                      const size_t size,
                      bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto y = hn::Set(d, hwy::ConvertScalarTo<DataType>(scalar));
  auto one = hn::Set(d, hwy::ConvertScalarTo<DataType>(1));
  auto zero = hn::Set(d, hwy::ConvertScalarTo<DataType>(0));
  auto op = [y, one, zero](const hn::Vec<D>& vec) {
    auto z = hn::Lt(vec, y);
    return hn::IfThenElse(z, one, zero);
  };
  map(std::move(op), out_array, x_array, size, aligned);
}

template <typename DataType>
void le_vector_scalar(const DataType* HWY_RESTRICT x_array,
                      float scalar,
                      DataType* HWY_RESTRICT out_array,
                      const size_t size,
                      bool aligned) {
  using D = hn::ScalableTag<DataType>;
  constexpr D d;
  auto y = hn::Set(d, hwy::ConvertScalarTo<DataType>(scalar));
  auto one = hn::Set(d, hwy::ConvertScalarTo<DataType>(1));
  auto zero = hn::Set(d, hwy::ConvertScalarTo<DataType>(0));
  auto op = [y, one, zero](const hn::Vec<D>& vec) {
    auto z = hn::Le(vec, y);
    return hn::IfThenElse(z, one, zero);
  };
  map(std::move(op), out_array, x_array, size, aligned);
}

}  // namespace HWY_NAMESPACE
}  // namespace adept
HWY_AFTER_NAMESPACE();

#endif  // SIMD_LOGIC_INL_H_