Date: 2022-03-31T04:46:08+00:00

Actual source code: kinit.kokkos.cxx

      1: #include <petsc/private/deviceimpl.h>
      2: #include <Kokkos_Core.hpp>

      4: PetscBool PetscKokkosInitialized = PETSC_FALSE;

      6: PetscErrorCode PetscKokkosFinalize_Private(void)
      7: {
      8:   Kokkos::finalize();
      9:   return 0;
     10: }

     12: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
     13: {
     14:   *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
     15:   return 0;
     16: }

     18: /* Initialize Kokkos if not yet */
     19: PetscErrorCode PetscKokkosInitializeCheck(void)
     20: {
     21:   if (!Kokkos::is_initialized()) {
     22:     auto args = Kokkos::InitArguments{}; /* use default constructor */

     24: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
     25:     /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
     26:     PetscDeviceContext dctx;

     28:     PetscDeviceContextGetCurrentContext(&dctx);
     29:     PetscMPIIntCast(dctx->device->deviceId,&args.device_id);
     30: #endif

     32:     args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);

     34:     /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
     35:        Otherwise, let's keep the default value (-1) of args.num_threads.
     36:     */
     37: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
     38:     args.num_threads = PetscNumOMPThreads;
     39: #endif

     41:     Kokkos::initialize(args);
     42:     PetscBeganKokkos = PETSC_TRUE;
     43:   }
     44:   PetscKokkosInitialized = PETSC_TRUE;
     45:   return 0;
     46: }