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 PetscDeviceContext dctx; 38 39 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 40 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 41 args.set_device_id(static_cast<int>(dctx->device->deviceId)); 42 #else 43 PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id)); 44 #endif 45 #endif 46 47 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 48 args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 49 #else 50 args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 51 #endif 52 53 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 54 Otherwise, let's keep the default value (-1) of args.num_threads. 55 */ 56 #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 57 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 58 args.set_num_threads(PetscNumOMPThreads); 59 #else 60 args.num_threads = PetscNumOMPThreads; 61 #endif 62 #endif 63 64 Kokkos::initialize(args); 65 #if defined(PETSC_HAVE_CUDA) 66 extern cudaStream_t PetscDefaultCudaStream; 67 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream)); 68 #elif defined(PETS_HAVE_HIP) 69 extern hipStream_t PetscDefaultHipStream; 70 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream)); 71 #else 72 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace()); 73 #endif 74 PetscBeganKokkos = PETSC_TRUE; 75 } 76 PetscKokkosInitialized = PETSC_TRUE; 77 PetscFunctionReturn(PETSC_SUCCESS); 78 } 79