TVM自定义修改代码示例

最新推荐文章于 2025-02-11 10:52:16 发布
最新推荐文章于 2025-02-11 10:52:16 发布
阅读量1.2k 收藏 1
点赞数
CC 4.0 BY-SA版权
分类专栏： AICompiler&AIFramework 人工智能训练与推理 智能推理 文章标签： python 开发语言 后端
吴建明
本文链接：https://blog.csdn.net/wujianing_110117/article/details/121280208
本文介绍了如何在TVM中添加一个新的设备类型"DPU"，涉及从补全Python接口到代码生成的全过程。首先，通过修改Python的target.py和runtime的runtime_ctypes.py文件，添加了DPU的设备类型。接着，深入C++代码，实现TargetCreate函数，维护IR Pass变换，并在codegen模块中添加DPU代码生成支持。最后，调整CMake配置以确保新设备的编译选项。文章还简要讨论了TVM代码生成流程，包括从Tensor Expression到代码生成的步骤。

摘要生成于 C知道 ，由 DeepSeek-R1 满血版支持， 前往体验 >

TVM自定义修改代码示例
一．TVM设备添加代码
因为要添加的设备是一种类似于GPU的加速卡，TVM中提供了对GPU编译器的各种支持，有openCl，OpenGL和CUDA等，这里选取比较熟悉的CUDA进行模仿生成。总体上看，TVM是一个多层的结构。
TVM在python这一层提供了相关的设备接口，然后使用tvm.build真正的编译，然后调用get_source函数来获得想要的源码（或者IR，如llvm选项提供的是LLVM的IR，或者PTX选项提供的就是NVPTX类型的IR）。
因此，添加新设备(device)推测的步骤就是：

  1. 补全相应的python接口

  2. 找到python和C交互的接口

  3. 正确维护中间代码的IR pass变换中新设备引入的特性

  4. 代码生成对新设备和新特性的支持

  5. 添加编译选项支持(非必须)
    以下就分别就这4个步骤进行介绍。

  6. 补全相应的python接口
    测试代码中使用的是字符串解析，但是从其它tutorial中发现，还存在一种tvm.target.cuda()的设备建立方式，这个很明显比字符串解析，相对找起来容易（字符串最终对应的也是这种方式）。按照这种方式找到了tvm/python/tvm/target.py文件中，这个类中定义了现在能支持的target。添加新的target叫做dpu。
    def dpu(model=‘unknown’, options=None):

    “”"Returns a dpu target.
    Parameters

    model: str
    The model of dpu device
    options : str or list of str
    Additional options
    “”"

    opts = _merge_opts([’-model=%s’ % model], options)
    return _api_internal._TargetCreate(“dpu”, *opts)
    每个设备都包括硬件自身的上下文信息和硬件上运行软件运行时，就是runtime，在TVM中相关的软件运行时信息在tvm/python/tvm/_ffi/runtime_ctypes.py文件中，添加对dpu的支持
    在class TVMContext的两个掩码MASK2STR和STR2MASK中分别添加：
    13: ‘dpu’,
    和
    ‘dpu’:13,

  7. 找到python和C交互的接口
    回到刚才的target.py文件中，核心的代码只有两句
    opts = _merge_opts([’-model=%s’ % model], options)
    return _api_internal._TargetCreate(“dpu”, *opts)
    第一句是将model和相关的options组合在一起，就是个字符串相关的拼接，没有特别多需要关注的内容，后边有一个_api_internel._TargetCreate的函数调用，从名字上看起来非常的重要，是创建真正的Target的，但是，在tvm/python文件中，无论如何都找不到该函数的实现。
    在这里插入图片描述

前边已经提到过TVM中使用的是python提供接口，真正的实现都是在C++中，因此，这里猜测是调用了C语言的实现。下面列一下TVM相关的文件夹
在这里插入图片描述

3rdparty是很多第三方库的实现
build 目录是建立的编译后的.so文件所在的位置
docs 是相关的文档
include C++代码的include的主目录
jvm 是java相关的文件夹
nnvm 是中间的nnvm算子所在的目录
python 是python文件所在的目录，所有与python相关的都在该目录中
rust apps conda docker golang web verilog都是特有领域中的内容，对一般项目没有影响
tests 是测试文件，中间包含了作者写的很多测试，是学习TVM的另一个手段
Tutorial是官网上相关的历程
vta 是TVM的软件栈
cmake包含了所有的编译配置文件，和CmakeLists.txt共同工作
src 是全部的C++代码
topi 是Tensor Operator Index Library，后续进行详细介绍
在src目录下搜索_TargetCreate，得到src/codegen/build_module.cc:116中有相关的内容
TVM_REGISTER_API("_TargetCreate")
.set_body([](TVMArgs args, TVMRetValue* ret) {
std::string target_name = args[0];
std::vectorstd::string options;
for (int i = 1; i < args.num_args; ++i) {
std::string arg = args[i];
options.push_back(arg);
}
ret = CreateTarget(target_name, options);
});
这段代码就是通过一种TVM_REGISTER_API的注册机制，注册_TargetCreate函数，真正的函数体是.set_body内执行的，实际上C++中tvm::CreateTarget函数。TVM_REGISTER_API的注册机制在TVM项目中非常普遍，其实现在项目中也有，不是主要的研究内容，不需要改，所以不另行赘述。
3. 正确维护中间代码的IR pass变换中新设备引入的特性
在src/codegen/build_module.cc文件中的tvm::CreateTarget函数中添加对dpu的支持
else if (target_name == “dpu”) {
t->device_type = kDLDPU;
}
这里边的kDLDPU是一个DLDeviceType类型值，实现是在3rdparty/dlpack/include/dlpack/dlpack.h中添加的
kDLDPU =13,
在include/tvm/runtime/device_api.h：200补充对kDLDPU的支持
case kDLDPU: return “dpu”;
Target部分添加完了，还需要补充运行时的内容。
运行时的内容在src/runtime/目录下，需要在module.cc中添加对dpu 的支持。
在RuntimeEnabled函数中，添加
else if (target == “dpu”) {
f_name = “device_api.dpu”;
}
这只是添加了一个名字的支持，需要新建一个dpu目录，里边存放DPUModuleNode、DPUWorkspace等支持，测试代码的getSource函数的真正实现，存放在这里边，主要模仿CUDA和openCl的实现进行。目前存放有dpu_common.h、dpu_device_api.cc、dpu_module.cc、dpu_module.h四个文件，大概1K行代码，实现逻辑不是很复杂。
4. 代码生成对新设备和新特性的支持
上边准备好了module部分，也就是运行时，但是这里第一步想要实现的是一个能在dpu编译器上运行的C代码。需要在codegen部分添加对dpu这个设备的支持。
codegen是在tvm.build（Python）中形成的，在其对应的C++实现上是codegen/build_module.cc文件，之前添加了名字的支持，现在还需要添加这个真正的Target调用点
Target DPU(const std::vectorstd::string& options ) {
return CreateTarget(“dpu”, options);
}
最主要的codegen对DPU的支持是新建CodeGenDPU类，这个类的实现在该目录的codegen_dpu.h和codegen_dpu.cc文件内。其它的函数可以不实现，有两个函数必须实现
runtime::Module BuildDPU(Array funcs) {
using tvm::runtime::Registry;
bool output_ssa = false;
CodeGenDPU cg;
cg.Init(output_ssa);
for (LoweredFunc f : funcs) {
cg.AddFunction(f);
}
std::string code = cg.Finish();
if (const auto f = Registry::Get(“tvm_callback_dpu_postproc”)) {
code = (*f)(code).operator std::string();
}
return DPUModuleCreate(code, “dpu”, ExtractFuncInfo(funcs), code);
}

TVM_REGISTER_API(“codegen.build_dpu”)
.set_body([](TVMArgs args, TVMRetValue* rv) {
*rv = BuildDPU(args[0]);
});
5. 添加编译选项支持
上边可以说是完成了从设备添加到代码生成的部分，但是如果只有上边，新添加的设备一直无法运行。但如果仅是对一个设备进行修改的话，这部分没有必要。后来排查发现是部分代码未编译进去导致的。所以开始修改cmake配置。
在上一个TVM调试文档中提到，编译需要打开LLVM和CUDA选项，这里新添加了dpu的设备，需要增加一个新的编译选项，在cmake/config.cmake中添加
1
2 #Build DPU
set(USE_DPU ON)
cmake目录下存在着modules和util目录，modules是指定了相关设备的目录等配置，util文件夹下的内容，用来寻找如CUDA等的配置。暂时只需要modules下添加DPU.cmake，这部分的配置代码相对比较简单，就是指定runtime对应的目录。
1
2
3
4
5
6

最低0.47元/天 解锁文章

200万优质内容无限畅学
如何学习TVM的代码？
qq_33287871的博客
02-14 877
问题： 对陈天奇团队的开源深度学习编译器TVM很感兴趣，特别是看到18年发的论文中提到的在FPGA上的部署。对于基础知识薄弱（如体系架构编译等方面）的学生，应该如何学习TVM的代码呢？ 如提问有误请指正，求轻喷。 相关链接： TVM: End-to-End Optimization Stack for Deep Learning 自动生成硬件优化内核：陈天奇等人发布深度学习编译器TVM dmlc/tvm 作者：蓝色 链接：https://www.zhihu.com/question/268423574/an
【TVM系列三】算子转换调用流程
Flying Cat
11-23 1394
文章同步更新在公众号 AIPlayer，欢迎扫码关注，共同进步 目录 一、前言 二、TVM代码结构 三、 前端流程 四、算子relay转换 五、Relay到TOPI实现的调用过程 六、总结 一、前言 本文将从源码分析一个算子在TVM中是如何从前端到后端的转换过程。首先来看一下keras模型编译推理的一个示例。 通过pip安装好keras和tensorflow，按照TVM官方文档的示例进行一些修改，因为环境配置的不同，python包的导入以及编译target进行了修改，示例使用c..
TVM自定义算子+测试
free1993的博客
09-17 2766
TVM自定义算子定义算子编译器后端算子的测试变量的声明调用底层算子部署到特定硬件上运行全部代码 定义算子 TVM分为编译器前端和编译器后端。编译器前端主要解决的事情是将其他框架（pytorch等）转换成自己的symbolic的表示方式，所有的基本操作都会转换成sym.基本算子的形式，然后通过注册的形式，使得sym.基本算子调用后端的算子实现。编译器后端的算子实现分为两部分，第一部分是compute的部分，compute定义了基本的计算规则，schedule定义了基本计算规则在不同的硬件上的实际执行的方式。
陈天奇等人提出TVM：深度学习自动优化代码生成器
HowardWood的博客
02-17 1721
TVM 是由华盛顿大学在读博士陈天奇等人提出的深度学习自动代码生成方法，去年 8 月机器之心曾对其进行过简要介绍。该技术能自动为大多数计算硬件生成可部署优化代码，其性能可与当前最优的供应商提供的优化计算库相比，且可以适应新型专用加速器后端。近日，这项研究的论文《TVM: End-to-End Optimization Stack for Deep Learning》终于完成，内容包含新方法的介绍与...
【TVM源码学习笔记】附录1 TVM python调用C++的机制
zx_ros的博客
02-15 4077
tvm relay ir在python端通过特定机制调用对应的C++接口。本文将讨论这些调用机制的详细实现
【TVM系列五】添加Relay自定义算子
最新发布
生活需要深度
02-11 912
这就需要实现一个函数，传递相应的参数给对应的relay算子，并且返回对应算子的Call Node（这个算子最终在Relay表达式的AST里面）。A：在函数中会创建输入与输出的TypeVarNode，然后创建TypeRelationNode将类型关系函数管理起来，并定义一个FuncTypeNode将这些定义好的对象作为输入，最终赋值给op_type。本文以实现一个axis_abs的自定义算子为例介绍如何在tvm中添加新的relay算子，该算子实现的功能是以输入的3维tensor取某一维度的指定切片取绝对值。
TVM开发三个示例分析
吴建明wujianming_110117
11-08 745
TVM开发三个示例分析 把自主生成的代码生成TVM 把自主生成的代码生成TVM 目录 简介 要生成C代码。 要生成任何其它图形表示。 实现一个C代码生成器 实现【CodegenC】 运算符代码生成 输入变量的代码生成 代码发送 实现【CSourceCodegen 】 实现【GenCFunc 】 实现【CreateCSourceModule 】 注册代码生成 为表示实现一个代码生成 实现【ExampleJsonCodeGen 】 实现自定义运行时 实现构造函数 实现【GetFunction 】 实
TVM代码生成codegen
吴建明wujianming_110117
05-24 797
TVM代码生成codegen 硬件后端提供程序（例如Intel，NVIDIA，ARM等），提供诸如cuBLAS或cuDNN之类的内核库以及许多常用的深度学习内核，或者提供框架例，如带有图形引擎的DNNL或TensorRT，使用户以某种方式描述模型，实现高性能。此外，新兴的深度学习加速器还具有自己的编译器，内核库或runtime框架。 当用户尝试在新的内核库或设备上工作时，必须学习新的编程接口。对统一编程接口的需求变得越来越重要，使所有用户和硬件后端提供程序都在同一页面上。 为了与广泛使用的深度学习框架共享编
【从零开始学深度学习编译器】七，万字长文入门TVM Pass
I good vegetable a!
06-08 1556
0x0. 前言 这篇文章基于TVM 0.8.0.dev版本。在【从零开始学深度学习编译器】五，TVM Relay以及Pass简介 这篇推文中已经简单介绍了Relay和Pass机制。但是那篇文章介绍得不是很细节，对Relay的抽象语法树结构以及Pass的代码实现都没有细读，因此这篇文章以Constant Folding为例来详细介绍TVM Relay Pass。 0x1. TVM Relay树结构 在介绍Pass之前需要理清一些关键概念。 AST 摘自wiki 在计算机科学中，抽象语法树（Abstract
TVM实现hardware backend
吴建明wujianming_110117
11-13 1009
TVM实现hardware backend 官方的矩阵相加的示例如下： 2个矩阵相加的实现 for (int i = 0; i < n; ++i) { C[i] = A[i] + B[i]; } 怎么优化？ 可以并行相加，如下 for (int bx = 0; bx < ceil(n / 64); ++bx) { for (int tx = 0; tx < 64; ++tx) { int i = bx * 64 + tx; if (i < n) { C[i] = A[i] + B[i
TVM自定义数据类型
吴建明wujianming_110117
03-14 318
TVM自定义数据类型 本文将介绍“自定义数据类型”框架，该框架可在TVM中使用自定义数据类型。 介绍 在设计加速器时，关键是如何近似地表示硬件中的实数。这个问题具有长期的行业标准解决方案：IEEE 754浮点标准。然而，当试图通过构建高度专业化的设计来最大限度地利用硬件时，使用通用IEEE 754浮点数是否有意义？知道工作负载的数字要求，是否可以构建更小，更快或更省电的数据类型？答案是肯定的！研究人员已经开始在学术和工业加速器设计中尝试新的数据类型。例如，Google的Tensor处理单元（TPU）使用bf
TVM yolov3优化代码修改（编译运行OK）
吴建明wujianming_110117
11-21 2966
TVM yolov3优化代码修改（编译运行OK） yolov3_quantize_sample.py 附https://github.com/makihiro/tvm_yolov3_sample代码： yolov3_quantize_sample.py import nnvm import nnvm.frontend.darknet import nnvm.testing.yolo_detection import nnvm.testing.darknet import matplotlib.
把你自己的代码生成TVM 
weixin_42164269的博客
02-13 2165
把你自己的代码生成要TVM 随着深度学习工作负载所针对的硬件设备的数量不断增加，用户在各种设备上实现高性能所需的知识也在不断增加。为了使数据科学家不必担心开发新模型时的性能，硬件后端提供程序要么为MKLDNN或cuDNN之类的库提供许多常用的深度学习运算符，要么提供诸如TensorRT的框架以使用户以某种方式描述其模型实现高性能。但是，用户尝试在新的库或设备上工作时必须学习新的编程界面。结果，...
自主数据类型：在TVM中启用自定义数据类型探索
吴建明wujianming_110117
12-16 380
自主数据类型：在TVM中启用自定义数据类型探索 介绍 在设计加速器时，一个重要的决定是如何在硬件中近似地表示实数。这个问题有一个长期的行业标准解决方案：IEEE 754浮点标准.1。然而，当试图通过构建高度专业化的设计来最大限度地利用硬件时，使用通用IEEE 754浮点有意义吗？如果知道工作负载的数字需求，是否可以构建一个更小、更快或更省电的数据类型？答案是肯定的！研究人员已经开始在学术和工业加速器设计中尝试新的数据类型。例如，Google的张量处理单元（TPU）使用bfloat类型：一个被截断为16位的单
初探TVM--使用Tensor Engine来编写算子
shaojie_wang的博客
02-22 1519
TVM的tensor expression
使用TVM构建一个简单的神经网络框架（仅仅使用TVM的算子）
TH_NUM的博客
06-07 888
tvm 对神经网络的算子、编译和底层硬件进行了分块抽象。 如何构建一个神经网络框架呢？ 下面我们用简单的MLP分类网络仅仅含有加法和乘法操作以及softmax。 y = softmax(Wx+b) 1.第一步构建计算图和反向梯度计算图 ...
(TVM开发代码学习)给Relay系统添加新算子
weixin_44114416的博客
12-28 2207
本文学习的是官网中的Adding an Operator to Relay — tvm 0.9.dev182+ge718f5a8a documentation (apache.org) 添加新的Relay算子需要以下步骤： 1、添加一个attribute node，声明在编译时已知的固定参数。 2、给算子写一个type relation，整合进Relay类的系统文件里。 3、在C++源码中用RELAY_REGISTER_OP宏定义来注册算子的数量、类型等信息。 4、定义算子的计算过程。 5、给算
TVM User Tutorial -- Introduction to TOPI
zachariah2000的博客
03-06 552
Author: Ehsan M. Kermani 这是 TVM Operator Inventory (TOPI) 的介绍性教程。 TOPI 提供了 numpy 风格的通用操作和调度，其抽象程度高于 TVM。 在本教程中，我们将了解 TOPI 如何使我们免于在 TVM 中编写模板代码。 from __future__ import absolute_import, print_function import tvm import tvm.testing from tvm import te from t
TVM_深度学习编译器
扬起你的笑脸~把影子藏在身后
03-15 1110
TVM所做的是要比传统compiler更偏上层的，你可以把它理解成source-to-source compiler，需要其他的后端(backend)来生成最后的指令。比如当编译的Target是Intel CPU时，翻译的顺序是Relay IR -> TVM IR/ Halide IR -> LLVM IR，之后交给LLVM生成最后的可执行程序。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值