1 #pragma once 2 3 #include <petsclog.h> // PetscLogGpuTimeBegin()/End() 4 #include <petscsys.h> // SETERRQ() 5 #include <petscdevice_cupm.h> // PETSC_USING_NVCC 6 7 #include <thrust/version.h> // THRUST_VERSION 8 #include <thrust/system_error.h> // thrust::system_error 9 #include <thrust/execution_policy.h> // thrust::cuda/hip::par 10 11 namespace Petsc 12 { 13 14 namespace device 15 { 16 17 namespace cupm 18 { 19 20 #if PetscDefined(USING_NVCC) 21 #if !defined(THRUST_VERSION) 22 #error "THRUST_VERSION not defined!" 23 #endif 24 #if THRUST_VERSION >= 101600 25 #define PETSC_THRUST_HAS_ASYNC 1 26 #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::cuda::par_nosync.on(s), __VA_ARGS__) 27 #else 28 #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::cuda::par.on(s), __VA_ARGS__) 29 #endif 30 #elif PetscDefined(USING_HCC) 31 #if !defined(THRUST_VERSION) 32 #error "THRUST_VERSION not defined!" 33 #endif 34 #if THRUST_VERSION >= 101600 35 #define PETSC_THRUST_HAS_ASYNC 1 36 #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::hip::par_nosync.on(s), __VA_ARGS__) 37 #else 38 #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::hip::par.on(s), __VA_ARGS__) 39 #endif 40 #else 41 #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(__VA_ARGS__) 42 #endif 43 44 #ifndef PETSC_THRUST_HAS_ASYNC 45 #define PETSC_THRUST_HAS_ASYNC 0 46 #endif 47 48 namespace detail 49 { 50 51 struct PetscLogGpuTimer { 52 PetscLogGpuTimer() noexcept 53 { 54 PetscFunctionBegin; 55 PetscCallAbort(PETSC_COMM_SELF, PetscLogGpuTimeBegin()); 56 PetscFunctionReturnVoid(); 57 } 58 59 ~PetscLogGpuTimer() noexcept 60 { 61 PetscFunctionBegin; 62 PetscCallAbort(PETSC_COMM_SELF, PetscLogGpuTimeEnd()); 63 PetscFunctionReturnVoid(); 64 } 65 }; 66 67 } // namespace detail 68 69 #define THRUST_CALL(...) \ 70 [&] { \ 71 const auto timer = ::Petsc::device::cupm::detail::PetscLogGpuTimer{}; \ 72 return PETSC_THRUST_CALL_PAR_ON(__VA_ARGS__); \ 73 }() 74 75 #define PetscCallThrust(...) \ 76 do { \ 77 try { \ 78 { \ 79 __VA_ARGS__; \ 80 } \ 81 } catch (const thrust::system_error &ex) { \ 82 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "Thrust error: %s", ex.what()); \ 83 } \ 84 } while (0) 85 86 } // namespace cupm 87 88 } // namespace device 89 90 } // namespace Petsc 91