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