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>
  5: #include <petscdevice_cupm.h>

  7: PetscBool    PetscKokkosInitialized = PETSC_FALSE; // Has Kokkos been initialized (either by PETSc or by users)?
  8: PetscScalar *PetscScalarPool        = nullptr;
  9: PetscInt     PetscScalarPoolSize    = 0;

 11: Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr;

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

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

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

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

 62:     /* To use PetscNumOMPThreads, one has to configure PETSc --with-openmp.
 63:        Otherwise, let's keep the default value (-1) of args.num_threads.
 64:     */
 65: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
 66:   #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 67:     args.set_num_threads(PetscNumOMPThreads);
 68:   #else
 69:     args.num_threads = PetscNumOMPThreads;
 70:   #endif
 71: #endif
 72:     PetscCallCXX(Kokkos::initialize(args));
 73:     PetscBeganKokkos = PETSC_TRUE;
 74:   }

 76:   if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by PETSc or by user, we need to init PetscKokkosExecutionSpacePtr
 77: #if (defined(KOKKOS_ENABLE_CUDA) && defined(PETSC_HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && defined(PETSC_HAVE_HIP))
 78:     PetscDeviceContext dctx;
 79:     PetscDeviceType    dtype;

 81:     PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); // it internally sets PetscDefaultCuda/HipStream
 82:     PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype));

 84:   #if defined(PETSC_HAVE_CUDA)
 85:     if (dtype == PETSC_DEVICE_CUDA) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
 86:   #elif defined(PETSC_HAVE_HIP)
 87:     if (dtype == PETSC_DEVICE_HIP) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
 88:   #endif
 89: #else
 90:     // In all other cases, we use Kokkos default
 91:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
 92: #endif
 93:   }

 95:   if (!PetscScalarPoolSize) { // A pool for a small count of PetscScalars
 96:     PetscScalarPoolSize = 1024;
 97:     PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_malloc(sizeof(PetscScalar) * PetscScalarPoolSize)));
 98:   }

100:   PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called
101:   PetscFunctionReturn(PETSC_SUCCESS);
102: }