1 #include <petscdevice.h> 2 #include <petsc/private/petscimpl.h> 3 #include <Kokkos_Core.hpp> 4 5 PetscBool PetscKokkosInitialized = PETSC_FALSE; 6 7 PetscErrorCode PetscKokkosFinalize_Private(void) 8 { 9 PetscFunctionBegin; 10 Kokkos::finalize(); 11 PetscFunctionReturn(0); 12 } 13 14 PetscErrorCode PetscKokkosIsInitialized_Private(PetscBool *isInitialized) 15 { 16 PetscFunctionBegin; 17 *isInitialized = Kokkos::is_initialized() ? PETSC_TRUE : PETSC_FALSE; 18 PetscFunctionReturn(0); 19 } 20 21 /* Initialize Kokkos if not yet */ 22 PetscErrorCode PetscKokkosInitializeCheck(void) 23 { 24 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) 25 PetscErrorCode ierr; 26 #endif 27 Kokkos::InitArguments args; 28 int devId = -1; 29 30 PetscFunctionBegin; 31 if (!Kokkos::is_initialized()) { 32 args.num_threads = -1; /* Kokkos default value of each parameter is -1 */ 33 args.num_numa = -1; 34 args.device_id = -1; 35 args.ndevices = -1; 36 args.skip_device = -1; 37 38 #if defined(PETSC_HAVE_KOKKOS_INIT_WARNINGS) 39 args.disable_warnings = false; 40 #else 41 args.disable_warnings = true; 42 #endif 43 44 #if defined(KOKKOS_ENABLE_CUDA) 45 cudaError_t cerr; 46 47 ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); 48 cerr = cudaGetDevice(&devId);CHKERRCUDA(cerr); 49 #elif defined(KOKKOS_ENABLE_HIP) /* Kokkos does not support CUDA and HIP at the same time */ 50 hipError_t herr; 51 52 ierr = PetscHIPInitializeCheck();CHKERRQ(ierr); 53 herr = hipGetDevice(&devId);CHKERRHIP(herr); 54 #endif 55 56 /* To use PetscNumOMPThreads, one has to configure petsc --with-openmp. 57 Otherwise, let's keep the default value (-1) of args.num_threads. 58 */ 59 #if defined(KOKKOS_ENABLE_OPENMP) && defined(PETSC_HAVE_OPENMP) 60 args.num_threads = PetscNumOMPThreads; 61 #endif 62 63 args.device_id = devId; 64 Kokkos::initialize(args); 65 PetscBeganKokkos = PETSC_TRUE; 66 } 67 PetscKokkosInitialized = PETSC_TRUE; 68 PetscFunctionReturn(0); 69 } 70