Actual source code: kinit.kokkos.cxx

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

  6: PetscBool    PetscKokkosInitialized = PETSC_FALSE;
  7: PetscScalar *PetscScalarPool        = nullptr;
  8: PetscInt     PetscScalarPoolSize    = 0;

 10: Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr;

 12: PetscErrorCode PetscKokkosFinalize_Private(void)
 13: {
 14:   PetscFunctionBegin;
 15:   PetscCallCXX(delete PetscKokkosExecutionSpacePtr);
 16:   PetscCallCXX(Kokkos::kokkos_free(PetscScalarPool));
 17:   PetscScalarPoolSize = 0;
 18:   if (PetscBeganKokkos) {
 19:     PetscCallCXX(Kokkos::finalize());
 20:     PetscBeganKokkos = PETSC_FALSE;
 21:   }
 22:   PetscFunctionReturn(PETSC_SUCCESS);
 23: }

 25: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
 26: {
 27:   PetscFunctionBegin;
 28:   *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
 29:   PetscFunctionReturn(PETSC_SUCCESS);
 30: }

 32: /* Initialize Kokkos if not yet */
 33: PetscErrorCode PetscKokkosInitializeCheck(void)
 34: {
 35:   PetscFunctionBegin;
 36:   if (!Kokkos::is_initialized()) {
 37: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 38:     auto args = Kokkos::InitializationSettings();
 39: #else
 40:     auto args = Kokkos::InitArguments{}; /* use default constructor */
 41: #endif

 43: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
 44:     /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
 45:     PetscDevice device;
 46:     PetscInt    deviceId;
 47:     PetscCall(PetscDeviceCreate(PETSC_DEVICE_DEFAULT(), PETSC_DECIDE, &device));
 48:     PetscCall(PetscDeviceGetDeviceId(device, &deviceId));
 49:     PetscCall(PetscDeviceDestroy(&device));
 50:   #if PETSC_PKG_KOKKOS_VERSION_GE(4, 0, 0)
 51:     // if device_id is not set, and no gpus have been found, kokkos will use CPU
 52:     if (deviceId >= 0) args.set_device_id(static_cast<int>(deviceId));
 53:   #elif PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 54:     args.set_device_id(static_cast<int>(deviceId));
 55:   #else
 56:     PetscCall(PetscMPIIntCast(deviceId, &args.device_id));
 57:   #endif
 58: #endif

 60: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 61:     args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
 62: #else
 63:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
 64: #endif

 66:     /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
 67:        Otherwise, let's keep the default value (-1) of args.num_threads.
 68:     */
 69: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
 70:   #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 71:     args.set_num_threads(PetscNumOMPThreads);
 72:   #else
 73:     args.num_threads = PetscNumOMPThreads;
 74:   #endif
 75: #endif
 76:     PetscCallCXX(Kokkos::initialize(args));
 77:     PetscBeganKokkos = PETSC_TRUE;
 78:   }
 79:   if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr
 80: #if defined(PETSC_HAVE_CUDA)
 81:     extern cudaStream_t PetscDefaultCudaStream;
 82:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
 83: #elif defined(PETS_HAVE_HIP)
 84:     extern hipStream_t PetscDefaultHipStream;
 85:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
 86: #else
 87:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
 88: #endif
 89:   }
 90:   if (!PetscScalarPoolSize) { // A pool for a small count of PetscScalars
 91:     PetscScalarPoolSize = 1024;
 92:     PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_malloc(sizeof(PetscScalar) * PetscScalarPoolSize)));
 93:   }

 95:   PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called
 96:   PetscFunctionReturn(PETSC_SUCCESS);
 97: }