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