Improves ATen CUDAEvent (#11293)

mruberry · facebook-github-bot · commit a2afad2b696e · 2018-09-06T21:39:44.000-07:00
Summary: After submitting PR #9726, PR #10581 created a different CUDAEvent class. The CUDAEvent proposed in #9726 was similar to the c10d::CUDAEvent class with additional testing and functionality. In particular, it was movable but not copyable. The CUDAEvent created by #10581 is refcounted and copyable. This PR retains the refcounting of the latter PR while fixing several bugs, adding tests, and extending the functionality to support testing and usage like in PR #8354. In particular, this PR: - Adds set_device() to CUDAContext - Adds three CUDAEvent tests to stream_test.cpp - Fixes three bugs: - Refcounting was broken. Destroying an of the RAIIs holding a particular CUDAEvent would destroy the event UNLESS it was the last RAII (the check was backwards). - Moving an event would cause a segfault. - Events were not destroyed on the device they were created on. See PR #9415 (pietern) - Adds the happened() and recordOnce() functions - Changes the record() functions to not be const - Adds additional assertions to verify correctness This PR does not: - Make c10d use the ATen CUDAEvent (this is appropriate for a separate PR) Whether events should be refcounted is an interesting question. It adds some atomic operations and makes event creation eager. Making events movable but not copyable (like the c10d events) avoids these costs and allows events to be lazily constructed. Lazy construction is preferable when working with containers (like std::array or std::vector) and because the event's device can be set automatically to the first stream it's recorded on. With eager construction the user is required to understand that events have a device and acquire the device of the stream the event will be recorded on upfront. This can be seen here: https://github.com/pytorch/pytorch/blob/542aadd9a7609892e207c1e15de08a975b697752/aten/src/ATen/native/cudnn/RNN.cpp#L1130-L1132 and that file is the only one which currently uses the ATen CUDAEvent. Refcounting does allow single writer multi-reader scenarios, although these scenarios can be also be supported by providing indirect access to the underlying CUDAEvent. I believe all current and planned usage scenarios do not require refcounting, and if desired I can update this PR to remove refcounting and make the ATen event movable but not copyable like the c10d event. I think not refcounting is preferable because it can improve performance, ease usability, and simplify the code (as seen with two of the above bugs). I have decided to separate this from PR #8354 since while it's required for PR #8354 the changes are, clearly, of independent interest. PR #8354 has a new dependency on this one, however. I am closing PR #9726 in favor of this PR. apaszke ezyang pietern Pull Request resolved: #11293 Differential Revision: D9665836 Pulled By: soumith fbshipit-source-id: a1513fa4f9761e2f304d126e402f6b6950e1c1d2
diff --git a/aten/src/ATen/cuda/CUDAContext.cpp b/aten/src/ATen/cuda/CUDAContext.cpp
@@ -16,6 +16,10 @@ int64_t current_device() {
   return cur_device;
 }
 
+void set_device(int64_t device) {
+  AT_CUDA_CHECK(cudaSetDevice((int)device));
+}
+
 cudaDeviceProp* getCurrentDeviceProperties() {
   return THCState_getCurrentDeviceProperties(at::globalContext().getTHCState());
 }
diff --git a/aten/src/ATen/cuda/CUDAContext.h b/aten/src/ATen/cuda/CUDAContext.h
@@ -39,6 +39,8 @@ AT_API int64_t getNumGPUs();
 
 AT_API int64_t current_device();
 
+AT_API void set_device(int64_t device);
+
 AT_API cudaDeviceProp* getCurrentDeviceProperties();
 
 AT_API cudaDeviceProp* getDeviceProperties(int64_t device);
diff --git a/aten/src/ATen/cuda/CUDAEvent.cpp b/aten/src/ATen/cuda/CUDAEvent.cpp
diff --git a/aten/src/ATen/cuda/CUDAEvent.h b/aten/src/ATen/cuda/CUDAEvent.h
@@ -1,78 +1,116 @@
 #pragma once
 
-#include <cstdint>
-#include <utility>
+#include "ATen/cuda/ATenCUDAGeneral.h"
+#include "ATen/cuda/CUDAStream.h"
+#include "ATen/cuda/CUDAContext.h"
+#include "ATen/cuda/Exceptions.h"
+#include "ATen/core/Error.h"
+#include "ATen/DeviceGuard.h"
 
 #include "cuda_runtime_api.h"
 
-#include <ATen/core/ATenGeneral.h>
-#include <ATen/Error.h>
+#include <cstdint>
+#include <utility>
+
+namespace at { namespace cuda {
 
 /*
-* A CUDA event interface with no CUDA build dependency.
+* CUDAEvents are movable not copyable wrappers around CUDA's events.
 *
-* Includes the CUDAEvent RAII class and a pointer-based event API.
+* CUDAEvents are constructed lazily when recorded on streams. The events
+* have a device, and this device is acquired from the first recording stream.
+* Later streams that record to the event must share this device, but streams
+* on any device can wait on the event.
 */
-
-struct CUDAEventInternals;
-
-namespace at {
-namespace cuda {
-
-struct CUDAStream;
-
-namespace detail {
-
-// Pointer-based API (for internal use)
-// Note: ATen/Context is preferred to work with streams safely
-AT_API CUDAEventInternals* CUDAEvent_create(unsigned int flags);
-AT_API void CUDAEvent_retain(CUDAEventInternals* internals);
-AT_API void CUDAEvent_uncheckedFree(CUDAEventInternals* internals);
-AT_API cudaEvent_t CUDAEvent_event(CUDAEventInternals* internals);
-AT_API int64_t CUDAEvent_device(CUDAEventInternals* internals);
-
-} // namespace detail
-
-struct CUDAEvent {
+struct AT_CUDA_API CUDAEvent {
   // Constants
   static constexpr unsigned int DEFAULT_FLAGS = cudaEventDisableTiming;
 
   // Constructors
-  CUDAEvent(unsigned int flags = DEFAULT_FLAGS)
-    : internals_(detail::CUDAEvent_create(flags)) {}
-
-  ~CUDAEvent() { detail::CUDAEvent_uncheckedFree(internals_); }
-
-  CUDAEvent(const CUDAEvent& other) {
-    detail::CUDAEvent_retain(other.internals_);
-    internals_ = other.internals_;
+  CUDAEvent(unsigned int flags = DEFAULT_FLAGS) 
+  : flags_{flags} { }
+
+  // Note: event destruction done on creating device to avoid creating a 
+  // CUDA context on other devices.
+  ~CUDAEvent() { 
+    try {
+      if (is_created_) {
+        at::DeviceGuard device_guard{(int)device_};
+        cudaEventDestroy(event_);
+      }
+    } catch (...) { /* No throw */ }
   }
 
-  CUDAEvent(CUDAEvent&& other) {
-    std::swap(internals_, other.internals_);
-  }
+  CUDAEvent(const CUDAEvent&) = delete;
+  CUDAEvent& operator=(const CUDAEvent&) = delete;
 
-  CUDAEvent& operator=(CUDAEvent other) noexcept {
-    std::swap(internals_, other.internals_);
+  CUDAEvent(CUDAEvent&& other) { moveHelper(std::move(other)); } 
+  CUDAEvent& operator=(CUDAEvent&& other) {
+    moveHelper(std::move(other));
     return *this;
   }
 
-  operator cudaEvent_t() const { return detail::CUDAEvent_event(internals_); }
+  operator cudaEvent_t() const { return event(); }
 
   // Less than operator (to allow use in sets)
   friend bool operator<(const CUDAEvent& left, const CUDAEvent& right) {
-    return left.internals_ < right.internals_;
+    return left.event_ < right.event_;
   }
 
-  int64_t device() const { return detail::CUDAEvent_device(internals_); }
-  cudaEvent_t event() const { return detail::CUDAEvent_event(internals_); }
-  CUDAEventInternals* internals() const { return internals_; }
+  bool isCreated() const { return is_created_; }
+  int64_t device() const { return device_; }
+  cudaEvent_t event() const { return event_; }
+  
+  bool happened() const { 
+    return (was_recorded_ && cudaEventQuery(event_) == cudaSuccess);
+  }
+
+  void record() { record(getCurrentCUDAStream()); }
+  
+  void recordOnce(const CUDAStream& stream) { 
+    if (!was_recorded_) record(stream);
+  }
+  
+  void record(const CUDAStream& stream) {
+    if (is_created_) {
+      AT_ASSERT(device_ == stream.device());
+    } else {
+      create(stream.device());
+    }
+
+    AT_CUDA_CHECK(cudaEventRecord(event_, stream));
+    was_recorded_ = true;
+  }
 
-  void record() const; // Record on the current stream
-  void record(const CUDAStream& stream) const;
+  void block (const CUDAStream& stream) {
+    if (is_created_) {
+      AT_CUDA_CHECK(cudaStreamWaitEvent(stream, event_, 0));
+    }
+  }
+  
 
 private:
-  CUDAEventInternals* internals_;
+  unsigned int flags_ = DEFAULT_FLAGS;
+  bool is_created_ = false;
+  bool was_recorded_ = false;
+  int64_t device_ = -1;
+  cudaEvent_t event_;
+
+  void moveHelper(CUDAEvent&& other) {
+    std::swap(flags_, other.flags_);
+    std::swap(is_created_, other.is_created_);
+    std::swap(was_recorded_, other.was_recorded_);
+    std::swap(device_, other.device_);
+    std::swap(event_, other.event_);
+  }
+
+  void create(const int64_t device) {
+    at::DeviceGuard device_guard{(int)device};
+    AT_CUDA_CHECK(cudaEventCreateWithFlags(&event_, flags_));
+
+    is_created_ = true;
+    device_ = device;
+  }
 };
 
 } // namespace cuda
diff --git a/aten/src/ATen/cuda/CUDAStream.cpp b/aten/src/ATen/cuda/CUDAStream.cpp
@@ -209,7 +209,8 @@ int64_t CUDAStream_device(CUDAStreamInternals* ptr) {
 }
 
 void CUDAStream_synchronize_with(CUDAStreamInternals* ptr, const CUDAEvent& event) {
-    AT_CUDA_CHECK(cudaStreamWaitEvent(ptr->stream, event, 0));
+    if (event.isCreated())
+      AT_CUDA_CHECK(cudaStreamWaitEvent(ptr->stream, event, 0));
 }
 
 } // namespace detail
diff --git a/aten/src/ATen/cuda/CUDAStream.h b/aten/src/ATen/cuda/CUDAStream.h
@@ -77,7 +77,7 @@ AT_CUDA_API int64_t CUDAStream_device(CUDAStreamInternals*);
 
 // RAII for a CUDA stream
 // Allows use as a cudaStream_t, copying, moving, and metadata access.
-struct CUDAStream {
+struct AT_CUDA_API CUDAStream {
 
   // Constructors
   CUDAStream() = default;
diff --git a/aten/src/ATen/test/stream_test.cpp b/aten/src/ATen/test/stream_test.cpp
@@ -3,6 +3,7 @@
 
 #include "ATen/cuda/CUDAContext.h"
 #include "ATen/cuda/CUDAGuard.h"
+#include "ATen/cuda/CUDAEvent.h"
 
 #include "cuda_runtime.h"
 
@@ -211,7 +212,6 @@ TEST_CASE("Streampool Round Robin") {
   REQUIRE(hasDuplicates);
 }
 
-// Note: to be expanded once CUDAEvent PR is accepted
 TEST_CASE("Multi-GPU") {
   if (at::cuda::getNumGPUs() < 2) return;
 
@@ -226,3 +226,44 @@ TEST_CASE("Multi-GPU") {
   at::DeviceGuard device_guard{1};
   REQUIRE(s1 == at::cuda::getCurrentCUDAStream());
 }
+
+TEST_CASE("CUDAEvent Syncs") {
+  const auto stream = at::cuda::createCUDAStream();
+  at::cuda::CUDAEvent event;
+
+  REQUIRE(!event.happened());
+
+  event.recordOnce(stream);
+
+  const auto wait_stream0 = at::cuda::createCUDAStream();
+  const auto wait_stream1 = at::cuda::createCUDAStream();
+
+  wait_stream0.synchronize_with(event);
+  wait_stream1.synchronize_with(event);
+
+  cudaStreamSynchronize(wait_stream0);
+  REQUIRE(event.happened());
+}
+
+TEST_CASE("Cross-Device Events") {
+  if (at::cuda::getNumGPUs() < 2) return;
+
+  const auto stream0 = at::cuda::createCUDAStream();
+  at::cuda::CUDAEvent event0;
+
+  at::cuda::set_device(1);
+  const auto stream1 = at::cuda::createCUDAStream();
+  at::cuda::CUDAEvent event1;
+
+  event0.record(stream0);
+  event1.record(stream1);
+  
+  event0 = std::move(event1);
+  
+  REQUIRE(event0.device() == 1);
+
+  stream0.synchronize_with(event0);
+  
+  cudaStreamSynchronize(stream0);
+  REQUIRE(event0.happened());
+}

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,10 @@ int64_t current_device() {`
`16`	`16`	`return cur_device;`
`17`	`17`	`}`
`18`	`18`
	`19`	`+void set_device(int64_t device) {`
	`20`	`+ AT_CUDA_CHECK(cudaSetDevice((int)device));`
	`21`	`+}`
	`22`	`+`
`19`	`23`	`cudaDeviceProp* getCurrentDeviceProperties() {`
`20`	`24`	`return THCState_getCurrentDeviceProperties(at::globalContext().getTHCState());`
`21`	`25`	`}`
Original file line number	Diff line number	Diff line change
`@@ -209,7 +209,8 @@ int64_t CUDAStream_device(CUDAStreamInternals* ptr) {`
`209`	`209`	`}`
`210`	`210`
`211`	`211`	`void CUDAStream_synchronize_with(CUDAStreamInternals* ptr, const CUDAEvent& event) {`
`212`		`- AT_CUDA_CHECK(cudaStreamWaitEvent(ptr->stream, event, 0));`
	`212`	`+ if (event.isCreated())`
	`213`	`+ AT_CUDA_CHECK(cudaStreamWaitEvent(ptr->stream, event, 0));`
`213`	`214`	`}`
`214`	`215`
`215`	`216`	`} // namespace detail`