Kokkos Remote Spaces adds distributed shared memory (DSM) support to Kokkos. This enables a global view on data for a convenient multi-GPU, multi-node, and multi-device programming.
A new memory space type, namely the DefaultRemoteMemorySpace type, represents a Kokkos memory space with remote access semantic. Kokkos View specialized to this memory space through template arguments expose this semantic to the programmer. The underlying implementation of remote memory accesses relies on PGAS as a backend layer.
Currently, four PGAS backends are supported namely SHMEM, NVSHMEM, ROCSHMEM and MPI One-sided (preview). SHMEM and MPI One-sided are host-only programming models and thus support distributed memory accesses on hosts only. This corresponds to Kokkos' execution spaces such as Serial or OpenMP. NVSHMEM or ROCSHMEM support device-initiated communication on compatible devices and consequently require the Kokkos CUDA or HIP execution space. The following diagram shows how Kokkos Remote Spaces fits into the landscape of Kokkos and PGAS libraries.
Kokkos::Experimental::DefaultRemoteMemorySpace;
Kokkos::Experimental::DefaultRemoteMemorySpace::fence();
Kokkos::Experimental::PartitionedLayout;
Kokkos::Experimental::get_local_range(View_t);
Kokkos::Experimental::get_range(View_t, rank);
Kokkos::View(...);}
Kokkos::subview(...);}
Kokkos::deep_copy(...);
Kokkos::local_deep_copy(...);
Kokkos::MemoryTraits::Atomic;
The following example illustrates the type definition of a Kokkos remote view (ViewRemote_3D_t). Further, it shows the instantiation of a remote view (view), in this case of a 3-dimensional array of 20 elements per dimension, and a subsequent instantiation of a subview that can span over multiple virtual address spaces (sub_view). It is worth pointing out that GlobalLayouts, per definition, distribute arrays by the left-most dimension. View data can be accesses similarly to Kokkos views.
We have included more examples in the source code distribution namely RandomAccess as well as CGSolve. CGSolve is an example of a representative code in scientific computing that implements the conjugate gradient method in Kokkos. Taking a closer look shows that CGSolve involves a sequence of kernels that compute the product of two vectors (dot product), perform a matrix-vector multiplication (gemm), and compute a vector sum (axpy). While data can be partitioned for the dot product and the vector sum computations, the matrix-vector multiplication accesses vector elements across all partitions (memory spaces). Because of this, it can be challenging to maintain resource utilization when scaling to multiple processes or GPUs. RandomAccess implements random access to memory.
The following charts provide a perspective on performance and development complexity. They show measured performance expressed as bandwidth (GB/sec) and quantify development complexity with Kokkos Remote Spaces compared to other programming models by comparing LOC (lines of code). We compare different executions on the Lassen supercomputer and against a reference implementation with MPI+Cuda. Here, the executions correspond to configurations with one, two, and four Nvidia V100 GPUs. It can be observed that executions on >2 GPUs result in fine-grained memory accesses over the inter-processor communication interface which is significantly slower (~4x) than the NVLink interface connecting 2 GPUs each (NVLink complex). Reducing the access granularity by presorting remote non-regular accesses to the distributed data structure and moving data in bulk using local deep copies would be the appropriate optimization strategy in this case.
In this chart, GI (global indexing) and LI (local indexing) mark two implementations of CGSolve where one implementation uses Kokkos views with global layout (GlobalLayoutLeft and where the other uses LayoutLeft and thus relies on the programmer to compute the PE index. In the latter case, this computation can be implemented with a binary left-shift which yields a slight performance advantage.
Kokkos Remote Spaces is built using CMake version 3.17 or later. Is a stand-alone project with dependencies on Kokkos and a selected PGAS backend library. The following steps document the build process. Note that building in the root directory in not allowed.
| Path | Description |
|---|---|
| Kokkos_ROOT | Path to the root of the Kokkos install |
| SHMEM_ROOT | Path to the root of a SHMEM installation if enabled |
| NVSHMEM_ROOT | Path to the root of a NVSHMEM installation if enabled |
| Variable | Default | Description |
|---|---|---|
| KRS_ENABLE_SHMEMSPACE | OFF | Enables the SHMEM backend |
| KRS_ENABLE_NVSHMEMSPACE | OFF | Enables the NVSHMEM backend |
| KRS_ENABLE_MPISPACE | OFF | Enables the MPI backend |
| KRS_ENABLE_APPLICATIONS | OFF | Enables building examples |
| KRS_ENABLE_TESTS | OFF | Enables building tests |
Building with SHMEM
$: cmake . -DKRS_ENABLE_SHMEMSPACE=ON
-DKokkos_ROOT=${KOKKOS_INSTALL_DIR}
-DSHMEM_ROOT=${PATH_TO_SHMEM}
-DCMAKE_CXX_COMPILER=oshc++
$: makeBuilding with NVSHMEM
$: cmake . -DKRS_ENABLE_NVSHMEMSPACE=ON
-DKokkos_ROOT=${KOKKOS_INSTALL_DIR}
-DNVSHMEM_ROOT=${PATH_TO_NVSHMEM}
-DCMAKE_CXX_COMPILER=nvcc_wrapper
$: makeBuilding with MPI
$: cmake . -DKRS_ENABLE_MPISPACE=ON
-DKokkos_ROOT=${KOKKOS_INSTALL_DIR}
-DCMAKE_CXX_COMPILER=mpicxx
$: makeApplications depend at least on Kokkos Remote Spaces and may depend on Kokkos Kernels or others. The following sample shows a cmake build file to generate the build scripts for "MyRemoteApp". It depends on Kokkos Remote Spaces and Kokkos Kernels.
#Example
cmake_minimum_required(VERSION 3.13)
project(MyRemoteApp LANGUAGES CXX)
find_package(KokkosKernels REQUIRED)
find_package(KokkosRemote REQUIRED)
add_executable(MatVec matvec.cpp)
target_link_libraries(MatVec PRIVATE \
Kokkos::kokkoskernels Kokkos::kokkosremotespaces)This cmake build fike can be used as
cmake .. -DKokkosKernels_ROOT=$KokkosKernels_INSTALL_PATH -DKokkosRemote_ROOT=$KokkosRemoteSpaces_INSTALL_PATHNote: Kokkos Remote Spaces is in an experimental development stage.
Jan Ciesko, jciesko@sandia.gov