TVM: Device/Target Interactions

任何新的运行时环境都必须实现三个主要方面:

DeviceAPI类为特定设备提供了一个句柄，以及用于与之交互的API。它定义了一个通用接口，用于查询设备参数(例如可用内存、线程数量等)和执行简单操作(例如从主机复制内存，或在设备的缓冲区之间复制内存)。
Target类包含函数将在其上运行的设备的描述。它同时暴露给目标代码生成器和优化pass。
目标代码生成器从IRModule构造一个由一个或多个PackedFunc组成的Module。

DeviceAPI

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

属性查询(Attribute queries)——GetAttr允许查询设备的各种参数，如设备名称、线程数等。可查询的参数定义在device_api.h中的DeviceAttrKind enum中。并非所有设备都支持该enum中定义的所有可查询参数。如果一个参数不能被查询(例如Vulkan上的kMaxClockRate)，或者一个参数不适用(例如CPU上的kWarpSize)，那么这些查询应该返回nullptr。
设置活动设备(Setting active device)—SetDevice应该将特定设备设置为活动设备。如果由特定目标的代码生成的PackedFunc需要在一个设备上执行，它应该在活动设备上运行。
内存管理(Memory management)——用于在设备上分配和释放内存的实用程序。

分配数据空间(Allocate data space)— AllocDataSpace和FreeDataSpace用于分配和释放设备上的空间。这些分配的空间可以提供给算子的输入和输出，并构成算子图的主要数据流。必须能够在主机和设备数据空间之间传输数据。返回一个不透明的void。虽然有些实现返回内存地址，但这不是必需的，而且void可能是一个不透明的句柄，只有生成它的设备后端才能解释它。void*被用作其他后端特定函数的参数，例如CopyDataFromTo。

分配工作空间(Allocate work space)—AllocWorkspace和FreeWorkspace分别分配和释放设备上的空间。与数据空间不同的是，它们用于存储算子定义中的中间值，并且不要求可在主机和设备间转移。如果DeviceAPI子类没有实现这些方法，它们将默认调用父类相应的DataSpace函数。

复制数据(Copy data)—CopyDataFromTo应该将数据从一个位置复制到另一个位置。复制的类型由dev_from和dev_to参数确定。实现应该支持在单个设备上将内存从CPU复制到设备、从设备复制到CPU、从一个缓冲区复制到另一个缓冲区。如果源或目标位置在CPU上，对应的void指向一个CPU地址，该地址可以传递给memcpy。如果源或目标位置在设备上，则相应的void是此前由AllocDataSpace或AllocWorkspace生成的。

这些拷贝在特定的TVMStreamHandle上排队执行。但是，实现不应该假设CPU缓冲区在CopyDataFromTo调用完成后仍然有效或可访问。

执行流管理(Execution stream management)——用于处理TVMStreamHandle的实用程序，它表示用于执行命令的并行执行流。

创建流(Create stream)——CreateStream和FreeStream应该为执行流分配/释放一个句柄。如果一个设备只实现一个命令队列，那么CreateStream应该返回nullptr。

设置活动流(Set active stream)—SetStream应该将一个流设置为活动流。在活动时，如果特定目标的代码生成器生成的PackedFunc需要在设备上执行，那么工作应该提交给活动流。

同步到CPU(Synchronize to CPU)——StreamSync应该同步一个执行流到CPU。对StreamSync的调用应该在StreamSync完成之前，返回提交的所有内存传输和计算。

流之间的同步(Synchronize between streams)——SyncStreamFromTo应该在源流和目标流之间引入同步栅栏。也就是说，在源流完成当前队列中的所有命令之前，目标流不能超越队列中的当前命令。

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

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

FooDeviceAPI* FooDeviceAPI::Global() {

  static FooDeviceAPI inst;

  return &inst;

}

将该函数注册到tvm注册表:

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

(1) 在c_runtime_api.h的TVMDeviceExtType enum中添加新DeviceAPI的条目。该值应该是一个大于DLDeviceType::kDLExtDev但小于DeviceAPIManager::kMaxDeviceAPI的未使用的值。

(2) 在device_api.h中的DeviceName中添加条件分支，将enum值转换为对应字符串。这个字符串表示应该匹配对应的TVM_REGISTER_GLOBAL注册名称。

(3) 在tmv .runtime. device的MASK2STR和STR2MASK字典中，为新的enum值添加条目。

Target Definition

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

在target_kind.cc中，添加TVM_REGISTER_TARGET_KIND的新声明，传递新目标的字符串名称，以及运行该目标的设备的TVMDeviceExtType或DLDeviceType enum值。通常，目标名称和设备名称是匹配的(例如，“cuda”目标运行在kDLCUDA设备上)。 但也有例外，比如多个不同的代码生成目标可以在同一个物理设备上运行(例如"llvm"和"c"目标都运行在kDLCPU设备类型上)。

特定目标类型的所有选项都是用add_attr_option函数添加的，带有可选的默认值。可以使用set_attrs_preprocessor添加预处理器，以定义任何参数，这些参数是动态地基于其他参数或从设备属性查询的。

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

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

Target Code Generators

代码生成器接受一个优化的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中可以通过名称访问。