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:     PetscDevice device;
 38:     PetscInt    deviceId;
 39:     PetscCall(PetscDeviceCreate(PETSC_DEVICE_DEFAULT(), PETSC_DECIDE, &device));
 40:     PetscCall(PetscDeviceGetDeviceId(device, &deviceId));
 41:     PetscCall(PetscDeviceDestroy(&device));
 42:   #if PETSC_PKG_KOKKOS_VERSION_GE(4, 0, 0)
 43:     // if device_id is not set, and no gpus have been found, kokkos will use CPU
 44:     if (deviceId >= 0) args.set_device_id(static_cast<int>(deviceId));
 45:   #elif PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 46:     args.set_device_id(static_cast<int>(deviceId));
 47:   #else
 48:     PetscCall(PetscMPIIntCast(deviceId, &args.device_id));
 49:   #endif
 50: #endif

 52: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 53:     args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
 54: #else
 55:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
 56: #endif

 58:     /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
 59:        Otherwise, let's keep the default value (-1) of args.num_threads.
 60:     */
 61: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
 62:   #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 63:     args.set_num_threads(PetscNumOMPThreads);
 64:   #else
 65:     args.num_threads = PetscNumOMPThreads;
 66:   #endif
 67: #endif
 68:     PetscCallCXX(Kokkos::initialize(args));
 69:     PetscBeganKokkos = PETSC_TRUE;
 70:   }
 71:   if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr
 72: #if defined(PETSC_HAVE_CUDA)
 73:     extern cudaStream_t PetscDefaultCudaStream;
 74:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
 75: #elif defined(PETS_HAVE_HIP)
 76:     extern hipStream_t PetscDefaultHipStream;
 77:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
 78: #else
 79:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
 80: #endif
 81:   }
 82:   PetscKokkosInitialized = PETSC_TRUE;
 83:   PetscFunctionReturn(PETSC_SUCCESS);
 84: }