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 /* To use PetscNumOMPThreads, one has to configure PETSc --with-openmp. 63 Otherwise, let's keep the default value (-1) of args.num_threads. 64 */ 65 #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 66 #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0) 67 args.set_num_threads(PetscNumOMPThreads); 68 #else 69 args.num_threads = PetscNumOMPThreads; 70 #endif 71 #endif 72 PetscCallCXX(Kokkos::initialize(args)); 73 PetscBeganKokkos = PETSC_TRUE; 74 } 75 76 if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by PETSc or by user, we need to init PetscKokkosExecutionSpacePtr 77 #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)) 78 PetscDeviceContext dctx; 79 PetscDeviceType dtype; 80 81 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); // it internally sets PetscDefaultCuda/HipStream 82 PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype)); 83 84 #if defined(PETSC_HAVE_CUDA) 85 if (dtype == PETSC_DEVICE_CUDA) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream)); 86 #elif defined(PETSC_HAVE_HIP) 87 if (dtype == PETSC_DEVICE_HIP) PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream)); 88 #elif defined(PETSC_HAVE_SYCL) 89 if (dtype == PETSC_DEVICE_SYCL) { 90 void *handle; 91 PetscCall(PetscDeviceContextGetStreamHandle(dctx, &handle)); // Kind of PetscDefaultSyclStream 92 PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(*(sycl::queue *)handle)); 93 } 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