AI System 人工智能系统 TVM深度学习编译器 DSL IR优化计算图编译优化内存内核调度优化 DAG 图优化 DFS TaiChi 函数注册机 Registry

本文链接：https://blog.csdn.net/xiaoxiaowenqiang/article/details/116420429

本文详细介绍了TVM深度学习编译器，包括DSL设计、IR优化、计算图调度、自定义优化pass编写等内容。TVM作为一个源到源的编译器，支持多种深度学习框架，通过将计算描述和调度分开，实现高效的目标平台代码生成。文章还涉及了TVM的Python与C++交互、函数注册机、Pass优化操作等核心技术。

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

DSL 领域专用语言 TVM深度学习编译器 AI System 人工智能系统

参考项目

TaiChi 三维动画渲染物理仿真引擎DSL
TVM 深度学习DSL 密集计算DSL
LLVM 模块化编译器
 编译器相关资料参考

强烈推荐 AI System 人工智能系统课件计算图编译优化调度优化

课程介绍课程概述和系统/AI基础
人工智能系统概述人工智能系统发展历史，神经网络基础，人工智能系统基础
深度神经网络计算框架基础反向传播和自动求导，张量，有向无环图，执行图论文和系统：PyTorch, TensorFlow
矩阵运算与计算机体系结构矩阵运算，CPU/SIMD, GPGPU, ASIC/TPU
论文和系统：Blas, TPU
分布式训练算法数据并行，模型并行，分布式SGD
论文和系统：PipeDream
分布式训练系统 MPI, parameter servers, all-reduce, RDMA
论文和系统: Horovod
异构计算集群调度与资源管理系统集群上运行DNN任务：容器，资源分配，调度
论文和系统：Kubeflow, OpenPAI, Gandiva
深度学习推导系统效率，延迟，吞吐量，部署
论文和系统：TensorRT, TensorFlowLite, ONNX
计算图的编译与优化 IR，子图模式匹配，矩阵乘和内存优化
论文和系统：XLA, MLIR, TVM, NNFusion
神经网络的压缩与稀疏化优化模型压缩，稀疏化，剪枝
自动机器学习系统超参调优，神经网络结构搜索（NAS）
论文和系统：Hyperband, SMAC, ENAX, AutoKeras, NNI
强化学习系统 RL理论，RL系统
论文和系统：AC3, RLlib, AlphaZero
安全与隐私联邦学习，安全，隐私
论文和系统：DeepFake
利用人工智能来优化计算机系统问题 AI用于传统系统问题，AI用于系统算法
论文和系统：Learned Indexes, Learned query path

DSL（Domain Specified Language）领域专用语言，它是用于解决特定领域问题的语言。

所谓领域专用语言（domain specific language / DSL），其基本思想是“求专不求全”，不像通用目的语言那样目标范围涵盖一切软件问题，而是专门针对某一特定问题的计算机语言。

与GPPL（通用目的编程语言）相比，DSL针对的目标是特定的领域。

DSL可以分为两种：内部DSL、外部DSL。

外部DSL是自我包含的语言，它们有自己特定语法、解析器和词法分析器等等，它往往是一种小型的编程语言，甚至不会像GPPL那样需要源文件。

与之相对的则是内部DSL。内部DSL其实更像是种别称，构建于其他语言之上，依托于另外一种语言的前端描述(词法语法语义解析部分)，生成自己的中间表示IR，基于中间IR，进行优化变化等操作，最后转换到最终输出。

TVM

TVM是用于高效内核代码构建的版本领域专用语言（Domain-Specialed-Language，DSL）。

TVM与LLVM的架构非常相似。TVM针对不同的深度学习框架和硬件平台，实现了统一的软件栈，以尽可能高效的方式，将不同框架下的深度学习模型部署到硬件平台上。

如果从编译器的视角来看待如何解决这个问题，各种框架写的网络可以根据特定的规则转化成某种统一的表示形式，在统一表示的基础上进行一些可重用的图优化，之后再用不同的后端来生成对应不同设备的代码，这就是目前各家都在尝试的设计思路了。

TensorFlow 的 XLA 会把高级代码抽象成 XLA HLO 的表示，做目标无关优化之后再用对应后端来生成更深一层的代码。

NVIDIA 的 TensorRT 的优化策略也是在图转化之后的统一表示上做，例如根据设定好的规则来做一些相邻计算单元的合并（Kernel Fusion）等等，但TensorRT最重要的Kernel部分的实现未开源。

TVM所做的是要比传统compiler更偏上层的，你可以把它理解成 source-to-source compiler（源到源编译器），需要其他的后端(backend)来生成最后的指令。比如当编译的Target是Intel CPU时，翻译的顺序是Relay IR/前端描述 -> TVM IR -> LLVM IR，之后交给LLVM生成最后的可执行程序。当编译的Target是NVIDIA GPU时，这个顺序就变成了 Relay IR/前端描述-> TVM IR -> CUDA program (比如生成一个文件, my_kernel.cu)，然后调用cuda的runtime compiler来生成CUDA可执行文件。

把你自己的代码生成TVM
tvm.d2l.ai 笔记
 tvm.schedule 自动代码生成调度
 tvm在线文档
 在TVM中添加新设备Codegen
tvm会议
 课程资料
 在android上安装和运行tvm
TVM代码走读（九）计算和调度

TVM 使用示例

Tensor Expression 支持常见的算术运算和数学运算，并涵盖了常见的算子。该语言没有指定循环结构和许多其他执行细节，并且它提供了为各种后端添加硬件感知优化的灵活性。根据解耦计算和调度的原则，用调度表示从张量表达式到底层代码的映射。

Tensor Expression的最初的想法来源于 Halide，核心在于把代码的计算和调度 Schedule 分开。

TVM提出 Schedules 的概念，指的是一种将计算描述（张量运算）降低到后端（底层）优化实现的特定规则。这也是TVM实现的核心。

其理念是对 Schedules 空间和用于遍历此空间的转换进行建模,从而提供生成低级代码的不同方法。

矩阵赋值

import tvm
from tvm import te

A = te.placeholder((10, 10))
B = tvm.compute((10, 10), lambda i, j: A[i, j])# B[i,j] = A[i, j]
s = te.create_schedule(B.op)      # 构建优化调度器
# 可定义优化调度
f = tvm.build(s, [A, B], "hello") # 编译生成代码
print(f.get_source())             # 打印生成的代码

生成的代码

// hello tvm backend!
void default_function( void* args,  void* arg_type_ids, int32_t num_args,  void* out_ret_value,  void* out_ret_tcode) {
   
  // ......
  float* placeholder = (float*)(((DLTensor*)arg0)[0].data);
  // ......
  float* compute = (float*)(((DLTensor*)arg1)[0].data);
  // ......
  for (int32_t i = 0; i < 10; ++i) {
   
    for (int32_t j = 0; j < 10; ++j) {
   
      compute[((i * 10) + j)] = placeholder[((i * 10) + j)]; // 矩阵赋值
    }
  }
}

矩阵乘法

import tvm
from tvm import te
# 定义变量占位符（萝卜坑位）
n = te.var("n")
m = te.var("m")
L = te.var("L")
A = te.placeholder((n, L), name="a", dtype=dtype)
B = te.placeholder((L, m), name="B", dtype=dtype)
# 规约运算轴
k = te.reduce_axis((0, L), name="k")
# 定义矩阵乘法运算
C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")

# schedule调度优化
s = te.create_schedule(C.op)# 定义调度器
y, x = s[C].op.axis         # 获取 外层for循环轴 i,j
k = s[C].op.reduce_axis[0]  # 获取内层 规约运算轴
yo, yi = s[C].split(y, 8)   # 循环轴拆分
xo, xi = s[C].split(x, 8)
s[C].reorder(yo, xo, k, yi, xi) # 多个循环轴 重新排序
# 编译
f = tvm.build(s, [A, B], "hello") # 编译生成代码
print(f.get_source())             # 打印生成的代码

相关TVM API

placeholder((n,), name=‘a’)
- def placeholder(shape, dtype=None, name=“placeholder”)
- 作用：类似 tf.placeholder，占位符，后续可指定输入数据。
compute(A.shape, lambda i: A[i] + B[i], name=‘c’)
- def compute(shape, fcompute, name=“compute”, tag=“”, attrs=None)
- 作用：定义计算图。 fcompute 是 lambda function of indices-> value,匿名函数。
create_schedule(C.op)
- def create_schedule(ops)
- 作用：为一系列Ops构建默认调度Schedule。
lower(s, [A, B, C], simple_mode=True)
*　def lower(sch, args, name=“default_function”, binds=None, simple_mode=False)
- 作用：输入schedule对象和输入输出tensor list。
- 如果 simple_mode 为True，则输出simple and compact statement（起始就是C语言形式的伪代码，表示op操作具体过程）。
- 如果 simple_mode 为False，我也不知道输出啥。
tvm.build(s, [A, B, C])
- def build(inputs, args=None, target=None, target_host=None, name=“default_function”, binds=None)
- 通过输入Schedule等参数，构建可运行的module。
mod.export_library(mod_fname) 与 tvm.module.load(mod_fname)
- def export_library(self, file_name, fcompile=None, **kwargs)
- def load(path, fmt=“”)
- 作用：将 tvm.module.Module 导出为本地文件，或将本地文件导入生成 tvm.module.Module 对象。
relay.frontend.from_mxnet(model, {‘data’: x.shape})
- def from_mxnet(symbol, shape=None, dtype=“float32”, arg_params=None, aux_params=None)
- 作用：将MXNet模型转换为TVM Relay模型。
with relay.build_config(opt_level=3)
- def build_config(opt_level=2, fallback_device=_nd.cpu(), required_pass=None, disabled_pass=None):
- 作用：为 relay.build 设置一些参数，比如图优化等。
tvm.context(target)
- def context(dev_type, dev_id=0)
- 作用：为给定的设备设置运行环境。
tvm.contrib.graph_runtime.create(graph, mod, ctx)
- def create(graph_json_str, libmod, ctx)
- 作用：通过 graph/module/ctx 创建实际可执行模块，之后导入输入数据、获取模型预测结果都通过这个方法返回的对象。
- 举例：rt.set_input, rt.run, rt.get_output 等。
relay.save_param_dict(params)
- def save_param_dict(params)
- 作用：保存模型参数。
remote = rpc.connect(‘172.31.0.149’, 9090)
- def connect(url, port, key=“”, session_timeout=0)
- 作用：建立远程rpc连接，方便交叉编译等功能。
- 该remote对象有很多功能：
- remote.upload(mod_fname)：上传文件
- remote.load_module(mod_fname)：导入模块
- remote.cpu()：建立远程context对象

前端描述分析

针对神经网络层级的描述接口 topi，底层大多使用 compute实现

 A = te.placeholder((batch, in_dim), name="A")
 B = te.placeholder((out_dim, in_dim), name="B")
 bnn_A = topi.nn.binarize_pack(A)  
 bnn_B = topi.nn.binarize_pack(B)
 # binary dense
 bnn_A1 = te.placeholder(bnn_A.shape, dtype=bnn_A.dtype)
 bnn_B1 = te.placeholder(bnn_B.shape, dtype=bnn_B.dtype)
 bnn_C = topi.nn.binary_dense(bnn_A1, bnn_B1)

张量描述语言 te.compute

n = te.var("n")
m = te.var("m")
A = te.placeholder((m, n), name='A')
B = te.placeholder((m, n), name='B')
C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name='C')
# 生成 ComputeOp

tvm IR构建语言 tvm DSL

def test_ir(A, B, C):
	ib = tvm.tir.ir_builder.create()
	n = C.shape[0]
	A = ib.buffer_ptr(A)
	B = ib.buffer_ptr(B)
	C = ib.buffer_ptr(C)
	i = ib.allocate("int32", (1,), name="i", scope="local")
	i[0] = 0
	
	with ib.for_range(0, n) as j:
	    C[j] = 0.0
	
	with ib.for_range(0, n, kind="vectorize") as j:
	    with ib.while_loop