由林振斌（华为）著

OpenReg 是一个使用 PyTorch 核心框架的“PrivateUse1”机制的树外后端实现的独立演示。该实现有两个主要目的：

  1. 参考实现：为第三方设备厂商通过 PrivateUse1 与 PyTorch 集成提供了一个实用的模板。
  2. CI 测试基础设施：为持续集成管道提供设备无关的测试能力。

使用说明

模块安装

cd {project}/test/cpp_extensions/open_registration_extension
python setup.py install

用例

import torch
import pytorch_openreg

if __name__ == "__main__":
   print(torch.ones(1, 2, device='openreg'))

架构概述

流程管理

OpenReg 通过启动 N 个独立的子进程实现虚拟设备隔离，每个子进程维护专用的请求/响应队列以进行进程间通信。父进程驱动程序将设备操作封装到命令包中，这些命令包：

  1. 通过请求队列发送到目标设备
  2. 异步处理，并通过响应队列返回结果

Parent-Subprocess Communication Flow

图：父进程与子进程通信流程

内存管理

设备内存分配在各个子进程中发生，以确保：

  1. 设备间严格的内存隔离
  2. 物理设备约束的逼真模拟

组件分解

_aten_impl.py

此模块承担双重职责：

  1. 悬钩注册：
    • 使用 _IMPL_REGISTRY 绑定 C++ 后端钩子（例如，getDevice、getStream）到设备驱动实现
  2. 回退机制：
    • 定义一个新的 torch.Library ，用于注册回退，当调用 PrivateUse1 的后端内核时将被调用。它包含处理所有类型本地函数的逻辑，计算输出元数据，分配它，并且仅调用设备守护进程以执行计算

_device_daemon.py

核心子系统

  1. 分配器：
    • HostAllocator : 管理父进程中的固定内存
    • 处理具有张量重建功能的设备内存
  2. 驱动程序（父进程）：
    • 维护设备上下文（活动设备/流）
    • 实现设备控制操作：
      • 设置设备/获取设备
      • 设备数量
      • 交换流
    • 通过基于队列的 IPC orchestrate 命令执行
  3. 执行器（子进程）:
    • 处理命令类型：
      • 内存操作（ malloc / free ）
      • 张量计算（ run_op ）
      • 数据传输（ send_data / recv_data ）
      • 流/事件管理（主要由于 CPU 同步特性而无需操作）

_meta_parser.py

关键特性：

  • 实现跨进程对象传输的序列化工具
  • OpenRegTensorMeta 类封装了以下完整的张量元数据：
    • 输出张量重建
    • 设备端计算准备

设计考虑

执行特性

  • 同步计算：CPU 运算符执行需要同步处理
  • 流/事件语义：由于同步执行模型，实现为空操作
  • 内存隔离：通过子进程分配强制执行严格的按设备内存边界

此架构能够在保持 PyTorch 兼容性的同时，通过标准后端接口实现设备集成的真实模拟。