1 #include <petscsys.h> 2 #include <petsc/private/petscimpl.h> 3 #include <Kokkos_Core.hpp> 4 5 PetscErrorCode PetscKokkosFinalize_Private(void) 6 { 7 PetscFunctionBegin; 8 Kokkos::finalize(); 9 PetscFunctionReturn(0); 10 } 11 12 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 13 { 14 PetscFunctionBegin; 15 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 16 PetscFunctionReturn(0); 17 } 18 19 /* Initialize Kokkos if not yet */ 20 PetscErrorCode PetscKokkosInitializeCheck(void) 21 { 22 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) 23 PetscErrorCode ierr; 24 #endif 25 Kokkos::InitArguments args; 26 int devId = -1; 27 28 PetscFunctionBegin; 29 if (!Kokkos::is_initialized()) { 30 #if defined(KOKKOS_ENABLE_CUDA) 31 ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 32 cudaGetDevice(&devId); 33 #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 34 ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 35 hipGetDevice(&devId); 36 #endif 37 args.device_id = devId; 38 Kokkos::initialize(args); 39 PetscBeganKokkos = PETSC_TRUE; 40 } 41 PetscFunctionReturn(0); 42 } 43