Wait for data CUDA stream instead of sync. (#9144)

---------

Co-authored-by: Philip Hyunsu Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Jiaming Yuan
2023-05-09 09:52:21 +08:00
committed by GitHub
parent a075aa24ba
commit 85988a3178
3 changed files with 24 additions and 22 deletions

View File

@@ -1352,14 +1352,12 @@ class CUDAStream {
cudaStream_t stream_;
public:
CUDAStream() {
dh::safe_cuda(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
}
~CUDAStream() {
dh::safe_cuda(cudaStreamDestroy(stream_));
}
CUDAStream() { dh::safe_cuda(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); }
~CUDAStream() { dh::safe_cuda(cudaStreamDestroy(stream_)); }
[[nodiscard]] CUDAStreamView View() const { return CUDAStreamView{stream_}; }
[[nodiscard]] cudaStream_t Handle() const { return stream_; }
CUDAStreamView View() const { return CUDAStreamView{stream_}; }
void Sync() { this->View().Sync(); }
};