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