news 2026/4/17 20:44:57

【TVM教程】设备/目标交互

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
【TVM教程】设备/目标交互

TVM 现已更新到 0.21.0 版本,TVM 中文文档已经和新版本对齐。

Apache TVM 是一个深度的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。

在线运行 TVM 学习教程

链接是:https://hyper.ai/notebooks/48919?utm_source=Distribute&utm_medium=Distribute-TVM&utm_campaign=Distribute-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”}。

DeviceAPI

DeviceAPI(设备 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。)

为了使 TVM 能够使用新的 DeviceAPI,需要执行以下注册步骤:

  1. 创建一个实例化 DeviceAPI 并返回其指针的函数:
FooDeviceAPI* FooDeviceAPI::Global() { static FooDeviceAPI inst; return &inst; }
  1. 在 TVM 注册表中注册:
TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("device_api.foo", FooDeviceAPI::Global); }
  1. 在 base.h 的TVMDeviceExtType枚举中为新的 DeviceAPI 添加条目。值需大于DLDeviceType::kDLExtDev,且小于DeviceAPIManager::kMaxDeviceAPI
  2. 在 device_api.h 的DeviceName中添加对应枚举 → 字符串映射,该字符串需与GlobalDef().def中一致。
  3. 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 名称,以及对应运行设备的TVMDeviceExtTypeDLDeviceType。通常情况下,target 名称和设备名称一致(如"cuda"运行于kDLCUDA),但也有例外(例如"llvm""c"目标都运行于kDLCPU)。

所有 target 选项通过add_attr_option添加,可带默认值。可以使用set_target_parser添加解析器,用于处理依赖其他参数或硬件属性的动态参数。

该参数解析器定义了如何从字符串格式构造 target。这由Target::Target(const String&)构造函数执行,该构造函数接受 JSON 格式字符串,通常通过 Python:

tvm.target.Target('{"kind": "cuda", "max_num_threads": 1024}')

在代码生成器中,可通过以下方式访问 target 属性:

Target 代码生成器

代码生成器将优化后的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中可通过名称访问。

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2026/4/18 3:30:51

学生数据统计不再头疼,学工一体化平台助力辅导员高效减负

✅作者简介&#xff1a;合肥自友科技 &#x1f4cc;核心产品&#xff1a;智慧校园平台(包括教工管理、学工管理、教务管理、考务管理、后勤管理、德育管理、资产管理、公寓管理、实习管理、就业管理、离校管理、科研平台、档案管理、学生平台等26个子平台) 。公司所有人员均有多…

作者头像 李华
网站建设 2026/4/18 3:37:43

<span class=“js_title_inner“>全栈“进化”公开课邀您参加|文心Moment大会走进大模型高效微调与极致推理全栈工程实践</span>

大模型效果越来越强&#xff0c;为什么真正上线却越来越难&#xff1f;从训练到推理、从算力到部署&#xff0c;工程能力正在成为决定成败的关键。从“模型可用”到“工程好用”&#xff0c;大模型正在经历一场真正的全栈进化。 1月22日下午13&#xff1a;00&#xff0c;全栈“…

作者头像 李华
网站建设 2026/4/18 3:30:46

Python入门篇【异常】

Python【异常】 文章目录Python【异常】一、Bug定义二、捕获异常2.1.基本捕获异常2.2.指定捕获异常提示&#xff1a;以下是本篇文章正文内容&#xff0c;下面案例可供参考 一、Bug定义 早期计算机采用大量继电器工作&#xff0c;马克二型计算机就是这样的。1945年9月9日&#…

作者头像 李华
网站建设 2026/4/18 3:36:23

Thread.sleep()与Object.wait()的区别解析

文章目录Thread.sleep() 与 Object.wait() 的区别解析一、引言&#xff1a;线程的基本操作二、Thread.sleep() 的详解1. 基本概念2. 示例代码3. 核心特点三、Object.wait() 的详解1. 基本概念2. 示例代码3. 核心特点四、Thread.sleep() 和 Object.wait() 的区别1. 调用方式2. 锁…

作者头像 李华
网站建设 2026/4/18 3:33:50

2025年热门的AI Agent方向及国内外代表产品

2025年作为Agent发展元年&#xff0c;Manus敲响了AI 应用大发展的号角&#xff0c;各种Agent应用层出不穷。这些智能体基本框架是怎样的&#xff0c;哪些更有前景&#xff1f; 最近在海外看到一张「2025 年 AI Agent 技术趋势全景图」&#xff0c;觉得非常精彩&#xff0c;分享…

作者头像 李华