1 #include <petscsys.h> 2 #include <petsc/private/petscimpl.h> 3 #include <Kokkos_Core.hpp> 4 5 /* These wrappers are used as C bindings for the Kokkos routines */ 6 PetscErrorCode PetscKokkosInitialize_Private(void) 7 { 8 Kokkos::InitArguments args; 9 int devId = -1; 10 11 PetscFunctionBegin; 12 #if defined(KOKKOS_ENABLE_CUDA) 13 cudaGetDevice(&devId); 14 #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 15 hipGetDevice(&devId); 16 #endif 17 args.device_id = devId; 18 Kokkos::initialize(args); 19 PetscFunctionReturn(0); 20 } 21 22 PetscErrorCode PetscKokkosFinalize_Private(void) 23 { 24 PetscFunctionBegin; 25 Kokkos::finalize(); 26 PetscFunctionReturn(0); 27 } 28 29 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 30 { 31 PetscFunctionBegin; 32 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 33 PetscFunctionReturn(0); 34 } 35 36 /* Initialize the device lazily just before creating the first device object. */ 37 PetscErrorCode PetscKokkosInitializeCheck(void) 38 { 39 PetscErrorCode ierr; 40 41 PetscFunctionBegin; 42 #if defined(KOKKOS_ENABLE_CUDA) 43 ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 44 #elif defined(KOKKOS_ENABLE_HIP) 45 ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 46 #else 47 ierr = PetscKokkosInitialize_Private();CHKERRQ(ierr); 48 #endif 49 PetscFunctionReturn(0); 50 }