tbb.hh

Go to the documentation of this file.

#pragma once
 
#include <DiFfRG/common/kokkos.hh>
 
// kept: several headers rely on this pulling in tbb::parallel_for as well.
#include <tbb/tbb.h>
 
#include <cstddef>
#include <cstdlib>
#include <functional>
 
namespace DiFfRG
{
  namespace internal
  {
    inline constexpr std::size_t reduction_serial_cutoff = 4096;
 
    inline constexpr std::size_t reduction_grainsize = 256;
 
    template <int dim>
    void unflatten_index(std::size_t flat, const device::array<size_t, dim> &grid_size,
                         device::array<size_t, dim> &idx)
    {
      for (int d = dim - 1; d >= 0; --d) {
        idx[d] = flat % grid_size[d];
        flat /= grid_size[d];
      }
    }
 
    template <int dim> void advance_index(const device::array<size_t, dim> &grid_size, device::array<size_t, dim> &idx)
    {
      for (int d = dim - 1; d >= 0; --d) {
        if (++idx[d] < grid_size[d]) return;
        idx[d] = 0;
      }
    }
 
    template <int dim, typename NT, typename FUN>
    NT serial_sum(const device::array<size_t, dim> &grid_size, const FUN &functor, const std::size_t begin,
                  const std::size_t end)
    {
      device::array<size_t, dim> idx{{}};
      unflatten_index<dim>(begin, grid_size, idx);
      NT value(0);
      for (std::size_t flat = begin; flat < end; ++flat) {
        value += functor(idx);
        advance_index<dim>(grid_size, idx);
      }
      return value;
    }
  } // namespace internal
 
  template <int dim, typename NT, typename FUN>
  NT TBBReduction(const device::array<size_t, dim> &grid_size, const FUN &functor)
  {
    std::size_t total = 1;
    for (int d = 0; d < dim; ++d)
      total *= grid_size[d];
    if (total == 0) return NT(0);
 
    if (total <= internal::reduction_serial_cutoff)
      return internal::serial_sum<dim, NT, FUN>(grid_size, functor, 0, total);
 
    return tbb::parallel_deterministic_reduce(
        tbb::blocked_range<std::size_t>(0, total, internal::reduction_grainsize), NT(0),
        [&](const tbb::blocked_range<std::size_t> &r, NT value) -> NT {
          return value + internal::serial_sum<dim, NT, FUN>(grid_size, functor, r.begin(), r.end());
        },
        std::plus<NT>(), tbb::simple_partitioner());
  }
} // namespace DiFfRG