File kokkos_internal.h
File List > code_source > templat > include > TempLat > parallel > devices > kokkos > kokkos_internal.h
Go to the documentation of this file
#ifndef TEMPLAT_PARALLEL_KOKKOS_INTERNAL_H
#define TEMPLAT_PARALLEL_KOKKOS_INTERNAL_H
/* This file is part of TempLat, available at https://cosmolattice.github.io/templat .
Copyright 2021-2026 The TempLat authors, see AUTHORS.md.
Released under the MIT license, see LICENSE.md. */
// File info: Main contributor(s): Franz R. Sattler, Year: 2025
#include "TempLat/parallel/devices/kokkos/kokkos.h"
namespace TempLat
{
// Need to forward-declare this - we cannot include the layout header here, as it would create a circular
// dependency.
template <size_t NDim> struct LayoutStruct;
namespace device_kokkos
{
// ------------------------------------------------
// Getting View types with stars
// ------------------------------------------------
template <size_t NDim, typename T> struct GetKokkosNDStarType {
using type = typename GetKokkosNDStarType<NDim - 1, T>::type *;
};
template <typename T> struct GetKokkosNDStarType<1, T> {
using type = T *;
};
// ------------------------------------------------
// View types
// ------------------------------------------------
using DefaultLayout = Kokkos::LayoutRight;
using DefaultExecutionSpace = Kokkos::DefaultExecutionSpace;
using DefaultHostExecutionSpace = Kokkos::DefaultHostExecutionSpace;
// ------------------------------------------------
// Getting the policies for iteration
// ------------------------------------------------
template <size_t NDim> auto getLocalKokkosPolicy(const LayoutStruct<NDim> &layout)
{
Kokkos::Array<uint64_t, NDim> start_iteration;
Kokkos::Array<uint64_t, NDim> stop_iteration;
const auto localSizes = layout.getSizesInMemory();
const size_t nGhosts = layout.getNGhosts();
// What's going on here: on GPU, it is beneficial to reverse the memory access pattern, for coalesced access.
// However, we do not want to impose this on the level of the memory layouts. In particular, this would
// require additional transpositions when going to Fourier space, which is not what we want. So we do the
// transposition within the thread dispatch, if we are on a GPU. Otherwise, for optimal cached memory access
// on CPU, we do not reverse the access pattern.
for (int d = 0; d < (int)NDim; ++d) {
const int _d = device_kokkos::reverse_access_pattern ? (int)NDim - 1 - d : d;
start_iteration[_d] = nGhosts;
stop_iteration[_d] = start_iteration[_d] + localSizes[d];
}
if constexpr (NDim == 1) {
return Kokkos::RangePolicy<DefaultExecutionSpace>(start_iteration[0], stop_iteration[0]);
} else {
return Kokkos::MDRangePolicy<DefaultExecutionSpace, Kokkos::Rank<NDim>>(start_iteration, stop_iteration);
}
}
template <size_t NDim> auto getLocalKokkosOuterPolicy(const LayoutStruct<NDim> &layout)
{
static_assert(NDim >= 2, "The outer policy needs a dimension left over for the inner loop.");
const auto localSizes = layout.getSizesInMemory();
const size_t nGhosts = layout.getNGhosts();
if constexpr (NDim == 2) {
return Kokkos::RangePolicy<DefaultExecutionSpace>(nGhosts, nGhosts + localSizes[0]);
} else {
Kokkos::Array<uint64_t, NDim - 1> start_iteration;
Kokkos::Array<uint64_t, NDim - 1> stop_iteration;
for (int d = 0; d < (int)NDim - 1; ++d) {
start_iteration[d] = nGhosts;
stop_iteration[d] = nGhosts + localSizes[d];
}
return Kokkos::MDRangePolicy<DefaultExecutionSpace, Kokkos::Rank<NDim - 1>>(start_iteration, stop_iteration);
}
}
template <size_t NDim, typename I>
auto getLocalKokkosPolicy(const device_kokkos::array<I, NDim> &starts, const device_kokkos::array<I, NDim> &stops)
{
Kokkos::Array<uint64_t, NDim> start_iteration;
Kokkos::Array<uint64_t, NDim> stop_iteration;
// What's going on here: on GPU, it is beneficial to reverse the memory access pattern, for coalesced access.
// However, we do not want to impose this on the level of the memory layouts. In particular, this would
// require additional transpositions when going to Fourier space, which is not what we want. So we do the
// transposition within the thread dispatch, if we are on a GPU. Otherwise, for optimal cached memory access
// on CPU, we do not reverse the access pattern.
for (int d = 0; d < (int)NDim; ++d) {
const int _d = device_kokkos::reverse_access_pattern ? (int)NDim - 1 - d : d;
start_iteration[_d] = starts[d];
stop_iteration[_d] = start_iteration[_d] + stops[d];
}
if constexpr (NDim == 1) {
return Kokkos::RangePolicy<DefaultExecutionSpace>(start_iteration[0], stop_iteration[0]);
} else {
return Kokkos::MDRangePolicy<DefaultExecutionSpace, Kokkos::Rank<NDim>>(start_iteration, stop_iteration);
}
}
// ------------------------------------------------
// Helpers for tuples and arrays
// ------------------------------------------------
// If you wonder about how performant this is, take a look at https://godbolt.org/z/d33cMaEG1
template <size_t i, typename Head, typename... Tail>
requires(i <= sizeof...(Tail))
DEVICE_INLINE_FUNCTION constexpr auto tuple_last(const device_kokkos::tuple<Head, Tail...> &t)
{
static_assert(i <= sizeof...(Tail), "Cannot take a longer tail than the tuple.");
if constexpr (sizeof...(Tail) + 1 == i)
return apply([](auto &head, auto &...tail) { return device_kokkos::tie(head, tail...); }, t);
else if constexpr (sizeof...(Tail) == i)
return apply([](auto & /*head*/, auto &...tail) { return device_kokkos::tie(tail...); }, t);
else
return apply(
[](auto & /*head*/, auto &...tail) { return device_kokkos::tuple_last<i>(device_kokkos::tie(tail...)); },
t);
}
template <int i, typename Head, typename... Tail>
DEVICE_INLINE_FUNCTION constexpr auto tuple_first(const device_kokkos::tuple<Head, Tail...> &t)
{
static_assert(i <= sizeof...(Tail), "Cannot take a longer sequence than the tuple.");
static_assert(i >= 0, "Cannot take a longer sequence than the tuple.");
if constexpr (i == 0)
return device_kokkos::tuple();
else if constexpr (i == 1)
return apply([](auto &head, auto &.../*tail*/) { return device_kokkos::tie(head); }, t);
else
return apply(
[](auto &head, auto &...tail) {
return device_kokkos::tuple_cat(device_kokkos::tie(head),
device_kokkos::tuple_first<i - 1>(device_kokkos::tie(tail...)));
},
t);
}
template <typename... Args, std::size_t... Is>
DEVICE_INLINE_FUNCTION auto reverse_tuple(const device_kokkos::tuple<Args...> &tuple, std::index_sequence<Is...>)
{
return device_kokkos::tie(device_kokkos::get<sizeof...(Args) - 1 - Is>(tuple)...);
}
template <typename... Args> DEVICE_INLINE_FUNCTION auto reverse_tuple(const device_kokkos::tuple<Args...> &tuple)
{
return reverse_tuple(tuple, std::make_index_sequence<sizeof...(Args)>());
}
template <typename Arg, size_t N, std::size_t... Is>
DEVICE_INLINE_FUNCTION auto reverse_array(const device_kokkos::array<Arg, N> &arr, std::index_sequence<Is...>)
{
return device_kokkos::array<Arg, N>{{get<N - 1 - Is>(arr)...}};
}
template <typename Arg, size_t N> DEVICE_INLINE_FUNCTION auto reverse_array(const device_kokkos::array<Arg, N> &arr)
{
return reverse_array(arr, std::make_index_sequence<N>());
}
// ------------------------------------------------
// Wrapping multidimensional device lambdas for use with Kokkos
// ------------------------------------------------
template <size_t NDim, typename FUN> struct KokkosNDLambdaWrapper {
KokkosNDLambdaWrapper(const FUN &_fun) : fun(_fun) {};
template <typename... Args>
requires(sizeof...(Args) == NDim)
DEVICE_INLINE_FUNCTION void operator()(const Args &...args) const
{
// What's going on here: on GPU, it is beneficial to reverse the memory access pattern, for coalesced access.
// However, we do not want to impose this on the level of the memory layouts. In particular, this would
// require additional transpositions when going to Fourier space, which is not what we want. So we do the
// transposition within the thread dispatch, if we are on a GPU. Otherwise, for optimal cached memory access
// on CPU, we do not reverse the access pattern.
if constexpr (device_kokkos::reverse_access_pattern)
fun(device_kokkos::reverse_array(device_kokkos::IdxArray<NDim>{{static_cast<Idx>(args)...}}));
else
fun({{static_cast<Idx>(args)...}});
}
FUN fun;
};
template <size_t NDim, typename FUN> struct KokkosNDLambdaWrapperInnerLoop {
KokkosNDLambdaWrapperInnerLoop(const FUN &_fun, Idx _innerStart, Idx _innerStop)
: fun(_fun), innerStart(_innerStart), innerStop(_innerStop) {};
template <typename... Args>
requires(sizeof...(Args) == NDim - 1)
DEVICE_INLINE_FUNCTION void operator()(const Args &...args) const
{
const Idx start = innerStart, stop = innerStop;
TEMPLAT_ASSUME_INDEPENDENT
for (Idx i = start; i < stop; ++i)
fun(device_kokkos::IdxArray<NDim>{{static_cast<Idx>(args)..., i}});
}
FUN fun;
Idx innerStart, innerStop;
};
template <size_t NDim, typename FUN> struct KokkosNDLambdaWrapperReduction {
KokkosNDLambdaWrapperReduction(const FUN &_fun) : fun(_fun) {};
template <typename... Args>
requires(sizeof...(Args) == NDim + 1)
DEVICE_INLINE_FUNCTION void operator()(Args &&...args) const
{
// What's going on here: on GPU, it is beneficial to reverse the memory access pattern, for coalesced access.
// However, we do not want to impose this on the level of the memory layouts. In particular, this would
// require additional transpositions when going to Fourier space, which is not what we want. So we do the
// transposition within the thread dispatch, if we are on a GPU. Otherwise, for optimal cached memory access
// on CPU, we do not reverse the access pattern.
auto tuple = device_kokkos::tie(args...);
if constexpr (device_kokkos::reverse_access_pattern)
fun(makeArray(device_kokkos::reverse_tuple(device_kokkos::tuple_first<NDim>(tuple))),
device_kokkos::get<NDim>(tuple)); // the last argument is the reduction result
else
fun(makeArray(device_kokkos::tuple_first<NDim>(tuple)),
device_kokkos::get<NDim>(tuple)); // the last argument is the reduction result
}
FUN fun;
template <typename... Args>
requires(sizeof...(Args) == NDim)
DEVICE_INLINE_FUNCTION auto makeArray(device_kokkos::tuple<Args...> &&tuple) const
{
return device_kokkos::apply(
[](const auto &...args) { return device_kokkos::IdxArray<NDim>{{static_cast<Idx>(args)...}}; },
std::move(tuple));
}
};
} // namespace device_kokkos
} // namespace TempLat
#endif // KOKKOS_H