Skip to content

Commit 3975bc6

Browse files
cwpearsoncwpearson
committed
Initial attempt at comm/space wrapping and noncontig views
1 parent fbaf029 commit 3975bc6

14 files changed

+287
-37
lines changed

‎perf_tests/test_2dhalo.cpp

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,22 +46,20 @@ void send_recv(benchmark::State &, MPI_Comm comm, const Space &space, int nx, in
4646
auto ym1_s = Kokkos::subview(v, make_pair(1, nx + 1), 1, Kokkos::ALL);
4747
auto ym1_r = Kokkos::subview(v, make_pair(1, nx + 1), 0, Kokkos::ALL);
4848

49-
std::vector<KokkosComm::Req> reqs;
50-
// std::cerr << get_rank(rx, ry) << " -> " << get_rank(xp1, ry) << "\n";
51-
reqs.push_back(KokkosComm::isend(space, xp1_s, get_rank(xp1, ry), 0, comm));
52-
reqs.push_back(KokkosComm::isend(space, xm1_s, get_rank(xm1, ry), 1, comm));
53-
reqs.push_back(KokkosComm::isend(space, yp1_s, get_rank(rx, yp1), 2, comm));
54-
reqs.push_back(KokkosComm::isend(space, ym1_s, get_rank(rx, ym1), 3, comm));
49+
KokkosComm::Handle<Space> h = KokkosComm::plan(space, comm, [=](KokkosComm::Handle<Space> &handle) {
50+
KokkosComm::isend(handle, xp1_s, get_rank(xp1, ry), 0);
51+
KokkosComm::isend(handle, xm1_s, get_rank(xm1, ry), 1);
52+
KokkosComm::isend(handle, yp1_s, get_rank(rx, yp1), 2);
53+
KokkosComm::isend(handle, ym1_s, get_rank(rx, ym1), 3);
54+
});
5555

5656
KokkosComm::recv(space, xm1_r, get_rank(xm1, ry), 0, comm);
5757
KokkosComm::recv(space, xp1_r, get_rank(xp1, ry), 1, comm);
5858
KokkosComm::recv(space, ym1_r, get_rank(rx, ym1), 2, comm);
5959
KokkosComm::recv(space, yp1_r, get_rank(rx, yp1), 3, comm);
6060

6161
// wait for comm
62-
for (KokkosComm::Req &req : reqs) {
63-
req.wait();
64-
}
62+
h.wait();
6563
}
6664

6765
void benchmark_2dhalo(benchmark::State &state) {

‎src/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ target_include_directories(KokkosComm INTERFACE
66
)
77
target_include_directories(KokkosComm INTERFACE
88
$<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/impl>
9+
$<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/mpi>
910
)
1011
target_include_directories(KokkosComm INTERFACE
1112
$<INSTALL_INTERFACE:include>

‎src/KokkosComm.hpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
#pragma once
1818

19+
#include "KokkosComm_point_to_point.hpp"
1920
#include "KokkosComm_collective.hpp"
2021
#include "KokkosComm_version.hpp"
2122
#include "KokkosComm_isend.hpp"
@@ -33,8 +34,6 @@ namespace KokkosComm {
3334

3435
using Impl::alltoall;
3536
using Impl::barrier;
36-
using Impl::irecv;
37-
using Impl::isend;
3837
using Impl::recv;
3938
using Impl::send;
4039

‎src/KokkosComm_config.hpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include "mpi/KokkosComm_mpi.hpp"
20+
#include "KokkosComm_concepts.hpp"
21+
22+
namespace KokkosComm {
23+
24+
using GenericTransport = ::KokkosComm::Mpi;
25+
using SpecialTransport = ::KokkosComm::Mpi;
26+
template <KokkosExecutionSpace ExecSpace>
27+
using Handle = Mpi::Handle<ExecSpace>;
28+
29+
} // namespace KokkosComm

‎src/KokkosComm_plan.hpp

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
namespace KokkosComm {
20+
21+
template <KokkosExecutionSpace ExecSpace, typename CommFunc>
22+
Mpi::Handle<ExecSpace> plan(const ExecSpace &space, MPI_Comm comm, CommFunc f) {
23+
return Mpi::Plan<ExecSpace, CommFunc>(space, comm, f).handle();
24+
}
25+
26+
} // namespace KokkosComm

‎src/KokkosComm_point_to_point.hpp

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include "KokkosComm_api.hpp"
20+
#include "KokkosComm_concepts.hpp"
21+
#include "KokkosComm_config.hpp"
22+
#include "KokkosComm_plan.hpp"
23+
24+
namespace KokkosComm {
25+
26+
template <typename Handle, KokkosView RecvView>
27+
void irecv(Handle &h, RecvView &rv, int src, int tag) {
28+
if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Irecv>) {
29+
SpecialTransport::irecv(h, rv, src, tag);
30+
} else {
31+
GenericTransport::irecv(h, rv, src, tag);
32+
}
33+
}
34+
35+
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
36+
KokkosComm::Handle<ExecSpace> irecv(const ExecSpace &space, const RecvView &rv, int dest, int tag, MPI_Comm comm) {
37+
using MyHandle = KokkosComm::Handle<ExecSpace>;
38+
return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::irecv(handle, rv, dest, tag); });
39+
}
40+
41+
template <typename Handle, KokkosView SendView>
42+
void isend(Handle &h, SendView &sv, int src, int tag) {
43+
if constexpr (Impl::api_avail_v<SpecialTransport, Impl::Api::Isend>) {
44+
SpecialTransport::isend(h, sv, src, tag);
45+
} else {
46+
GenericTransport::isend(h, sv, src, tag);
47+
}
48+
}
49+
50+
template <KokkosExecutionSpace ExecSpace, KokkosView SendView>
51+
KokkosComm::Handle<ExecSpace> isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
52+
using MyHandle = KokkosComm::Handle<ExecSpace>;
53+
return KokkosComm::plan(space, comm, [=](MyHandle &handle) { KokkosComm::isend(handle, sv, dest, tag); });
54+
}
55+
56+
} // namespace KokkosComm

‎src/impl/KokkosComm_api.hpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
namespace KokkosComm::Impl {
20+
21+
enum class Api { Irecv, Isend };
22+
23+
// catch-all: no transports implement any APIs
24+
template <typename Transport, Impl::Api API>
25+
struct api_avail : public std::false_type {};
26+
27+
template <typename Transport, Impl::Api API>
28+
constexpr bool api_avail_v = api_avail<Transport, API>::value;
29+
30+
} // namespace KokkosComm::Impl

‎src/impl/KokkosComm_contiguous.hpp

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
//@HEADER
2+
// ************************************************************************
3+
//
4+
// Kokkos v. 4.0
5+
// Copyright (2022) National Technology & Engineering
6+
// Solutions of Sandia, LLC (NTESS).
7+
//
8+
// Under the terms of Contract DE-NA0003525 with NTESS,
9+
// the U.S. Government retains certain rights in this software.
10+
//
11+
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12+
// See https://kokkos.org/LICENSE for license information.
13+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14+
//
15+
//@HEADER
16+
17+
#pragma once
18+
19+
#include <string>
20+
21+
#include <Kokkos_Core.hpp>
22+
23+
#include "KokkosComm_concepts.hpp"
24+
25+
namespace KokkosComm::Impl {
26+
27+
template <KokkosView View>
28+
struct contiguous_view {
29+
using type = Kokkos::View<typename View::non_const_data_type, Kokkos::LayoutRight, typename View::memory_space>;
30+
};
31+
32+
template <KokkosView View>
33+
using contiguous_view_t = contiguous_view<View>::type;
34+
35+
template <KokkosView View, KokkosExecutionSpace Space>
36+
auto allocate_contiguous_for(const Space &space, const std::string &label, View &v) {
37+
using non_const_packed_view_type = contiguous_view_t<View>;
38+
39+
if constexpr (KokkosComm::rank<View>() == 1) {
40+
return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0));
41+
} else if constexpr (KokkosComm::rank<View>() == 2) {
42+
return non_const_packed_view_type(Kokkos::view_alloc(space, Kokkos::WithoutInitializing, label), v.extent(0),
43+
v.extent(1));
44+
} else {
45+
static_assert(std::is_void_v<View>, "allocate_contiguous_for for views > rank 2 not implemented");
46+
}
47+
}
48+
49+
template <KokkosExecutionSpace Space, KokkosView DstView, KokkosView SrcView>
50+
auto resize_contiguous_for(const Space &space, DstView &out, const SrcView &in) {
51+
static_assert(DstView::rank == SrcView::rank, "");
52+
53+
if constexpr (KokkosComm::rank<DstView>() == 1) {
54+
Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0));
55+
} else if constexpr (KokkosComm::rank<DstView>() == 2) {
56+
Kokkos::realloc(Kokkos::view_alloc(space, Kokkos::WithoutInitializing), out, in.extent(0), in.extent(1));
57+
} else {
58+
static_assert(std::is_void_v<DstView>, "realloc_contiguous_for for views > rank 2 not implemented");
59+
}
60+
}
61+
62+
} // namespace KokkosComm::Impl

‎src/impl/KokkosComm_include_mpi.hpp

Lines changed: 0 additions & 12 deletions
This file was deleted.

‎src/impl/KokkosComm_types.hpp

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,4 +107,50 @@ MPI_Datatype mpi_type() {
107107

108108
template <typename Scalar>
109109
inline MPI_Datatype mpi_type_v = mpi_type<Scalar>();
110+
111+
template <KokkosView View>
112+
MPI_Datatype view_mpi_type(const View &view) {
113+
#define USE_CACHE
114+
115+
#if defined(USE_CACHE)
116+
constexpr int P = 2 * View::rank;
117+
using Key = std::array<int, P>;
118+
119+
auto key_from = [](const View &v) -> Key {
120+
Key key;
121+
for (size_t d = 0; d < View::rank; d++) {
122+
key[d] = v.extent(d);
123+
key[d + 1] = v.stride(d);
124+
}
125+
return key;
126+
};
127+
128+
static std::map<Key, MPI_Datatype> cache;
129+
130+
Key key = key_from(view);
131+
if (cache.count(key) > 0) {
132+
return cache[key];
133+
}
134+
#endif
135+
136+
using value_type = typename View::non_const_value_type;
137+
MPI_Datatype type = mpi_type_v<value_type>;
138+
139+
// This doesn't work for 1D contiguous views into reduce because it
140+
// represents the whole 1D view as 1 Hvector, rather than N elements.
141+
// FIXME: is there a more generic way to handle this, maybe by treating
142+
// the last dimension specially under certain circumstances?
143+
for (size_t d = 0; d < KokkosComm::rank<View>(); ++d) {
144+
MPI_Datatype newtype;
145+
MPI_Type_create_hvector(KokkosComm::extent(view, d) /*count*/, 1 /*block length*/,
146+
KokkosComm::stride(view, d) * sizeof(value_type), type, &newtype);
147+
type = newtype;
148+
}
149+
MPI_Type_commit(&type);
150+
#if defined(USE_CACHE)
151+
cache[key] = type;
152+
#endif
153+
return type;
154+
}
155+
110156
}; // namespace KokkosComm::Impl

‎unit_tests/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ target_link_libraries(test-mpi MPI::MPI_CXX)
4545
add_executable(test-main test_main.cpp
4646
test_gtest_mpi.cpp
4747
test_isendirecv.cpp
48-
test_isendrecv.cpp
48+
#test_isendrecv.cpp # FIXME: disabled during isend development
4949
test_sendrecv.cpp
5050
test_barrier.cpp
5151
test_alltoall.cpp

‎unit_tests/test_isendirecv.cpp

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,13 +49,14 @@ void test_1d(const View1D &a) {
4949
int dst = 1;
5050
Kokkos::parallel_for(
5151
a.extent(0), KOKKOS_LAMBDA(const int i) { a(i) = i; });
52-
KokkosComm::Req req = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD);
53-
req.wait();
52+
53+
KokkosComm::Handle h = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD);
54+
h.wait();
5455
} else if (1 == rank) {
5556
int src = 0;
56-
MPI_Request req;
57-
KokkosComm::irecv(a, src, 0, MPI_COMM_WORLD, req);
58-
MPI_Wait(&req, MPI_STATUS_IGNORE);
57+
58+
KokkosComm::Handle h = KokkosComm::irecv(Kokkos::DefaultExecutionSpace(), a, src, 0, MPI_COMM_WORLD);
59+
h.wait();
5960
int errs;
6061
Kokkos::parallel_reduce(
6162
a.extent(0), KOKKOS_LAMBDA(const int &i, int &lsum) { lsum += a(i) != Scalar(i); }, errs);
@@ -82,12 +83,16 @@ void test_2d(const View2D &a) {
8283
int dst = 1;
8384
Kokkos::parallel_for(
8485
policy, KOKKOS_LAMBDA(int i, int j) { a(i, j) = i * a.extent(0) + j; });
85-
KokkosComm::Req req = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD);
86-
req.wait();
86+
std::cerr << __FILE__ << ":" << __LINE__ << " isend...\n";
87+
KokkosComm::Handle h = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD);
88+
h.wait();
89+
std::cerr << __FILE__ << ":" << __LINE__ << " isend done\n";
8790
} else if (1 == rank) {
88-
int src = 0;
89-
KokkosComm::Req req = KokkosComm::irecv(a, src, 0, MPI_COMM_WORLD);
90-
req.wait();
91+
int src = 0;
92+
std::cerr << __FILE__ << ":" << __LINE__ << " irecv...\n";
93+
KokkosComm::Handle h = KokkosComm::irecv(Kokkos::DefaultExecutionSpace(), a, src, 0, MPI_COMM_WORLD);
94+
h.wait();
95+
std::cerr << __FILE__ << ":" << __LINE__ << " irecv done\n";
9196
int errs;
9297
Kokkos::parallel_reduce(
9398
policy, KOKKOS_LAMBDA(int i, int j, int &lsum) { lsum += a(i, j) != Scalar(i * a.extent(0) + j); }, errs);
@@ -100,9 +105,19 @@ TYPED_TEST(IsendIrecv, 1D_contig) {
100105
test_1d(a);
101106
}
102107

108+
TYPED_TEST(IsendIrecv, 1D_noncontig) {
109+
auto a = ViewBuilder<typename TestFixture::Scalar, 1>::view(noncontig{}, "a", 1013);
110+
test_1d(a);
111+
}
112+
103113
TYPED_TEST(IsendIrecv, 2D_contig) {
104114
auto a = ViewBuilder<typename TestFixture::Scalar, 2>::view(contig{}, "a", 137, 17);
105115
test_2d(a);
106116
}
107117

118+
TYPED_TEST(IsendIrecv, 2D_noncontig) {
119+
auto a = ViewBuilder<typename TestFixture::Scalar, 2>::view(noncontig{}, "a", 137, 17);
120+
test_2d(a);
121+
}
122+
108123
} // namespace

0 commit comments

Comments
 (0)