blob: e70579b55c0aeebe3c9c2c970048aaf97f81c8ea (
plain)
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
|
#include "tensorflow/stream_executor/cuda/cuda_stream.h"
#include "tensorflow/stream_executor/lib/status.h"
namespace perftools {
namespace gputools {
namespace cuda {
bool CUDAStream::Init() {
return CUDADriver::CreateStream(parent_->cuda_context(), &cuda_stream_);
}
void CUDAStream::Destroy() {
{
mutex_lock lock{mu_};
if (completed_event_ != nullptr) {
port::Status status =
CUDADriver::DestroyEvent(parent_->cuda_context(), &completed_event_);
if (!status.ok()) {
LOG(ERROR) << status.error_message();
}
}
}
CUDADriver::DestroyStream(parent_->cuda_context(), &cuda_stream_);
}
bool CUDAStream::IsIdle() const {
return CUDADriver::IsStreamIdle(parent_->cuda_context(), cuda_stream_);
}
bool CUDAStream::GetOrCreateCompletedEvent(CUevent *completed_event) {
mutex_lock lock{mu_};
if (completed_event_ != nullptr) {
*completed_event = completed_event_;
return true;
}
if (!CUDADriver::CreateEvent(parent_->cuda_context(), &completed_event_,
CUDADriver::EventFlags::kDisableTiming)
.ok()) {
return false;
}
*completed_event = completed_event_;
return true;
}
} // namespace cuda
} // namespace gputools
} // namespace perftools
|
