1c2b86a48SJunchao Zhang #include <petscsys.h> 2c2b86a48SJunchao Zhang #include <petsc/private/petscimpl.h> 3c2b86a48SJunchao Zhang #include <Kokkos_Core.hpp> 4c2b86a48SJunchao Zhang 5c2b86a48SJunchao Zhang PetscErrorCode PetscKokkosFinalize_Private(void) 6c2b86a48SJunchao Zhang { 7c2b86a48SJunchao Zhang PetscFunctionBegin; 8c2b86a48SJunchao Zhang Kokkos::finalize(); 9c2b86a48SJunchao Zhang PetscFunctionReturn(0); 10c2b86a48SJunchao Zhang } 11c2b86a48SJunchao Zhang 12c2b86a48SJunchao Zhang PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 13c2b86a48SJunchao Zhang { 14c2b86a48SJunchao Zhang PetscFunctionBegin; 15c2b86a48SJunchao Zhang *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 16c2b86a48SJunchao Zhang PetscFunctionReturn(0); 17c2b86a48SJunchao Zhang } 18375e5adfSJunchao Zhang 19*34766dafSJunchao Zhang /* Initialize Kokkos if not yet */ 20375e5adfSJunchao Zhang PetscErrorCode PetscKokkosInitializeCheck(void) 21375e5adfSJunchao Zhang { 22375e5adfSJunchao Zhang PetscErrorCode ierr; 23*34766dafSJunchao Zhang Kokkos::InitArguments args; 24*34766dafSJunchao Zhang int devId = -1; 25375e5adfSJunchao Zhang 26375e5adfSJunchao Zhang PetscFunctionBegin; 27*34766dafSJunchao Zhang if (!Kokkos::is_initialized()) { 28375e5adfSJunchao Zhang #if defined(KOKKOS_ENABLE_CUDA) 29375e5adfSJunchao Zhang ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 30*34766dafSJunchao Zhang cudaGetDevice(&devId); 31*34766dafSJunchao Zhang #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 32375e5adfSJunchao Zhang ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 33*34766dafSJunchao Zhang hipGetDevice(&devId); 34375e5adfSJunchao Zhang #endif 35*34766dafSJunchao Zhang args.device_id = devId; 36*34766dafSJunchao Zhang Kokkos::initialize(args); 3759e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 3859e55d94SJunchao Zhang } 39375e5adfSJunchao Zhang PetscFunctionReturn(0); 40375e5adfSJunchao Zhang }