mpi: isend profiling region, space handling

kokkos · May 15, 2024 · 0017caa · 0017caa
1 parent 0981355
commit 0017caa
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/impl/KokkosComm_isend.hpp b/src/impl/KokkosComm_isend.hpp
@@ -33,6 +33,7 @@ namespace KokkosComm::Impl {
 
 template <KokkosView SendView>
 void isend(const SendView &sv, int dest, int tag, MPI_Comm comm, MPI_Request &req) {
+  Kokkos::Tools::pushRegion("KokkosComm::Impl::isend");
   using KCT = typename KokkosComm::Traits<SendView>;
 
   if (KCT::is_contiguous(sv)) {
@@ -41,6 +42,7 @@ void isend(const SendView &sv, int dest, int tag, MPI_Comm comm, MPI_Request &re
   } else {
     throw std::runtime_error("only contiguous views supported for low-level isend");
   }
+  Kokkos::Tools::popRegion();
 }
 
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
@@ -79,6 +81,7 @@ KokkosComm::Req isend(const ExecSpace &space, const SendView &sv, int dest, int
     req.keep_until_wait(args.view);
   } else {
     using SendScalar = typename SendView::value_type;
+    space.fence();  // can't issue isend until work in space is complete
     mpi_isend_fn(KCT::data_handle(sv), KCT::span(sv), mpi_type_v<SendScalar>, dest, tag, comm, &req.mpi_req());
     if (KCT::is_reference_counted()) {
       req.keep_until_wait(sv);