1 #include <petsc/private/deviceimpl.h> 2 #include <Kokkos_Core.hpp> 3 4 PetscBool PetscKokkosInitialized = PETSC_FALSE; 5 6 PetscErrorCode PetscKokkosFinalize_Private(void) 7 { 8 PetscFunctionBegin; 9 Kokkos::finalize(); 10 PetscFunctionReturn(0); 11 } 12 13 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 14 { 15 PetscFunctionBegin; 16 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 17 PetscFunctionReturn(0); 18 } 19 20 /* Initialize Kokkos if not yet */ 21 PetscErrorCode PetscKokkosInitializeCheck(void) 22 { 23 PetscFunctionBegin; 24 if (!Kokkos::is_initialized()) { 25 #if PETSC_PKG_KOKKOS_VERSION_GE(3,6,99) 26 auto args = Kokkos::InitializationSettings(); 27 #else 28 auto args = Kokkos::InitArguments{}; /* use default constructor */ 29 #endif 30 31 #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL)) 32 /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 33 PetscDeviceContext dctx; 34 35 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 36 #if PETSC_PKG_KOKKOS_VERSION_GE(3,6,99) 37 args.set_device_id(static_cast<int>(dctx->device->deviceId)); 38 #else 39 PetscCall(PetscMPIIntCast(dctx->device->deviceId,&args.device_id)); 40 #endif 41 #endif 42 43 #if PETSC_PKG_KOKKOS_VERSION_GE(3,6,99) 44 args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 45 #else 46 args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 47 #endif 48 49 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 50 Otherwise, let's keep the default value (-1) of args.num_threads. 51 */ 52 #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 53 #if PETSC_PKG_KOKKOS_VERSION_GE(3,6,99) 54 args.set_num_threads(PetscNumOMPThreads); 55 #else 56 args.num_threads = PetscNumOMPThreads; 57 #endif 58 #endif 59 60 Kokkos::initialize(args); 61 PetscBeganKokkos = PETSC_TRUE; 62 } 63 PetscKokkosInitialized = PETSC_TRUE; 64 PetscFunctionReturn(0); 65 } 66