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