1 #include <petsc/private/deviceimpl.h> 2 #include <petsc/private/kokkosimpl.hpp> 3 #include <petscpkg_version.h> 4 #include <petsc_kokkos.hpp> 5 #include <petscdevice_cupm.h> 6 7 PetscBool PetscKokkosInitialized = PETSC_FALSE; // Has Kokkos been initialized (either by petsc or by users)? 8 PetscScalar *PetscScalarPool = nullptr; 9 PetscInt PetscScalarPoolSize = 0; 10 11 Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr; 12 13 PetscErrorCode PetscKokkosFinalize_Private(void) 14 { 15 PetscFunctionBegin; 16 PetscCallCXX(delete PetscKokkosExecutionSpacePtr); 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) && defined(PETSC_HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && defined(PETSC_HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && defined(PETSC_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 81 if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr 82 #if (defined(KOKKOS_ENABLE_CUDA) && defined(PETSC_HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && defined(PETSC_HAVE_HIP)) 83 PetscDeviceContext dctx; 84 PetscDeviceType dtype; 85 86 PetscDeviceContextGetCurrentContext(&dctx); // it internally sets PetscDefaultCuda/HipStream 87 PetscDeviceContextGetDeviceType(dctx, &dtype); 88 89 #if defined(PETSC_HAVE_CUDA) 90 if (dtype == PETSC_DEVICE_CUDA) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream)); 91 #elif defined(PETSC_HAVE_HIP) 92 if (dtype == PETSC_DEVICE_HIP) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream)); 93 #endif 94 #else 95 // In all other cases, we use Kokkos default 96 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace()); 97 #endif 98 } 99 100 if (!PetscScalarPoolSize) { // A pool for a small count of PetscScalars 101 PetscScalarPoolSize = 1024; 102 PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_malloc(sizeof(PetscScalar) * PetscScalarPoolSize))); 103 } 104 105 PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called 106 PetscFunctionReturn(PETSC_SUCCESS); 107 } 108