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

src/backends/cpu/vector_ops-inl.hpp
264 строки9 KB
kolkir
Refactor simd algos to use map reduce
16 мар 2025, 23:09
16 мар 2025, 23:095f536e1
#if defined(SIMD_VECTOR_OPS_INL_H_) == defined(HWY_TARGET_TOGGLE)
#ifdef SIMD_VECTOR_OPS_INL_H_
#undef SIMD_VECTOR_OPS_INL_H_
#else
#define SIMD_VECTOR_OPS_INL_H_
#endif

#include <hwy/highway.h>
#include <array>

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

template <typename DataType>
struct AlignedLoader {
  static constexpr hn::ScalableTag<DataType> d{};
  static auto load(const auto* src, size_t) { return hn::Load(d, src); };
  static void store(const auto& src, auto* dst, size_t) { hn::Store(src, d, dst); }
};

template <typename DataType>
struct UnAlignedLoader {
  static constexpr hn::ScalableTag<DataType> d{};
  static auto load(const auto* src, size_t) { return hn::LoadU(d, src); };
  static void store(const auto& src, auto* dst, size_t) { hn::StoreU(src, d, dst); }
};

template <typename DataType>
struct RemainingLoader {
  static constexpr hn::ScalableTag<DataType> d{};
  static auto load(const auto* src, size_t remaining) { return hn::LoadN(d, src, remaining); };
  static void store(const auto& src, auto* dst, size_t remaining) {
    hn::StoreN(src, d, dst, remaining);
  }
};

template <typename DataType, template <typename> typename Loader, typename Func>
void map_(const Func& func, DataType* output_data, const DataType* input_data, size_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);
  size_t i = 0;
  if (size >= lanes) {
    for (; i <= (size - lanes); i += lanes) {
      auto vec = func(Loader<DataType>::load(input_data + i, 0));
      Loader<DataType>::store(vec, output_data + i, 0);
    }
  }

  const auto remaining = size - i;
  if (remaining > 0) {
    auto vec = func(RemainingLoader<DataType>::load(input_data + i, remaining));
    RemainingLoader<DataType>::store(vec, output_data + i, remaining);
  }
}

template <typename DataType, typename Func>
void map(const Func& func,
         DataType* output_data,
         const DataType* input_data,
         size_t size,
         bool aligned) {
  if (aligned)
    map_<DataType, AlignedLoader>(func, output_data, input_data, size);
  else
    map_<DataType, UnAlignedLoader>(func, output_data, input_data, size);
}

template <typename DataType, template <typename> typename Loader, typename Func>
void map2_(const Func& func,
           DataType* output_data,
           const DataType* input_data1,
           const DataType* input_data2,
           size_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);
  size_t i = 0;
  if (size >= lanes) {
    for (; i <= (size - lanes); i += lanes) {
      auto res = func(Loader<DataType>::load(input_data1 + i, 0),
                      Loader<DataType>::load(input_data2 + i, 0));
      Loader<DataType>::store(res, output_data + i, 0);
    }
  }

  const auto remaining = size - i;
  if (remaining > 0) {
    auto res = func(RemainingLoader<DataType>::load(input_data1 + i, remaining),
                    RemainingLoader<DataType>::load(input_data2 + i, remaining));
    RemainingLoader<DataType>::store(res, output_data + i, remaining);
  }
}

template <typename DataType, typename Func>
void map2(const Func& func,
          DataType* output_data,
          const DataType* input_data1,
          const DataType* input_data2,
          size_t size,
          bool aligned) {
  if (aligned)
    map2_<DataType, AlignedLoader>(func, output_data, input_data1, input_data2, size);
  else
    map2_<DataType, UnAlignedLoader>(func, output_data, input_data1, input_data2, size);
}

template <typename DataType, template <typename> typename Loader, typename Vec, typename Func>
inline DataType vec_reduce_all_(const Func& func, Vec acc_vec, size_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);
  DataType acc_arr[lanes];
  Loader<DataType>::store(acc_vec, acc_arr, size);
  std::array<DataType, lanes> acc_arr_next = {0};
  for (size_t i = 1; i < size; ++i) {
    acc_arr_next[0] = acc_arr[i];
    auto acc_vec_next = Loader<DataType>::load(acc_arr_next.data(), size);
    acc_vec = func(acc_vec, acc_vec_next);
  }
  Loader<DataType>::store(acc_vec, acc_arr, size);
  return acc_arr[0];
}

template <typename DataType, template <typename> typename Loader, typename Func>
DataType reduce_all_(const Func& func, const DataType* input_data, size_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);

  if (size < lanes)
    return vec_reduce_all_<DataType, RemainingLoader>(
        func, Loader<DataType>::load(input_data, size), size);

  auto acc = Loader<DataType>::load(input_data, 0);
  size_t i = lanes;
  if (size >= lanes) {
    for (; i <= (size - lanes); i += lanes) {
      acc = func(acc, Loader<DataType>::load(input_data + i, 0));
    }
  }

  const auto remaining = size - i;
  if (remaining > 0) {
    acc = func(acc, RemainingLoader<DataType>::load(input_data + i, remaining));
  }
  return vec_reduce_all_<DataType, Loader>(func, acc, lanes);
}

template <typename DataType, typename Func>
DataType reduce_all(const Func& func, const DataType* input_data, size_t size, bool aligned) {
  if (aligned)
    return reduce_all_<DataType, AlignedLoader>(func, input_data, size);
  else
    return reduce_all_<DataType, UnAlignedLoader>(func, input_data, size);
}

template <typename DataType,
          template <typename>
          typename Loader,
          typename MapFunc,
          typename ReduceFunc>
DataType map_reduce_all_(const MapFunc& map_func,
                         const ReduceFunc& red_func,
                         const DataType* input_data,
                         int64_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);

  if (size < lanes)
    return vec_reduce_all_<DataType, RemainingLoader>(
        red_func, map_func(Loader<DataType>::load(input_data, size)), size);

  auto acc = Loader<DataType>::load(input_data, 0);
  size_t i = lanes;
  if (size >= lanes) {
    for (; i <= (size - lanes); i += lanes) {
      auto data = Loader<DataType>::load(input_data + i, 0);
      auto map_res = map_func(data);
      acc = red_func(acc, map_res);
    }
  }

  const auto remaining = size - i;
  if (remaining > 0) {
    auto data = RemainingLoader<DataType>::load(input_data + i, remaining);
    auto map_res = map_func(data);
    acc = red_func(acc, map_res);
  }

  return vec_reduce_all_<DataType, Loader>(red_func, acc, lanes);
}

template <typename DataType, typename MapFunc, typename ReduceFunc>
DataType map_reduce_all(const MapFunc& map_func,
                        const ReduceFunc& red_func,
                        const DataType* input_data,
                        size_t size,
                        bool aligned) {
  if (aligned)
    return map_reduce_all_<DataType, AlignedLoader>(map_func, red_func, input_data, size);
  else
    return map_reduce_all_<DataType, UnAlignedLoader>(map_func, red_func, input_data, size);
}

template <typename DataType,
          template <typename>
          typename Loader,
          typename MapFunc,
          typename ReduceFunc>
DataType map2_reduce_all_(const MapFunc& map_func,
                          const ReduceFunc& red_func,
                          const DataType* input_data1,
                          const DataType* input_data2,
                          size_t size) {
  constexpr hn::ScalableTag<DataType> d;
  const size_t lanes = hn::Lanes(d);
  if (size < lanes) {
    auto data1 = RemainingLoader<DataType>::load(input_data1, size);
    auto data2 = RemainingLoader<DataType>::load(input_data2, size);
    auto map_res = map_func(data1, data2);
    return vec_reduce_all_<DataType, RemainingLoader>(red_func, map_res, size);
  }

  auto acc =
      map_func(Loader<DataType>::load(input_data1, 0), Loader<DataType>::load(input_data2, 0));
  size_t i = lanes;
  if (size >= lanes) {
    for (; i <= (size - lanes); i += lanes) {
      auto data1 = Loader<DataType>::load(input_data1 + i, 0);
      auto data2 = Loader<DataType>::load(input_data2 + i, 0);
      auto map_res = map_func(data1, data2);
      acc = red_func(acc, map_res);
    }
  }

  const auto remaining = size - i;
  if (remaining > 0) {
    auto data1 = RemainingLoader<DataType>::load(input_data1 + i, remaining);
    auto data2 = RemainingLoader<DataType>::load(input_data2 + i, remaining);
    auto map_res = map_func(data1, data2);
    acc = red_func(acc, map_res);
  }
  return vec_reduce_all_<DataType, Loader>(red_func, acc, lanes);
}

template <typename DataType, typename MapFunc, typename ReduceFunc>
DataType map2_reduce_all(const MapFunc& map_func,
                         const ReduceFunc& red_func,
                         const DataType* input_data1,
                         const DataType* input_data2,
                         size_t size,
                         bool aligned) {
  if (aligned)
    return map2_reduce_all_<DataType, AlignedLoader>(map_func, red_func, input_data1, input_data2,
                                                     size);
  else
    return map2_reduce_all_<DataType, UnAlignedLoader>(map_func, red_func, input_data1, input_data2,
                                                       size);
}

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

#endif  // SIMD_VECTOR_OPS_INL_H_