不太行啊

[kevinJhz]

不太行啊

Originally posted by @kevinJhz in #1 (review)

[kevinJhz]

EP 设备抽象接口文档

代码路径

ep 模块的接口定义在 oneflow/core/ep/include 目录中，CUDA 相关的实现在 oneflow/core/ep/cuda 目录中， CPU 相关的实现在 oneflow/core/ep/cpu

接口说明

Event

Event 类是事件的抽象，对标 CUDA Event 。

class Event {
 public:
  OF_DISALLOW_COPY_AND_MOVE(Event);
  Event() = default;
  virtual ~Event() = default;

  // 查询事件是否完成，对标 [cudaEventQuery](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html#group__CUDART__EVENT_1g2bf738909b4a059023537eaa29d8a5b7), 如果对应的事件已经完成，返回 Maybe<bool>(true)，如果尚未完成返回 Maybe<bool>(false)，如果查询失败返回对应的错误，如果事件尚未 Record ，认为事件已经完成。
  virtual Maybe<bool> QueryDone() = 0;
  // 等待事件完成，对标 [cudaEventSynchronize](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html#group__CUDART__EVENT_1g949aa42b30ae9e622f6ba0787129ff22) , 成功返回 Maybe<void>::Ok(), 否则返回对应的错误
  virtual Maybe<void> Sync() = 0;
};

Stream

Stream 是流或者队列的抽象，对标 CUDA Stream ，Stream 不是线程安全的。

class Stream {
 public:
  OF_DISALLOW_COPY_AND_MOVE(Stream);
  Stream() = default;
  virtual ~Stream() = default;

  // 返回对应的设备类型
  virtual DeviceType device_type() const = 0;
  // 返回对应的 Device
  virtual Device* device() const = 0;
  // 同步 Stream，等待该 Stream 上所有添加的任务都已经完成后返回，同步失败则返回对应的错误，对标 [cudaStreamSynchronize](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g82b5784f674c17c6df64affe618bf45e)
  virtual Maybe<void> Sync() = 0;
  // 记录一个事件，对标 [cudaEventRecord](https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html#group__CUDART__EVENT_1gf4fcb74343aa689f4159791967868446)
  virtual void RecordEvent(Event* event) = 0;

  // 初始化调用线程，该方法的实现可以包含在当前线程使用该 Stream 所需要的必备设置或者对提高性能有帮助的设置，比如为CUDA设置当前设备(cudaSetDevice)，设置 CPU/内存的亲和性，设置将调试信息添加到当前 Stream 等
  virtual Maybe<void> OnExecutionContextSetup() { return Maybe<void>::Ok(); }
  // OnExecutionContextSetup 操作的逆过程
  virtual Maybe<void> OnExecutionContextTeardown() { return Maybe<void>::Ok(); }

  // 将 Stream 转为具体 Stream 类的便捷方法
  template<typename T>
  T* As() {
    return static_cast<T*>(this);
  }
};

AllocationOptions

内存分配选项，用于对设备内存分配进行详细的配置，目前有两个配置项，pinned device 适用于 CPU 内存，代表该内存作为 CPU 与特定非 CPU 设备交互的内存，比如将CPU数据拷贝到某个非CPU设备，pinned device 只能配置一次，如果尝试覆盖会报错；numa node affinity 也适用 CPU 内存，表示应该内存尽量分配在指定的 numa node 节点上，该选项非强制选择，后端可以自行决定是否处理该配置，同时该配置也允许覆盖。

class AllocationOptions {
 public:
  
  // 查询是否配置了对应的 pinned device
  bool HasPinnedDevice() const;

  // pinned device 的设备类型，如果没有配置 pinned device，将会报错
  DeviceType GetPinnedDeviceType() const;

  // pinned device 的设备 index，如果没有配置 pinned device，将会报错
  size_t GetPinnedDeviceIndex() const;

  // 配置 pinned device, 如果当前 pinned device已经配置，将会报错
  void SetPinnedDevice(DeviceType device_type, size_t device_index);

  // 清除 pinned device 配置
  void ClearPinnedDevice();

  // 查询是否配置了 numa 亲和性
  bool HasNumaNodeAffinity() const;

  // 返回配置的 numa 亲和节点
  size_t GetNumaNodeAffinity();

  // 配置 numa 亲和节点
  void SetNumaNodeAffinity(size_t numa_node);

  // 清楚 nuam 亲和节点配置
  void ClearNumaNodeAffinity();
};

Device

Device 是设备的抽象

class Device {
 public:
  OF_DISALLOW_COPY_AND_MOVE(Device);
  Device() = default;
  virtual ~Device() = default;

  // 将当前设备设置为当前线程活跃设备，对标 cudaSetDevice 
  virtual void SetAsActiveDevice() = 0;

  // 设备属性
  
  // 设备类型
  virtual DeviceType device_type() const = 0;
  // 设置 index
  virtual size_t device_index() const = 0;
  // 设备对应的 DeviceManager 指针
  virtual DeviceManager* device_manager() const = 0;

  // Stream 管理
  // 创建一个 Stream，调用者必须在使用结束后调用 DestroyStream
  virtual Stream* CreateStream() = 0;
  // 销毁一个 Stream
  virtual void DestroyStream(Stream* stream) = 0;

  // 事件管理，如果事件的创建销毁开销比较大，Device 实现应考虑复用 Event
  
  // 创建一个 Event，用户必须在使用结束后销毁该 Event 
  virtual Event* CreateEvent();
  // 销毁一个 Event
  virtual void DestroyEvent(Event* event);
  // 一次性创建多个 Event
  virtual void CreateEvents(Event** events, size_t count) = 0;
  // 一次性销毁多个 Event
  virtual void DestroyEvents(Event** events, size_t count) = 0;

  // 内存管理
  
  // 在当前设备上分配内存
  virtual Maybe<void> Alloc(const AllocationOptions& options, void** ptr, size_t size) = 0;
  // 释放当前设备上分配的内存
  virtual void Free(const AllocationOptions& options, void* ptr) = 0;
  // 分配适用于和当前设备交互的 CPU 内存，如 cudaMallocHost
  virtual Maybe<void> AllocPinned(const AllocationOptions& options, void** ptr, size_t size) = 0;
  // 释放适用于和当前设备交互的 CPU 内存，如果 cudaFreeHost
  virtual void FreePinned(const AllocationOptions& options, void* ptr) = 0;
};

ActiveDeviceGuard

使用 RAII 维护当前线程当前设备类型的活跃设备

class ActiveDeviceGuard {
 public:
  OF_DISALLOW_COPY_AND_MOVE(ActiveDeviceGuard);
  // 构造时保存当前线程相同设备类型的活跃设备，并将活跃设备切换为 device
  explicit ActiveDeviceGuard(Device* device);
  // 析构时将活跃线程切换回保存的设备
  ~ActiveDeviceGuard();

 private:
  size_t saved_active_device_;
  DeviceManager* device_manager_;
};

DeviceManager

设备管理类，每个设备类型对应一个 DeviceManager

class DeviceManager {
 public:
  OF_DISALLOW_COPY_AND_MOVE(DeviceManager);
  DeviceManager() = default;
  virtual ~DeviceManager() = default;

  // 返回对应的 DeviceManagerRegistry 对象指针。
  virtual DeviceManagerRegistry* registry() const = 0;
  // 获取 device_index 对应的设备的 shared_ptr，多次使用同一个 device_index 调用应返回相同的对象，应实现为线程安全的。
  virtual std::shared_ptr<Device> GetDevice(size_t device_index) = 0;
  // 获取该类型设备的数量，primary_device_index 指示在执行底层查询前应设置的活跃设备，例如用于在 CUDA 环境中避免创建不必要的 CUDA Context
  virtual size_t GetDeviceCount(size_t primary_device_index) = 0;
  // 使用默认的活跃设备查询设备数量
  virtual size_t GetDeviceCount() = 0;
  // 查询当前活跃的设备 index 
  virtual size_t GetActiveDeviceIndex() = 0;
  // 设置 device_index 对应的设备为当前活跃设备
  virtual void SetActiveDeviceByIndex(size_t device_index) = 0;
};

DeviceManagerFactory

DeviceManager 的工厂，用于创建 DeviceManager

class DeviceManagerFactory {
 public:
  OF_DISALLOW_COPY_AND_MOVE(DeviceManagerFactory);
  DeviceManagerFactory() = default;
  virtual ~DeviceManagerFactory() = default;

  // 创建一个 DeviceManager
  virtual std::unique_ptr<DeviceManager> NewDeviceManager(DeviceManagerRegistry* registry) = 0;
  // 工厂对应的设备类型
  virtual DeviceType device_type() const = 0;
  // 工厂对应的设备名称，和 device_tag 接近
  virtual std::string device_type_name() const = 0;
  // 打印驱动、运行时库版本等信息
  virtual void DumpVersionInfo() const {}
};

DeviceManagerRegistry

DeviceManagerRegistry 的注册中心，一般通过 Global<ep::DeviceManagerRegistry>::Get()访问其单例对象，也是 EP 设备抽象的根入口

class DeviceManagerRegistry {
 public:
  OF_DISALLOW_COPY_AND_MOVE(DeviceManagerRegistry);
  DeviceManagerRegistry();
  ~DeviceManagerRegistry();

  // 获得 device_type 对应的 DeviceManager 对象指针
  DeviceManager* GetDeviceManager(DeviceType device_type);
  // 获得 device_type 和 device_index 对应的设备指针，等价于 GetDeviceManager 和 GetDeviceManager::GetDevice 的组合
  std::shared_ptr<Device> GetDevice(DeviceType device_type, size_t device_index);

  // 注册一个 DeviceManagerFactory
  static void RegisterDeviceManagerFactory(std::unique_ptr<DeviceManagerFactory>&& factory);
  // 打印全部的设备对应的驱动运行时库版本等信息
  static void DumpVersionInfo();
  // device_type 到 device_name 的转换
  static std::string GetDeviceTypeNameByDeviceType(DeviceType device_type);
  // device_name 到 device_type 的转换
  static DeviceType GetDeviceTypeByDeviceTypeName(const std::string& device_type_name);

 private:
  class Impl;
  std::unique_ptr<Impl> impl_;
};