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
PetscKokkosFinalize_Private(void)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
PetscKokkosIsInitialized_Private(PetscBool * isInitialized)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 */
PetscKokkosInitializeCheck(void)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