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