mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
e7cccc23b9
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59560 `at::cuda::CUDAStream` has the `query` and `synchronize` methods, but `c10::Stream` does not, and I couldn't find any generic way to accomplish this. Hence I added helpers to do this to the DeviceGuardImpl interface, and then defined these methods on `c10::Stream`. (I had to do it out-of-line to circumvent a circular dependency). ghstack-source-id: 130932249 Test Plan: CI Reviewed By: ezyang Differential Revision: D28931377 fbshipit-source-id: cd0c19cf021e305d0c0cf9af364afb445d010248
28 lines
880 B
C++
28 lines
880 B
C++
#include <c10/core/Stream.h>
|
|
#include <c10/core/impl/VirtualGuardImpl.h>
|
|
|
|
namespace c10 {
|
|
|
|
// Return whether all asynchronous work previously enqueued on this stream
|
|
// has completed running on the device.
|
|
bool Stream::query() const {
|
|
impl::VirtualGuardImpl impl{device_.type()};
|
|
return impl.queryStream(*this);
|
|
}
|
|
|
|
// Wait (by blocking the calling thread) until all asynchronous work enqueued
|
|
// on this stream has completed running on the device.
|
|
void Stream::synchronize() const {
|
|
impl::VirtualGuardImpl impl{device_.type()};
|
|
impl.synchronizeStream(*this);
|
|
}
|
|
|
|
// Not very parsable, but I don't know a good compact syntax for streams.
|
|
// Feel free to change this into something more compact if needed.
|
|
std::ostream& operator<<(std::ostream& stream, const Stream& s) {
|
|
stream << "stream " << s.id() << " on device " << s.device();
|
|
return stream;
|
|
}
|
|
|
|
} // namespace c10
|