Skip to content

Commit

Permalink
Merge pull request #1152 from aprokop/kokkos_4.3_std_algorithms
Browse files Browse the repository at this point in the history 
Switch to using Kokkos StdAlgorithms post 4.3
  • Loading branch information
@aprokop
aprokop authored Sep 26, 2024
2 parents 62bf8c2 + 580530d commit 1944c3e
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 221 deletions.
108 changes: 4 additions & 104 deletions src/kokkos_ext/ArborX_DetailsKokkosExtStdAlgorithms.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,116 +13,16 @@
#define ARBORX_DETAILS_KOKKOS_EXT_STD_ALGORITHMS_HPP

#include <ArborX_DetailsKokkosExtAccessibilityTraits.hpp>
#include <ArborX_Exception.hpp>

#include <Kokkos_Core.hpp>
#include <Kokkos_StdAlgorithms.hpp>

namespace ArborX::Details::KokkosExt
{

template <typename ExecutionSpace, typename SrcView, typename DstView,
typename InitValueType>
void exclusive_scan(ExecutionSpace const &space, SrcView const &src,
DstView const &dst, InitValueType init)
{
static_assert(Kokkos::is_execution_space<ExecutionSpace>::value);
static_assert(Kokkos::is_view<SrcView>::value);
static_assert(Kokkos::is_view<DstView>::value);
static_assert(
is_accessible_from<typename SrcView::memory_space, ExecutionSpace>::value,
"Source view must be accessible from the execution space");
static_assert(
is_accessible_from<typename DstView::memory_space, ExecutionSpace>::value,
"Destination view must be accessible from the execution "
"space");
static_assert(std::is_same_v<typename SrcView::value_type,
typename DstView::non_const_value_type>,
"exclusive_scan requires non-const destination type");
static_assert(unsigned(DstView::rank) == unsigned(SrcView::rank) &&
unsigned(DstView::rank) == unsigned(1),
"exclusive_scan requires Views of rank 1");

using ValueType = typename DstView::value_type;

auto const n = src.extent(0);
ARBORX_ASSERT(n == dst.extent(0));
Kokkos::parallel_scan(
"ArborX::Algorithms::exclusive_scan",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, n),
KOKKOS_LAMBDA(int i, ValueType &update, bool final_pass) {
auto const tmp = src(i);
if (final_pass)
dst(i) = update + init;
update += tmp;
});
}

template <typename ExecutionSpace, typename ViewType>
typename ViewType::non_const_value_type
reduce(ExecutionSpace const &space, ViewType const &v,
typename ViewType::non_const_value_type init)
{
static_assert(Kokkos::is_execution_space<ExecutionSpace>::value);
static_assert(Kokkos::is_view<ViewType>::value);
static_assert(is_accessible_from<typename ViewType::memory_space,
ExecutionSpace>::value,
"Source view must be accessible from the execution space");
static_assert(ViewType::rank == 1, "accumulate requires a View of rank 1");

// NOTE: Passing the argument init directly to the parallel_reduce() while
// using a lambda does not yield the expected result because Kokkos will
// supply a default init method that sets the reduction result to zero.
// Rather than going through the hassle of defining a custom functor for
// the reduction, introduce here a temporary variable and add it to init
// before returning.
typename ViewType::non_const_value_type tmp;
Kokkos::parallel_reduce(
"ArborX::Algorithms::accumulate",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, v.extent(0)),
KOKKOS_LAMBDA(int i, typename ViewType::non_const_value_type &update) {
update += v(i);
},
tmp);
init += tmp;
return init;
}

template <typename ExecutionSpace, typename SrcView, typename DstView>
void adjacent_difference(ExecutionSpace const &space, SrcView const &src,
DstView const &dst)
{
static_assert(Kokkos::is_execution_space<ExecutionSpace>::value);
static_assert(Kokkos::is_view<SrcView>::value);
static_assert(Kokkos::is_view<DstView>::value);
static_assert(
is_accessible_from<typename SrcView::memory_space, ExecutionSpace>::value,
"Source view must be accessible from the execution space");
static_assert(
is_accessible_from<typename DstView::memory_space, ExecutionSpace>::value,
"Destination view must be accessible from the execution space");
static_assert(SrcView::rank == 1 && DstView::rank == 1,
"adjacent_difference operates on rank-1 views");
static_assert(
std::is_same_v<typename DstView::value_type,
typename DstView::non_const_value_type>,
"adjacent_difference requires non-const destination value type");
static_assert(std::is_same_v<typename SrcView::non_const_value_type,
typename DstView::value_type>,
"adjacent_difference requires same value type for source and "
"destination");

auto const n = src.extent(0);
ARBORX_ASSERT(n == dst.extent(0));
ARBORX_ASSERT(src != dst);
Kokkos::parallel_for(
"ArborX::Algorithms::adjacent_difference",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
if (i > 0)
dst(i) = src(i) - src(i - 1);
else
dst(i) = src(i);
});
}
using Kokkos::Experimental::adjacent_difference;
using Kokkos::Experimental::exclusive_scan;
using Kokkos::Experimental::reduce;

template <typename ExecutionSpace, typename ViewType>
void iota(ExecutionSpace const &space, ViewType const &v,
Expand Down
117 changes: 0 additions & 117 deletions test/tstDetailsKokkosExtStdAlgorithms.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,120 +50,3 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(iota, DeviceType, ARBORX_DEVICE_TYPES)
Kokkos::deep_copy(w_host, w);
BOOST_TEST(w_ref == w_host, tt::per_element());
}

BOOST_AUTO_TEST_CASE_TEMPLATE(exclusive_scan, DeviceType, ARBORX_DEVICE_TYPES)
{
namespace KokkosExt = ArborX::Details::KokkosExt;
using ExecutionSpace = typename DeviceType::execution_space;

ExecutionSpace space{};

int const n = 10;
Kokkos::View<int *, DeviceType> x("x", n);
std::vector<int> x_ref(n, 1);
x_ref.back() = 0;
auto x_host = Kokkos::create_mirror_view(x);
for (int i = 0; i < n; ++i)
x_host(i) = x_ref[i];
Kokkos::deep_copy(x, x_host);

Kokkos::View<int *, DeviceType> y("y", n);
KokkosExt::exclusive_scan(space, x, y, 0);

std::vector<int> y_ref(n);
std::iota(y_ref.begin(), y_ref.end(), 0);
auto y_host = Kokkos::create_mirror_view(y);
Kokkos::deep_copy(y_host, y);
Kokkos::deep_copy(x_host, x);
BOOST_TEST(y_host == y_ref, tt::per_element());
BOOST_TEST(x_host == x_ref, tt::per_element());
// in-place
KokkosExt::exclusive_scan(space, x, x, 0);
Kokkos::deep_copy(x_host, x);
BOOST_TEST(x_host == y_ref, tt::per_element());
int const m = 11;
BOOST_TEST(n != m);
Kokkos::View<int *, DeviceType> z("z", m);
BOOST_CHECK_THROW(KokkosExt::exclusive_scan(space, x, z, 0),
ArborX::SearchException);
Kokkos::View<double[3], DeviceType> v("v");
auto v_host = Kokkos::create_mirror_view(v);
v_host(0) = 1.;
v_host(1) = 1.;
v_host(2) = 0.;
Kokkos::deep_copy(v, v_host);
// in-place with init value
KokkosExt::exclusive_scan(space, v, v, 5.);
Kokkos::deep_copy(v_host, v);
std::vector<double> v_ref = {5., 6., 7.};
BOOST_TEST(v_host == v_ref, tt::per_element());
Kokkos::View<double *, DeviceType> w("w", 4);
BOOST_CHECK_THROW(KokkosExt::exclusive_scan(space, v, w, 0),
ArborX::SearchException);
v_host(0) = 1.;
v_host(1) = 0.;
v_host(2) = 0.;
Kokkos::deep_copy(v, v_host);
Kokkos::resize(w, 3);
KokkosExt::exclusive_scan(space, v, w, 0);
auto w_host = Kokkos::create_mirror_view(w);
Kokkos::deep_copy(w_host, w);
std::vector<double> w_ref = {0., 1., 1.};
BOOST_TEST(w_host == w_ref, tt::per_element());
}

BOOST_AUTO_TEST_CASE_TEMPLATE(reduce, DeviceType, ARBORX_DEVICE_TYPES)
{
using ExecutionSpace = typename DeviceType::execution_space;
ExecutionSpace space{};

namespace KokkosExt = ArborX::Details::KokkosExt;

Kokkos::View<int[6], DeviceType> v("v");
Kokkos::deep_copy(v, 5);
BOOST_TEST(KokkosExt::reduce(space, v, 3) == 33);

Kokkos::View<int *, DeviceType> w("w", 5);
KokkosExt::iota(space, w, 2);
BOOST_TEST(KokkosExt::reduce(space, w, 4) == 24);
}

BOOST_AUTO_TEST_CASE_TEMPLATE(adjacent_difference, DeviceType,
ARBORX_DEVICE_TYPES)
{
using ExecutionSpace = typename DeviceType::execution_space;
ExecutionSpace space{};

using ArborX::Details::KokkosExt::adjacent_difference;

Kokkos::View<int[5], DeviceType> v("v");
auto v_host = Kokkos::create_mirror_view(v);
v_host(0) = 2;
v_host(1) = 4;
v_host(2) = 6;
v_host(3) = 8;
v_host(4) = 10;
Kokkos::deep_copy(v, v_host);
// In-place operation is not allowed
BOOST_CHECK_THROW(adjacent_difference(space, v, v), ArborX::SearchException);
auto w = Kokkos::create_mirror(DeviceType(), v);
BOOST_CHECK_NO_THROW(adjacent_difference(space, v, w));
auto w_host = Kokkos::create_mirror_view(w);
Kokkos::deep_copy(w_host, w);
std::vector<int> w_ref(5, 2);
BOOST_TEST(w_host == w_ref, tt::per_element());

Kokkos::View<float *, DeviceType> x("x", 10);
Kokkos::deep_copy(x, 3.14);
BOOST_CHECK_THROW(adjacent_difference(space, x, x), ArborX::SearchException);
Kokkos::View<float[10], DeviceType> y("y");
BOOST_CHECK_NO_THROW(adjacent_difference(space, x, y));
std::vector<float> y_ref(10);
y_ref[0] = 3.14;
auto y_host = Kokkos::create_mirror_view(y);
Kokkos::deep_copy(y_host, y);
BOOST_TEST(y_host == y_ref, tt::per_element());

Kokkos::resize(x, 5);
BOOST_CHECK_THROW(adjacent_difference(space, y, x), ArborX::SearchException);
}

0 comments on commit 1944c3e

Please sign in to comment.