深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

Posted on 2023年6月10日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

topi 是 tvm 的一个张量算子库，提供了很多常见的算子操作，例如 conv2D，transpose，等等。

tvm 在编译计算图的时候，会先将计算图从 relay ir 翻译成 tir，再将 tir 翻译成目标设备代码（比如llvm，比如c，比如cuda），但是 relay ir 和 tir 之间还有一种中间语言叫 tensor expression，这是一种专门为张量计算而设计的语言。relay ir 的绝大多数部分 op 其实都是通过 tensor expression 这种语言来实现的

https://tvm.apache.org/docs/reference/langref/relay_expr.html

An operator is a primitive operation, such as add or conv2d, not defined in the Relay language. Operators are declared in the global operator registry in C++. Many common operators are backed by TVM’s Tensor Operator Inventory.

topi 和 te 的关系：

te( tensor expression ) 是一门函数式编程语言，面向张量计算的
topi 是基于 te 实现的张量算子库

类似于编程语言和标准库之间的关系

1. topi 算子列表

源码目录：src/topi

src/topi 只是把算子通过 TVM_REGISTER_GLOBAL 注册到 tvm 的全局函数中，并没有太多细节

函数的实现，全部是以 inline 函数的形式，实现在 include/tvm/topi 中，由于算子都是通过 tensor expression 来实现的，因此所有算子的实现最终都会调用 tvm::te::compute 来进行计算

最常见的有4类操作：broadcast，elemwise，nn，reduction

broadcast	elemwise	nn	reduction
topi.add topi.subtract topi.multiply topi.divide topi.floor_divide topi.mod topi.floor_mod topi.maximum topi.minimum topi.power topi.left_shift topi.logical_and topi.logical_or topi.logical_xor topi.bitwise_and topi.bitwise_or topi.bitwise_xor topi.right_shift topi.greater topi.less topi.equal topi.not_equal topi.greater_equal topi.less_equal topi.broadcast_to	topi.acos topi.acosh topi.asin topi.asinh topi.atanh topi.exp topi.fast_exp topi.erf topi.fast_erf topi.tan topi.cos topi.cosh topi.sin topi.sinh topi.tanh topi.fast_tanh topi.atan topi.sigmoid topi.sqrt topi.rsqrt topi.log topi.log2 topi.log10 topi.identity topi.negative topi.clip topi.cast topi.reinterpret topi.elemwise_sum topi.sign topi.full topi.full_like topi.logical_not topi.bitwise_not	topi.nn.relu topi.nn.leaky_relu topi.nn.prelu topi.nn.pad topi.nn.space_to_batch_nd topi.nn.batch_to_space_nd topi.nn.nll_loss topi.nn.dense topi.nn.bias_add topi.nn.dilate topi.nn.flatten topi.nn.scale_shift_nchw topi.nn.scale_shift_nhwc topi.nn.pool_grad topi.nn.global_pool topi.nn.adaptive_pool topi.nn.adaptive_pool3d topi.nn.pool1d topi.nn.pool2d topi.nn.pool3d topi.nn.softmax topi.nn.log_softmax topi.nn.lrn topi.nn.binarize_pack topi.nn.binary_dense	topi.sum topi.min topi.max topi.argmin topi.argmax topi.prod topi.all topi.any

2. topi 编程示例

由于 topi 实现的所有算子，最终都会注册到 tvm 的全局函数中，并已 relay.op 的形式存在

因此，我们在描述计算图的时候，可以直接调用相关的算子，相当于我们平时写 c 代码的时候，有很多编译器的 builtin 函数一样

比如 topi.add 对应的 relay.op 是 relay.add，其他类推

a1 = relay.var("a1", shape=(1,), dtype="float32")
c1 = relay.const(10, 'float32')
c2 = relay.add(c1, a1)

深入浅出 tvm – (13) Relay Pass 之常量折叠

Posted on 2023年6月10日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

和算子融合一样，常量折叠是编译领域里最常见的一个优化，简单来说，就是把常量表达式前置计算，在编译阶段就计算好，然后以常量的形式翻译成底层机器码，以提高执行效率，减少计算量

实际上大部分的编译器，常量折叠一般包含2种优化技术：常量折叠和常量传播

1. 基本概念

1.1. 常量折叠

constant folding，常量折叠，编译器优化技术之一，通过对编译时常量或常量表达式进行计算来简化代码。以下面的代码为例：

i = 320 * 200 * 32;

上面的代码中，编译器通常会在编译过程中直接对表达式进行求值，计算出320 * 200 * 32的结果，而不会生成2个乘法指令。

还有一些更复杂（但不清楚tvm是否支持，后面验证下）。比如，在执行一些复杂表达式的计算时，我们可以将表达式内部一些常量运算合并，最终起到简化的效果，如下

优化前（左边）每个表达式的运算量是 8 flop，优化后（右边）的运算量是 2 flop，运算效率极大提升了

不过，实际上 tvm 并没有做的那么高级，tvm 的常亮折叠只用来处理一些比较简单的场景

1.2. 常量传播

constant propagation，常量传播，同样也是编译器最常见的优化技术之一，在编译的过程中，对常量依赖做一些基本的推演和提前计算，再使用常量折叠技术来简化代码。如下：

int x = 14;
int y = 7 - x / 2;
return y * (28 / x + 2);

//常量传播

int x = 14;
int y = 7 - 14 / 2;
return y * (28 / 14 + 2);

//常量折叠

int x = 14;
int y = 0;
return 0;

深入浅出 tvm – (12) Relay Pass 之布局转换

Posted on 2023年6月7日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

深入浅出 tvm – (11) Relay Pass 之算子融合

Posted on 2023年6月5日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

算子融合，FuseOps，是编译领域最常见的一个优化技术，在这里也算属于 relay 里面最复杂的一类优化了，整个优化的核心逻辑 1k+ 行

代码：src/relay/transforms/fuse_ops.cc

算子融合的目的最终是要解决 AI 处理器的内存墙、并行墙的问题，提升 Tensor 数据的访存局部性。目前算子融合的技术路线有比较多，这里不涉及，我们只需要知道 tvm 是基于支配树来实现算子融合的就行了

1. 基本概念

1.1. 算子融合

算子融合，即将多个算子组合在一起放到同一个核中，通过算子融合的方式，不需要将中间结果保存到全局内存，进而减少执行所需要的时间。

tvm 中将算子分为7种类型：

kElemWise：2个 tensor 之间按元素逐个操作的算子，实际上所有四则运算都是这种类型，https://deeplizard.com/learn/video/QscEWm0QTRY
kBroadcast：见上述链接，到操作两个不同形状的 tensor 时
kInjective：一对一映射函数，比如 add / sqrt / exp 等操作算子（operator）
kCommReduce：多到少的映射，输入到输出具有降维性质，如：sum / max / min等操作操作算子（operator）
kOutEWiseFusable：这是计算比较复杂的，如：conv2d / bn / relu等操作算子（operator）
kTuple：xx
kOpaque：无法被融合的算符，比如 sort

根据以上对算符的不同类型，TVM提供了三种融合规则（我看论文是这么写的，但是现在不止3种了？）：

原文：https://layman-n-ish.github.io/pdfs/TVM_Review_Report.pdf

从融合算子的内部视角看，这种融合实际上是数据计算pipeline化，即两次计算中间数据不再经历store-load过程，而是直接给到下一个计算单元完成计算。

深入浅出 tvm – (10) Relay Pass 之剔除无用函数

Posted on 2023年6月2日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

我们从一个最简单的 Pass 开始入手，RemoveUnusedFunctions，带大家一步一步了解 Pass 优化的具体过程

这个 Pass 看名字顾名思义，就是要把 Module 里执行不到的函数删除掉

代码：src/relay/transforms/dead_code.cc

1. 无用函数的定义

无用函数是指那些根本不会调用到的函数

我们用 tvm.relay 手写一个最简单的网络，来跟踪优化的过程和效果

如下，有2个函数，f1 和 f2，其中 f1 就是无用的函数，因为 main 函数进来后，根本不会调用到这个函数

import tvm
from tvm import relay
from tvm.contrib import relay_viz
 
a1 = relay.var("a1", shape=(1,), dtype="float32")
a2 = relay.var("a2", shape=(1,), dtype="float32")
add_op = relay.add(a1, a2)
f1 = relay.Function([a1, a2], add_op)

d1 = relay.var("d1", shape=(1, 32, 56, 56), dtype="float32")
w1 = relay.var("w1", shape=(32, 32, 3, 3), dtype="float32")
b1 = relay.var("b1", shape=(32,), dtype="float32")
conv = relay.nn.conv2d(d1, w1, strides=(1, 1), padding=(1, 1))
bias = relay.nn.bias_add(conv, b1)
relu = relay.nn.relu(bias)
f2 = relay.Function([d1, w1, b1], relu)

mod = tvm.IRModule({'add_func': f1, 'main': f2})
mod = relay.transform.InferType()(mod)

mod = relay.transform.RemoveUnusedFunctions()(mod)

优化完之后，右边的图就被去掉了（图片的可视化，参考之前的 Relay 计算图）

深入浅出 tvm – (9) Relay Pass 优化

Posted on 2023年6月1日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

Pass 又称为 transform，是编译领域最常见的一种优化技术，tvm 里的 pass 设计，主要是参考了 llvm 的设计思想

pass 本质上是一种图到图的转换，pass不改变图计算的结果

我们知道 tvm 计算图有两种 level 的表示 ir，一种是 relay ir，一种是 tir，不同的 ir 有不同的 pass 优化逻辑，因此 tvm 中的 pass 也有两种：

relay 层的 pass。代码在 relay/transform，包括很多图结构的优化，比如算符融合，常量折叠，等等，属于偏前端的优化
tir 层的 pass，偏向底层的优化，比如 prefetch 注入，unrollLoop 等

这里我们只讲 relay 相关的 pass，tir 的后续再讲

1. 基本过程

前面我们将 relay.build 的时候就知道 build 过程中的 OptimizeImpl 阶段就是在执行 Pass 优化，我们来看下 OptimizeImpl 函数的实现

src/relay/backend/build_module.cc

这个函数很长，但是非常简单，我省略了一部分代码

pass_seqs 就是一个数组，函数根据 config，选择相应的 pass，然后放到这个数组里面，最后执行 pass 优化的时候，直接一个 for 循环执行 pass_func() 就完了，由于每个 pass_func 的输入和输出都是 IRModule，串行执行即可

  IRModule OptimizeImpl(IRModule relay_module) {
    ICHECK(relay_module.defined()) << "The IRModule must be defined for the Relay compiler.";

    backend::BindParamsInModule(relay_module, params_);

    Array pass_seqs =
        GetPassPrefix(/*is_homogenous=*/config_->primitive_targets.size() == 1, /*is_vm=*/false);
    transform::PassContext pass_ctx = PassContext::Current();
    // ...    

    relay_module = transform::InferType()(relay_module);

    // Inline the functions that have been lifted by the module scope.
    //
    // TODO(@zhiics) Note that we need to be careful about the subgraphs with
    // global function calls. We should make sure that these callees are also
    // inline functions. However, this should be very unlikely for accelerators
    // and vendor-provided libraries. So we don't handle for now.
    relay_module = transform::Inline()(relay_module);
    relay_module = transform::InferType()(relay_module);
    relay_module = transform::LabelOps()(relay_module);
    relay_module = transform::AnnotateMemoryScope(config_)(relay_module);

    ICHECK(relay_module.defined());

    return relay_module;
  }

从这个函数看，relay层有3类pass：

GetPassPrefix() 函数会首先返回一系列，最常见的、公共的pass，大概20来个左右，可以细看函数
homogeneous 相关的 pass
auto schedule 相关的 pass
meta schedule 相关的 pass

2. 常见的 Pass 列表

grep -r ^Pass src/relay/transforms/ | grep “{“| sed ‘s/(/ /g’ | awk ‘{print $2}’ | sort

基本可得到所有的 Pass 函数声明，如下我只写了重点的部分

AlterOpLayout：替换操作符的布局或用其他表达式替换基本操作符
CanonicalizeOps：将特殊算子规范化为基本算子
CombineParallelConv2D：合并 conv2d 操作
CombineParallelDense：合并 dense 操作
ConvertLayout：布局转换
DeadCodeElimination：删除没用的代码
DefuseOps：FuseOps的逆操作
EliminateCommonSubexpr：删除常见的子表达式
FastMath：将昂贵的非线性函数转换为快速但近似的对应函数。
FoldConstant：常量表达式折叠
FuseOps：将 expr 中的操作符融合为更大的操作符
InferType：类型推断
Inline：执行内联操作
LazyGradientInit：减少梯度张量的内存使用
MergeCompilerRegions：合并编译区域
SimplifyExpr：简化表达式，比如合并连续的 reshapes 操作
SimplifyInference：简化推理阶段的 data-flow
SplitArgs：将具有大量参数的函数切割成更小的块
ToBasicBlockNormalForm：将表达式转换为基本块的形式
ToMixedPrecision：自动混合精度重写

成功，源于对美学的执著追求

月度归档： 2023 年 6 月

深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

1. topi 算子列表

2. topi 编程示例

深入浅出 tvm – (13) Relay Pass 之常量折叠

1. 基本概念

1.1. 常量折叠

1.2. 常量传播

深入浅出 tvm – (12) Relay Pass 之布局转换

深入浅出 tvm – (11) Relay Pass 之算子融合

1. 基本概念

1.1. 算子融合

深入浅出 tvm – (10) Relay Pass 之剔除无用函数

1. 无用函数的定义

深入浅出 tvm – (9) Relay Pass 优化

1. 基本过程

2. 常见的 Pass 列表

2023 年 6 月
一	二	三	四	五	六	日
			1	2	3	4
5	6	7	8	9	10	11
12	13	14	15	16	17	18
19	20	21	22	23	24	25
26	27	28	29	30