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 1934766dafSJunchao Zhang /* Initialize Kokkos if not yet */ 20375e5adfSJunchao Zhang PetscErrorCode PetscKokkosInitializeCheck(void) 21375e5adfSJunchao Zhang { 22*55c8bb78SStefano Zampini #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) 23375e5adfSJunchao Zhang PetscErrorCode ierr; 24*55c8bb78SStefano Zampini #endif 2534766dafSJunchao Zhang Kokkos::InitArguments args; 2634766dafSJunchao Zhang int devId = -1; 27375e5adfSJunchao Zhang 28375e5adfSJunchao Zhang PetscFunctionBegin; 2934766dafSJunchao Zhang if (!Kokkos::is_initialized()) { 30375e5adfSJunchao Zhang #if defined(KOKKOS_ENABLE_CUDA) 31375e5adfSJunchao Zhang ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 3234766dafSJunchao Zhang cudaGetDevice(&devId); 3334766dafSJunchao Zhang #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 34375e5adfSJunchao Zhang ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 3534766dafSJunchao Zhang hipGetDevice(&devId); 36375e5adfSJunchao Zhang #endif 3734766dafSJunchao Zhang args.device_id = devId; 3834766dafSJunchao Zhang Kokkos::initialize(args); 3959e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 4059e55d94SJunchao Zhang } 41375e5adfSJunchao Zhang PetscFunctionReturn(0); 42375e5adfSJunchao Zhang } 43