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