1a4af0ceeSJacob Faibussowitsch #include <petsc/private/deviceimpl.h> 2c2b86a48SJunchao Zhang #include <Kokkos_Core.hpp> 3c2b86a48SJunchao Zhang 445639126SStefano Zampini PetscBool PetscKokkosInitialized = PETSC_FALSE; 545639126SStefano Zampini 6*9371c9d4SSatish Balay PetscErrorCode PetscKokkosFinalize_Private(void) { 7c2b86a48SJunchao Zhang PetscFunctionBegin; 8c2b86a48SJunchao Zhang Kokkos::finalize(); 9c2b86a48SJunchao Zhang PetscFunctionReturn(0); 10c2b86a48SJunchao Zhang } 11c2b86a48SJunchao Zhang 12*9371c9d4SSatish Balay PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) { 13c2b86a48SJunchao Zhang PetscFunctionBegin; 14c2b86a48SJunchao Zhang *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 15c2b86a48SJunchao Zhang PetscFunctionReturn(0); 16c2b86a48SJunchao Zhang } 17375e5adfSJunchao Zhang 1834766dafSJunchao Zhang /* Initialize Kokkos if not yet */ 19*9371c9d4SSatish Balay PetscErrorCode PetscKokkosInitializeCheck(void) { 20375e5adfSJunchao Zhang PetscFunctionBegin; 2134766dafSJunchao Zhang if (!Kokkos::is_initialized()) { 22c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 23c66e0907SJunchao Zhang auto args = Kokkos::InitializationSettings(); 24c66e0907SJunchao Zhang #else 2562825ce1SJacob Faibussowitsch auto args = Kokkos::InitArguments{}; /* use default constructor */ 26c66e0907SJunchao Zhang #endif 27b84ac304SJunchao Zhang 2862825ce1SJacob Faibussowitsch #if (defined(KOKKOS_ENABLE_CUDA) && PetscDefined(HAVE_CUDA)) || (defined(KOKKOS_ENABLE_HIP) && PetscDefined(HAVE_HIP)) || (defined(KOKKOS_ENABLE_SYCL) && PetscDefined(HAVE_SYCL)) 2962825ce1SJacob Faibussowitsch /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 30a4af0ceeSJacob Faibussowitsch PetscDeviceContext dctx; 31b9d03b0cSStefano Zampini 329566063dSJacob Faibussowitsch PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 33c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 34c66e0907SJunchao Zhang args.set_device_id(static_cast<int>(dctx->device->deviceId)); 35c66e0907SJunchao Zhang #else 369566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(dctx->device->deviceId, &args.device_id)); 37375e5adfSJunchao Zhang #endif 38c66e0907SJunchao Zhang #endif 3911f0be55SJunchao Zhang 40c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 41c66e0907SJunchao Zhang args.set_disable_warnings(!PetscDefined(HAVE_KOKKOS_INIT_WARNINGS)); 42c66e0907SJunchao Zhang #else 4362825ce1SJacob Faibussowitsch args.disable_warnings = !PetscDefined(HAVE_KOKKOS_INIT_WARNINGS); 44c66e0907SJunchao Zhang #endif 4562825ce1SJacob Faibussowitsch 4611f0be55SJunchao Zhang /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 4711f0be55SJunchao Zhang Otherwise, let's keep the default value (-1) of args.num_threads. 4811f0be55SJunchao Zhang */ 4962825ce1SJacob Faibussowitsch #if defined(KOKKOS_ENABLE_OPENMP) && PetscDefined(HAVE_OPENMP) 50c66e0907SJunchao Zhang #if PETSC_PKG_KOKKOS_VERSION_GE(3, 6, 99) 51c66e0907SJunchao Zhang args.set_num_threads(PetscNumOMPThreads); 52c66e0907SJunchao Zhang #else 5311f0be55SJunchao Zhang args.num_threads = PetscNumOMPThreads; 5411f0be55SJunchao Zhang #endif 55c66e0907SJunchao Zhang #endif 5611f0be55SJunchao Zhang 5734766dafSJunchao Zhang Kokkos::initialize(args); 5859e55d94SJunchao Zhang PetscBeganKokkos = PETSC_TRUE; 5959e55d94SJunchao Zhang } 6045639126SStefano Zampini PetscKokkosInitialized = PETSC_TRUE; 61375e5adfSJunchao Zhang PetscFunctionReturn(0); 62375e5adfSJunchao Zhang } 63