kokkos.hh

Go to the documentation of this file.

#pragma once
 
#include <DiFfRG/common/tuples.hh>
#include <Kokkos_Core.hpp>
#include <type_traits>
 
#ifdef KOKKOS_ENABLE_CUDA
#include <cuda/std/array>
#include <cuda/std/tuple>
#include <cuda/std/utility>
#else
#include <array>
#include <tuple>
#include <utility>
#endif
 
namespace DiFfRG
{
  struct TBB_ExecutionSpace {
    using memory_space = Kokkos::DefaultHostExecutionSpace::memory_space;
    void fence() const {}
  };
 
  class ExecutionSpaces
  {
  public:
    using GPU_exec_space = Kokkos::DefaultExecutionSpace;
    using GPU_memory_space = GPU_exec_space::memory_space;
 
    using Threads_exec_space = Kokkos::DefaultHostExecutionSpace;
    using Threads_memory_space = Threads_exec_space::memory_space;
 
    using TBB_exec_space = TBB_ExecutionSpace;
    using TBB_memory_space = TBB_exec_space::memory_space;
  };
 
  using GPU_memory = ExecutionSpaces::GPU_memory_space;
  using Threads_memory = ExecutionSpaces::Threads_memory_space;
  using TBB_memory = ExecutionSpaces::TBB_memory_space;
 
  using CPU_memory = Kokkos::DefaultHostExecutionSpace::memory_space;
 
  using GPU_exec = ExecutionSpaces::GPU_exec_space;
  using Threads_exec = ExecutionSpaces::Threads_exec_space;
  using TBB_exec = ExecutionSpaces::TBB_exec_space;
 
  // Ensure that CPU memory space is the same as Threads memory space and TBB memory space
  // We assume that this is true, and when switching to a different memory space, it is always unique.
  static_assert(std::is_same_v<CPU_memory, Threads_memory>,
                "CPU memory space must be the same as Threads memory space");
  static_assert(std::is_same_v<CPU_memory, TBB_memory>, "CPU memory space must be the same as TBB memory space");
 
  template <typename MemorySpace>
  using other_memory_space_t = std::conditional_t<std::is_same_v<MemorySpace, GPU_memory>, CPU_memory, GPU_memory>;
 
  template <class Scalar, class SavedScalar, class Space> struct SumPlus {
  public:
    // Required
    using reducer = SumPlus<Scalar, SavedScalar, Space>;
    using value_type = std::remove_cv_t<Scalar>;
    using saved_type = std::remove_cv_t<SavedScalar>;
    static_assert(!std::is_pointer_v<value_type> && !std::is_array_v<value_type>);
 
    using result_view_type = Kokkos::View<value_type, Space>;
 
  private:
    result_view_type value;
    bool references_scalar_v;
 
    const saved_type plus_value;
 
  public:
    KOKKOS_INLINE_FUNCTION
    SumPlus(value_type &value_, const saved_type &plus_value_)
        : value(&value_), references_scalar_v(true), plus_value(plus_value_)
    {
    }
 
    KOKKOS_INLINE_FUNCTION
    SumPlus(const result_view_type &value_, const saved_type &plus_value_)
        : value(value_), references_scalar_v(false), plus_value(plus_value_)
    {
    }
 
    // Required
    KOKKOS_INLINE_FUNCTION
    void join(value_type &dest, const value_type &src) const { dest += src; }
 
    KOKKOS_INLINE_FUNCTION
    void init(value_type &val) const { val = Kokkos::reduction_identity<value_type>::sum(); }
 
    KOKKOS_INLINE_FUNCTION
    value_type &reference() const { return *value.data(); }
 
    KOKKOS_INLINE_FUNCTION
    result_view_type view() const { return value; }
 
    KOKKOS_INLINE_FUNCTION
    bool references_scalar() const { return references_scalar_v; }
 
    KOKKOS_INLINE_FUNCTION
    void final(value_type &update) const { update += plus_value; }
  };
 
  namespace device
  {
#ifdef KOKKOS_ENABLE_CUDA
    template <typename... T> using tuple = cuda::std::tuple<T...>;
    template <typename T, std::size_t N> using array = cuda::std::array<T, N>;
    using cuda::std::apply;
    using cuda::std::forward;
    using cuda::std::forward_as_tuple;
    using cuda::std::get;
    using cuda::std::index_sequence;
    using cuda::std::integer_sequence;
    using cuda::std::make_integer_sequence;
    using cuda::std::make_tuple;
    using cuda::std::tie;
    using cuda::std::tuple_cat;
    using cuda::std::tuple_element;
#else
    template <typename... T> using tuple = std::tuple<T...>;
    template <typename T, std::size_t N> using array = std::array<T, N>;
    using std::apply;
    using std::forward;
    using std::forward_as_tuple;
    using std::get;
    using std::index_sequence;
    using std::integer_sequence;
    using std::make_integer_sequence;
    using std::make_tuple;
    using std::tie;
    using std::tuple_cat;
    using std::tuple_element;
#endif
  } // namespace device
 
  template <int dim, typename T> struct GetKokkosNDStarType {
    using type = typename GetKokkosNDStarType<dim - 1, T>::type *;
  };
  template <typename T> struct GetKokkosNDStarType<1, T> {
    using type = T *;
  };
 
  // ------------------------------------------------
  // Getting View types
  // ------------------------------------------------
 
  template <int dim, typename T, typename ExecutionSpace>
  using KokkosNDView = Kokkos::View<typename GetKokkosNDStarType<dim, T>::type, // Get the star syntax for
                                                                                // dimensionality recursively with
                                    ExecutionSpace                              // Choice between GPU and CPU
                                    >;
  template <int dim, typename T, typename ExecutionSpace>
  using KokkosNDViewRestrict =
      Kokkos::View<typename GetKokkosNDStarType<dim, T>::type, // Get the star syntax for dimensionality recursively
                                                               // with a helper
                   ExecutionSpace,                             // Choice between GPU and CPU
                   Kokkos::MemoryTraits<Kokkos::Restrict>      // No-alias hint for compiler optimization
                   >;
 
  template <int dim, typename T, typename ExecutionSpace>
  using KokkosNDViewUnmanaged =
      Kokkos::View<typename GetKokkosNDStarType<dim, T>::type, // Get the star syntax for dimensionality recursively
                                                               // with a helper
                   ExecutionSpace,                             // Choice between GPU and CPU
                   Kokkos::MemoryTraits<Kokkos::Unmanaged>     // No allocation: Attach to existing memory
                   >;
 
  template <int dim, typename T, typename ExecutionSpace>
  auto make_kokkos_nd_view(const std::string &label, const device::array<size_t, dim> &extents)
  {
    return device::apply([&](const auto &...args) { return KokkosNDView<dim, T, ExecutionSpace>(label, args...); },
                         extents);
  }
 
  template <int dim, typename T, typename ExecutionSpace>
  auto make_kokkos_nd_view_restrict(const std::string &label, const device::array<size_t, dim> &extents)
  {
    return device::apply(
        [&](const auto &...args) { return KokkosNDViewRestrict<dim, T, ExecutionSpace>(label, args...); }, extents);
  }
 
  // ------------------------------------------------
  // Getting ranges to iterate over
  // ------------------------------------------------
  template <int dim, typename ExecutionSpace> struct KokkosNDRangeHelper {
    using type = Kokkos::MDRangePolicy<Kokkos::Rank<dim>, ExecutionSpace>;
  };
  template <typename ExecutionSpace> struct KokkosNDRangeHelper<1, ExecutionSpace> {
    using type = Kokkos::RangePolicy<ExecutionSpace>;
  };
  template <int dim, typename ExecutionSpace> using KokkosNDRange = KokkosNDRangeHelper<dim, ExecutionSpace>::type;
 
  template <int dim>
  device::array<size_t, dim> compute_tile_hints(const device::array<size_t, dim> &extents, size_t max_threads = 256)
  {
    device::array<size_t, dim> tile;
    for (int i = 0; i < dim; ++i)
      tile[i] = 1;
 
    size_t budget = max_threads;
    // Fill from innermost dimension outward
    for (int i = dim - 1; i >= 0; --i) {
      tile[i] = std::min(extents[i], budget);
      budget /= tile[i];
      if (budget == 0) break;
    }
    return tile;
  }
 
  template <int dim, typename ExecutionSpace>
  auto make_kokkos_nd_range(ExecutionSpace &space, const device::array<size_t, dim> start,
                         const device::array<size_t, dim> end)
  {
    if constexpr (dim == 1) {
      return KokkosNDRange<dim, ExecutionSpace>(space, start[0], end[0]);
    } else {
      Kokkos::Array<size_t, dim> start_view;
      Kokkos::Array<size_t, dim> end_view;
      for (size_t i = 0; i < dim; ++i) {
        start_view[i] = start[i];
        end_view[i] = end[i];
      }
      return KokkosNDRange<dim, ExecutionSpace>(space, start_view, end_view);
    }
  }
 
  template <int dim, typename ExecutionSpace>
  auto make_kokkos_nd_range(ExecutionSpace &space, const device::array<size_t, dim> start,
                         const device::array<size_t, dim> end, const device::array<size_t, dim> tile)
  {
    if constexpr (dim == 1) {
      return KokkosNDRange<dim, ExecutionSpace>(space, start[0], end[0]);
    } else {
      Kokkos::Array<size_t, dim> start_view;
      Kokkos::Array<size_t, dim> end_view;
      Kokkos::Array<size_t, dim> tile_view;
      for (size_t i = 0; i < dim; ++i) {
        start_view[i] = start[i];
        end_view[i] = end[i];
        tile_view[i] = tile[i];
      }
      return KokkosNDRange<dim, ExecutionSpace>(space, start_view, end_view, tile_view);
    }
  }
 
  template <int dim, typename TeamType>
  KOKKOS_FORCEINLINE_FUNCTION auto make_kokkos_nd_thread_range(const TeamType &team, const device::array<size_t, dim> end)
  {
    if constexpr (dim == 1) {
      return Kokkos::TeamThreadRange(team, end[0]);
    } else {
      return device::apply([&](const auto &...args) { return Kokkos::TeamThreadMDRange(team, args...); }, end);
    }
  }
 
  // ------------------------------------------------
  // Wrap Kokkos lambdas
  // ------------------------------------------------
 
  template <int dim, typename FUN> struct KokkosNDLambdaWrapper {
    KOKKOS_FUNCTION
    KokkosNDLambdaWrapper(const FUN &_fun) : fun(_fun) {};
 
    template <typename... Args>
      requires(sizeof...(Args) == dim)
    KOKKOS_FORCEINLINE_FUNCTION void operator()(Args &&...args) const
    {
      fun({{std::forward<Args>(args)...}});
    }
 
    FUN fun;
  };
 
  template <int dim, typename FUN> struct KokkosNDLambdaWrapperReduction {
    KOKKOS_FUNCTION
    KokkosNDLambdaWrapperReduction(const FUN &_fun) : fun(_fun) {};
 
    template <typename... Args>
      requires(sizeof...(Args) == dim + 1)
    KOKKOS_FORCEINLINE_FUNCTION void operator()(Args &&...args) const
    {
      impl(device::make_integer_sequence<size_t, dim>{}, device::forward<Args>(args)...);
    }
 
    FUN fun;
 
  private:
    template <size_t... Is, typename... Args>
    KOKKOS_FORCEINLINE_FUNCTION void impl(device::integer_sequence<size_t, Is...>, Args &&...args) const
    {
      auto tuple = device::tie(args...);
      fun(device::array<size_t, dim>{{static_cast<size_t>(device::get<Is>(tuple))...}}, device::get<dim>(tuple));
    }
  };
} // namespace DiFfRG
 
#include <autodiff/forward/real.hpp>
 
namespace Kokkos
{ // reduction identity must be defined in Kokkos namespace
  template <size_t N, class T> struct reduction_identity<autodiff::Real<N, T>> {
    KOKKOS_FORCEINLINE_FUNCTION static autodiff::Real<N, T> sum() { return autodiff::Real<N, T>(); }
  };
} // namespace Kokkos