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 } 35375e5adfSJunchao Zhang 36375e5adfSJunchao Zhang /* Initialize the device lazily just before creating the first device object. */ 37375e5adfSJunchao Zhang PetscErrorCode PetscKokkosInitializeCheck(void) 38375e5adfSJunchao Zhang { 39375e5adfSJunchao Zhang PetscErrorCode ierr; 40375e5adfSJunchao Zhang 41375e5adfSJunchao Zhang PetscFunctionBegin; 42375e5adfSJunchao Zhang #if defined(KOKKOS_ENABLE_CUDA) 43375e5adfSJunchao Zhang ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 44375e5adfSJunchao Zhang #elif defined(KOKKOS_ENABLE_HIP) 45375e5adfSJunchao Zhang ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 46375e5adfSJunchao Zhang #endif 47*59e55d94SJunchao Zhang if (!Kokkos::is_initialized()) { 48*59e55d94SJunchao Zhang ierr = PetscKokkosInitialize_Private();CHKERRQ(ierr); 49*59e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 50*59e55d94SJunchao Zhang } 51375e5adfSJunchao Zhang PetscFunctionReturn(0); 52375e5adfSJunchao Zhang }