DeviceAPI#

DeviceAPI 表示对特定硬件设备 API 的句柄。（例如，CUDADeviceAPI 处理通过 CUDA 框架进行的所有交互。）大多数 DeviceAPI 方法接受 device_id 参数，用于指定应访问哪个设备。在 Python 中，通常使用 tvm.runtime.device() 函数来访问这些设备，该函数返回通过特定 API 访问的特定设备的句柄。（例如，tvm.runtime.device('cuda',0) 提供对物理设备 0 的访问，该设备通过 CUDA API 进行访问。）

• 属性查询 - GetAttr 允许查询不同的设备特定参数，例如设备名称、线程数量等。可查询的参数定义在 device_api.h 中的 enum DeviceAttrKind 枚举中。并非所有设备都支持所有可查询的参数。如果某个参数无法查询（例如 Vulkan 上的 kMaxClockRate），或者某个参数不适用（例如 CPU 上的 kWarpSize），则这些查询应返回 nullptr。

• 设置活动设备 - SetDevice 应将特定设备设置为活动状态。如果由目标特定代码生成器生成的 PackedFunc 需要在设备上执行，则它应在活动设备上运行。

• 内存管理 - 提供在设备上分配和释放内存的工具。

  • 分配数据空间 - AllocDataSpace 和 FreeDataSpace 用于在设备上分配和释放空间。这些分配可以作为算子的输入和输出，并构成算子图的主要数据流。必须能够将数据从主机传输到数据空间或从数据空间传输到主机。返回值是不透明的 void*。虽然某些实现返回内存地址，但这并非必需，void* 可能是仅由生成它的设备后端解释的不透明句柄。void* 用作其他后端特定函数（例如 CopyDataFromTo）的参数。

  • 分配工作空间 - AllocWorkspace 和 FreeWorkspace 用于在设备上分配和释放空间。与数据空间不同，这些空间用于存储算子定义中的中间值，并且不需要能够在主机设备之间传输。如果 DeviceAPI 的子类未实现这些方法，它们将默认调用相应的 DataSpace 函数。

  • 复制数据 - CopyDataFromTo 应将数据从一个位置复制到另一个位置。复制的类型由 dev_from 和 dev_to 参数决定。实现应支持从 CPU 复制到设备、从设备复制到 CPU，以及在单个设备上的缓冲区之间复制数据。如果源或目标位置在 CPU 上，则相应的 void* 指向可以传递给 memcpy 的 CPU 地址。如果源或目标位置在设备上，则相应的 void* 先前由 AllocDataSpace 或 AllocWorkspace 生成。

    这些复制操作会被排队在特定的 TVMStreamHandle 上执行。然而，实现不应假设在调用 CopyDataFromTo 完成后，CPU 缓冲区仍然有效或可访问。

• 执行流管理 - 用于处理 TVMStreamHandle 的工具，TVMStreamHandle 表示用于执行命令的并行执行流。

  • 创建流 - CreateStream 和 FreeStream 应分配/释放一个执行流的句柄。如果设备仅实现单个命令队列，则 CreateStream 应返回 nullptr。

  • 设置活动流 - SetStream 应将流设置为活动状态。在活动状态下，如果由目标特定代码生成器生成的 PackedFunc 需要在设备上执行，则工作应提交到活动流中。

  • 与 CPU 同步 - StreamSync 应将执行流与 CPU 同步。在 StreamSync 调用之前提交的所有内存传输和计算完成后，StreamSync 的调用才会返回。

  • 流之间的同步 - SyncStreamFromTo 应在源流和目标流之间引入同步屏障。也就是说，在源流完成当前排队的所有命令之前，目标流不得继续执行当前排队的命令。

为了能够被 TVM 框架使用，新的 DeviceAPI 应按照以下步骤进行注册。

1. 创建函数来实例化新的 DeviceAPI，并返回指向它的指针:

   FooDeviceAPI* FooDeviceAPI::Global() {
     static FooDeviceAPI inst;
     return &inst;
   }
   

2. 将该函数注册到 tvm 注册表中:

   TVM_REGISTER_GLOBAL("device_api.foo").set_body_typed(FooDeviceAPI::Global);
   

1. 在 c_runtime_api.h 中的 TVMDeviceExtType 枚举中添加新 DeviceAPI 的条目。该值应为大于 DLDeviceType::kDLExtDev 但小于 DeviceAPIManager::kMaxDeviceAPI 的未使用值。

2. 在 device_api.h 中的 DeviceName 函数中添加 case，将枚举值转换为字符串表示形式。该字符串表示形式应与 TVM_REGISTER_GLOBAL 中给出的名称匹配。

3. 为新的枚举值在 tvm.runtime.Device 的 MASK2STR 和 STR2MASK 字典中添加条目。

目标定义#

Target 对象是关于物理设备、其硬件/驱动程序限制及其功能的属性查找表。Target 在优化和代码生成阶段均可访问。虽然所有运行时目标都使用相同的 Target 类，但每个运行时目标可能需要添加特定于目标的选项。

在 target_kind.cc 中，添加新的 TVM_REGISTER_TARGET_KIND 声明，传递新目标的字符串名称，以及该目标应运行的设备的 TVMDeviceExtType 或 DLDeviceType 枚举值。通常，目标名称和设备名称会匹配。（例如，"cuda" 目标在 kDLCUDA 设备上运行。）但也有例外情况，例如当多个不同的代码生成目标可以在同一物理设备上运行时。（例如，"llvm" 和 "c" 目标都在 kDLCPU 设备类型上运行。）

特定目标类型的所有选项都通过 add_attr_option 函数添加，并可选择提供默认值。可以通过 set_target_parser 添加 Target 解析器，以处理基于其他参数动态生成或从设备属性查询的参数。

此参数定义了解析器，可以解包目标的字符串描述。这是在 C++ 的 Target::Target(const String&) 构造函数中完成的，该构造函数接受 JSON 格式的字符串，通常通过 tvm.target.Target Python 对象调用。例如，tvm.target.Target('{"kind": "cuda", "max_num_threads": 1024}') 将创建 cuda 目标，同时覆盖默认的最大线程数。

在代码生成器中，可以使用 C++ 中的 target->GetAttr<T>(param_name) 或 Python 中的 target.attrs 字典访问目标属性。

目标代码生成器#

代码生成器接收优化后的 IRModule 并将其转换为可执行的表示形式。每个代码生成器都必须注册才能被 TVM 框架使用。这是通过注册名为 "target.build.foo" 的函数来完成的，其中 foo 与上述 TVM_REGISTER_TARGET_KIND 定义中使用的名称相同。:

tvm::runtime::Module GeneratorFooCode(IRModule mod, Target target);
TVM_REGISTER_GLOBAL("target.build.foo").set_body_typed(GeneratorFooCode);

代码生成器接受两个参数。第一个是要编译的 IRModule，第二个是描述代码应运行的设备的 Target。由于执行编译的环境不一定与执行代码的环境相同，代码生成器不应在设备本身上执行任何属性查找，而应访问存储在 Target 中的参数。

输入 IRModule 中的每个函数都应通过名称在输出 runtime::Module 中可访问。