The KOKKOS accelerator package doc page explains:

• what hardware and software the accelerated package requires
• how to build SPARTA with the accelerated package
• how to run with the accelerated package either via command-line switches or modifying the input script
• speed-ups to expect
• guidelines for best performance
• restrictions

5.3 KOKKOS package

Kokkos is a templated C++ library that provides abstractions to allow a single implementation of an application kernel (e.g. a collision style) to run efficiently on different kinds of hardware, such as GPUs or many-core CPUs. Kokkos maps the C++ kernel onto different backend languages such as CUDA, OpenMP, or Pthreads. The Kokkos library also provides data abstractions to adjust (at compile time) the memory layout of data structures like 2d and 3d arrays to optimize performance on different hardware. For more information on Kokkos, see Github. Kokkos is part of Trilinos. The Kokkos library was written primarily by Carter Edwards, Christian Trott, and Dan Sunderland (all Sandia).

The SPARTA KOKKOS package contains versions of collide, fix, and compute styles that use data structures and macros provided by the Kokkos library, which is included with SPARTA in /lib/kokkos. The KOKKOS package was developed primarily by Stan Moore (Sandia) with contributions of various styles by others, including Dan Ibanez (Sandia), Tim Fuller (Sandia), and Sam Mish (Sandia). For more information on developing using Kokkos abstractions see the Kokkos programmers' guide at /lib/kokkos/doc/Kokkos_PG.pdf.

The KOKKOS package supports multiple execution backends (per MPI task): Serial (MPI-only for CPUs), OpenMP (threading for many-core CPUs), CUDA (NVIDIA GPUs), HIP (AMD GPUs), and SYCL (Intel GPUs and other SYCL-capable devices). You choose the backend at build time to produce an executable compatible with specific hardware.

NOTE: The KOKKOS package must be built using CMake. GNU Makefile builds are not supported for the KOKKOS package.

NOTE: Kokkos support within SPARTA must be built with a C++20 compatible compiler. For a list of compilers that have been tested with the Kokkos library, see the Kokkos README.

Building SPARTA with the KOKKOS package with CMake:

To build with the KOKKOS package, start with one of the provided preset files in /cmake/presets/. Preset files encode all the settings needed for a particular hardware target. The available KOKKOS presets are:

You may need to override -D Kokkos_ARCH_TYPE=ON to match your specific hardware. For example:

• for Sandy Bridge CPUs, set -D Kokkos_ARCH_SNB=ON
• for Broadwell CPUs, set -D Kokkos_ARCH_BDW=ON
• for K80 GPUs, set -D Kokkos_ARCH_KEPLER37=ON
• for V100 GPUs and Power9 CPUs, set -D Kokkos_ARCH_VOLTA70=ON -D Kokkos_ARCH_POWER9=ON
• for A100 GPUs, set -D Kokkos_ARCH_AMPERE80=ON
• for H100 GPUs, set -D Kokkos_ARCH_HOPPER90=ON

See the Advanced Kokkos Options section below for a complete listing of all Kokkos architecture options.

NOTE: If you are migrating from GNU Makefile builds, see the table of preset files above and pick the one that matches your hardware. The Kokkos KOKKOS_DEVICES and KOKKOS_ARCH Makefile variables map to CMake options as follows: KOKKOS_DEVICES=OpenMP becomes -DKokkos_ENABLE_OPENMP=ON, KOKKOS_DEVICES=Cuda becomes -DKokkos_ENABLE_CUDA=ON, KOKKOS_ARCH=Volta70 becomes -DKokkos_ARCH_VOLTA70=ON, etc.

Compile for CPU-only (MPI only, no threading):

Use a C++20 compatible compiler. Then do the following:

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_mpi_only.cmake /path/to/sparta/cmake
make -j 4 

The resulting executable will be named spa_kokkos_mpi_only.

Compile for CPU-only (MPI plus OpenMP threading):

NOTE: To build with Kokkos support for OpenMP threading, your compiler must support the OpenMP interface. You should have one or more multi-core CPUs so that multiple threads can be launched by each MPI task running on a CPU.

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_omp.cmake /path/to/sparta/cmake
make -j 4 

The resulting executable will be named spa_kokkos_omp.

To select a specific CPU architecture (e.g. Haswell), add:

cmake -C /path/to/sparta/cmake/presets/kokkos_omp.cmake -DKokkos_ARCH_HSW=ON /path/to/sparta/cmake 

Compile for NVIDIA GPUs using CUDA (with OpenMPI or MPICH):

NOTE: To build with Kokkos support for NVIDIA GPUs, NVIDIA CUDA software version 11.0 or later must be installed on your system.

The kokkos_cuda.cmake preset defaults to NVIDIA Hopper (H100) GPUs. To target a different GPU, override the architecture flag. Common NVIDIA GPU models and their Kokkos arch flags are:

• V100 (Volta) -> -DKokkos_ARCH_VOLTA70=ON
• A100 (Ampere) -> -DKokkos_ARCH_AMPERE80=ON
• H100/H200 (Hopper) -> -DKokkos_ARCH_HOPPER90=ON (default in kokkos_cuda.cmake)
• L40S/RTX 4090 (Ada) -> -DKokkos_ARCH_ADA89=ON
• GB200 (Blackwell) -> -DKokkos_ARCH_BLACKWELL100=ON

Build for the default Hopper (H100) GPU:

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_cuda.cmake /path/to/sparta/cmake
make -j 4 

Build for A100 GPUs (override the default Hopper arch):

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_cuda.cmake   -DKokkos_ARCH_HOPPER90=OFF -DKokkos_ARCH_AMPERE80=ON   /path/to/sparta/cmake
make -j 4 

Build for V100 GPUs:

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_cuda.cmake   -DKokkos_ARCH_HOPPER90=OFF -DKokkos_ARCH_VOLTA70=ON   /path/to/sparta/cmake
make -j 4 

The resulting executable will be named spa_kokkos_cuda.

Compile for AMD GPUs using HIP:

NOTE: To build with Kokkos support for AMD GPUs, ROCm software version 5.2 or later must be installed on your system.

The kokkos_hip.cmake preset targets AMD MI250X GPUs (GFX90A). Common AMD GPU models and their Kokkos arch flags are:

• MI250X/MI250/MI210 -> -DKokkos_ARCH_VEGA90A=ON (default in kokkos_hip.cmake)
• MI300X/MI300A -> -DKokkos_ARCH_AMD_GFX942=ON
• RX 7900 XTX (RDNA3) -> -DKokkos_ARCH_AMD_GFX1100=ON

Build for the default MI250X GPU:

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_hip.cmake /path/to/sparta/cmake
make -j 4 

Build for MI300X GPUs (override the default MI250X arch):

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_hip.cmake   -DKokkos_ARCH_VEGA90A=OFF -DKokkos_ARCH_AMD_GFX942=ON   /path/to/sparta/cmake
make -j 4 

For AMD MI300A APU systems using Cray MPICH (e.g. El Capitan):

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/elcapitan_kokkos.cmake /path/to/sparta/cmake
make -j 4 

The resulting executable will be named spa_kokkos_hip or spa_elcapitan_kokkos respectively.

Compile for Intel GPUs using SYCL:

NOTE: To build with Kokkos support for Intel GPUs via SYCL, Intel's oneAPI toolkit (version 2024.0 or later) must be installed on your system. Use the icpx compiler from the Intel oneAPI toolkit. The GPU architecture is detected automatically at runtime; no arch flag override is needed.

The kokkos_sycl.cmake preset targets Intel Ponte Vecchio (PVC) GPUs and uses Intel MKL for FFTs:

mkdir build
cd build
cmake -C /path/to/sparta/cmake/presets/kokkos_sycl.cmake /path/to/sparta/cmake
make -j 4 

The resulting executable will be named spa_kokkos_sycl.

Running SPARTA with the KOKKOS package:

All Kokkos operations occur within the context of an individual MPI task running on a single node of the machine. The total number of MPI tasks used by SPARTA (one or multiple per compute node) is set in the usual manner via the mpirun or mpiexec commands, and is independent of Kokkos. The mpirun or mpiexec command sets the total number of MPI tasks used by SPARTA (one or multiple per compute node) and the number of MPI tasks used per node. E.g. the mpirun command in OpenMPI does this via its -np and -npernode switches. Ditto for MPICH via -np and -ppn.

Running on a multi-core CPU:

Here is a quick overview of how to use the KOKKOS package for CPU acceleration, assuming one or more 16-core nodes.

mpirun -np 16 spa_kokkos_mpi_only -k on -sf kk -in in.collide        # 1 node, 16 MPI tasks/node, no multi-threading
mpirun -np 2 -ppn 1 spa_kokkos_omp -k on t 16 -sf kk -in in.collide  # 2 nodes, 1 MPI task/node, 16 threads/task
mpirun -np 2 spa_kokkos_omp -k on t 8 -sf kk -in in.collide          # 1 node,  2 MPI tasks/node, 8 threads/task
mpirun -np 32 -ppn 4 spa_kokkos_omp -k on t 4 -sf kk -in in.collide  # 8 nodes, 4 MPI tasks/node, 4 threads/task 

To run using the KOKKOS package, use the "-k on", "-sf kk" and "-pk kokkos" command-line switches in your mpirun command. You must use the "-k on" command-line switch to enable the KOKKOS package. It takes additional arguments for hardware settings appropriate to your system. Those arguments are documented here. For OpenMP use:

-k on t Nt 

The "t Nt" option specifies how many OpenMP threads per MPI task to use with a node. The default is Nt = 1, which is MPI-only mode. Note that the product of MPI tasks * OpenMP threads/task should not exceed the physical number of cores (on a node), otherwise performance will suffer. If hyperthreading is enabled, then the product of MPI tasks * OpenMP threads/task should not exceed the physical number of cores * hardware threads. The "-k on" switch also issues a "package kokkos" command (with no additional arguments) which sets various KOKKOS options to default values, as discussed on the package command doc page.

The "-sf kk" command-line switch will automatically append the "/kk" suffix to styles that support it. In this manner no modification to the input script is needed. Alternatively, one can run with the KOKKOS package by editing the input script as described below.

NOTE: When using a single OpenMP thread, the Kokkos Serial backend will give better performance than the OpenMP backend because some of the overhead to make the code thread-safe is removed.

NOTE: The default for the package kokkos command is to use "threaded" communication. However, when running on CPUs, it will typically be faster to use "classic" non-threaded communication. Use the "-pk kokkos" command-line switch to change the default package kokkos options. See its doc page for details and default settings. Experimenting with its options can provide a speed-up for specific calculations. For example:

mpirun -np 16 spa_kokkos_mpi_only -k on -sf kk -pk kokkos comm classic -in in.collide       # non-threaded comm 

For OpenMP, the KOKKOS package uses data duplication (i.e. thread-private arrays) by default to avoid thread-level write conflicts in some compute styles. Data duplication is typically fastest for small numbers of threads (i.e. 8 or less) but does increase memory footprint and is not scalable to large numbers of threads. An alternative to data duplication is to use thread-level atomics, which don't require duplication. When using the Kokkos Serial backend or the OpenMP backend with a single thread, no duplication or atomics are used. For CUDA, the KOKKOS package always uses atomics in these computes when necessary. The use of atomics instead of duplication can be forced by compiling with the "-DSPARTA_KOKKOS_USE_ATOMICS" compile switch.

Core and Thread Affinity:

When using multi-threading, it is important for performance to bind both MPI tasks to physical cores, and threads to physical cores, so they do not migrate during a simulation.

If you are not certain MPI tasks are being bound (check the defaults for your MPI installation), binding can be forced with these flags:

OpenMPI 1.8: mpirun -np 2 -bind-to socket -map-by socket ./spa_openmpi ...
Mvapich2 2.0: mpiexec -np 2 -bind-to socket -map-by socket ./spa_mvapich ... 

For binding threads with KOKKOS OpenMP, use thread affinity environment variables to force binding. With OpenMP 3.1 (gcc 4.7 or later, intel 12 or later) setting the environment variable OMP_PROC_BIND=true should be sufficient. In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads. For binding threads with the KOKKOS pthreads option, compile SPARTA the KOKKOS HWLOC=yes option as described below.

Running on GPUs:

Use the "-k" command-line switch to specify the number of GPUs per node, and the number of threads per MPI task. Typically the -np setting of the mpirun command should set the number of MPI tasks/node to be equal to the # of physical GPUs on the node. You can assign multiple MPI tasks to the same GPU with the KOKKOS package, but this is usually only faster if significant portions of the input script have not been ported to use Kokkos. Using CUDA MPS is recommended in this scenario. As above for multi-core CPUs (and no GPU), if N is the number of physical cores/node, then the number of MPI tasks/node should not exceed N.

-k on g Ng 

Here are examples of how to use the KOKKOS package for GPUs, assuming one or more nodes, each with two GPUs.

NVIDIA GPUs (CUDA):

mpirun -np 2 spa_kokkos_cuda -k on g 2 -sf kk -in in.collide          # 1 node,   2 MPI tasks/node, 2 GPUs/node
mpirun -np 32 -ppn 2 spa_kokkos_cuda -k on g 2 -sf kk -in in.collide  # 16 nodes, 2 MPI tasks/node, 2 GPUs/node (32 GPUs total) 

NOTE: Use the "-pk kokkos" command-line switch to change the default package kokkos options. See its doc page for details and default settings. For example:

mpirun -np 2 spa_kokkos_cuda -k on g 2 -sf kk -pk kokkos gpu/aware off -in in.collide      # set gpu/aware MPI support off 

AMD GPUs (HIP):

mpirun -np 2 spa_kokkos_hip -k on g 2 -sf kk -in in.collide          # 1 node,   2 MPI tasks/node, 2 GPUs/node
mpirun -np 16 -ppn 8 spa_kokkos_hip -k on g 8 -sf kk -in in.collide  # 2 nodes,  8 MPI tasks/node, 8 GPUs/node 

For Cray MPICH systems (e.g. El Capitan), set the GTL (GPU Transfer Library) environment variable before running:

export MPICH_GPU_SUPPORT_ENABLED=1
mpirun -np 8 spa_elcapitan_kokkos -k on g 8 -sf kk -in in.collide 

Intel GPUs (SYCL):

mpirun -np 4 spa_kokkos_sycl -k on g 4 -sf kk -in in.collide          # 1 node, 4 MPI tasks/node, 4 GPUs/node 

NOTE: Using OpenMP threading and CUDA/HIP/SYCL together is currently not possible with the SPARTA KOKKOS package.

NOTE: For good performance of the KOKKOS package on GPUs, you must have Kepler generation GPUs (or later). The Kokkos library exploits texture cache options not supported by Tesla generation GPUs (or older).

NOTE: When using a GPU, you will achieve the best performance if your input script does not use fix or compute styles which are not yet Kokkos-enabled. This allows data to stay on the GPU for multiple timesteps, without being copied back to the host CPU. Invoking a non-Kokkos fix or compute, or performing I/O for stat or dump output will cause data to be copied back to the CPU incurring a performance penalty.

NOTE: Most non-Kokkos styles degrade this way, costing performance but still producing correct results. A few, however, are rejected outright and will stop the run, rather than falling back to the host. Those restrictions are documented on the page for the style that imposes them.

NOTE: A run may define any number of instances of the styles that the device kernels capture -- surf_collide, surf_react, and the per-event tally computes such as compute boundary, compute surf, compute isurf/grid, compute react/boundary, compute react/surf and compute react/isurf/grid. Each is held in a runtime-sized device buffer, so there is no compile-time cap. The one remaining exception is that at most two compute surf instances may tally for a single fix emit/surf; exceeding that stops the run with an explanatory message.

NOTE: Building with -DSPARTA_KOKKOS_FIXED_LISTS instead holds those styles in fixed-size arrays inside the kernel functor, which restores the former caps: at most two instances of each surf_react style, at most four surf react models in total, at most two active instances of each per-event tally compute listed above, and at most four gas-phase tally computes. The two layouts produce identical results; which is faster is a hardware-dependent question, so both are kept so they can be compared by rebuilding.

Run with the KOKKOS package by editing an input script:

Alternatively the effect of the "-sf" or "-pk" switches can be duplicated by adding the package kokkos or suffix kk commands to your input script.

The discussion above for building SPARTA with the KOKKOS package, the mpirun/mpiexec command, and setting appropriate thread are the same.

You must still use the "-k on" command-line switch to enable the KOKKOS package, and specify its additional arguments for hardware options appropriate to your system, as documented above.

You can use the suffix kk command, or you can explicitly add a "kk" suffix to individual styles in your input script, e.g.

collide vss/kk air ar.vss 

You only need to use the package kokkos command if you wish to change any of its option defaults, as set by the "-k on" command-line switch.

Speed-ups to expect:

The performance of KOKKOS running in different modes is a function of your hardware, which KOKKOS-enable styles are used, and the problem size.

Generally speaking, when running on CPUs only, with a single thread per MPI task, the performance difference of a KOKKOS style and (un-accelerated) styles (MPI-only mode) is typically small (less than 20%).

See the Benchmark page of the SPARTA web site for performance of the KOKKOS package on different hardware.

Advanced Kokkos options:

There are other allowed options when building with the KOKKOS package. A few options are listed here; for a full list of all options, please refer to the Kokkos documentation. As above, these options can be set as variables on the command line or in a CMake presets file. For default CMake values, see cmake -LH | grep -i kokkos.

The CMake option Kokkos_ENABLE_OPTION sets the parallelization method used for Kokkos code (within SPARTA). For example, the CMake option Kokkos_ENABLE_SERIAL=ON means that no threading will be used. The CMake option Kokkos_ENABLE_OPENMP=ON means that OpenMP threading will be used. The CMake option Kokkos_ENABLE_CUDA=ON means an NVIDIA GPU running CUDA will be used.

As described above, the CMake option Kokkos_ARCH_TYPE=ON enables compiler switches needed when compiling for a specific hardware:

The CMake option Kokkos_ENABLE_CUDA_OPTION enables additional options for CUDA. For example, the CMake option Kokkos_ENABLE_CUDA_UVM=ON enables the use of CUDA "Unified Virtual Memory" (UVM) in Kokkos. UVM allows to one to use the host CPU memory to supplement the memory used on the GPU (with some performance penalty) and thus enables running larger problems that would otherwise not fit into the RAM on the GPU.

The CMake option Kokkos_ENABLE_DEBUG=ON is useful when developing a Kokkos-enabled style within SPARTA. This option enables printing of run-time debugging information that can be useful and also enables runtime bounds checking on Kokkos data structures, but may slow down performance.

Restrictions:

Currently, there are no precision options with the KOKKOS package. All compilation and computation is performed in double precision.