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