1a4af0ceeSJacob Faibussowitsch #include <petsc/private/deviceimpl.h> 2*0e6b6b59SJacob Faibussowitsch #include <petscpkg_version.h> 3c2b86a48SJunchao Zhang #include <Kokkos_Core.hpp> 4c2b86a48SJunchao Zhang 545639126SStefano Zampini PetscBool PetscKokkosInitialized = PETSC_FALSE; 645639126SStefano Zampini 79371c9d4SSatish Balay PetscErrorCode PetscKokkosFinalize_Private(void) { 8c2b86a48SJunchao Zhang PetscFunctionBegin; 9c2b86a48SJunchao Zhang Kokkos::finalize(); 10c2b86a48SJunchao Zhang PetscFunctionReturn(0); 11c2b86a48SJunchao Zhang } 12c2b86a48SJunchao Zhang 139371c9d4SSatish Balay PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) { 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 */ 209371c9d4SSatish Balay PetscErrorCode PetscKokkosInitializeCheck(void) { 21375e5adfSJunchao Zhang PetscFunctionBegin; 2234766dafSJunchao Zhang if (!Kokkos::is_initialized()) { 23c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 24c66e0907SJunchao Zhang auto args = Kokkos::InitializationSettings(); 25c66e0907SJunchao Zhang #else 2662825ce1SJacob Faibussowitsch auto args = Kokkos::InitArguments{}; /* use default constructor */ 27c66e0907SJunchao Zhang #endif 28b84ac304SJunchao Zhang 2962825ce1SJacob Faibussowitsch #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL)) 3062825ce1SJacob Faibussowitsch /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 31a4af0ceeSJacob Faibussowitsch PetscDeviceContext dctx; 32b9d03b0cSStefano Zampini 339566063dSJacob Faibussowitsch PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 34c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 35c66e0907SJunchao Zhang args.set_device_id(static_cast<int>(dctx->device->deviceId)); 36c66e0907SJunchao Zhang #else 379566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id)); 38375e5adfSJunchao Zhang #endif 39c66e0907SJunchao Zhang #endif 4011f0be55SJunchao Zhang 41c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 42c66e0907SJunchao Zhang args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 43c66e0907SJunchao Zhang #else 4462825ce1SJacob Faibussowitsch args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 45c66e0907SJunchao Zhang #endif 4662825ce1SJacob Faibussowitsch 4711f0be55SJunchao Zhang /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 4811f0be55SJunchao Zhang Otherwise, let's keep the default value (-1) of args.num_threads. 4911f0be55SJunchao Zhang */ 5062825ce1SJacob Faibussowitsch #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 51c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 52c66e0907SJunchao Zhang args.set_num_threads(PetscNumOMPThreads); 53c66e0907SJunchao Zhang #else 5411f0be55SJunchao Zhang args.num_threads = PetscNumOMPThreads; 5511f0be55SJunchao Zhang #endif 56c66e0907SJunchao Zhang #endif 5711f0be55SJunchao Zhang 5834766dafSJunchao Zhang Kokkos::initialize(args); 5959e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 6059e55d94SJunchao Zhang } 6145639126SStefano Zampini PetscKokkosInitialized = PETSC_TRUE; 62375e5adfSJunchao Zhang PetscFunctionReturn(0); 63375e5adfSJunchao Zhang } 64