reduce_sum.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_PRIM_FUNCTOR_REDUCE_SUM_HPP
#define STAN_MATH_PRIM_FUNCTOR_REDUCE_SUM_HPP
 
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/functor/apply.hpp>
 
#include <tbb/task_arena.h>
#include <tbb/parallel_reduce.h>
#include <tbb/blocked_range.h>
 
#include <algorithm>
#include <tuple>
#include <vector>
 
namespace stan {
namespace math {
 
namespace internal {
 
template <typename ReduceFunction, typename Enable, typename ReturnType,
          typename Vec, typename... Args>
struct reduce_sum_impl;
 
template <typename ReduceFunction, typename ReturnType, typename Vec,
          typename... Args>
struct reduce_sum_impl<ReduceFunction, require_arithmetic_t<ReturnType>,
                       ReturnType, Vec, Args...> {
  struct recursive_reducer {
    Vec vmapped_;
    std::stringstream msgs_;
    std::tuple<Args...> args_tuple_;
    return_type_t<Vec, Args...> sum_{0.0};
 
    recursive_reducer(Vec&& vmapped, std::ostream* msgs, Args&&... args)
        : vmapped_(std::forward<Vec>(vmapped)),
          args_tuple_(std::forward<Args>(args)...) {}
 
    recursive_reducer(recursive_reducer& other, tbb::split)
        : vmapped_(other.vmapped_), args_tuple_(other.args_tuple_) {}
 
    inline void operator()(const tbb::blocked_range<size_t>& r) {
      if (r.empty()) {
        return;
      }
 
      std::decay_t<Vec> sub_slice;
      sub_slice.reserve(r.size());
      for (size_t i = r.begin(); i < r.end(); ++i) {
        sub_slice.emplace_back(vmapped_[i]);
      }
 
      sum_ += math::apply(
          [&](auto&&... args) {
            return ReduceFunction()(sub_slice, r.begin(), r.end() - 1, &msgs_,
                                    args...);
          },
          args_tuple_);
    }
 
    inline void join(const recursive_reducer& rhs) {
      sum_ += rhs.sum_;
      msgs_ << rhs.msgs_.str();
    }
  };
 
  inline ReturnType operator()(Vec&& vmapped, bool auto_partitioning,
                               int grainsize, std::ostream* msgs,
                               Args&&... args) const {
    const std::size_t num_terms = vmapped.size();
    if (vmapped.empty()) {
      return 0.0;
    }
    recursive_reducer worker(std::forward<Vec>(vmapped), msgs,
                             std::forward<Args>(args)...);
 
    if (auto_partitioning) {
      tbb::parallel_reduce(
          tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker);
    } else {
      tbb::simple_partitioner partitioner;
      tbb::parallel_deterministic_reduce(
          tbb::blocked_range<std::size_t>(0, num_terms, grainsize), worker,
          partitioner);
    }
    if (msgs) {
      *msgs << worker.msgs_.str();
    }
 
    return worker.sum_;
  }
};
 
}  // namespace internal
 
template <typename ReduceFunction, typename Vec,
          typename = require_vector_like_t<Vec>, typename... Args>
inline auto reduce_sum(Vec&& vmapped, int grainsize, std::ostream* msgs,
                       Args&&... args) {
  using return_type = return_type_t<Vec, Args...>;
 
  check_positive("reduce_sum", "grainsize", grainsize);
 
#ifdef STAN_THREADS
  return internal::reduce_sum_impl<ReduceFunction, void, return_type, Vec,
                                   ref_type_t<Args&&>...>()(
      std::forward<Vec>(vmapped), true, grainsize, msgs,
      std::forward<Args>(args)...);
#else
  if (vmapped.empty()) {
    return return_type(0.0);
  }
 
  return ReduceFunction()(std::forward<Vec>(vmapped), 0, vmapped.size() - 1,
                          msgs, std::forward<Args>(args)...);
#endif
}
 
}  // namespace math
}  // namespace stan
 
#endif