1 #include <petsc/private/deviceimpl.h> 2 #include <Kokkos_Core.hpp> 3 4 PetscBool PetscKokkosInitialized = PETSC_FALSE; 5 6 PetscErrorCode PetscKokkosFinalize_Private(void) 7 { 8 PetscFunctionBegin; 9 Kokkos::finalize(); 10 PetscFunctionReturn(0); 11 } 12 13 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 14 { 15 PetscFunctionBegin; 16 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 17 PetscFunctionReturn(0); 18 } 19 20 #define PETSC_AND_KOKKOS_HAVE(CUPM) (defined(KOKKOS_ENABLE_##CUPM) && PetscDefined(HAVE_##CUPM)) 21 22 /* Initialize Kokkos if not yet */ 23 PetscErrorCode PetscKokkosInitializeCheck(void) 24 { 25 Kokkos::InitArguments args; 26 27 PetscFunctionBegin; 28 if (!Kokkos::is_initialized()) { 29 args.num_threads = -1; /* Kokkos default value of each parameter is -1 */ 30 args.num_numa = -1; 31 args.device_id = -1; 32 args.ndevices = -1; 33 args.skip_device = -1; 34 35 #if defined(PETSC_HAVE_KOKKOS_INIT_WARNINGS) 36 args.disable_warnings = false; 37 #else 38 args.disable_warnings = true; 39 #endif 40 41 #if PETSC_AND_KOKKOS_HAVE(CUDA) || PETSC_AND_KOKKOS_HAVE(HIP) || PETSC_AND_KOKKOS_HAVE(SYCL) 42 PetscDeviceContext dctx; 43 PetscErrorCode ierr; 44 45 ierr = PetscDeviceContextGetCurrentContext(&dctx);CHKERRQ(ierr); 46 ierr = PetscMPIIntCast(dctx->device->deviceId,&args.device_id);CHKERRQ(ierr); 47 #endif 48 49 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 50 Otherwise, let's keep the default value (-1) of args.num_threads. 51 */ 52 #if defined(KOKKOS_ENABLE_OPENMP) && defined(PETSC_HAVE_OPENMP) 53 args.num_threads = PetscNumOMPThreads; 54 #endif 55 56 Kokkos::initialize(args); 57 PetscBeganKokkos = PETSC_TRUE; 58 } 59 PetscKokkosInitialized = PETSC_TRUE; 60 PetscFunctionReturn(0); 61 } 62