Actual source code: cupmstream.hpp

  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_, &current_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