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 PetscKokkosExecutionSpacePtr = nullptr; 18 PetscCallCXX(Kokkos::kokkos_free(PetscScalarPool)); 19 PetscScalarPoolSize = 0; 20 if (PetscBeganKokkos) { 21 PetscCallCXX(Kokkos::finalize()); 22 PetscBeganKokkos = PETSC_FALSE; 23 } 24 PetscFunctionReturn(PETSC_SUCCESS); 25 } 26 27 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 28 { 29 PetscFunctionBegin; 30 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 31 PetscFunctionReturn(PETSC_SUCCESS); 32 } 33 34 /* Initialize Kokkos if not yet */ 35 PetscErrorCode PetscKokkosInitializeCheck(void) 36 { 37 PetscFunctionBegin; 38 if (!Kokkos::is_initialized()) { 39 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 40 auto args = Kokkos::InitializationSettings(); 41 #else 42 auto args = Kokkos::InitArguments{}; /* use default constructor */ 43 #endif 44 45 #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)) 46 /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 47 PetscDevice device; 48 PetscInt deviceId; 49 PetscCall(PetscDeviceCreate(PETSC_DEVICE_DEFAULT(), PETSC_DECIDE, &device)); 50 PetscCall(PetscDeviceGetDeviceId(device, &deviceId)); 51 PetscCall(PetscDeviceDestroy(&device)); 52 #if PETSC_PKG_KOKKOS_VERSION_GE(4, 0, 0) 53 // if device_id is not set, and no gpus have been found, kokkos will use CPU 54 if (deviceId >= 0) args.set_device_id(static_cast<int>(deviceId)); 55 #elif PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 56 args.set_device_id(static_cast<int>(deviceId)); 57 #else 58 PetscCall(PetscMPIIntCast(deviceId, &args.device_id)); 59 #endif 60 #endif 61 62 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 63 args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 64 #else 65 args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 66 #endif 67 68 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 69 Otherwise, let's keep the default value (-1) of args.num_threads. 70 */ 71 #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 72 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 73 args.set_num_threads(PetscNumOMPThreads); 74 #else 75 args.num_threads = PetscNumOMPThreads; 76 #endif 77 #endif 78 PetscCallCXX(Kokkos::initialize(args)); 79 PetscBeganKokkos = PETSC_TRUE; 80 } 81 82 if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr 83 #if (defined(KOKKOS_ENABLE_CUDA) && defined(PETSC_HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && defined(PETSC_HAVE_HIP)) 84 PetscDeviceContext dctx; 85 PetscDeviceType dtype; 86 87 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); // it internally sets PetscDefaultCuda/HipStream 88 PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype)); 89 90 #if defined(PETSC_HAVE_CUDA) 91 if (dtype == PETSC_DEVICE_CUDA) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream)); 92 #elif defined(PETSC_HAVE_HIP) 93 if (dtype == PETSC_DEVICE_HIP) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream)); 94 #endif 95 #else 96 // In all other cases, we use Kokkos default 97 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace()); 98 #endif 99 } 100 101 if (!PetscScalarPoolSize) { // A pool for a small count of PetscScalars 102 PetscScalarPoolSize = 1024; 103 PetscCallCXX(PetscScalarPool = static_cast<PetscScalar *>(Kokkos::kokkos_malloc(sizeof(PetscScalar) * PetscScalarPoolSize))); 104 } 105 106 PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called 107 PetscFunctionReturn(PETSC_SUCCESS); 108 } 109