深入浅出 tvm – (16) Tensor Expression 详解

Posted on 2023年8月26日2023年11月26日 by remaper in 深入浅出 tvm

https://homes.cs.washington.edu/~arvind/papers/tvm.pdf

TVM addresses these challenges with three key modules.

(1) We introduce a tensor expression languageto build operators and provide program transformation primitives that generate different versions of the program with various optimizations.

tensor expression 是 tvm 引入的一门面向张量计算的函数式编程语言。前面我们知道 topi 就是基于 te 实现了一系列常见的张量算子库

在 tvm 里面，每个算子都可以表达为一个 te 表达式，te 接受张量作为输入，并返回一个新的张量

这个语言本身是非常简单的：

Type：因为处理的输入输出，都是 tensor，所以只有一种数据类型（不像 relay ir，有 Var、Constant、Function、Call 等各种类型）
Op：操作符也很简单，最常见的就2个 placeholder_op，compute_op

te 的代码主要在 include/tvm/te/ src/te/ 这2个目录，全部加起来才1.3w行，所以说是并不复杂

这篇文章，我们从编译器的视角来理解一下 tensor expression

参考资料：

1. te 的设计目标

te 可能应该有其他更多的设计目标

但从我目前了解到的来看，te 最大的作用应该是简化tensor 算子的编写难度，并支持进一步的调度和优化

比如向量加法

1）TVMScript 版本

在没有 te 之前，比如只有 tir 的时候，用户只能用 TVM script 来编写tensor 算子。如下：

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 我们通过 T.handle 进行数据交换，类似于内存指针
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 通过 handle 创建 Buffer
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # block 是针对计算的抽象
            with T.block("B"):
                # 定义一个空间（可并行）block 迭代器，并且将它的值绑定成 i
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0

ir_module = MyModule

2）tensor expression 版本

但是如果用 te 来描述向量加法的话，只需要2行，并且通过script() 打印出来会发现和 TVM script 是一样的

from tvm import te

A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")

func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())

Pathways – Asynchronous Distributed Dataflow for ML

Posted on 2023年7月1日2023年11月1日 by remaper in 大模型训练

论文：https://arxiv.org/pdf/2203.12533.pdf

google 认为 pathways 是下一代AI架构：https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/

其他参考：

1. 论文背景、动机、贡献

在探讨现有 AI 模型的局限时，Jeff Dean 曾经说过，今天的人工智能系统总是从头开始学习新问题，我们为数千个单独的任务开发了数千个单独的模型，这种方式需要大量的数据和算力，是非常低效的

理想的发展模式应该是训练1个模型来干多个事情

而导致这个现象的根本原因，是因为过去10年的机器学习，总是算法&硬件&系统一起演进的，这种 co-evolution 的方式有一定的隐患：系统过于针对当前的任务，而不能很好的适应未来的需求

未来 ML 的需求是什么？

模型会越来越大：
1. 多模态：复杂，参数多。纯粹的数据并行已经不行了，开始引入流水线并行、模型并行等新的并行训练方式
2. 计算稀疏：大模型特有的一个问题
硬件异构：SPMD 不适用了，因为我们需要在不同的“硬件”上运行不同的代码，需要转向 MPMD
训推一体：提升硬件利用率

为了实现这个愿景，论文提出了一种叫 [Pathways] 的通用AI架构。Pathways 旨在用一个架构同时处理多项任务，并且拥有快速学习新任务、更好地理解世界的能力

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

Posted on 2023年6月10日2023年9月21日 by remaper in 深入浅出 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日2023年9月21日 by remaper in 深入浅出 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日2023年9月21日 by remaper in 深入浅出 tvm

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

Posted on 2023年6月5日2023年9月21日 by remaper in 深入浅出 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日2023年9月21日 by remaper in 深入浅出 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日2023年9月21日 by remaper in 深入浅出 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：自动混合精度重写

深入浅出 tvm – (8) Relay IR 之 Let binding

Posted on 2023年5月29日2023年9月21日 by remaper in 深入浅出 tvm

Let Binding 这个概念也来自函数式编程，为了理解引入 Let Binding 的原因，我们先看下面这个例子

这是一个经典的计算图表达方式，唯一与传统不太一致的，就是使用 Function 来代替 Graph 的概念。在这个例子中，需要注意的是 add 的两个参数都是 1%，在 DAG 中出现了一个共享节点，我们假设 add 表达式的代码如下

class Add : public Expr {
  Expr lhs;
  Expr rhs;
};

在上面的例子中，也就意味着 lhs 成员和 rhs 成员指向了同一个表达式。在传统的计算图模式中，我们会按照拓扑序求值，1% 只会计算一次。但是从做编译的角度来说，实现的 codegen 大概率是这样

auto Add::Codegen() {
  auto lhs_val = lhs->Codegen();
  auto rhs_val = rhs->Codegen();
  return add(lhs_val, rhs_val);
}

也就意味着 log(%x)+log(%x) ，所以这里就出现了歧义，而为了能够保持和传统计算图的兼容，TVM 在打印 IR 到文本格式时，会为每个节点打印一行，并分配一个临时 id （也就是上面的%1, %2)，所以从文本格式看可能发现不了什么。为了解决歧义，TVM 引入了 Let Binding, 看下面这个例子

Let 表达式为 Let (var, value, body) , 其将value求值赋给var，然后返回body的求值结果，其对应的 IR/AST 称为 A-normal form，如上图，上面那个 Let 将 log(%x) 作为 value 求值赋给其 var（也就是 %v1），被 add 引用。通过这种方式，解决了歧义问题。

更简单地说，Let Binding 将表达式 Expr 绑定到局部的不可变 var 中。从代码实现角度，如果将传统计算图的方式称为 Graph Binding，与 Let Binding 对立的话。在 Graph Binding中，add 的 lhs 和 rhs 都直接指向了 log(%x) 这个计算表达式 Expr，而在 Let Binding 中，add 的 lhs 和 rhs 都指向了一个局部变量表达式，而这个局部变量表达式的值，由 Let 表达式负责赋值。我们可以看 Let 表达式的代码

class LetNode : public ExprNode {
 protected:
  // LetNode uses own deleter to indirectly call non-recursive destructor
  Object::FDeleter saved_deleter_;
  static void Deleter_(Object* ptr);

 public:
  /*! \brief The variable we bind to */
  Var var;
  /*! \brief The value we bind var to */
  Expr value;
  /*! \brief The body of the let binding */
  Expr body;

事实上，Let Binding 的一个核心用途就是指定计算的Scope，我们看下面这个没有使用 Let Binding 的例子

如果看文本格式，我们会觉得 %1 应该在 if 表达式外计算，但看 AST/IR, 我们又迷惑了，在下面这个例子中，这个迷惑就更明显了

fn (%x) {
  %1 = log(%x)
  %2 = fn(%y) {
    add(%y, %1)
  }
  %2
}

我们应该在闭包的外面还是里面做 %1 的计算呢？而 Let Binding 解决了这个问题，我们只要加上 Let 就能解决这个歧义。

2024年 4月
一	二	三	四	五	六	日
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

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

1. 系统架构

1. te 的设计目标

1. 论文背景、动机、贡献

1. topi 算子列表

2. topi 编程示例

1. 基本概念

1.1. 常量折叠

1.2. 常量传播

1. 基本概念

1.1. 算子融合

1. 无用函数的定义

1. 基本过程

2. 常见的 Pass 列表