Skip to content

Commit c7702ae

Browse files
dssgabrieldssgabriel
committed
refactor(mpi): propagate Communicator changes to MPI backend
Signed-off-by: Gabriel Dos Santos <gabriel.dossantos@cea.fr>
1 parent c2d87fe commit c7702ae

File tree

10 files changed

+183
-120
lines changed

10 files changed

+183
-120
lines changed

‎src/KokkosComm/collective.hpp‎

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -31,19 +31,24 @@
3131
namespace KokkosComm::Experimental {
3232

3333
/// Copy the `v` view from the `root` rank to all ranks' `v` view.
34-
template <KokkosView View, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
35-
CommunicationSpace CommSpace = DefaultCommunicationSpace>
36-
auto broadcast(Handle<ExecSpace, CommSpace>& h, View v, int root) -> Request<CommSpace> {
34+
template <
35+
KokkosView View,
36+
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
37+
CommunicationSpace CommSpace = DefaultCommunicationSpace>
38+
auto broadcast(Communicator<CommSpace, ExecSpace>& h, View v, int root) -> Request<CommSpace> {
3739
return Impl::Broadcast<View, ExecSpace, CommSpace>::execute(h, v, root);
3840
}
3941

4042
/// Copy the `sv` view from each rank to the `rv` view, receiving data from rank `i` at offset
4143
/// `i * KokkosComm::span(sv)`.
4244
///
4345
/// Note: this assumes the span of the `rv` view to be `h.size() * KokkosComm::span(sv)`.
44-
template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
45-
CommunicationSpace CommSpace = DefaultCommunicationSpace>
46-
auto allgather(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv) -> Request<CommSpace> {
46+
template <
47+
KokkosView SendView,
48+
KokkosView RecvView,
49+
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
50+
CommunicationSpace CommSpace = DefaultCommunicationSpace>
51+
auto allgather(Communicator<CommSpace, ExecSpace>& h, const SendView sv, RecvView rv) -> Request<CommSpace> {
4752
return Impl::AllGather<SendView, RecvView, ExecSpace, CommSpace>::execute(h, sv, rv);
4853
}
4954

@@ -53,27 +58,37 @@ auto allgather(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv)
5358
/// Data to receive from source rank `j` is placed into the `rv` view at offset `j * count`.
5459
///
5560
/// Note: this assumes the span of both `sv` and `rv` views to be `h.size * count`.
56-
template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
57-
CommunicationSpace CommSpace = DefaultCommunicationSpace>
58-
auto alltoall(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv, int count) -> Request<CommSpace> {
61+
template <
62+
KokkosView SendView,
63+
KokkosView RecvView,
64+
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
65+
CommunicationSpace CommSpace = DefaultCommunicationSpace>
66+
auto alltoall(Communicator<CommSpace, ExecSpace>& h, const SendView sv, RecvView rv, int count) -> Request<CommSpace> {
5967
return Impl::AllToAll<SendView, RecvView, ExecSpace, CommSpace>::execute(h, sv, rv, count);
6068
}
6169

6270
/// Reduce the `sv` view using the `RedOp` operation and copy the result to all ranks' `rv` view.
63-
template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp,
64-
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
65-
CommunicationSpace CommSpace = DefaultCommunicationSpace>
66-
auto allreduce(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv, RedOp) -> Request<CommSpace> {
71+
template <
72+
KokkosView SendView,
73+
KokkosView RecvView,
74+
ReductionOperator RedOp,
75+
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
76+
CommunicationSpace CommSpace = DefaultCommunicationSpace>
77+
auto allreduce(Communicator<CommSpace, ExecSpace>& h, const SendView sv, RecvView rv, RedOp) -> Request<CommSpace> {
6778
return Impl::AllReduce<SendView, RecvView, RedOp, ExecSpace, CommSpace>::execute(h, sv, rv);
6879
}
6980

7081
/// Reduce the `sv` view using the `RedOp` operation and copy the result to the `root` rank's `rv` view.
7182
///
7283
/// The `rv` view is only used on the `root` rank and ignored for all other ranks.
73-
template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp,
74-
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
75-
CommunicationSpace CommSpace = DefaultCommunicationSpace>
76-
auto reduce(Handle<ExecSpace, CommSpace>& h, const SendView sv, RecvView rv, int root, RedOp) -> Request<CommSpace> {
84+
template <
85+
KokkosView SendView,
86+
KokkosView RecvView,
87+
ReductionOperator RedOp,
88+
KokkosExecutionSpace ExecSpace = Kokkos::DefaultExecutionSpace,
89+
CommunicationSpace CommSpace = DefaultCommunicationSpace>
90+
auto reduce(Communicator<CommSpace, ExecSpace>& h, const SendView sv, RecvView rv, int root, RedOp)
91+
-> Request<CommSpace> {
7792
return Impl::Reduce<SendView, RecvView, RedOp, ExecSpace, CommSpace>::execute(h, sv, rv, root);
7893
}
7994

‎src/KokkosComm/mpi/allgather.hpp‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,19 @@ auto iallgather(const ExecSpace& space, const SView sv, RView rv, MPI_Comm comm)
2525
static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::iallgather: View value types must be identical");
2626
Kokkos::Tools::pushRegion("KokkosComm::mpi::iallgather");
2727

28-
fail_if(!is_contiguous(sv) || !is_contiguous(rv),
29-
"KokkosComm::mpi::iallgather: unimplemented for non-contiguous views");
28+
fail_if(
29+
!is_contiguous(sv) || !is_contiguous(rv), "KokkosComm::mpi::iallgather: unimplemented for non-contiguous views"
30+
);
3031

3132
// Sync: Work in space may have been used to produce view data.
3233
space.fence("fence before non-blocking all-gather");
3334

3435
Request<MpiSpace> req;
3536
// All ranks send/recv same count
36-
MPI_Iallgather(data_handle(sv), span(sv), datatype<MpiSpace, ST>, data_handle(rv), span(sv), datatype<MpiSpace, RT>,
37-
comm, req.request_ptr());
37+
MPI_Iallgather(
38+
data_handle(sv), span(sv), datatype_for<MpiSpace>(sv), data_handle(rv), span(sv), datatype_for<MpiSpace>(rv),
39+
comm, req.request_ptr()
40+
);
3841
req.extend_view_lifetime(sv);
3942
req.extend_view_lifetime(rv);
4043

@@ -53,8 +56,10 @@ void allgather(const SendView& sv, const RecvView& rv, MPI_Comm comm) {
5356
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(rv), "low-level allgather requires contiguous recv view");
5457

5558
const int count = KokkosComm::span(sv); // all ranks send/recv same count
56-
MPI_Allgather(KokkosComm::data_handle(sv), count, datatype<MpiSpace, SendScalar>(), KokkosComm::data_handle(rv),
57-
count, datatype<MpiSpace, RecvScalar>(), comm);
59+
MPI_Allgather(
60+
KokkosComm::data_handle(sv), count, datatype<MpiSpace, SendScalar>(), KokkosComm::data_handle(rv), count,
61+
datatype<MpiSpace, RecvScalar>(), comm
62+
);
5863

5964
Kokkos::Tools::popRegion();
6065
}
@@ -69,8 +74,10 @@ void allgather(const ExecSpace& space, const RecvView& rv, const size_t recvCoun
6974
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(rv), "low-level allgather requires contiguous recv view");
7075

7176
space.fence("fence before allgather"); // work in space may have been used to produce send view data
72-
MPI_Allgather(MPI_IN_PLACE, 0 /*ignored*/, MPI_DATATYPE_NULL /*ignored*/, KokkosComm::data_handle(rv), recvCount,
73-
datatype<MpiSpace, RecvScalar>(), comm);
77+
MPI_Allgather(
78+
MPI_IN_PLACE, 0 /*ignored*/, MPI_DATATYPE_NULL /*ignored*/, KokkosComm::data_handle(rv), recvCount,
79+
datatype<MpiSpace, RecvScalar>(), comm
80+
);
7481

7582
Kokkos::Tools::popRegion();
7683
}
@@ -79,8 +86,10 @@ template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvVi
7986
void allgather(const ExecSpace& space, const SendView& sv, const RecvView& rv, MPI_Comm comm) {
8087
Kokkos::Tools::pushRegion("KokkosComm::Mpi::allgather");
8188

82-
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
83-
"allgather for non-contiguous views not implemented");
89+
KokkosComm::mpi::fail_if(
90+
!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
91+
"allgather for non-contiguous views not implemented"
92+
);
8493

8594
space.fence("fence before allgather"); // work in space may have been used to produce send view data
8695
allgather(sv, rv, comm);
@@ -93,8 +102,8 @@ namespace Experimental::Impl {
93102

94103
template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace>
95104
struct AllGather<SendView, RecvView, ExecSpace, MpiSpace> {
96-
static auto execute(Handle<ExecSpace, MpiSpace>& h, const SendView sv, RecvView rv) -> Request<MpiSpace> {
97-
return mpi::iallgather(h.space(), sv, rv, h.comm());
105+
static auto execute(Communicator<MpiSpace, ExecSpace>& h, const SendView sv, RecvView rv) -> Request<MpiSpace> {
106+
return mpi::iallgather(h.exec(), sv, rv, h.comm());
98107
}
99108
};
100109

‎src/KokkosComm/mpi/allreduce.hpp‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,15 +34,17 @@ auto iallreduce(const ExecSpace& space, const SView sv, RView rv, MPI_Op op, MPI
3434
not std::is_same_v<typename SView::execution_space, Kokkos::HIP> and
3535
not std::is_same_v<typename RView::execution_space, Kokkos::HIP>,
3636
#endif
37-
"KokkosComm::mpi::iallreduce: Unsupported with Open MPI + Kokkos CUDA/HIP backend");
37+
"KokkosComm::mpi::iallreduce: Unsupported with Open MPI + Kokkos CUDA/HIP backend"
38+
);
3839
#endif
3940
using ST = typename SView::non_const_value_type;
4041
using RT = typename RView::non_const_value_type;
4142
static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::iallreduce: View value types must be identical");
4243
Kokkos::Tools::pushRegion("KokkosComm::mpi::iallreduce");
4344

44-
fail_if(!is_contiguous(sv) || !is_contiguous(rv),
45-
"KokkosComm::mpi::iallreduce: unimplemented for non-contiguous views");
45+
fail_if(
46+
!is_contiguous(sv) || !is_contiguous(rv), "KokkosComm::mpi::iallreduce: unimplemented for non-contiguous views"
47+
);
4648

4749
// Sync: Work in space may have been used to produce view data.
4850
space.fence("fence before non-blocking all-gather");
@@ -63,16 +65,19 @@ void allreduce(SendView const& sv, RecvView const& rv, MPI_Op op, MPI_Comm comm)
6365

6466
using SendScalar = typename SendView::value_type;
6567
using RecvScalar = typename RecvView::value_type;
66-
static_assert(std::is_same_v<std::remove_cv_t<SendScalar>, std::remove_cv_t<RecvScalar> >,
67-
"Send and receive views have different value types");
68+
static_assert(
69+
std::is_same_v<std::remove_cv_t<SendScalar>, std::remove_cv_t<RecvScalar> >,
70+
"Send and receive views have different value types"
71+
);
6872

6973
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(sv), "low-level allreduce requires contiguous send view");
7074
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(rv), "low-level allreduce requires contiguous recv view");
7175
KokkosComm::mpi::fail_if(sv.size() != rv.size(), "allreduce requires send and receive views to have the same size");
7276

7377
int const count = sv.size();
74-
MPI_Allreduce(KokkosComm::data_handle(sv), KokkosComm::data_handle(rv), count, datatype<MpiSpace, SendScalar>(), op,
75-
comm);
78+
MPI_Allreduce(
79+
KokkosComm::data_handle(sv), KokkosComm::data_handle(rv), count, datatype<MpiSpace, SendScalar>(), op, comm
80+
);
7681

7782
Kokkos::Tools::popRegion();
7883
}
@@ -95,8 +100,10 @@ template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvVi
95100
void allreduce(ExecSpace const& space, SendView const& sv, RecvView const& rv, MPI_Op op, MPI_Comm comm) {
96101
Kokkos::Tools::pushRegion("KokkosComm::mpi::allreduce");
97102

98-
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
99-
"allreduce for non-contiguous views not implemented");
103+
KokkosComm::mpi::fail_if(
104+
!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
105+
"allreduce for non-contiguous views not implemented"
106+
);
100107

101108
space.fence("fence before allreduce"); // work in space may have been used to produce send view data
102109
allreduce(sv, rv, op, comm);
@@ -121,8 +128,8 @@ namespace Experimental::Impl {
121128

122129
template <KokkosView SendView, KokkosView RecvView, ReductionOperator RedOp, KokkosExecutionSpace ExecSpace>
123130
struct AllReduce<SendView, RecvView, RedOp, ExecSpace, MpiSpace> {
124-
static auto execute(Handle<ExecSpace, MpiSpace>& h, const SendView& sv, RecvView rv) -> Request<MpiSpace> {
125-
return mpi::iallreduce(h.space(), sv, rv, reduction_op<MpiSpace, RedOp>(), h.mpi_comm());
131+
static auto execute(Communicator<MpiSpace, ExecSpace>& h, const SendView& sv, RecvView rv) -> Request<MpiSpace> {
132+
return mpi::iallreduce(h.exec(), sv, rv, reduction_op<MpiSpace, RedOp>(), h.comm());
126133
}
127134
}; // namespace Experimental::Impl
128135

‎src/KokkosComm/mpi/alltoall.hpp‎

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,19 @@ auto ialltoall(const ExecSpace& space, const SView sv, RView rv, int count, MPI_
2525
static_assert(std::is_same_v<ST, RT>, "KokkosComm::mpi::ialltoall: View value types must be identical");
2626
Kokkos::Tools::pushRegion("KokkosComm::mpi::ialltoall");
2727

28-
fail_if(!is_contiguous(sv) || !is_contiguous(rv),
29-
"KokkosComm::mpi::ialltoall: unimplemented for non-contiguous views");
28+
fail_if(
29+
!is_contiguous(sv) || !is_contiguous(rv), "KokkosComm::mpi::ialltoall: unimplemented for non-contiguous views"
30+
);
3031

3132
// Sync: Work in space may have been used to produce view data.
3233
space.fence("fence before non-blocking all-gather");
3334

3435
Request<MpiSpace> req;
3536
// All ranks send/recv same count
36-
MPI_Ialltoall(data_handle(sv), count, datatype<MpiSpace, ST>(), data_handle(rv), count, datatype<MpiSpace, RT>(),
37-
comm, req.request_ptr());
37+
MPI_Ialltoall(
38+
data_handle(sv), count, datatype<MpiSpace, ST>(), data_handle(rv), count, datatype<MpiSpace, RT>(), comm,
39+
req.request_ptr()
40+
);
3841
req.extend_view_lifetime(sv);
3942
req.extend_view_lifetime(rv);
4043

@@ -43,8 +46,14 @@ auto ialltoall(const ExecSpace& space, const SView sv, RView rv, int count, MPI_
4346
}
4447

4548
template <KokkosExecutionSpace ExecSpace, KokkosView SendView, KokkosView RecvView>
46-
void alltoall(const ExecSpace& space, const SendView& sv, const size_t sendCount, const RecvView& rv,
47-
const size_t recvCount, MPI_Comm comm) {
49+
void alltoall(
50+
const ExecSpace& space,
51+
const SendView& sv,
52+
const size_t sendCount,
53+
const RecvView& rv,
54+
const size_t recvCount,
55+
MPI_Comm comm
56+
) {
4857
Kokkos::Tools::pushRegion("KokkosComm::mpi::alltoall");
4958

5059
using SendScalar = typename SendView::value_type;
@@ -53,8 +62,10 @@ void alltoall(const ExecSpace& space, const SendView& sv, const size_t sendCount
5362
// Make sure views are ready
5463
space.fence("KokkosComm::mpi::alltoall");
5564

56-
KokkosComm::mpi::fail_if(!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
57-
"alltoall for non-contiguous views not implemented");
65+
KokkosComm::mpi::fail_if(
66+
!KokkosComm::is_contiguous(sv) || !KokkosComm::is_contiguous(rv),
67+
"alltoall for non-contiguous views not implemented"
68+
);
5869

5970
int size;
6071
MPI_Comm_size(comm, &size);
@@ -72,8 +83,10 @@ void alltoall(const ExecSpace& space, const SendView& sv, const size_t sendCount
7283
KokkosComm::mpi::fail_if(true, ss.str().data());
7384
}
7485

75-
MPI_Alltoall(KokkosComm::data_handle(sv), sendCount, datatype<MpiSpace, SendScalar>(), KokkosComm::data_handle(rv),
76-
recvCount, datatype<MpiSpace, RecvScalar>(), comm);
86+
MPI_Alltoall(
87+
KokkosComm::data_handle(sv), sendCount, datatype<MpiSpace, SendScalar>(), KokkosComm::data_handle(rv), recvCount,
88+
datatype<MpiSpace, RecvScalar>(), comm
89+
);
7790

7891
Kokkos::Tools::popRegion();
7992
}
@@ -100,8 +113,10 @@ void alltoall(const ExecSpace& space, const RecvView& rv, const size_t recvCount
100113
KokkosComm::mpi::fail_if(true, ss.str().data());
101114
}
102115

103-
MPI_Alltoall(MPI_IN_PLACE, 0 /*ignored*/, MPI_BYTE /*ignored*/, KokkosComm::data_handle(rv), recvCount,
104-
datatype<MpiSpace, RecvScalar>(), comm);
116+
MPI_Alltoall(
117+
MPI_IN_PLACE, 0 /*ignored*/, MPI_BYTE /*ignored*/, KokkosComm::data_handle(rv), recvCount,
118+
datatype<MpiSpace, RecvScalar>(), comm
119+
);
105120

106121
Kokkos::Tools::popRegion();
107122
}
@@ -111,8 +126,9 @@ namespace Experimental::Impl {
111126

112127
template <KokkosView SendView, KokkosView RecvView, KokkosExecutionSpace ExecSpace>
113128
struct AllToAll<SendView, RecvView, ExecSpace, MpiSpace> {
114-
static auto execute(Handle<ExecSpace, MpiSpace>& h, const SendView sv, RecvView rv, int count) -> Request<MpiSpace> {
115-
return mpi::ialltoall(h.space(), sv, rv, count, h.mpi_comm());
129+
static auto execute(Communicator<MpiSpace, ExecSpace>& h, const SendView sv, RecvView rv, int count)
130+
-> Request<MpiSpace> {
131+
return mpi::ialltoall(h.exec(), sv, rv, count, h.comm());
116132
}
117133
};
118134

‎src/KokkosComm/mpi/broadcast.hpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ auto ibroadcast(const ExecSpace& space, View& v, int root, MPI_Comm comm) -> Req
2828
space.fence("fence before non-blocking broadcast");
2929

3030
Request<MpiSpace> req;
31-
MPI_Ibcast(data_handle(v), span(v), datatype<MpiSpace, T>, root, comm, req.request_ptr());
31+
MPI_Ibcast(data_handle(v), span(v), datatype_for<MpiSpace>(v), root, comm, req.request_ptr());
3232
req.extend_view_lifetime(v);
3333

3434
Kokkos::Tools::popRegion();
@@ -63,8 +63,8 @@ namespace Experimental::Impl {
6363

6464
template <KokkosView View, KokkosExecutionSpace ExecSpace>
6565
struct Broadcast<View, ExecSpace, MpiSpace> {
66-
static auto execute(Handle<ExecSpace, MpiSpace>& h, View& v, int root) -> Request<MpiSpace> {
67-
return KokkosComm::mpi::ibroadcast(h.space(), v, root, h.comm());
66+
static auto execute(Communicator<MpiSpace, ExecSpace>& h, View& v, int root) -> Request<MpiSpace> {
67+
return KokkosComm::mpi::ibroadcast(h.exec(), v, root, h.comm());
6868
}
6969
};
7070

‎src/KokkosComm/mpi/communicator.hpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ class Communicator<MpiSpace, Ex> {
4646
/// A process that passes `MPI_UNDEFINED` as `color` will not join a new communicator and `nullopt` is returned.
4747
[[nodiscard]] static auto split(
4848
const communicator_type comm, int color, int key, const execution_space& exec = Kokkos::DefaultExecutionSpace{}
49-
) noexcept -> std::optional<Communicator<execution_space, communication_space>> {
49+
) noexcept -> std::optional<Communicator<communication_space, execution_space>> {
5050
communicator_type new_comm;
5151
MPI_Comm_split(comm, color, key, &new_comm);
5252
// Something may have failed, but `color` may be `MPI_UNDEFINED` and we are now outside the split communicator
@@ -55,7 +55,7 @@ class Communicator<MpiSpace, Ex> {
5555
}
5656
return Communicator<communication_space, execution_space>(new_comm, exec, true);
5757
}
58-
[[nodiscard]] auto split(int color, int key) -> std::optional<Communicator<execution_space, communication_space>> {
58+
[[nodiscard]] auto split(int color, int key) -> std::optional<Communicator<communication_space, execution_space>> {
5959
return Communicator::split(comm_, color, key, exec_);
6060
}
6161

@@ -73,7 +73,7 @@ class Communicator<MpiSpace, Ex> {
7373
}
7474
return Communicator<communication_space, execution_space>(new_comm, exec, true);
7575
}
76-
[[nodiscard]] auto duplicate() -> std::optional<Communicator<execution_space, communication_space>> {
76+
[[nodiscard]] auto duplicate() -> std::optional<Communicator<communication_space, execution_space>> {
7777
return Communicator::duplicate(comm_, exec_);
7878
}
7979

‎src/KokkosComm/mpi/irecv.hpp‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,21 +20,23 @@ namespace Impl {
2020
// Recv implementation for Mpi
2121
template <KokkosExecutionSpace ExecSpace, KokkosView RecvView>
2222
struct Recv<RecvView, ExecSpace, MpiSpace> {
23-
static Request<MpiSpace> execute(Handle<ExecSpace, MpiSpace>& h, const RecvView& rv, int src) {
23+
static Request<MpiSpace> execute(Communicator<MpiSpace, ExecSpace>& h, const RecvView& rv, int src) {
2424
using Packer = typename mpi::Impl::PackTraits<RecvView>::packer_type;
2525

26-
const ExecSpace& space = h.space();
26+
const ExecSpace& space = h.exec();
2727

2828
Request<MpiSpace> req;
2929
if (KokkosComm::is_contiguous(rv)) {
3030
space.fence("fence before irecv");
31-
MPI_Irecv(KokkosComm::data_handle(rv), KokkosComm::span(rv), datatype<MpiSpace, typename RecvView::value_type>(),
32-
src, POINTTOPOINT_TAG, h.mpi_comm(), req.request_ptr());
31+
MPI_Irecv(
32+
KokkosComm::data_handle(rv), KokkosComm::span(rv), datatype<MpiSpace, typename RecvView::value_type>(), src,
33+
POINTTOPOINT_TAG, h.comm(), req.request_ptr()
34+
);
3335
req.extend_view_lifetime(rv);
3436
} else {
3537
auto args = Packer::allocate_packed_for(space, "TODO", rv);
3638
space.fence("fence before irecv");
37-
MPI_Irecv(args.view.data(), args.count, args.datatype, src, POINTTOPOINT_TAG, h.mpi_comm(), req.request_ptr());
39+
MPI_Irecv(args.view.data(), args.count, args.datatype, src, POINTTOPOINT_TAG, h.comm(), req.request_ptr());
3840
// implicitly extends args.view and rv lifetime due to lambda capture
3941
req.add_callback([space, rv, args]() {
4042
Packer::unpack_into(space, rv, args.view);

0 commit comments

Comments
 (0)