Skip to content

Commit

Permalink
Don't restrict index type in builtin reducers
Browse files Browse the repository at this point in the history
  • Loading branch information
@masterleinad
masterleinad committed May 15, 2023
1 parent e94b5dd commit df5681d
Show file tree
Hide file tree
Showing 2 changed files with 212 additions and 3 deletions.
3 changes: 0 additions & 3 deletions core/src/Kokkos_Parallel_Reduce.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -416,7 +416,6 @@ struct MinLoc {
using index_type = std::remove_cv_t<Index>;
static_assert(!std::is_pointer_v<scalar_type> &&
!std::is_array_v<scalar_type>);
static_assert(std::is_integral_v<index_type>);

public:
// Required
Expand Down Expand Up @@ -472,7 +471,6 @@ struct MaxLoc {
using index_type = std::remove_cv_t<Index>;
static_assert(!std::is_pointer_v<scalar_type> &&
!std::is_array_v<scalar_type>);
static_assert(std::is_integral_v<index_type>);

public:
// Required
Expand Down Expand Up @@ -597,7 +595,6 @@ struct MinMaxLoc {
using index_type = std::remove_cv_t<Index>;
static_assert(!std::is_pointer_v<scalar_type> &&
!std::is_array_v<scalar_type>);
static_assert(std::is_integral_v<index_type>);

public:
// Required
Expand Down
212 changes: 212 additions & 0 deletions core/unit_test/TestReducers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,17 @@

//--------------------------------------------------------------------------

namespace Test {
struct MyPair : Kokkos::pair<int, int> {};
} // namespace Test

template <>
struct Kokkos::reduction_identity<Test::MyPair> {
KOKKOS_FUNCTION static Test::MyPair min() {
return Test::MyPair{{INT_MAX, INT_MAX}};
}
};

namespace Test {

struct ReducerTag {};
Expand Down Expand Up @@ -74,6 +85,20 @@ struct TestReducers {
}
};

struct MinLocFunctor2D {
Kokkos::View<const Scalar**, ExecSpace> values;

KOKKOS_INLINE_FUNCTION
void operator()(
const int& i, const int& j,
typename Kokkos::MinLoc<Scalar, MyPair>::value_type& value) const {
if (values(i, j) < value.val) {
value.val = values(i, j);
value.loc = {{i, j}};
}
}
};

struct MaxLocFunctor {
Kokkos::View<const Scalar*, ExecSpace> values;

Expand All @@ -88,6 +113,20 @@ struct TestReducers {
}
};

struct MaxLocFunctor2D {
Kokkos::View<const Scalar**, ExecSpace> values;

KOKKOS_INLINE_FUNCTION
void operator()(
const int& i, const int& j,
typename Kokkos::MaxLoc<Scalar, MyPair>::value_type& value) const {
if (values(i, j) > value.val) {
value.val = values(i, j);
value.loc = {{i, j}};
}
}
};

struct MinMaxLocFunctor {
Kokkos::View<const Scalar*, ExecSpace> values;

Expand All @@ -107,6 +146,25 @@ struct TestReducers {
}
};

struct MinMaxLocFunctor2D {
Kokkos::View<const Scalar**, ExecSpace> values;

KOKKOS_INLINE_FUNCTION
void operator()(
const int& i, const int& j,
typename Kokkos::MinMaxLoc<Scalar, MyPair>::value_type& value) const {
if (values(i, j) > value.max_val) {
value.max_val = values(i, j);
value.max_loc = {{i, j}};
}

if (values(i, j) < value.min_val) {
value.min_val = values(i, j);
value.min_loc = {{i, j}};
}
}
};

struct BAndFunctor {
Kokkos::View<const Scalar*, ExecSpace> values;

Expand Down Expand Up @@ -598,6 +656,44 @@ struct TestReducers {
}
}

static void test_minloc_2d(int N) {
using reducer_type = Kokkos::MinLoc<Scalar, MyPair>;
using value_type = typename reducer_type::value_type;

Kokkos::View<Scalar**, ExecSpace> values("Values", N, N);
auto h_values = Kokkos::create_mirror_view(values);
Scalar reference_min = std::numeric_limits<Scalar>::max();
MyPair reference_loc = {{-1, -1}};

for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
h_values(i, j) = (Scalar)(rand() % 100000 + 2);

if (h_values(i, j) < reference_min) {
reference_min = h_values(i, j);
reference_loc = {{i, j}};
} else if (h_values(i, j) == reference_min) {
// Make min unique.
h_values(i, j) += Scalar(1);
}
}
Kokkos::deep_copy(values, h_values);

MinLocFunctor2D f;
f.values = values;

{
value_type min_scalar;
reducer_type reducer_scalar(min_scalar);

Kokkos::parallel_reduce(
Kokkos::MDRangePolicy<Kokkos::Rank<2>, ExecSpace>({0, 0}, {N, N}), f,
reducer_scalar);
ASSERT_EQ(min_scalar.val, reference_min);
ASSERT_EQ(min_scalar.loc, reference_loc);
}
}

static void test_maxloc(int N) {
using value_type = typename Kokkos::MaxLoc<Scalar, int>::value_type;

Expand Down Expand Up @@ -661,6 +757,44 @@ struct TestReducers {
}
}

static void test_maxloc_2d(int N) {
using reducer_type = Kokkos::MaxLoc<Scalar, MyPair>;
using value_type = typename reducer_type::value_type;

Kokkos::View<Scalar**, ExecSpace> values("Values", N, N);
auto h_values = Kokkos::create_mirror_view(values);
Scalar reference_max = std::numeric_limits<Scalar>::min();
MyPair reference_loc = {{-1, -1}};

for (int i = 0; i < N; ++i)
for (int j = 0; j < N; ++j) {
h_values(i, j) = (Scalar)(rand() % 100000 + 2);

if (h_values(i, j) > reference_max) {
reference_max = h_values(i, j);
reference_loc = {{i, j}};
} else if (h_values(i, j) == reference_max) {
// Make max unique.
h_values(i, j) -= Scalar(1);
}
}
Kokkos::deep_copy(values, h_values);

MaxLocFunctor2D f;
f.values = values;

{
value_type max_scalar;
reducer_type reducer_scalar(max_scalar);

Kokkos::parallel_reduce(
Kokkos::MDRangePolicy<Kokkos::Rank<2>, ExecSpace>({0, 0}, {N, N}), f,
reducer_scalar);
ASSERT_EQ(max_scalar.val, reference_max);
ASSERT_EQ(max_scalar.loc, reference_loc);
}
}

static void test_minmaxloc(int N) {
using value_type = typename Kokkos::MinMaxLoc<Scalar, int>::value_type;

Expand Down Expand Up @@ -777,6 +911,78 @@ struct TestReducers {
}
}

static void test_minmaxloc_2d(int N) {
using reducer_type = Kokkos::MinMaxLoc<Scalar, MyPair>;
using value_type = typename reducer_type::value_type;

Kokkos::View<Scalar**, ExecSpace> values("Values", N, N);
auto h_values = Kokkos::create_mirror_view(values);
Scalar reference_max = std::numeric_limits<Scalar>::min();
Scalar reference_min = std::numeric_limits<Scalar>::max();
MyPair reference_minloc = {{-1, -1}};
MyPair reference_maxloc = {{-1, -1}};

for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
h_values(i, j) = (Scalar)(rand() % 100000 + 2);
}

for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
if (h_values(i, j) > reference_max) {
reference_max = h_values(i, j);
reference_maxloc = {{i, j}};
} else if (h_values(i, j) == reference_max) {
// Make max unique.
h_values(i, j) -= Scalar(1);
}
}

for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
if (h_values(i, j) < reference_min) {
reference_min = h_values(i, j);
reference_minloc = {{i, j}};
} else if (h_values(i, j) == reference_min) {
// Make min unique.
h_values(i, j) += Scalar(1);
}
}

Kokkos::deep_copy(values, h_values);

MinMaxLocFunctor2D f;
f.values = values;
{
value_type minmax_scalar;
reducer_type reducer_scalar(minmax_scalar);

Kokkos::parallel_reduce(
Kokkos::MDRangePolicy<Kokkos::Rank<2>, ExecSpace>({0, 0}, {N, N}), f,
reducer_scalar);

ASSERT_EQ(minmax_scalar.min_val, reference_min);
for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
if ((minmax_scalar.min_loc == MyPair{{i, j}}) &&
(h_values(i, j) == reference_min)) {
reference_minloc = {{i, j}};
}
}
ASSERT_EQ(minmax_scalar.min_loc, reference_minloc);

ASSERT_EQ(minmax_scalar.max_val, reference_max);
for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
if ((minmax_scalar.max_loc == MyPair{{i, j}}) &&
(h_values(i, j) == reference_max)) {
reference_maxloc = {{i, j}};
}
}
ASSERT_EQ(minmax_scalar.max_loc, reference_maxloc);
}
}

static void test_BAnd(int N) {
Kokkos::View<Scalar*, ExecSpace> values("Values", N);
auto h_values = Kokkos::create_mirror_view(values);
Expand Down Expand Up @@ -985,11 +1191,13 @@ struct TestReducers {
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_minloc(10003);
test_minloc_2d(100);
#endif
test_max(10007);
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_maxloc(10007);
test_maxloc_2d(100);
#endif
#if defined(KOKKOS_ENABLE_OPENMPTARGET) && defined(KOKKOS_COMPILER_CLANG) && \
(KOKKOS_COMPILER_CLANG < 1300)
Expand All @@ -998,6 +1206,7 @@ struct TestReducers {
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_minmaxloc(10007);
test_minmaxloc_2d(100);
#endif
#endif
}
Expand All @@ -1012,11 +1221,13 @@ struct TestReducers {
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_minloc(10003);
test_minloc_2d(100);
#endif
test_max(10007);
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_maxloc(10007);
test_maxloc_2d(100);
#endif
#if defined(KOKKOS_ENABLE_OPENMPTARGET) && defined(KOKKOS_COMPILER_CLANG) && \
(KOKKOS_COMPILER_CLANG < 1300)
Expand All @@ -1025,6 +1236,7 @@ struct TestReducers {
#if !defined(KOKKOS_ENABLE_OPENACC)
// FIXME_OPENACC - OpenACC (V3.3) does not support custom reductions.
test_minmaxloc(10007);
test_minmaxloc_2d(100);
#endif
#endif
test_BAnd(35);
Expand Down

0 comments on commit df5681d

Please sign in to comment.