1: #pragma once

  3: #include <petscdevice_cupm.h>

  5: #include <thrust/version.h>          // THRUST_VERSION
  6: #include <thrust/system_error.h>     // thrust::system_error
  7: #include <thrust/execution_policy.h> // thrust::cuda/hip::par

  9: namespace Petsc
 10: {

 12: namespace device
 13: {

 15: namespace cupm
 16: {

 18: #if PetscDefined(USING_NVCC)
 19:   #if !defined(THRUST_VERSION)
 20:     #error "THRUST_VERSION not defined!"
 21:   #endif
 22:   #if THRUST_VERSION >= 101600
 23:     #define PETSC_THRUST_HAS_ASYNC                 1
 24:     #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::cuda::par_nosync.on(s), __VA_ARGS__)
 25:   #else
 26:     #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::cuda::par.on(s), __VA_ARGS__)
 27:   #endif
 28: #elif PetscDefined(USING_HCC)
 29:   #if !defined(THRUST_VERSION)
 30:     #error "THRUST_VERSION not defined!"
 31:   #endif
 32:   #if THRUST_VERSION >= 101600
 33:     #define PETSC_THRUST_HAS_ASYNC                 1
 34:     #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::hip::par_nosync.on(s), __VA_ARGS__)
 35:   #else
 36:     #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(thrust::hip::par.on(s), __VA_ARGS__)
 37:   #endif
 38: #else
 39:   #define PETSC_THRUST_CALL_PAR_ON(func, s, ...) func(__VA_ARGS__)
 40: #endif

 42: #if !PetscDefined(THRUST_HAS_ASYNC)
 43:   #define PETSC_THRUST_HAS_ASYNC 0
 44: #endif

 46: namespace detail
 47: {

 49: struct PetscLogGpuTimer {
 50:   PetscLogGpuTimer() noexcept
 51:   {
 52:     PetscFunctionBegin;
 53:     PetscCallAbort(PETSC_COMM_SELF, PetscLogGpuTimeBegin());
 54:     PetscFunctionReturnVoid();
 55:   }

 57:   ~PetscLogGpuTimer() noexcept
 58:   {
 59:     PetscFunctionBegin;
 60:     PetscCallAbort(PETSC_COMM_SELF, PetscLogGpuTimeEnd());
 61:     PetscFunctionReturnVoid();
 62:   }
 63: };

 65: } // namespace detail

 67: #define THRUST_CALL(...) \
 68:   [&] { \
 69:     const auto timer = ::Petsc::device::cupm::detail::PetscLogGpuTimer{}; \
 70:     return PETSC_THRUST_CALL_PAR_ON(__VA_ARGS__); \
 71:   }()

 73: #define PetscCallThrust(...) \
 74:   do { \
 75:     try { \
 76:       { \
 77:         __VA_ARGS__; \
 78:       } \
 79:     } catch (const thrust::system_error &ex) { \
 80:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "Thrust error: %s", ex.what()); \
 81:     } \
 82:   } while (0)

 84: } // namespace cupm

 86: } // namespace device

 88: } // namespace Petsc