Actual source code: kinit.kokkos.cxx
1: #include <petsc/private/deviceimpl.h>
2: #include <petscpkg_version.h>
3: #include <petsc_kokkos.hpp>
5: PetscBool PetscKokkosInitialized = PETSC_FALSE;
7: Kokkos::DefaultExecutionSpace *PetscKokkosExecutionSpacePtr = nullptr;
9: PetscErrorCode PetscKokkosFinalize_Private(void)
10: {
11: PetscFunctionBegin;
12: PetscCallCXX(delete PetscKokkosExecutionSpacePtr);
13: Kokkos::finalize();
14: PetscFunctionReturn(PETSC_SUCCESS);
15: }
17: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
18: {
19: PetscFunctionBegin;
20: *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
21: PetscFunctionReturn(PETSC_SUCCESS);
22: }
24: /* Initialize Kokkos if not yet */
25: PetscErrorCode PetscKokkosInitializeCheck(void)
26: {
27: PetscFunctionBegin;
28: if (!Kokkos::is_initialized()) {
29: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
30: auto args = Kokkos::InitializationSettings();
31: #else
32: auto args = Kokkos::InitArguments{}; /* use default constructor */
33: #endif
35: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
36: /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
37: PetscDevice device;
38: PetscInt deviceId;
39: PetscCall(PetscDeviceCreate(PETSC_DEVICE_DEFAULT(), PETSC_DECIDE, &device));
40: PetscCall(PetscDeviceGetDeviceId(device, &deviceId));
41: PetscCall(PetscDeviceDestroy(&device));
42: #if PETSC_PKG_KOKKOS_VERSION_GE(4, 0, 0)
43: // if device_id is not set, and no gpus have been found, kokkos will use CPU
44: if (deviceId >= 0) args.set_device_id(static_cast<int>(deviceId));
45: #elif PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
46: args.set_device_id(static_cast<int>(deviceId));
47: #else
48: PetscCall(PetscMPIIntCast(deviceId, &args.device_id));
49: #endif
50: #endif
52: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
53: args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
54: #else
55: args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
56: #endif
58: /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
59: Otherwise, let's keep the default value (-1) of args.num_threads.
60: */
61: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
62: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
63: args.set_num_threads(PetscNumOMPThreads);
64: #else
65: args.num_threads = PetscNumOMPThreads;
66: #endif
67: #endif
68: PetscCallCXX(Kokkos::initialize(args));
69: PetscBeganKokkos = PETSC_TRUE;
70: }
71: if (!PetscKokkosExecutionSpacePtr) { // No matter Kokkos is init'ed by petsc or by user, we need to init PetscKokkosExecutionSpacePtr
72: #if defined(PETSC_HAVE_CUDA)
73: extern cudaStream_t PetscDefaultCudaStream;
74: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
75: #elif defined(PETS_HAVE_HIP)
76: extern hipStream_t PetscDefaultHipStream;
77: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
78: #else
79: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
80: #endif
81: }
82: PetscKokkosInitialized = PETSC_TRUE;
83: PetscFunctionReturn(PETSC_SUCCESS);
84: }