Program Listing for File reduce_cpp_.hpp

↰ Return to documentation for file (include/ripple/algorithm/kernel/reduce_cpp_.hpp)

#ifndef RIPPLE_ALGORITHM_KERNEL_REDUCE_CPP__HPP
#define RIPPLE_ALGORITHM_KERNEL_REDUCE_CPP__HPP

#include <ripple/iterator/iterator_traits.hpp>

namespace ripple::kernel::cpu {
namespace detail {

template <size_t Dims>
class Reducer;

template <>
class Reducer<0> {
  static constexpr auto dim = ripple::dimx();

 public:
  template <
    typename Iterator1,
    typename Iterator2,
    typename Pred,
    typename... Args>
  ripple_all static auto reduce(
    Iterator1&& it, Iterator2&& result, Pred&& pred, Args&&... args) noexcept
    -> void {
    static_assert(is_iterator_v<Iterator1>, "Reduction requires an iterator!");
    static_assert(is_iterator_v<Iterator2>, "Reduction requires an iterator!");

    constexpr size_t sub_iters = 4;

    // Not reducing the first element, hence -1.
    const auto elements  = it.size(dim) - 1;
    const auto iters     = elements / sub_iters;
    const auto rem_iters = elements % sub_iters;

    auto other = it;
    for (size_t i = 0; i < iters; ++i) {
      unrolled_for<sub_iters>([&](auto j) {
        const size_t idx = i * sub_iters + j;
        other            = it.offset(dim, idx + 1);
        pred.inplace(result, other);
      });
    }

    for (size_t i = 0; i < rem_iters; ++i) {
      const size_t idx = iters * sub_iters + i;
      other            = it.offset(dim, idx + 1);
      pred.inplace(result, other);
    }
  }
};

template <size_t Dim>
class Reducer {
  using NextReducer = Reducer<Dim - 1>;

  static constexpr auto dim = Dim == 1 ? ripple::dimy() : ripple::dimz();

 public:
  template <
    typename Iterator1,
    typename Iterator2,
    typename Pred,
    typename... Args>
  ripple_all static auto reduce(
    Iterator1&& it, Iterator2&& result, Pred&& pred, Args&&... args) noexcept
    -> void {
    static_assert(is_iterator_v<Iterator1>, "Reduction requires an iterator!");
    static_assert(is_iterator_v<Iterator2>, "Reduction requires an iterator!");

    // Reduce the first row/plane:
    NextReducer::reduce(
      it, result, ripple_forward(pred), ripple_forward(args)...);
    for (size_t i = 0; i < it.size(dim) - 1; ++i) {
      auto next = it.offset(dim, i + 1);
      // Next dimension doesn't do the first element:
      pred.inplace(result, next);

      NextReducer::reduce(
        next, result, ripple_forward(pred), ripple_forward(args)...);
    }
  }
};

} // namespace detail

template <typename T, size_t Dims, typename Pred, typename... Args>
auto reduce(
  const HostBlock<T, Dims>& block, Pred&& pred, Args&&... args) noexcept {
  using Reducer = detail::Reducer<Dims - 1>;

  // Store the inital value, and then reduce into the first iterated value.
  auto it         = block.begin();
  auto init_value = *it;
  Reducer::reduce(it, it, ripple_forward(pred), ripple_forward(args)...);

  // Get results and reset the first element, since it has been reduced into.
  auto result = *it;
  *it         = init_value;
  return result;
}

} // namespace ripple::kernel::cpu

#endif // RIPPLE_ALGORITHM_KERNEL_REDUCE_CPP__HPP