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 { 9 PetscFunctionBegin; 10 Kokkos::finalize(); 11 PetscFunctionReturn(PETSC_SUCCESS); 12 } 13 14 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 15 { 16 PetscFunctionBegin; 17 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 18 PetscFunctionReturn(PETSC_SUCCESS); 19 } 20 21 /* Initialize Kokkos if not yet */ 22 PetscErrorCode PetscKokkosInitializeCheck(void) 23 { 24 PetscFunctionBegin; 25 if (!Kokkos::is_initialized()) { 26 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 27 auto args = Kokkos::InitializationSettings(); 28 #else 29 auto args = Kokkos::InitArguments{}; /* use default constructor */ 30 #endif 31 32 #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL)) 33 /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 34 PetscDevice device; 35 PetscInt deviceId; 36 PetscCall(PetscDeviceCreate(PETSC_DEVICE_DEFAULT(), PETSC_DECIDE, &device)); 37 PetscCall(PetscDeviceGetDeviceId(device, &deviceId)); 38 PetscCall(PetscDeviceDestroy(&device)); 39 #if PETSC_PKG_KOKKOS_VERSION_GE(4, 0, 0) 40 // if device_id is not set, and no gpus have been found, kokkos will use CPU 41 if (deviceId >= 0) args.set_device_id(static_cast<int>(deviceId)); 42 #elif PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 43 args.set_device_id(static_cast<int>(deviceId)); 44 #else 45 PetscCall(PetscMPIIntCast(deviceId, &args.device_id)); 46 #endif 47 #endif 48 49 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 50 args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 51 #else 52 args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 53 #endif 54 55 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 56 Otherwise, let's keep the default value (-1) of args.num_threads. 57 */ 58 #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 59 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 60 args.set_num_threads(PetscNumOMPThreads); 61 #else 62 args.num_threads = PetscNumOMPThreads; 63 #endif 64 #endif 65 PetscCallCXX(Kokkos::initialize(args)); 66 PetscBeganKokkos = PETSC_TRUE; 67 } 68 PetscKokkosInitialized = PETSC_TRUE; 69 PetscFunctionReturn(PETSC_SUCCESS); 70 } 71