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) AllocDataSpaceFreeDataSpace用于分配和释放设备上的空间。这些分配的空间可以提供给算子的输入和输出,并构成算子图的主要数据流。必须能够在主机和设备数据空间之间传输数据。返回一个不透明的void。虽然有些实现返回内存地址,但这不是必需的,而且void可能是一个不透明的句柄,只有生成它的设备后端才能解释它。void*被用作其他后端特定函数的参数,例如CopyDataFromTo。
  • 分配工作空间(Allocate work space)AllocWorkspaceFreeWorkspace分别分配和释放设备上的空间。与数据空间不同的是,它们用于存储算子定义中的中间值,并且不要求可在主机和设备间转移。如果DeviceAPI子类没有实现这些方法,它们将默认调用父类相应的DataSpace函数。
  • 复制数据(Copy data)CopyDataFromTo应该将数据从一个位置复制到另一个位置。复制的类型由dev_from和dev_to参数确定实现应该支持在单个设备上将内存从CPU复制到设备、从设备复制到CPU、从一个缓冲区复制到另一个缓冲区。如果源或目标位置在CPU上,对应的void指向一个CPU地址,该地址可以传递给memcpy。如果源或目标位置在设备上,则相应的void是此前由AllocDataSpaceAllocWorkspace生成的。
    这些拷贝在特定的TVMStreamHandle上排队执行。但是,实现不应该假设CPU缓冲区在CopyDataFromTo调用完成后仍然有效或可访问。
  • 执行流管理(Execution stream management)——用于处理TVMStreamHandle的实用程序,它表示用于执行命令的并行执行流。
  • 创建流(Create stream)——CreateStreamFreeStream应该为执行流分配/释放一个句柄。如果一个设备只实现一个命令队列,那么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;
}
  1. 将该函数注册到tvm注册表:
TVM_REGISTER_GLOBAL("device_api.foo").set_body_typed(FooDeviceAPI::Global);

(1) 在c_runtime_api.hTVMDeviceExtType 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的新声明,传递新目标的字符串名称,以及运行该目标的设备的TVMDeviceExtTypeDLDeviceType 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中可以通过名称访问

posted @ 2022-07-31 12:22  牛犁heart  阅读(443)  评论(0编辑  收藏  举报