Actual source code: kinit.kokkos.cxx

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

  5: PetscBool PetscKokkosInitialized = PETSC_FALSE;

  7: Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr;

  9: PetscErrorCode PetscKokkosFinalize_Private(void)
 10: {
 11:   PetscFunctionBegin;
 12:   PetscCallCXX(delete PetscKokkosExecutionSpacePtr);
 13:   Kokkos::finalize();
 14:   PetscFunctionReturn(PETSC_SUCCESS);
 15: }

 17: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
 18: {
 19:   PetscFunctionBegin;
 20:   *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
 21:   PetscFunctionReturn(PETSC_SUCCESS);
 22: }

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

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

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

 47: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 48:     args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
 49: #else
 50:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
 51: #endif

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

 64:     Kokkos::initialize(args);
 65: #if defined(PETSC_HAVE_CUDA)
 66:     extern cudaStream_t PetscDefaultCudaStream;
 67:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
 68: #elif defined(PETS_HAVE_HIP)
 69:     extern hipStream_t PetscDefaultHipStream;
 70:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
 71: #else
 72:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
 73: #endif
 74:     PetscBeganKokkos = PETSC_TRUE;
 75:   }
 76:   PetscKokkosInitialized = PETSC_TRUE;
 77:   PetscFunctionReturn(PETSC_SUCCESS);
 78: }