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: }