MPSEvent Class — pytorch Architecture

Source Code

aten/src/ATen/mps/MPSEvent.h lines 13–66

class MPSEvent {
 public:
  explicit MPSEvent(id_t ID, MPSStream* stream, bool enable_timing);
  ~MPSEvent();

  // records an event on the stream
  void record(bool needsLock, bool syncEvent = false);
  // makes all future work submitted to the stream wait for this event.
  bool wait(bool needsLock, bool syncEvent = false);
  // schedules a notifyListener callback for the event.
  bool notify(bool needsLock, MTLSharedEventNotificationBlock block);
  // checks if events are already signaled.
  bool query() const;
  // blocks the CPU thread until all the GPU work that were scheduled
  // prior to recording this event are completed.
  bool synchronize();
  // resets this event with new parameters in case it gets reused from the event
  // pool
  void reset(MPSStream* stream, bool enable_timing);
  // returns the unique ID of the event instance
  id_t getID() const {
    return m_id;
  }
  // returns the completion timestamp of the event
  uint64_t getCompletionTime() const {
    return m_completion_time;
  }
  // if already recorded, waits for cpu_sync_cv to be signaled
  void waitForCpuSync();

 private:
  id_t m_id;
  // enables measuring the completion time of the notifyListener of this event
  bool m_enable_timing;
  uint64_t m_signalCounter = 0;
  MPSStream* m_stream = nullptr;
  MTLSharedEvent_t m_event = nullptr;
  MTLSharedEventListener* m_listener = nullptr;
  // used to sync the events created on this Stream with CPU
  std::mutex m_cpu_sync_mutex{};
  std::condition_variable m_cpu_sync_cv{};
  // CondVar predicate to sync the events created on this Stream with CPU
  bool m_cpu_sync_completed = false;
  // used to compute elapsed time
  uint64_t m_completion_time = 0;

  void recordLocked(bool syncEvent);
  bool waitLocked(bool syncEvent);
  bool notifyLocked(MTLSharedEventNotificationBlock block);
  void notifyCpuSync();
  static uint64_t getTime() {
    return clock_gettime_nsec_np(CLOCK_MONOTONIC_RAW);
  }
};