Add missing fences in all_to_all

dssgabriel · Jun 12, 2024 · ff1c526 · ff1c526
1 parent 2e34aef
commit ff1c526
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/impl/KokkosComm_alltoall.hpp b/src/impl/KokkosComm_alltoall.hpp
@@ -55,6 +55,9 @@ void alltoall(const ExecSpace &space, const SendView &sv, const size_t sendCount
   static_assert(KokkosComm::rank<SendView>() <= 1, "alltoall for SendView::rank > 1 not supported");
   static_assert(KokkosComm::rank<RecvView>() <= 1, "alltoall for RecvView::rank > 1 not supported");
 
+  // Make sure views are ready
+  space.fence("KokkosComm::Impl::alltoall");
+
   if (KokkosComm::PackTraits<SendView>::needs_pack(sv) || KokkosComm::PackTraits<RecvView>::needs_pack(rv)) {
     throw std::runtime_error("alltoall for non-contiguous views not implemented");
   } else {
@@ -90,6 +93,9 @@ void alltoall(const ExecSpace &space, const RecvView &rv, const size_t recvCount
 
   static_assert(RecvView::rank <= 1, "alltoall for RecvView::rank > 1 not supported");
 
+  // Make sure views are ready
+  space.fence("KokkosComm::Impl::alltoall");
+
   if (KokkosComm::PackTraits<RecvView>::needs_pack(rv)) {
     throw std::runtime_error("alltoall for non-contiguous views not implemented");
   } else {