Actual source code: kinit.kokkos.cxx
1: #include <petsc/private/deviceimpl.h>
2: #include <petscpkg_version.h>
3: #include <Kokkos_Core.hpp>
5: PetscBool PetscKokkosInitialized = PETSC_FALSE;
7: PetscErrorCode PetscKokkosFinalize_Private(void)
8: {
9: PetscFunctionBegin;
10: Kokkos::finalize();
11: PetscFunctionReturn(PETSC_SUCCESS);
12: }
14: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
15: {
16: PetscFunctionBegin;
17: *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
18: PetscFunctionReturn(PETSC_SUCCESS);
19: }
21: /* Initialize Kokkos if not yet */
22: PetscErrorCode PetscKokkosInitializeCheck(void)
23: {
24: PetscFunctionBegin;
25: if (!Kokkos::is_initialized()) {
26: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99)
27: auto args = Kokkos::InitializationSettings();
28: #else
29: auto args = Kokkos::InitArguments{}; /* use default constructor */
30: #endif
32: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
33: /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
34: PetscDeviceContext dctx;
36: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
37: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99)
38: args.set_device_id(static_cast<int>(dctx->device->deviceId));
39: #else
40: PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id));
41: #endif
42: #endif
44: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99)
45: args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS));
46: #else
47: args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
48: #endif
50: /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
51: Otherwise, let's keep the default value (-1) of args.num_threads.
52: */
53: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
54: #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99)
55: args.set_num_threads(PetscNumOMPThreads);
56: #else
57: args.num_threads = PetscNumOMPThreads;
58: #endif
59: #endif
61: Kokkos::initialize(args);
62: PetscBeganKokkos = PETSC_TRUE;
63: }
64: PetscKokkosInitialized = PETSC_TRUE;
65: PetscFunctionReturn(PETSC_SUCCESS);
66: }