1c2b86a48SJunchao Zhang #include <petscsys.h> 2c2b86a48SJunchao Zhang #include <petsc/private/petscimpl.h> 3c2b86a48SJunchao Zhang #include <Kokkos_Core.hpp> 4c2b86a48SJunchao Zhang 5*45639126SStefano Zampini PetscBool PetscKokkosInitialized = PETSC_FALSE; 6*45639126SStefano Zampini 7c2b86a48SJunchao Zhang PetscErrorCode PetscKokkosFinalize_Private(void) 8c2b86a48SJunchao Zhang { 9c2b86a48SJunchao Zhang PetscFunctionBegin; 10c2b86a48SJunchao Zhang Kokkos::finalize(); 11c2b86a48SJunchao Zhang PetscFunctionReturn(0); 12c2b86a48SJunchao Zhang } 13c2b86a48SJunchao Zhang 14c2b86a48SJunchao Zhang PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 15c2b86a48SJunchao Zhang { 16c2b86a48SJunchao Zhang PetscFunctionBegin; 17c2b86a48SJunchao Zhang *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 18c2b86a48SJunchao Zhang PetscFunctionReturn(0); 19c2b86a48SJunchao Zhang } 20375e5adfSJunchao Zhang 2134766dafSJunchao Zhang /* Initialize Kokkos if not yet */ 22375e5adfSJunchao Zhang PetscErrorCode PetscKokkosInitializeCheck(void) 23375e5adfSJunchao Zhang { 2455c8bb78SStefano Zampini #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) 25375e5adfSJunchao Zhang PetscErrorCode ierr; 2655c8bb78SStefano Zampini #endif 2734766dafSJunchao Zhang Kokkos::InitArguments args; 2834766dafSJunchao Zhang int devId = -1; 29375e5adfSJunchao Zhang 30375e5adfSJunchao Zhang PetscFunctionBegin; 3134766dafSJunchao Zhang if (!Kokkos::is_initialized()) { 32375e5adfSJunchao Zhang #if defined(KOKKOS_ENABLE_CUDA) 33375e5adfSJunchao Zhang ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 3434766dafSJunchao Zhang cudaGetDevice(&devId); 3534766dafSJunchao Zhang #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 36375e5adfSJunchao Zhang ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 3734766dafSJunchao Zhang hipGetDevice(&devId); 38375e5adfSJunchao Zhang #endif 3934766dafSJunchao Zhang args.device_id = devId; 4034766dafSJunchao Zhang Kokkos::initialize(args); 4159e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 4259e55d94SJunchao Zhang } 43*45639126SStefano Zampini PetscKokkosInitialized = PETSC_TRUE; 44375e5adfSJunchao Zhang PetscFunctionReturn(0); 45375e5adfSJunchao Zhang } 46