1: #pragma once

  3: #include <petsclog.h>
  4: #include <petscsys.h>
  5: #include <petscdevice_cupm.h>

  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

 11: namespace Petsc
 12: {

 14: namespace device
 15: {

 17: namespace cupm
 18: {

 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

 44: #ifndef PETSC_THRUST_HAS_ASYNC
 45:   #define PETSC_THRUST_HAS_ASYNC 0
 46: #endif

 48: namespace detail
 49: {

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

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

 67: } // namespace detail

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

 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)

 86: } // namespace cupm

 88: } // namespace device

 90: } // namespace Petsc