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