设备/目标交互

本文档旨在帮助有兴趣了解 TVM 框架如何与特定设备 API 交互的开发者，或者希望为新的 API 或新硬件实现支持的开发者。

对于任何新的运行时环境，必须实现三个主要方面。

  • DeviceAPI 类提供了特定设备的句柄，以及用于与其交互的 API。它为查询设备参数（例如，可用内存、线程数等）和执行简单操作（例如，从主机复制内存，或在设备上的缓冲区之间复制内存）定义了一个通用接口。

  • Target 类包含将要运行函数的设备的描述。它同时暴露给目标代码生成器和优化 pass。

  • 目标代码生成器 从 IRModule 构建一个 Module，该模块由一个或多个 PackedFunc 组成。

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 上执行。但是，实现不应假设 CPU 缓冲区在调用 CopyDataFromTo 完成后仍然有效或可访问。

  • 执行流管理 - 用于处理 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 中访问。