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) 42 /* Kokkos does not support CUDA and HIP at the same time (but we do :)) */ 43 PetscDeviceContext dctx; 44 PetscErrorCode ierr; 45 46 ierr = PetscDeviceContextGetCurrentContext(&dctx);CHKERRQ(ierr); 47 ierr = PetscMPIIntCast(dctx->device->deviceId,&args.device_id);CHKERRQ(ierr); 48 #endif 49 50 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 51 Otherwise, let's keep the default value (-1) of args.num_threads. 52 */ 53 #if defined(KOKKOS_ENABLE_OPENMP) && defined(PETSC_HAVE_OPENMP) 54 args.num_threads = PetscNumOMPThreads; 55 #endif 56 57 Kokkos::initialize(args); 58 PetscBeganKokkos = PETSC_TRUE; 59 } 60 PetscKokkosInitialized = PETSC_TRUE; 61 PetscFunctionReturn(0); 62 } 63