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: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
 63:     args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
 64: #else
 65:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
 66: #endif

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

 82:   if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr
 83: #if (defined(KOKKOS_ENABLE_CUDA) && defined(PETSC_HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && defined(PETSC_HAVE_HIP))
 84:     PetscDeviceContext dctx;
 85:     PetscDeviceType    dtype;

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

 90:   #if defined(PETSC_HAVE_CUDA)
 91:     if (dtype == PETSC_DEVICE_CUDA) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
 92:   #elif defined(PETSC_HAVE_HIP)
 93:     if (dtype == PETSC_DEVICE_HIP) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
 94:   #endif
 95: #else
 96:     // In all other cases, we use Kokkos default
 97:     PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
 98: #endif
 99:   }

101:   if (!PetscScalarPoolSize) { // A pool for a small count of PetscScalars
102:     PetscScalarPoolSize = 1024;
103:     PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_malloc(sizeof(PetscScalar) * PetscScalarPoolSize)));
104:   }

106:   PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called
107:   PetscFunctionReturn(PETSC_SUCCESS);
108: }