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