inclusive_scan.hpp

// SPDX-FileCopyrightText: Intel Corporation
//
// SPDX-License-Identifier: BSD-3-Clause
 
#pragma once
 
#include <dr/mp/algorithms/inclusive_exclusive_scan_impl.hpp>
 
namespace dr::mp {
 
template <dr::distributed_contiguous_range R,
          dr::distributed_contiguous_range O, typename BinaryOp, typename T>
auto inclusive_scan(R &&r, O &&o, BinaryOp &&binary_op, T init) {
  return __detail::inclusive_exclusive_scan_impl_<false>(
      std::forward<R>(r), rng::begin(std::forward<O>(o)),
      std::forward<BinaryOp>(binary_op), std::optional(init));
}
 
template <dr::distributed_contiguous_range R,
          dr::distributed_contiguous_range O, typename BinaryOp>
auto inclusive_scan(R &&r, O &&o, BinaryOp &&binary_op) {
  return __detail::inclusive_exclusive_scan_impl_<false>(
      std::forward<R>(r), rng::begin(std::forward<O>(o)),
      std::forward<BinaryOp>(binary_op));
}
 
template <dr::distributed_contiguous_range R,
          dr::distributed_contiguous_range O>
auto inclusive_scan(R &&r, O &&o) {
  return dr::mp::inclusive_scan(std::forward<R>(r), std::forward<O>(o),
                                std::plus<rng::range_value_t<R>>());
}
 
// Distributed iterator versions
 
template <dr::distributed_iterator Iter, dr::distributed_iterator OutputIter,
          typename BinaryOp, typename T>
OutputIter inclusive_scan(Iter first, Iter last, OutputIter d_first,
                          BinaryOp &&binary_op, T init) {
 
  return dr::mp::inclusive_scan(rng::subrange(first, last), d_first,
                                std::forward<BinaryOp>(binary_op), init);
}
 
template <dr::distributed_iterator Iter, dr::distributed_iterator OutputIter,
          typename BinaryOp>
OutputIter inclusive_scan(Iter first, Iter last, OutputIter d_first,
                          BinaryOp &&binary_op) {
 
  auto dist = rng::distance(first, last);
  auto d_last = d_first;
  rng::advance(d_last, dist);
  dr::mp::inclusive_scan(rng::subrange(first, last),
                         rng::subrange(d_first, d_last),
                         std::forward<BinaryOp>(binary_op));
 
  return d_last;
}
 
template <dr::distributed_iterator Iter, dr::distributed_iterator OutputIter>
OutputIter inclusive_scan(Iter first, Iter last, OutputIter d_first) {
  auto dist = rng::distance(first, last);
  auto d_last = d_first;
  rng::advance(d_last, dist);
  dr::mp::inclusive_scan(rng::subrange(first, last),
                         rng::subrange(d_first, d_last));
 
  return d_last;
}
 
} // namespace dr::mp