Automated PR - 2026-07-07

2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
@@ -0,0 +1,114 @@
+/**
+ * @file event.hpp
+ * @brief CUDA stream and event synchronization utilities.
+ *
+ * This header provides wrapper types and helper functions for managing
+ * CUDA events and stream synchronization in PyTorch/ATen environment.
+ * These utilities are used to coordinate asynchronous operations across
+ * multiple CUDA streams.
+ */
+
+#pragma once
+
+#include <ATen/cuda/CUDAContext.h>
+#include <memory>
+
+#include "cuda/exceptions.cuh"
+
+namespace ltx_kernels {
+
+/**
+ * @struct EventHandle
+ * @brief RAII wrapper for a CUDA event with automatic recording.
+ *
+ * EventHandle encapsulates a torch::Event and automatically records it
+ * on the specified (or current) CUDA stream upon construction. This
+ * provides a convenient way to capture the completion point of stream
+ * operations for synchronization purposes.
+ *
+ * ## Usage Example
+ *
+ * ```cpp
+ * // Record event on current stream
+ * EventHandle ev1;
+ *
+ * // Record event on specific stream
+ * EventHandle ev2(my_stream);
+ *
+ * // Make current stream wait for the event
+ * ev1.current_stream_wait();
+ * ```
+ */
+struct EventHandle {
+  /// Shared pointer to the underlying torch::Event
+  std::shared_ptr<torch::Event> event;
+
+  /**
+   * @brief Constructs an EventHandle and records on the current CUDA stream.
+   *
+   * The event captures the completion point of all operations submitted
+   * to the current stream before this constructor is called.
+   */
+  EventHandle() {
+    event = std::make_shared<torch::Event>(torch::kCUDA);
+    event->record(at::cuda::getCurrentCUDAStream());
+  }
+
+  /**
+   * @brief Constructs an EventHandle and records on the specified stream.
+   *
+   * @param stream The CUDA stream to record the event on
+   */
+  explicit EventHandle(const at::cuda::CUDAStream &stream) {
+    event = std::make_shared<torch::Event>(torch::kCUDA);
+    event->record(stream);
+  }
+
+  /// Copy constructor (shares the underlying event)
+  EventHandle(const EventHandle &other) = default;
+
+  /**
+   * @brief Makes the current CUDA stream wait for this event.
+   *
+   * After this call returns, operations submitted to the current stream
+   * will not execute until the event has been reached on its recording stream.
+   */
+  void current_stream_wait() const { at::cuda::getCurrentCUDAStream().unwrap().wait(*event); }
+};
+
+/**
+ * @brief Creates and records a CUDA event on the specified stream.
+ *
+ * @param s The CUDA stream to record on
+ * @return A torch::Event that has been recorded on stream s
+ */
+inline torch::Event create_event(const at::cuda::CUDAStream &s) {
+  auto event = torch::Event(torch::kCUDA);
+  event.record(s);
+  return event;
+}
+
+/**
+ * @brief Makes stream s_0 wait for stream s_1's current position.
+ *
+ * After this call, operations on s_0 will not execute until all operations
+ * currently queued on s_1 have completed.
+ *
+ * @param s_0 The stream that will wait
+ * @param s_1 The stream to wait for
+ * @pre s_0 and s_1 must be different streams
+ */
+inline void stream_wait(const at::cuda::CUDAStream &s_0, const at::cuda::CUDAStream &s_1) {
+  EP_HOST_ASSERT(s_0.id() != s_1.id());
+  s_0.unwrap().wait(create_event(s_1));
+}
+
+/**
+ * @brief Makes a stream wait for a previously recorded event.
+ *
+ * @param s The stream that will wait
+ * @param event The event to wait for
+ */
+inline void stream_wait(const at::cuda::CUDAStream &s, const EventHandle &event) { s.unwrap().wait(*event.event); }
+
+} // namespace ltx_kernels