1: #ifndef PETSC_CUPMSTREAM_HPP 2: #define PETSC_CUPMSTREAM_HPP 4: #include <petsc/private/cupminterface.hpp> 6: #include "../segmentedmempool.hpp" 7: #include "cupmevent.hpp" 9: #if defined(__cplusplus) 10: namespace Petsc 11: { 13: namespace device 14: { 16: namespace cupm 17: { 19: // A bare wrapper around a cupmStream_t. The reason it exists is because we need to uniquely 20: // identify separate cupm streams. This is so that the memory pool can accelerate allocation 21: // calls as it can just pass back a pointer to memory that was used on the same 22: // stream. Otherwise it must either serialize with another stream or allocate a new chunk. 23: // Address of the objects does not suffice since cupmStreams are very likely internally reused. 25: template <DeviceType T> 26: class CUPMStream : public StreamBase<CUPMStream<T>>, impl::Interface<T> { 27: using crtp_base_type = StreamBase<CUPMStream<T>>; 28: friend crtp_base_type; 30: public: 31: PETSC_CUPM_INHERIT_INTERFACE_TYPEDEFS_USING(interface_type, T); 33: using stream_type = cupmStream_t; 34: using id_type = typename crtp_base_type::id_type; 35: using event_type = CUPMEvent<T>; 36: using flag_type = unsigned int; 38: CUPMStream() noexcept = default; 40: PETSC_NODISCARD PetscErrorCode destroy() noexcept; 41: PETSC_NODISCARD PetscErrorCode create(flag_type) noexcept; 42: PETSC_NODISCARD PetscErrorCode change_type(PetscStreamType) noexcept; 44: private: 45: stream_type stream_{}; 46: id_type id_ = new_id_(); 48: PETSC_NODISCARD static id_type new_id_() noexcept; 50: // CRTP implementations 51: PETSC_NODISCARD stream_type get_stream_() const noexcept; 52: PETSC_NODISCARD id_type get_id_() const noexcept; 53: PETSC_NODISCARD PetscErrorCode record_event_(event_type &) const noexcept; 54: PETSC_NODISCARD PetscErrorCode wait_for_(event_type &) const noexcept; 55: }; 57: template <DeviceType T> 58: inline PetscErrorCode CUPMStream<T>::destroy() noexcept 59: { 60: if (stream_) { 61: cupmStreamDestroy(stream_); 62: stream_ = cupmStream_t{}; 63: id_ = 0; 64: } 65: return 0; 66: } 68: template <DeviceType T> 69: inline PetscErrorCode CUPMStream<T>::create(flag_type flags) noexcept 70: { 71: if (stream_) { 72: if (PetscDefined(USE_DEBUG)) { 73: flag_type current_flags; 75: cupmStreamGetFlags(stream_, ¤t_flags); 77: } 78: return 0; 79: } 80: cupmStreamCreateWithFlags(&stream_, flags); 81: id_ = new_id_(); 82: return 0; 83: } 85: template <DeviceType T> 86: inline PetscErrorCode CUPMStream<T>::change_type(PetscStreamType newtype) noexcept 87: { 88: if (newtype == PETSC_STREAM_GLOBAL_BLOCKING) { 89: destroy(); 90: } else { 91: const flag_type preferred = newtype == PETSC_STREAM_DEFAULT_BLOCKING ? cupmStreamDefault : cupmStreamNonBlocking; 93: if (stream_) { 94: flag_type flag; 96: cupmStreamGetFlags(stream_, &flag); 97: if ((flag != preferred) || (cupmStreamQuery(stream_) != cupmSuccess)) destroy(); 98: } 99: create(preferred); 100: } 101: return 0; 102: } 104: template <DeviceType T> 105: inline typename CUPMStream<T>::id_type CUPMStream<T>::new_id_() noexcept 106: { 107: static id_type id = 0; 108: return id++; 109: } 111: // CRTP implementations 112: template <DeviceType T> 113: inline typename CUPMStream<T>::stream_type CUPMStream<T>::get_stream_() const noexcept 114: { 115: return stream_; 116: } 118: template <DeviceType T> 119: inline typename CUPMStream<T>::id_type CUPMStream<T>::get_id_() const noexcept 120: { 121: return id_; 122: } 124: template <DeviceType T> 125: inline PetscErrorCode CUPMStream<T>::record_event_(event_type &event) const noexcept 126: { 127: event.record(stream_); 128: return 0; 129: } 131: template <DeviceType T> 132: inline PetscErrorCode CUPMStream<T>::wait_for_(event_type &event) const noexcept 133: { 134: cupmStreamWaitEvent(stream_, event.get(), 0); 135: return 0; 136: } 138: } // namespace cupm 140: } // namespace device 142: } // namespace Petsc 143: #endif // __cplusplus 145: #endif // PETSC_CUPMSTREAM_HPP