Line data    Source code

       1            1 : // Distributed under the MIT License.
       2              : // See LICENSE.txt for details.
       3              : 
       4              : /// \file
       5              : /// Macros for code that can run on GPUs. This header only includes the light
       6              : /// Kokkos headers that define them, so it can be included everywhere. Include
       7              : /// `Utilities/Kokkos/KokkosCore.hpp` instead to use Kokkos types and parallel
       8              : /// loops such as `Kokkos::View` and `Kokkos::parallel_for`.
       9              : 
      10              : #pragma once
      11              : 
      12              : // The build system defines `SPECTRE_KOKKOS` when configured with
      13              : // `-D SPECTRE_KOKKOS=ON`.
      14              : #ifdef SPECTRE_KOKKOS
      15              : // `KOKKOS_FUNCTION`, `KOKKOS_IF_ON_HOST`, `KOKKOS_IF_ON_DEVICE`, ...
      16              : #include <Kokkos_Macros.hpp>
      17              : // `Kokkos::abort`
      18              : #include <Kokkos_Abort.hpp>
      19              : 
      20              : /*!
      21              :  * \brief Defined while compiling device code for a GPU backend.
      22              :  *
      23              :  * GPU compilers compile each source file twice: once for the host and once for
      24              :  * the device. This macro is defined only in the device pass, so code that is
      25              :  * never called on the device can be excluded from device compilation. For
      26              :  * explicit instantiations use `SPECTRE_DEVICE_AND_HOST_TYPES` instead.
      27              :  *
      28              :  * Only use this for code at namespace scope. Inside functions use Kokkos'
      29              :  * `KOKKOS_IF_ON_HOST` and `KOKKOS_IF_ON_DEVICE` instead.
      30              :  */
      31              : #if defined(__CUDA_ARCH__) or defined(__HIP_DEVICE_COMPILE__) or \
      32              :     defined(__SYCL_DEVICE_ONLY__)
      33              : #define SPECTRE_KOKKOS_DEVICE_PASS 1
      34              : #endif
      35              : 
      36              : #else  // SPECTRE_KOKKOS
      37              : #define KOKKOS_FUNCTION
      38              : #define KOKKOS_INLINE_FUNCTION inline
      39              : // Without Kokkos all code runs on the host. The argument is wrapped in
      40              : // parentheses, e.g. `KOKKOS_IF_ON_HOST((io_call(...);))`, like in Kokkos.
      41              : #define SPECTRE_KOKKOS_STRIP_PARENS(...) __VA_ARGS__
      42              : #define KOKKOS_IF_ON_HOST(CODE) {SPECTRE_KOKKOS_STRIP_PARENS CODE}
      43              : #define KOKKOS_IF_ON_DEVICE(CODE) \
      44              :   {                               \
      45              :   }
      46              : #endif  // SPECTRE_KOKKOS
      47              : 
      48              : /*!
      49              :  * \brief Put before a `KOKKOS_FUNCTION` template whose instantiations for
      50              :  * some types can only run on the host.
      51              :  *
      52              :  * Examples are templates that call a callable passed by the caller (e.g. the
      53              :  * function of a root finder), or header-defined templates that are also
      54              :  * instantiated for `DataVector`. nvcc warns about these host-only
      55              :  * instantiations, even if they are only used on the host. This disables that
      56              :  * check for the function that follows. The compiler then also doesn't
      57              :  * diagnose device code that uses a host-only instantiation, so callables used
      58              :  * on the device must be `KOKKOS_FUNCTION`s or lambdas defined inside a kernel.
      59              :  */
      60              : #ifdef __NVCC__
      61              : #define SPECTRE_KOKKOS_DISABLE_EXEC_CHECK _Pragma("nv_exec_check_disable")
      62              : #else
      63            1 : #define SPECTRE_KOKKOS_DISABLE_EXEC_CHECK
      64              : #endif
      65              : 
      66              : /*!
      67              :  * \brief The types to instantiate a function for with `GENERATE_INSTANTIATIONS`
      68              :  * when only some of them are used on GPUs.
      69              :  *
      70              :  * Functions marked `KOKKOS_FUNCTION` can run on the CPU and, in GPU builds, on
      71              :  * a GPU. They are instantiated for the `DEVICE_TYPES` on the GPU and for the
      72              :  * `HOST_TYPES` on the CPU. List types that only exist on the CPU, such as
      73              :  * `DataVector`, only in `HOST_TYPES`. The `HOST_TYPES` should include the
      74              :  * `DEVICE_TYPES`, so the function can be called for them on the CPU as well.
      75              :  * For example,
      76              :  *
      77              :  * \code
      78              :  * GENERATE_INSTANTIATIONS(
      79              :  *     INSTANTIATE, (1, 2, 3),
      80              :  *     SPECTRE_DEVICE_AND_HOST_TYPES((double), (double, DataVector)))
      81              :  * \endcode
      82              :  *
      83              :  * instantiates the function for `double` on the GPU, and for `double` and
      84              :  * `DataVector` on the CPU. In builds without a GPU, this is simply the list of
      85              :  * `HOST_TYPES`.
      86              :  */
      87              : #ifdef SPECTRE_KOKKOS_DEVICE_PASS
      88              : #define SPECTRE_DEVICE_AND_HOST_TYPES(DEVICE_TYPES, HOST_TYPES) DEVICE_TYPES
      89              : #else
      90            1 : #define SPECTRE_DEVICE_AND_HOST_TYPES(DEVICE_TYPES, HOST_TYPES) HOST_TYPES
      91              : #endif