feat(mpi): implement host staging

CMakeLists.txt

@@ -6,8 +6,7 @@ cmake_minimum_required(VERSION 3.23)
 
 project(
   KokkosComm
-  LANGUAGES
-    CXX
+  LANGUAGES CXX
   VERSION 0.2.0
   DESCRIPTION "Experimental MPI interfaces (and more!) for the Kokkos C++ Performance Portability Programming ecosystem"
   HOMEPAGE_URL "https://kokkos.org/kokkos-comm/"
@@ -21,12 +20,14 @@ option(KokkosComm_ENABLE_TESTS "Build KokkosComm tests" OFF)
 option(KokkosComm_ENABLE_MPI "Build KokkosComm with MPI transport" ON)
 option(KokkosComm_ENABLE_NCCL "Build KokkosComm with NCCL transport" OFF)
 option(KokkosComm_ABORT_ON_ERROR "Runtime error checks trigger a global abort" OFF)
+option(KokkosComm_ENABLE_GPU_AWARE_MPI "Provided MPI is GPU-aware" OFF)
 
 # Resolve options
 set(KOKKOSCOMM_ENABLE_PERFTESTS ${KokkosComm_ENABLE_PERFTESTS} CACHE BOOL "" FORCE)
 set(KOKKOSCOMM_ENABLE_TESTS ${KokkosComm_ENABLE_TESTS} CACHE BOOL "" FORCE)
 set(KOKKOSCOMM_ENABLE_MPI ${KokkosComm_ENABLE_MPI} CACHE BOOL "" FORCE)
 set(KOKKOSCOMM_ENABLE_NCCL ${KokkosComm_ENABLE_NCCL} CACHE BOOL "" FORCE)
+set(KOKKOSCOMM_ENABLE_GPU_AWARE_MPI ${KokkosComm_ENABLE_GPU_AWARE_MPI} CACHE BOOL "" FORCE)
 
 find_package(Kokkos REQUIRED)
 if(KOKKOSCOMM_ENABLE_MPI)

src/KokkosComm/CMakeLists.txt

@@ -6,37 +6,37 @@ target_sources(
   KokkosComm
   INTERFACE
     FILE_SET kokkoscomm_public_headers
-    TYPE HEADERS
-    BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-    FILES
-      KokkosComm.hpp
-      collective.hpp
-      concepts.hpp
-      fwd.hpp
-      point_to_point.hpp
-      traits.hpp
-      datatype.hpp
-      reduction_op.hpp
+      TYPE HEADERS
+      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+      FILES
+        KokkosComm.hpp
+        collective.hpp
+        concepts.hpp
+        fwd.hpp
+        point_to_point.hpp
+        traits.hpp
+        datatype.hpp
+        reduction_op.hpp
 )
 
 # Implementation detail headers
 target_sources(
   KokkosComm
   INTERFACE
     FILE_SET kokkoscomm_impl_headers
-    TYPE HEADERS
-    BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-    FILES impl/contiguous.hpp
+      TYPE HEADERS
+      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+      FILES impl/contiguous.hpp impl/host_staging.hpp
 )
 
 # Configuration header
 target_sources(
   KokkosComm
   INTERFACE
     FILE_SET kokkoscomm_config_headers
-    TYPE HEADERS
-    BASE_DIRS ${CMAKE_BINARY_DIR}/src
-    FILES ${PROJECT_BINARY_DIR}/src/KokkosComm/config.hpp
+      TYPE HEADERS
+      BASE_DIRS ${CMAKE_BINARY_DIR}/src
+      FILES ${PROJECT_BINARY_DIR}/src/KokkosComm/config.hpp
 )
 
 if(KOKKOSCOMM_ENABLE_MPI)
@@ -45,40 +45,40 @@ if(KOKKOSCOMM_ENABLE_MPI)
     KokkosComm
     INTERFACE
       FILE_SET kokkoscomm_mpi_headers
-      TYPE HEADERS
-      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-      FILES
-        # Structures
-        mpi/mpi_space.hpp
-        mpi/comm_mode.hpp
-        mpi/handle.hpp
-        mpi/req.hpp
-        mpi/comm_mode.hpp
-        mpi/channel.hpp
-        # P2P
-        mpi/irecv.hpp
-        mpi/isend.hpp
-        mpi/recv.hpp
-        mpi/send.hpp
-        # Collectives
-        mpi/broadcast.hpp
-        mpi/allgather.hpp
-        mpi/allreduce.hpp
-        mpi/alltoall.hpp
-        mpi/reduce.hpp
-        mpi/scan.hpp
-        # Other/utilities
-        mpi/barrier.hpp
+        TYPE HEADERS
+        BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+        FILES
+          # Structures
+          mpi/mpi_space.hpp
+          mpi/comm_mode.hpp
+          mpi/handle.hpp
+          mpi/req.hpp
+          mpi/comm_mode.hpp
+          mpi/channel.hpp
+          # P2P
+          mpi/irecv.hpp
+          mpi/isend.hpp
+          mpi/recv.hpp
+          mpi/send.hpp
+          # Collectives
+          mpi/broadcast.hpp
+          mpi/allgather.hpp
+          mpi/allreduce.hpp
+          mpi/alltoall.hpp
+          mpi/reduce.hpp
+          mpi/scan.hpp
+          # Other/utilities
+          mpi/barrier.hpp
   )
 
   # Implementation detail MPI headers
   target_sources(
     KokkosComm
     INTERFACE
       FILE_SET kokkoscomm_mpi_impl_headers
-      TYPE HEADERS
-      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-      FILES mpi/impl/pack_traits.hpp mpi/impl/packer.hpp mpi/impl/tags.hpp mpi/impl/error_handling.hpp
+        TYPE HEADERS
+        BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+        FILES mpi/impl/pack_traits.hpp mpi/impl/packer.hpp mpi/impl/tags.hpp mpi/impl/error_handling.hpp
   )
 endif()
 
@@ -88,32 +88,32 @@ if(KOKKOSCOMM_ENABLE_NCCL)
     KokkosComm
     INTERFACE
       FILE_SET kokkoscomm_nccl_headers
-      TYPE HEADERS
-      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-      FILES
-        # Structures
-        nccl/nccl_space.hpp
-        nccl/handle.hpp
-        nccl/req.hpp
-        # P2P
-        nccl/send.hpp
-        nccl/recv.hpp
-        # Collectives
-        nccl/broadcast.hpp
-        nccl/alltoall.hpp
-        nccl/allgather.hpp
-        nccl/allreduce.hpp
-        nccl/reduce.hpp
+        TYPE HEADERS
+        BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+        FILES
+          # Structures
+          nccl/nccl_space.hpp
+          nccl/handle.hpp
+          nccl/req.hpp
+          # P2P
+          nccl/send.hpp
+          nccl/recv.hpp
+          # Collectives
+          nccl/broadcast.hpp
+          nccl/alltoall.hpp
+          nccl/allgather.hpp
+          nccl/allreduce.hpp
+          nccl/reduce.hpp
   )
 
   # Implementation detail NCCL headers
   target_sources(
     KokkosComm
     INTERFACE
       FILE_SET kokkoscomm_nccl_impl_headers
-      TYPE HEADERS
-      BASE_DIRS ${PROJECT_SOURCE_DIR}/src
-      FILES nccl/impl/pack_traits.hpp nccl/impl/packer.hpp nccl/impl/types.hpp
+        TYPE HEADERS
+        BASE_DIRS ${PROJECT_SOURCE_DIR}/src
+        FILES nccl/impl/pack_traits.hpp nccl/impl/packer.hpp nccl/impl/types.hpp
   )
 endif()
 
@@ -127,10 +127,7 @@ include(CheckCXXCompilerFlag)
 macro(kokkoscomm_check_and_add_compile_options)
   set(target ${ARGV0})
   set(flag ${ARGV1})
-  check_cxx_compiler_flag(
-    ${flag}
-    HAS_${flag}
-  )
+  check_cxx_compiler_flag(${flag} HAS_${flag})
   if(HAS_${flag})
     target_compile_options(${target} INTERFACE ${flag})
   endif()
@@ -140,12 +137,7 @@ endmacro()
 add_library(KokkosCommFlags INTERFACE)
 add_library(KokkosComm::KokkosCommFlags ALIAS KokkosCommFlags)
 target_compile_features(KokkosCommFlags INTERFACE cxx_std_20)
-set_target_properties(
-  KokkosCommFlags
-  PROPERTIES
-    CXX_EXTENSIONS
-      OFF
-)
+set_target_properties(KokkosCommFlags PROPERTIES CXX_EXTENSIONS OFF)
 
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wall)
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wextra)
@@ -158,43 +150,24 @@ kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wmissing-include-dirs)
 kokkoscomm_check_and_add_compile_options(KokkosCommFlags -Wno-gnu-zero-variadic-macro-arguments)
 
 # Linking
-target_link_libraries(
-  KokkosComm
-  INTERFACE
-    KokkosComm::KokkosCommFlags
-    Kokkos::kokkos
-)
+target_link_libraries(KokkosComm INTERFACE KokkosComm::KokkosCommFlags Kokkos::kokkos)
 if(KOKKOSCOMM_ENABLE_MPI)
   target_link_libraries(KokkosComm INTERFACE MPI::MPI_CXX)
 endif()
 if(KOKKOSCOMM_ENABLE_NCCL)
   target_link_libraries(KokkosComm INTERFACE NCCL::NCCL)
 endif()
-target_link_libraries(
-  KokkosComm
-  INTERFACE
-    KokkosComm::KokkosCommFlags
-    Kokkos::kokkos
-)
+target_link_libraries(KokkosComm INTERFACE KokkosComm::KokkosCommFlags Kokkos::kokkos)
 
 # Install library
 install(
-  TARGETS
-    KokkosComm
-    KokkosCommFlags
+  TARGETS KokkosComm KokkosCommFlags
   EXPORT KokkosCommTargets
-  FILE_SET
-  kokkoscomm_public_headers
-  FILE_SET
-  kokkoscomm_impl_headers
-  FILE_SET
-  kokkoscomm_mpi_headers
-  FILE_SET
-  kokkoscomm_mpi_impl_headers
-  FILE_SET
-  kokkoscomm_nccl_headers
-  FILE_SET
-  kokkoscomm_nccl_impl_headers
-  FILE_SET
-  kokkoscomm_config_headers
+  FILE_SET kokkoscomm_public_headers
+  FILE_SET kokkoscomm_impl_headers
+  FILE_SET kokkoscomm_mpi_headers
+  FILE_SET kokkoscomm_mpi_impl_headers
+  FILE_SET kokkoscomm_nccl_headers
+  FILE_SET kokkoscomm_nccl_impl_headers
+  FILE_SET kokkoscomm_config_headers
 )

src/KokkosComm/impl/host_staging.hpp

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#pragma once
+
+#include <Kokkos_Core.hpp>
+
+#include <KokkosComm/concepts.hpp>
+
+namespace KokkosComm::Impl {
+
+template <KokkosView V>
+inline constexpr bool needs_staging_v =
+    !Kokkos::SpaceAccessibility<Kokkos::HostSpace, typename V::memory_space>::accessible;
+
+/// Stage view on the host for non-GPU-aware communications.
+/// No-op if `view` is device-accessible from the host.
+template <KokkosView V>
+auto stage_for(const V& view) {
+  return Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, view);
+}
+
+/// Copy back to device (e.g. for receive operations).
+/// No-op if `dst` is device-accessible from the host.
+template <KokkosExecutionSpace E, KokkosView Dst, KokkosView Src>
+auto copy_back(const E& space, Dst& dst, const Src& src) -> void {
+  if constexpr (needs_staging_v<Dst>) {
+    Kokkos::deep_copy(space, dst, src);
+  }
+}
+
+}  // namespace KokkosComm::Impl