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: PetscDeviceContext dctx;
39: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
40: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
41: args.set_device_id(static_cast<int>(dctx->device->deviceId));
42: #else
43: PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id));
44: #endif
45: #endif
47: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
48: args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
49: #else
50: args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
51: #endif
53: /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
54: Otherwise, let's keep the default value (-1) of args.num_threads.
55: */
56: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
57: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 7, 0)
58: args.set_num_threads(PetscNumOMPThreads);
59: #else
60: args.num_threads = PetscNumOMPThreads;
61: #endif
62: #endif
64: Kokkos::initialize(args);
65: #if defined(PETSC_HAVE_CUDA)
66: extern cudaStream_t PetscDefaultCudaStream;
67: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultCudaStream));
68: #elif defined(PETS_HAVE_HIP)
69: extern hipStream_t PetscDefaultHipStream;
70: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace(PetscDefaultHipStream));
71: #else
72: PetscCallCXX(PetscKokkosExecutionSpacePtr = new Kokkos::DefaultExecutionSpace());
73: #endif
74: PetscBeganKokkos = PETSC_TRUE;
75: }
76: PetscKokkosInitialized = PETSC_TRUE;
77: PetscFunctionReturn(PETSC_SUCCESS);
78: }