TVM 现已更新到 0.21.0 版本TVM 中文文档已经和新版本对齐。Apache TVM 是一个深度的深度学习编译框架适用于 CPU、GPU 和各种机器学习加速芯片。在线运行 TVM 学习教程链接是https://hyper.ai/notebooks/48919?utm_sourceDistributeutm_mediumDistribute-TVMutm_campaignDistribute-TVM-260126本文档面向希望了解 TVM 框架如何与特定设备 API 进行交互的开发者或希望为新的 API 或新硬件添加支持的开发者。对于任何新的运行时环境需要实现三个主要部分DeviceAPI tvm-target-specific-device-api{.interpreted-text role“ref”} 类提供对特定设备的句柄以及用于与其交互的 API。它定义了一套通用接口用于查询设备参数例如可用内存、线程数量等以及执行简单操作例如从主机复制内存或在设备缓冲区之间复制数据。Target tvm-target-specific-target{.interpreted-text role“ref”} 类包含将要运行函数的设备描述。它同时暴露给目标代码生成器和优化 Pass。目标代码生成器 tvm-target-specific-codegen{.interpreted-text role“ref”} 从 IRModule 构建一个由一个或多个PackedFunc tvm-runtime-system-packed-func{.interpreted-text role“ref”} 组成的Module tvm-runtime-system-module{.interpreted-text role“ref”}。DeviceAPIDeviceAPI设备 API表示对特定硬件设备 API 的访问句柄。例如CUDADeviceAPI处理所有通过 CUDA 框架的交互。大多数DeviceAPI方法都接受一个device_id参数用于指定访问哪个设备。在 Python 中通常使用tvm.runtime.device{.interpreted-text role“py:func”} 函数访问特定设备该函数返回指定 API 所访问设备的句柄。例如tvm.runtime.device(cuda, 0)表示访问通过 CUDA API 访问的物理设备0。属性查询—GetAttr用于查询不同的设备特定参数例如设备名称、线程数量等。可查询的参数定义在enum DeviceAttrKind文件位置 device_api.h。 并非所有参数都适用于所有设备。如果某个参数无法查询例如 Vulkan 上的kMaxClockRate或不适用例如 CPU 上的kWarpSize应返回nullptr。设置活动设备—SetDevice应将某个设备设置为当前活动设备。如果目标代码生成器生成的PackedFunc需要在设备上执行该执行应发生在当前活动设备上。内存管理— 用于在设备上分配和释放内存的工具函数。分配数据空间—AllocDataSpace和FreeDataSpace用于在设备上分配和释放数据存储空间。这些空间可作为算子输入和输出并构成算子图的主要数据流。必须支持主机与数据空间之间的数据传输。返回值为不透明指针void*。某些实现返回真实地址但这不是必须的该指针也可能是仅可由设备后端解释的句柄。该void*将作为参数传递给其他后端函数例如CopyDataFromTo。分配工作空间—AllocWorkspace和FreeWorkspace用于分配和释放工作区。这些区域用于算子内部中间值存储不要求可与主机传输。如果子类未实现则默认调用对应的数据空间分配函数。数据复制—CopyDataFromTo应在不同位置之间复制数据。复制类型由dev_from和dev_to决定。实现应该支持将内存从CPU复制到设备从设备复制到CPU以及在单个设备上从一个缓冲区复制到另一个缓冲区。如果源或目标位于 CPU则指针为可直接用于memcpy的主机地址如果位于设备则指针必定由AllocDataSpace或AllocWorkspace生成。这些复制会排队在某个TVMStreamHandle流中执行。但是实现不应假设 CPU 缓冲区在函数返回后仍然有效或可访问。执行流管理— 管理TVMStreamHandle执行命令的并行流。创建流—CreateStream/FreeStream负责分配和释放执行流。如果设备只有单一指令队列则CreateStream应返回nullptr。设置活动流—SetStream用于将某个流设置为当前活跃流。目标代码生成器生成的函数执行时应提交到该流。同步到 CPU—StreamSync应同步流使之在执行完成前阻塞返回。流间同步—SyncStreamFromTo应在两个流之间插入同步屏障使目标流在源流执行完当前排队命令前无法继续执行。为了使 TVM 能够使用新的 DeviceAPI需要执行以下注册步骤创建一个实例化 DeviceAPI 并返回其指针的函数FooDeviceAPI* FooDeviceAPI::Global() { static FooDeviceAPI inst; return inst; }在 TVM 注册表中注册TVM_FFI_STATIC_INIT_BLOCK() { namespace refl tvm::ffi::reflection; refl::GlobalDef().def(device_api.foo, FooDeviceAPI::Global); }在 base.h 的TVMDeviceExtType枚举中为新的 DeviceAPI 添加条目。值需大于DLDeviceType::kDLExtDev且小于DeviceAPIManager::kMaxDeviceAPI。在 device_api.h 的DeviceName中添加对应枚举 → 字符串映射该字符串需与GlobalDef().def中一致。在tvm.runtime.Device的_DEVICE_TYPE_TO_NAME与_DEVICE_NAME_TO_TYPE字典中添加对应映射。Target 定义Target对象是有关物理设备、其硬件/驱动限制和能力的属性查询表。Target可在优化阶段和代码生成阶段使用。虽然所有运行时共享相同的Target类但不同运行时可能需要额外的 target 特定属性。在 target_kind.cc 中使用TVM_REGISTER_TARGET_KIND注册新的 target需传入 target 名称以及对应运行设备的TVMDeviceExtType或DLDeviceType。通常情况下target 名称和设备名称一致如cuda运行于kDLCUDA但也有例外例如llvm与c目标都运行于kDLCPU。所有 target 选项通过add_attr_option添加可带默认值。可以使用set_target_parser添加解析器用于处理依赖其他参数或硬件属性的动态参数。该参数解析器定义了如何从字符串格式构造 target。这由Target::Target(const String)构造函数执行该构造函数接受 JSON 格式字符串通常通过 Pythontvm.target.Target({kind: cuda, max_num_threads: 1024})在代码生成器中可通过以下方式访问 target 属性Ctarget-GetAttrT(param_name)Pythontarget.attrsTarget 代码生成器代码生成器将优化后的IRModule转换为可执行表示。每个代码生成器必须注册到 TVM 框架中其名称为target.build.foo其中foo与先前TVM_REGISTER_TARGET_KIND中的名称一致。示例tvm::runtime::Module GeneratorFooCode(IRModule mod, Target target); TVM_FFI_STATIC_INIT_BLOCK() { namespace refl tvm::ffi::reflection; refl::GlobalDef().def(target.build.foo, GeneratorFooCode); }代码生成器有两个参数。第一个是要编译的IRModule第二个是描述代码应该运行在哪个设备上的目标Target。由于编译环境不一定与执行环境相同因此代码生成器不应直接向设备查询属性而应始终使用Target中的属性。输入IRModule中的每个函数都应在输出的runtime::Module中可通过名称访问。