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 PetscErrorCode ierr; 23 Kokkos::InitArguments args; 24 int devId = -1; 25 26 PetscFunctionBegin; 27 if (!Kokkos::is_initialized()) { 28 #if defined(KOKKOS_ENABLE_CUDA) 29 ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 30 cudaGetDevice(&devId); 31 #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 32 ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 33 hipGetDevice(&devId); 34 #endif 35 args.device_id = devId; 36 Kokkos::initialize(args); 37 PetscBeganKokkos = PETSC_TRUE; 38 } 39 PetscFunctionReturn(0); 40 }