Pytorch c10 模块详解

using StreamId = int64_t;

class C10_API Stream final {
 private:
  Device device_;
  StreamId id_;

 public:
  enum Unsafe { UNSAFE };
  enum Default { DEFAULT };

  /// Unsafely construct a stream from a Device and a StreamId.  In
  /// general, only specific implementations of streams for a
  /// backend should manufacture Stream directly in this way; other users
  /// should use the provided APIs to get a stream.  In particular,
  /// we don't require backends to give any guarantees about non-zero
  /// StreamIds; they are welcome to allocate in whatever way they like.
  explicit Stream(Unsafe, Device device, StreamId id)
      : device_(device), id_(id) {}

  /// Construct the default stream of a Device.  The default stream is
  /// NOT the same as the current stream; default stream is a fixed stream
  /// that never changes, whereas the current stream may be changed by
  /// StreamGuard.
  explicit Stream(Default, Device device) : device_(device), id_(0) {}

  bool operator==(const Stream& other) const noexcept {
    return this->device_ == other.device_ && this->id_ == other.id_;
  }
  bool operator!=(const Stream& other) const noexcept {
    return !(*this == other);
  }

  Device device() const noexcept {
    return device_;
  }
  DeviceType device_type() const noexcept {
    return device_.type();
  }
  DeviceIndex device_index() const noexcept {
    return device_.index();
  }
  StreamId id() const noexcept {
    return id_;
  }

  // Enqueues a wait instruction in the stream's work queue.
  // This instruction is a no-op unless the event is marked
  // for recording. In that case the stream stops processing
  // until the event is recorded.
  template <typename T>
  void wait(const T& event) const {
    event.block(*this);
  }

  // Return whether all asynchronous work previously enqueued on this stream
  // has completed running on the device.
  bool query() const;

  // Wait (by blocking the calling thread) until all asynchronous work enqueued
  // on this stream has completed running on the device.
  void synchronize() const;

  // The purpose of this function is to more conveniently permit binding
  // of Stream to and from Python.  Without packing, I have to setup a whole
  // class with two fields (device and stream id); with packing I can just
  // store a single uint64_t.
  //
  // The particular way we pack streams into a uint64_t is considered an
  // implementation detail and should not be relied upon.
  uint64_t hash() const noexcept {
    // Concat these together into a 64-bit integer
    uint64_t bits = static_cast<uint64_t>(device_type()) << 56 |
        static_cast<uint64_t>(device_index()) << 48 |
        // Remove the sign extension part of the 64-bit address because
        // the id might be used to hold a pointer.
        (static_cast<uint64_t>(id()) & ((1ull << 48) - 1));
    return bits;
  }

  struct StreamData3 pack3() const {
    return {id(), device_index(), device_type()};
  }

  static Stream unpack3(
      StreamId stream_id,
      DeviceIndex device_index,
      DeviceType device_type) {
    TORCH_CHECK(isValidDeviceType(device_type));
    return Stream(UNSAFE, Device(device_type, device_index), stream_id);
  }

  // I decided NOT to provide setters on this class, because really,
  // why would you change the device of a stream?  Just construct
  // it correctly from the beginning dude.
};

C10_API std::ostream& operator<<(std::ostream& stream, const Stream& s);

} // namespace c10

class C10_CUDA_API CUDAStream {
 public:
  enum Unchecked { UNCHECKED };

  /// Construct a CUDAStream from a Stream.  This construction is checked,
  /// and will raise an error if the Stream is not, in fact, a CUDA stream.
  explicit CUDAStream(Stream stream) : stream_(stream) {
    TORCH_CHECK(stream_.device_type() == DeviceType::CUDA);
  }
  // 代码省略
  private:
    Stream stream_; // 公共抽象父类对象，定义在 c10/core/Stream.h
}

/**
 * Get the default CUDA stream, for the passed CUDA device, or for the
 * current device if no device index is passed.  The default stream is
 * where most computation occurs when you aren't explicitly using
 * streams.
 */
C10_API CUDAStream getDefaultCUDAStream(DeviceIndex device_index = -1);

/**
 * Get a new stream from the CUDA stream pool.  You can think of this
 * as "creating" a new stream, but no such creation actually happens;
 * instead, streams are preallocated from the pool and returned in a
 * round-robin fashion.
 *
 * You can request a stream from the high priority pool by setting
 * isHighPriority to true, or a stream for a specific device by setting device
 * (defaulting to the current CUDA stream.)
 */
C10_API CUDAStream
getStreamFromPool(const bool isHighPriority = false, DeviceIndex device = -1);
// no default priority to disambiguate overloads
C10_API CUDAStream
getStreamFromPool(const int priority, DeviceIndex device = -1);

// 目的：确保当前流的所有 CUDA 操作执行完毕，提供一种同步机制。
void synchronize() const {
    // 确保在调用 CUDA 同步操作前，将当前设备切换为流所属的设备
    DeviceGuard guard{stream_.device()};
    // 调用 c10::cuda::stream_synchronize() 对当前流进行同步，
    // 即阻塞直到流中所有操作执行完毕，确保 CUDA 操作已完成
    c10::cuda::stream_synchronize(stream());
}

// InlineEvent 是一个模板结构体，用于封装事件（Event）的操作，其行为依赖于具体的后端实现（通过模板参数 T 提供）。
// 注意：构造函数被删除，必须传入设备类型和可选的 EventFlag 进行初始化。
template <typename T>
struct InlineEvent final {
  // 禁用默认构造函数
  InlineEvent() = delete;

  // 构造函数：需要传入设备类型和可选的标志参数（默认为 PYTORCH_DEFAULT）
  InlineEvent(
      const DeviceType _device_type,
      const EventFlag _flag = EventFlag::PYTORCH_DEFAULT)
      : backend_{_device_type},            // 利用设备类型初始化后端实现对象
        device_type_{_device_type},         // 记录事件所属的设备类型
        flag_{_flag} {}                     // 记录事件标志

  // 禁用拷贝构造和拷贝赋值，避免误拷贝事件对象
  InlineEvent(const InlineEvent&) = delete;
  InlineEvent& operator=(const InlineEvent&) = delete;

  // 移动构造函数，允许事件对象在移动语义下转移所有权
  InlineEvent(InlineEvent&& other) noexcept
      : event_(other.event_),                        // 传递底层事件指针
        backend_(std::move(other.backend_)),         // 移动后端实现对象
        device_type_(other.device_type_),            // 拷贝设备类型
        device_index_(other.device_index_),          // 拷贝设备索引
        flag_(other.flag_),                          // 拷贝事件标志
        was_marked_for_recording_(other.was_marked_for_recording_) {
    // 移动后将原对象的事件指针置空，避免重复释放
    other.event_ = nullptr;
  }

  // 移动赋值运算符，利用 swap 实现
  InlineEvent& operator=(InlineEvent&& other) noexcept {
    swap(other);
    return *this;
  }

  // 交换两个 InlineEvent 对象的所有成员
  void swap(InlineEvent& other) noexcept {
    std::swap(event_, other.event_);
    std::swap(backend_, other.backend_);
    std::swap(device_type_, other.device_type_);
    std::swap(device_index_, other.device_index_);
    std::swap(flag_, other.flag_);
    std::swap(was_marked_for_recording_, other.was_marked_for_recording_);
  }

  // 析构函数：若 event_ 不为空，则通过后端对象释放资源
  ~InlineEvent() noexcept {
    if (event_)
      backend_.destroyEvent(event_, device_index_);
  }

  // 返回事件所属设备的类型
  DeviceType device_type() const noexcept {
    return device_type_;
  }
  
  // 返回事件所属设备的索引
  DeviceIndex device_index() const noexcept {
    return device_index_;
  }
  
  // 返回事件的标志
  EventFlag flag() const noexcept {
    return flag_;
  }
  
  // 查询事件是否已被记录
  bool was_marked_for_recording() const noexcept {
    return was_marked_for_recording_;
  }

  // 如果尚未记录，则记录一次当前流
  void recordOnce(const Stream& stream) {
    if (!was_marked_for_recording_)
      record(stream);
  }

  // 记录事件到指定流
  void record(const Stream& stream) {
    // 检查传入流的设备类型与事件所属设备类型是否匹配
    TORCH_CHECK(
        stream.device_type() == device_type_,
        "Event device type ",
        DeviceTypeName(device_type_),
        " does not match recording stream's device type ",
        DeviceTypeName(stream.device_type()),
        ".");
    
    // 调用后端接口进行事件记录
    // 参数：事件指针的地址、流、当前设备索引、标志参数
    backend_.record(&event_, stream, device_index_, flag_);
    
    // 标记事件已经被记录
    was_marked_for_recording_ = true;
    
    // 更新设备索引为当前流的设备索引
    device_index_ = stream.device_index();
  }

  // 阻塞等待指定流直到该事件之前的所有操作完成
  void block(const Stream& stream) const {
    if (!was_marked_for_recording_)
      return;

    // 检查传入流的设备类型与事件所属设备类型是否一致
    TORCH_CHECK(
        stream.device_type() == device_type_,
        "Event device type ",
        DeviceTypeName(device_type_),
        " does not match blocking stream's device type ",
        DeviceTypeName(stream.device_type()),
        ".");

    // 调用后端接口执行阻塞等待操作
    backend_.block(event_, stream);
  }

  // 查询事件状态：如果事件未被记录，则返回 true，否则调用后端接口查询
  bool query() const {
    if (!was_marked_for_recording_)
      return true;
    return backend_.queryEvent(event_);
  }

  // 返回底层事件指针
  void* eventId() const {
    return event_;
  }

  // 计算两个事件之间的经过时间（通常以毫秒为单位）
  double elapsedTime(const InlineEvent& other) const {
    TORCH_CHECK(
        other.was_marked_for_recording(),
        "other was not marked for recording.");
    TORCH_CHECK(
        was_marked_for_recording(), "self was not marked for recording.");
    TORCH_CHECK(
        other.device_type() == device_type_,
        "Event device type ",
        DeviceTypeName(device_type_),
        " does not match other's device type ",
        DeviceTypeName(other.device_type()),
        ".");
    // 调用后端接口计算事件间耗时
    return backend_.elapsedTime(event_, other.event_, device_index_);
  }

  // 阻塞等待当前事件完成
  void synchronize() const {
    if (!was_marked_for_recording_)
      return;
    backend_.synchronizeEvent(event_);
  }

 private:
  // 底层事件指针，通常指向后端资源（例如 CUDA 事件）
  void* event_ = nullptr;
  // 后端实现对象，通过模板参数 T 提供
  T backend_;
  // 事件所属设备的类型（例如 CPU、CUDA 等）
  DeviceType device_type_;
  // 事件所属设备的索引，初始为 -1 表示未设置
  DeviceIndex device_index_ = -1;
  // 事件的标志，默认为 PYTORCH_DEFAULT
  EventFlag flag_ = EventFlag::PYTORCH_DEFAULT;
  // 标识是否已经记录过事件
  bool was_marked_for_recording_ = false;
};