Repository navigation
Expand file tree
/
Copy pathwaitfor.hpp
More file actions
139 lines (115 loc) · 3.92 KB
/
Copy pathwaitfor.hpp
File metadata and controls
139 lines (115 loc) · 3.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
// Copyright Oliver Kowalke 2017.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
#ifndef BOOST_FIBERS_CUDA_WAITFOR_H
#define BOOST_FIBERS_CUDA_WAITFOR_H
#include <initializer_list>
#include <mutex>
#include <iostream>
#include <set>
#include <tuple>
#include <vector>
#include <boost/assert.hpp>
#include <boost/config.hpp>
#include <cuda.h>
#include <boost/fiber/detail/config.hpp>
#include <boost/fiber/detail/is_all_same.hpp>
#include <boost/fiber/condition_variable.hpp>
#include <boost/fiber/mutex.hpp>
#ifdef BOOST_HAS_ABI_HEADERS
# include BOOST_ABI_PREFIX
#endif
namespace boost {
namespace fibers {
namespace cuda {
namespace detail {
template< typename Rendezvous >
static void trampoline( cudaStream_t st, cudaError_t status, void * vp) {
Rendezvous * data = static_cast< Rendezvous * >( vp);
data->notify( st, status);
}
class single_stream_rendezvous {
public:
single_stream_rendezvous( cudaStream_t st) {
unsigned int flags = 0;
cudaError_t status = ::cudaStreamAddCallback( st, trampoline< single_stream_rendezvous >, this, flags);
if ( cudaSuccess != status) {
st_ = st;
status_ = status;
done_ = true;
}
}
void notify( cudaStream_t st, cudaError_t status) noexcept {
std::unique_lock< mutex > lk{ mtx_ };
st_ = st;
status_ = status;
done_ = true;
lk.unlock();
cv_.notify_one();
}
std::tuple< cudaStream_t, cudaError_t > wait() {
std::unique_lock< mutex > lk{ mtx_ };
cv_.wait( lk, [this]{ return done_; });
return std::make_tuple( st_, status_);
}
private:
mutex mtx_{};
condition_variable cv_{};
cudaStream_t st_{};
cudaError_t status_{ cudaErrorUnknown };
bool done_{ false };
};
class many_streams_rendezvous {
public:
many_streams_rendezvous( std::initializer_list< cudaStream_t > l) :
stx_{ l } {
results_.reserve( stx_.size() );
for ( cudaStream_t st : stx_) {
unsigned int flags = 0;
cudaError_t status = ::cudaStreamAddCallback( st, trampoline< many_streams_rendezvous >, this, flags);
if ( cudaSuccess != status) {
std::unique_lock< mutex > lk{ mtx_ };
stx_.erase( st);
results_.push_back( std::make_tuple( st, status) );
}
}
}
void notify( cudaStream_t st, cudaError_t status) noexcept {
std::unique_lock< mutex > lk{ mtx_ };
stx_.erase( st);
results_.push_back( std::make_tuple( st, status) );
if ( stx_.empty() ) {
lk.unlock();
cv_.notify_one();
}
}
std::vector< std::tuple< cudaStream_t, cudaError_t > > wait() {
std::unique_lock< mutex > lk{ mtx_ };
cv_.wait( lk, [this]{ return stx_.empty(); });
return results_;
}
private:
mutex mtx_{};
condition_variable cv_{};
std::set< cudaStream_t > stx_;
std::vector< std::tuple< cudaStream_t, cudaError_t > > results_;
};
}
void waitfor_all();
inline
std::tuple< cudaStream_t, cudaError_t > waitfor_all( cudaStream_t st) {
detail::single_stream_rendezvous rendezvous( st);
return rendezvous.wait();
}
template< typename ... STP >
std::vector< std::tuple< cudaStream_t, cudaError_t > > waitfor_all( cudaStream_t st0, STP ... stx) {
static_assert( boost::fibers::detail::is_all_same< cudaStream_t, STP ...>::value, "all arguments must be of type `CUstream*`.");
detail::many_streams_rendezvous rendezvous{ st0, stx ... };
return rendezvous.wait();
}
}}}
#ifdef BOOST_HAS_ABI_HEADERS
# include BOOST_ABI_SUFFIX
#endif
#endif // BOOST_FIBERS_CUDA_WAITFOR_H