Actual source code: kinit.kokkos.cxx
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>
7: PetscBool PetscKokkosInitialized = PETSC_FALSE; // Has Kokkos been initialized (either by petsc or by users)?
8: PetscScalar *PetscScalarPool = nullptr;
9: PetscInt PetscScalarPoolSize = 0;
11: Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr;
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: }
27: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
28: {
29: PetscFunctionBegin;
30: *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
31: PetscFunctionReturn(PETSC_SUCCESS);
32: }
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
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
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
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: }
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;
87: PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); // it internally sets PetscDefaultCuda/HipStream
88: PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype));
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: }
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: }
106: PetscKokkosInitialized = PETSC_TRUE; // PetscKokkosInitializeCheck() was called
107: PetscFunctionReturn(PETSC_SUCCESS);
108: }