Actual source code: kinit.kokkos.cxx

  1: #include <petsc/private/deviceimpl.h>
  2: #include <petscpkg_version.h>
  3: #include <Kokkos_Core.hpp>

  5: PetscBool PetscKokkosInitialized = PETSC_FALSE;

  7: PetscErrorCode PetscKokkosFinalize_Private(void)
  8: {
  9:   PetscFunctionBegin;
 10:   Kokkos::finalize();
 11:   PetscFunctionReturn(PETSC_SUCCESS);
 12: }

 14: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
 15: {
 16:   PetscFunctionBegin;
 17:   *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
 18:   PetscFunctionReturn(PETSC_SUCCESS);
 19: }

 21: /* Initialize Kokkos if not yet */
 22: PetscErrorCode PetscKokkosInitializeCheck(void)
 23: {
 24:   PetscFunctionBegin;
 25:   if (!Kokkos::is_initialized()) {
 26: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 27:     auto args = Kokkos::InitializationSettings();
 28: #else
 29:     auto args             = Kokkos::InitArguments{}; /* use default constructor */
 30: #endif

 32: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
 33:     /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
 34:     PetscDeviceContext dctx;

 36:     PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
 37:   #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 38:     args.set_device_id(static_cast<int>(dctx->device->deviceId));
 39:   #else
 40:     PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id));
 41:   #endif
 42: #endif

 44: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 45:     args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
 46: #else
 47:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
 48: #endif

 50:     /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
 51:        Otherwise, let's keep the default value (-1) of args.num_threads.
 52:     */
 53: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
 54:   #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 55:     args.set_num_threads(PetscNumOMPThreads);
 56:   #else
 57:     args.num_threads = PetscNumOMPThreads;
 58:   #endif
 59: #endif

 61:     Kokkos::initialize(args);
 62:     PetscBeganKokkos = PETSC_TRUE;
 63:   }
 64:   PetscKokkosInitialized = PETSC_TRUE;
 65:   PetscFunctionReturn(PETSC_SUCCESS);
 66: }