1: #include <petsc/private/deviceimpl.h>
2: #include <Kokkos_Core.hpp>
4: PetscBool PetscKokkosInitialized = PETSC_FALSE;
6: PetscErrorCode PetscKokkosFinalize_Private(void)
7: {
8: Kokkos::finalize();
9: return 0;
10: }
12: PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized)
13: {
14: *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE;
15: return 0;
16: }
18: /* Initialize Kokkos if not yet */
19: PetscErrorCode PetscKokkosInitializeCheck(void)
20: {
21: if (!Kokkos::is_initialized()) {
22: auto args = Kokkos::InitArguments{}; /* use default constructor */
24: #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL))
25: /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */
26: PetscDeviceContext dctx;
28: PetscDeviceContextGetCurrentContext(&dctx);
29: PetscMPIIntCast(dctx->device->deviceId,&args.device_id);
30: #endif
32: args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS);
34: /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp.
35: Otherwise, let's keep the default value (-1) of args.num_threads.
36: */
37: #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP)
38: args.num_threads = PetscNumOMPThreads;
39: #endif
41: Kokkos::initialize(args);
42: PetscBeganKokkos = PETSC_TRUE;
43: }
44: PetscKokkosInitialized = PETSC_TRUE;
45: return 0;
46: }