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

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(i[0] < num_iter):
	        C[j] += A[j] + B[j]
	        i[0] += 1
	
	return ib.get()

dtype = "float32"
A = te.placeholder((n,), name="A", dtype=dtype)
B = te.placeholder((n,), name="B", dtype=dtype)

C = te.extern(
    (n,),
    [A, B],
    lambda ins, outs: test_ir(ins[0], ins[1], outs[0]),
    name="while_vectorize",
    dtype=dtype,
)
s = te.create_schedule(C.op)

tvm Python前端描述 script / hybrid

import tvm
from tvm import tir
from tvm.script import ty
@tvm.script.tir
def flattened_elementwise_func(a: ty.handle, c: ty.handle) -> None:
    A = tir.match_buffer(a, (16, 16), "float32")
    C = tir.match_buffer(c, (16, 16), "float32")
    for i in tir.serial(0, 16):
        B_new = tir.allocate([16], "float32", "global")
        for j in tir.serial(0, 16):
            B_new[j] = tir.load("float32", A.data, ((i * 16) + j)) + 1.0
        for j in tir.serial(0, 16):
            C.data[((i * 16) + j)] = tir.load("float32", B_new, j) * 2.0


from tvm.te.hybrid import script
@script
def outer_product(n, m, a, b):
    """This is a simple outer product.
    Actually this function is not required to be documented.
    I write this docstring to test skipping docstring functionality.
    """
    c = output_tensor((n, m), a.dtype)
    for i in range(n):
        for j in range(m):
            assert i < n and j < m, "index out of range!"
            c[i, j] = a[i] * b[j]
    return c

OP介绍

tvm/python/tvm/te/operation.py
tvm/src/te/operation/compute_op.cc
tvm/src/te/operation/extern_op.cc
tvm/src/te/operation/hybrid_op.cc
tvm/src/te/operation/placeholder_op.cc

tvm 调用分析

TVM代码走读（九） 计算和调度

import tvm
from tvm import te
import numpy as np
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')

变量 te.var

Var 就类似于tensorflow中variable，创建了一个变量。其调用链为python/tvm/te/operation.py -> python/tvm/tir/expr.py -> src/tir/ir/expr.cc。var继承了PrimExpr类，建立var的时候创建了VarNode。VarNode中保存了变量的类型，名字等信息。

tensor 占位符 te.placeholder

Placeholder也类似tensorflow中的占位符，实际上最终创建了一个PlaceholderOp，保存了名字，shape，dtype信息。

compute 计算图构建

compute这个算子。Python调用位于python/tvm/te/operation.py中。这个主要是实现lamba函数算子的转换。

def compute(shape, fcompute, name=“compute”, tag="", attrs=None):
fcompute是对应着lamba表达式。首先从fcompute的__code__中提取出变量名称和数目信息，然后对应每个输入变量和对应的shape信息一起创建IterVar。IterVar有点像for循环中的循环变量，这里做了IterVar的抽象。然后将IterVar传入fcompute创建了函数体。然后就是根据fcompute的类型建立ComputeOp或者TensorComputeOp。这里我们追踪一下ComputeOp的实现。在src/te/operation/compute_op.cc中。也是构建一个ComputeOpNoe。并记录下数据，表达式信息。

def compute(shape, fcompute, name="compute", tag="", attrs=None):
    if isinstance(body, _tensor.TensorIntrinCall):
        for i, s in enumerate(shape[out_ndim:]):
            var_name = "ax" + str(i)
            dim_var.append(tvm.tir.IterVar((0, s), var_name, 4))
        op_node = _ffi_api.TensorComputeOp(name,
                                           tag,
                                           dim_var,
                                           body.reduce_axis,
                                           out_ndim,
                                           body.intrin,
                                           body.tensors,
                                           body.regions,
                                           body.scalar_inputs)
    else:
        if not isinstance(body, (list, tuple)):
            body = [body]
        body = convert(body)
        op_node = _ffi_api.ComputeOp(
            name, tag, attrs, dim_var, body)

调用链 python/tvm/te/operation.py [def compute() 函数 ] —> _ffi_api.TensorComputeOp() —> src/te/operation/compute_op.cc ComputeOp::ComputeOp()

ComputeOp::ComputeOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
                     Array<IterVar> axis, Array<PrimExpr> body) {
  if (!attrs.defined()) {
    attrs = Map<String, ObjectRef>();
  }
  auto n = make_object<ComputeOpNode>();
  n->name = std::move(name);
  n->tag = std::move(tag);
  n->attrs = std::move(attrs);
  n->axis = std::move(axis);
  n->body = std::move(body);
  if (n->body[0]->IsInstance<tir::ReduceNode>()) {
    const tir::ReduceNode* reduce = n->body[0].as<tir::ReduceNode>();
    n->reduce_axis = reduce->axis;
  }
  VerifyComputeOp(n.get());
  data_ = std::move(n);
}

body 中记录了 计算图信息，可以通过遍历的方式从中拿到输入节点，这里的思路是，存在于等式右边的tensor取索引变量(即为LOAD载入数据类型的节点 ProducerLoadNode )

Array<Tensor> ComputeOpNode::InputTensors() const {
  Array<Tensor> ret;
  std::unordered_set<Tensor> visited;// 集合 元素不重复
  for (auto& e : body) {
    // 对于 IR树 进行 后序遍历
    tir::PostOrderVisit(e, [&ret, &visited](const ObjectRef& n) {
      if (auto* pload = n.as<tir::ProducerLoadNode>()) {
        // 等式右边的tensor取索引节点 ProducerLoadNode
        Tensor t = Downcast<Tensor>(pload->producer);
        if (!visited.count(t)) {// 未记录过
          ret.push_back(t);     // 记录该输入tensor
          visited.insert(t);
        }
      }
    });
  }
  return ret;
}

te.compute可以总结为如下几步:

根据传入的fcompute,翻译成对应的表达式传入.
生成ComputeOpNode,记录计算节点.
根据计算节点,返回计算节点输出对应的Tensor.

之后 在 build中，根据 ComputeOp 的输入输出Tensor和其他op之间的关系 构建整个计算图

代码结构

Src：C++ code for operator compilation and deployment runtimes.
用于运算符编译和部署运行时的C++代码。

src/relay：Implementation of Relay, a new IR for deep learning framework superseding nnvm below.
Relay的实现，一种用于深度学习框架的新的IR（中间表示），它取代了下面的nnvm。我之后对tvm的讲解以及修改都是基于relay。

python：Python frontend that wraps C++ functions and objects implemented in src.
Python前端，它封装了在src中实现的c++函数和对象。

topi：Compute definitions and backend schedules for standard neural network operators.
标准神经网络operators的计算定义和后端调度。

nnvm：C++ code and Python frontend for graph optimization and compilation. After the introduction of Relay, it remains in the codebase for backward compatibility.
用于图形优化和编译的C++代码和Python前端。在引入Relay之后，为了向后兼容，它仍然保留在代码库中。

include：Header files in include are public APIs that share across modules.
include中的头文件是是跨模块共享的公共API。

TVM python与C相互调用分析

深入理解TVM：Object家族

深入理解TVM：Python/C++互调（上）

深入理解TVM：Python/C++互调（中）

深入理解TVM：Python/C++互调（下）

最底层 Object 类为基础，构建了PackedFunc交互接口以及IR遍历修改优化等结构。

可使用 runtime.DumpTypeTable 接口打印 以 Object 类为基础的类的派生关系。
后端 c++实现 src/runtime/object.cc

TVM_REGISTER_GLOBAL("runtime.DumpTypeTable").set_body_typed([](int min_child_count) {
  TypeContext::Global()->Dump(min_child_count);
});

已经使用 TVM_REGISTER_GLOBAL 注册到 共享函数库中了，可以从前端直接调用。

在前端 python/tvm/runtime/object.py 中添加 获取接口

def DumpTypeTable(min_child_count):
    return _ffi_api.DumpTypeTable(self, fmt)

在 模块初始化 tvm/ python / tvm / runtime / init.py 中 暴露接口

from .object import DumpTypeTable```

使用该接口获取 以 Object 类为基础的类的派生关系 表

import tvm

print(tvm.runtime.DumpTypeTable(0))

Object类在TVM里非常重要，基本上所有其他的类都直接或者间接继承自这个类;
ObjectPtr可以看作是指向Object的一个模拟指针的封装类, ObjectRef也可以看作是指向Object的一个封装类;

TVM里有个不成文的约定，所有以Node为结尾的类名都是继承自Object，不以Node结尾的类名都是继承自ObjectRef，例如：

// include/tvm/runtime/module.h
class Module : public ObjectRef {};
class ModuleNode : public Object {};

主要有几个静态变量约定需要遵守，包括：

_type_key：类的全局唯一的字符串标识符
_type_index：类的全局唯一的uint32_t标识符
_type_child_slots：为子类预留的index个数
_type_final：表示是不是没有子类了，可以通过TVM_DECLARE_FINAL_OBJECT_INFO这个宏来设置
_type_child_slots_can_overflow：标识是不是可以超过_type_child_slots定义的数量;

type_index应用:
可以用于IsInstance这个辅助函数的加速，内部实现就是直接判断子类的_type_index 的值是不是在父类预留的 _type_child_slots 范围之内

示例代码如下：

class ObjBase : public Object {
 public:
  // dynamically allocate slow
  static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
  static constexpr const uint32_t _type_child_slots = 1000;
  static constexpr const char* _type_key = "test.ObjBase";
  TVM_DECLARE_BASE_OBJECT_INFO(ObjBase, Object);
};
// 定义过类之后要进行注册
TVM_REGISTER_OBJECT_TYPE(ObjBase);

PackedFunc 这个类来说，它是python和c++互调的桥梁，此类实现代码在
include/tvm/runtime/packed_func.h 文件中，
这里面还有一个TypedPackedFunc类，它只是PackedFunc的一个wrapper，主要增加了类型检查的功能，开发TVM的c++代码要尽可能的使用这个类，其中用到了下面这几个关键的数据结构：

TVMValue ： union 数据类型，主要是为了储存c++和其它语言交互时所支持的几种类型的数据；
TVMArgs ：类，为了封装传给PackedFunc的所有参数，包含 参数数组 TVMValue* values ，参数数量 num_args；
TVMPODValue_ ： 内部使用的基类，plain old data 数据类型，实现核心是强制类型转换运算符重载；
TVMArgValue ：继承自前面的TVMPODValue_类，用作表示PackedFunc的一个参数，它和TVMPODValue_的区别是扩充了一些数据类型的支持，比如string、PackedFunc、TypedPackedFunc等，其中对后两个的支持是在c++代码中能够调用python函数的根本原因；
TVMRetValue ：继承自TVMPODValue_类，主要作用是作为存放调用PackedFunc返回值的容器；
TVMArgsSetter ：值设置类型，一个用于给TVMValue对象赋值的辅助类，主要通过重载函数调用运算符来实现；

PackedFunc

class PackedFunc {
 public:
  using FType = function<void(TVMArgs args, TVMRetValue* rv)>;
  PackedFunc() {}
  explicit PackedFunc(FType body) : body_(body) {}

​  template <typename... Args> // 不定数量参数类型 模板参数
  // 函数调用运算符重载
  inline TVMRetValue operator()(Args&&... args) const {
    const int kNumArgs = sizeof...(Args);// 参数数量
    const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
    TVMValue values[kArraySize];// 函数参数数组
    int type_codes[kArraySize]; // 参数类型数组
    detail::for_each(TVMArgsSetter(values, type_codes), 
      std::forward<Args>(args)...);// 利用 TVMArgsSetter 设置参数 到内部容器
    TVMRetValue rv;
    // 传入参数，调用函数体
    body_(TVMArgs(values, type_codes, kNumArgs), &rv);
    return rv;
  }
  
  inline void CallPacked(TVMArgs args, TVMRetValue* rv) const {
    body_(args, rv);
  }
​
 private:
  FType body_; // 函数实体 
};

c++端的函数注册，python端对c++端的函数调用都来源于c++端的注册函数，最主要的一个函数注册宏是 TVM_REGISTER_GLOBAL，code base里大概用了1300多次，除了这个注册宏，TVM里还有许多其它的注册宏；

注册的函数可以是普通函数，也可以是labda表达式，注册接口有三个：set_body、set_body_typed、set_body_method，第一个使用的是PackedFunc，后面两个使用的是TypedPackedFunc，PackedFunc在这个系列的上篇讲过了，TypedPackedFunc是PackedFunc的一个wrapper；

1. 使用set_body接口注册lambda表达式：
   注册参数固定为 (TVMArgs args, TVMRetValue* rv) 的函数

// src/topi/nn.cc
TVM_REGISTER_GLOBAL("topi.nn.relu")
    .set_body([](TVMArgs args, TVMRetValue* rv) {
  *rv = relu<float>(args[0]);
});

2. 使用set_body_typed接口注册lambda表达式：
   注册参数类型不固定的函数

// src/te/schedule/graph.cc
TVM_REGISTER_GLOBAL("schedule.PostDFSOrder")
    .set_body_typed([](
          const Array<Operation>& roots, 
          const ReadGraph& g) {
      return PostDFSOrder(roots, g);
    });

3. 使用set_body_method接口注册类内函数：
   注册类内任意类型函数

// src/ir/module.cc
TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar")
    .set_body_method<IRModule>(&IRModuleNode::GetGlobalVar);

TVM_REGISTER_GLOBAL宏定义

这个宏定义的本质就是在注册文件定义了一个static的引用变量，引用到注册机内部new出来的一个新的 Registry对象 ：

// include/tvm/runtime/registry.h
#define TVM_REGISTER_GLOBAL(OpName)               \
  static ::tvm::runtime::Registry& __mk_TVMxxx =  \
      ::tvm::runtime::Registry::Register(OpName)

上述的 三个 注册函数体接口 set_body** 都是 Register 类的方法。

先来看最核心的Manager类，它是Registry的内部类，用来存储注册的对象，先看下代码：

// src/runtime/registry.cc
struct Registry::Manager {
  // Manager是个单例, 只允许有一个 对象实例，函数前面 带 static 标识，静态数据 
  static Manager* Global() {
    static Manager* inst = new Manager();
    return inst;
  }
  std::mutex mutex; // 数据结构里面带锁，可以保证线程安全
                    // 修改数据时上锁 std::lock_guard lock(m->mutex);                   
  unordered_map<std::string, Registry*> fmap; // 函数存储 容器 键值对 map
  // 使用unordered_map来存储注册信息，注册对象是 Registry指针, 键为函数 名字符串
};

函数注册接口类

// include/tvm/runtime/registry.h
class Registry {
 public:
  // 三种注册函数体的接口函数
  Registry& set_body(PackedFunc f);    // 注册固定接口的 PackedFunc函数类型
  Registry& set_body_typed(FLambda f); // 注册任意接口类型函数，内部打包成 PackedFunc
  Registry& set_body_method(R (T::*f)(Args...)); // 注册 任意类T的具有任意参数类型的函数f
​
  static Registry& Register(const std::string& name);// 构建一个 包裹函数的祖册器类
  static const PackedFunc* Get(const std::string& name);// 获取注册的函数实体
  static std::vector ListNames(); // 获取已经注册在容器内的 函数名列表
​
 protected:
  std::string name_;     // 函数名
  PackedFunc func_;      // 保存 注册进来的函数体 (实际数据)
  friend struct Manager; // 友元结构体 函数存储 容器，内 fmap 为容器池
};

通过 Register 的 Get 方法 可获取到 实际的 函数入口

const PackedFunc* Registry::Get(const std::string& name) {
  Manager* m = Manager::Global(); // 静态对象（只有一个)
    // 各种  共有一个 存储容器 Manager
  std::lock_guard<std::mutex> lock(m->mutex); // 数据访问前， 上锁
  auto it = m->fmap.find(name); // 按名字查找 函数包裹 Registry 实例
  if (it == m->fmap.end()) return nullptr;
  return &(it->second->func_);// 返回 Registry 实例 中实际存储 函数的 数据func_
}

f1, f2, …, fn 函数实体
R1, R2， …, RN 函数包裹 类

共用 Manager 存储函数 fmap

这种注册机思想在 深度学习框架、业务引擎等框架内很常用

前后端调用

TVM使用python的ctypes模块来调用c++代码提供的API，ctypes是python内建的可以用于调用C/C++动态链接库函数的功能模块.

对于动态链接库提供的API，需要使用符合c语言编译和链接约定的API，因为python的ctype只和c兼容，而c++编译器会对函数和变量名进行 name mangling，所以需要使用__cplusplus宏和extern "C"来得到符合c语言编译和链接约定的API.

// TVM给python提供的接口主要都在这个文件：
// include/tvm/runtime/c_runtime_api.h，
// 下面主要展示了__cplusplus和extern "C"的用法，
// 以及几个关键的API。
#ifdef __cplusplus
extern "C" {
#endif
​
int TVMFuncListGlobalNames(...);
int TVMFuncGetGlobal(...);
int TVMFuncCall(...);
    
#ifdef __cplusplus
}  // TVM_EXTERN_C
#endif

加载TVM动态库

TVM的python代码从python/tvm/init.py中开始真正执行，即：

from ._ffi.base import TVMError, version
这句简单的import代码，会执行python/tvm/_ffi/init.py：

from .base import register_error
from .registry import register_func
from .registry import _init_api, get_global_func

上面的第一句，会导致python/tvm/_ffi/base.py中的下面代码被执行：

def _load_lib():
    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL)
    return lib, os.path.basename(lib_path[0])

_LIB, _LIB_NAME = _load_lib()
上面的lib_path[0]是TVM动态链接库的全路径名称，我是在linux系统做的试验，链接库的名称是/xxx/libtvm.so（不同的系统动态库的名字会有所不同，windows系统是.dll，苹果系统是.dylib，linux系统是.so），在_load_lib函数执行完成后，_LIB和_LIB_NAME都完成了初始化，其中_LIB是一个ctypes.CDLL类型的变量，可以认为它是能够操作TVM动态链接库的export symbols的一个全局句柄，_LIB_NAME是libtvm.so这个字符串。这样后续在python中，我们就能通过_LIB这个桥梁不断的和c++的部分进行交互。

python中来获取c++API的底层函数是_get_global_func：

#python/tvm/_ffi/_ctypes/packed_func.py
def _get_global_func(func_name):
    handle = ctypes.c_void_p() # void * 类型
    _LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle))
    return _make_packed_func(handle, False)

这里面handle是一个相当于void类型的指针变量，因为从ctypes的官方文档中可以查到，c_void_p对应的primitive C compatible data type是：

​ctype type c type python type
c_void_p void * int or None

_get_global_func中调用了TVMFuncGetGlobal这个API，看下这个API的实现就可以发现，handle最终保存了一个c++代码在堆中new出来的PackedFunc对象指针：

// src/runtime/registry.cc
int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
  const tvm::runtime::PackedFunc* fp 
      = tvm::runtime::Registry::Get(name);// 从函数注册仓库中获取对应的C++函数入口 PackedFunc 
  *out = new tvm::runtime::PackedFunc(*fp);
}

和c++PackedFunc的关联工作这时候才完成一半，在_get_global_func的最后调用了_make_packed_func这个函数：

python 端保存 c++端 函数的 PackedFunc 对象实例

#python/tvm/_ffi/_ctypes/packed_func.py
def _make_packed_func(handle, is_global):
    obj = PackedFunc.__new__(PackedFuncBase)
    obj.is_global = is_global
    obj.handle = handle
    return obj

可以看到_make_packed_func函数中创建了一个定义在python/tvm/runtime/packed_func.py 中的python PackedFunc对象，PackedFunc其实是一个空实现，它继承自 PackedFuncBase类，PackedFuncBase 类中定义了一个__call__函数：

对象呗调用时 会执行这个 函数行为

# python/tvm/_ffi/_ctypes/packed_func.py
class PackedFuncBase(object):
  def __call__(self, *args):
    # 打包 python 端 参数  成 packed_func 参数
    # _make_tvm_args 打包函数，包含了 python类型转C tvm类型
    values, tcodes, num_args = _make_tvm_args(args, temp_args)
    ret_val = TVMValue() # 函数返回值 类型
    ret_tcode = ctypes.c_int()
    _LIB.TVMFuncCall(
        self.handle, # 存储的 c++端 packed_func 指针
        values, # 打包好的 tvmargs
        tcodes, # 参数类型
        ctypes.c_int(num_args), # 参数数量
        ctypes.byref(ret_val),
        ctypes.byref(ret_tcode),
    )
    return ret_val

从上面可以看出，python的__call__函数调用了C的TVMFuncCall这个API，把前面保存有c++ PackedFunc对象地址的handle以及相关的函数参数传了进去，TVMFuncCall的主体代码如下：

// src/runtime/c_runtime_api.cc
int TVMFuncCall(TVMFunctionHandle handle, TVMValue* args, ...)
  (*static_cast<const PackedFunc*>(handle))
      .CallPacked(TVMArgs(args, arg_type_codes, num_args), &rv);
}

这样就完成了把c++中的PackedFunc映射到了python中的PackedFunc，在python代码中只需要调用python中创建好的PackedFunc对象，就会通过上面分析的过程来一步步调到c++的代码中。

注册的函数既包括c++中注册的函数，也包括python中注册的函数，其中主要是c++中注册的函数，通过 list_global_func_names 函数（实际上调用的 TVMFuncListGlobalNames 这个c++API）可以得到c++中注册的所有函数，目前有1500多个

先看_init_api这个函数，这个函数是把注册函数关联到各个模块的关键：

# python/tvm/_ffi/registry.py
def _init_api(prefix, module_name):
    target_module = sys.modules[module_name]
​
    for name in list_global_func_names():
        if not name.startswith(prefix):
            continue
        fname = name[len(prefix) + 1 :]
        f = get_global_func(name)
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = "TVM PackedFunc %s. " % fname
        setattr(target_module, ff.__name__, ff)

有三个最主要的点：

line3：sys.modules是一个全局字典，每当程序员导入新的模块，sys.modules将自动记录该模块。 当第二次再导入该模块时，python会直接到字典中查找，从而加快了程序运行的速度。
line9：get_global_func等同于上面章节仔细说了的_get_global_func这个函数，这个函数返回一个python端的PackedFunc对象，它的handle成员存储了c++中new出来的PackedFunc对象（以注册函数作为构造参数）的地址，python端的PackedFunc对象的__call__函数调用了c++的TVMFuncCall这个API，handle作为这个API的参数之一，c++端再把handle转成c++的PackedFunc对象来执行，这样就完成了从python端PackedFunc对象的执行到c++端PackedFunc对象的执行的映射。
line13：把前面代码构造的python端PackedFunc对象作为属性设置到相应的模块上.

然后各个模块中对_init_api来全局调用一次，就完成了关联，我在代码中找了几个作为示例，如下所示：

# python/tvm/runtime/_ffi_api.py
tvm._ffi._init_api("runtime", __name__)
​
# python/tvm/relay/op/op.py
tvm._ffi._init_api("relay.op", __name__)
​
# python/tvm/relay/backend/_backend.py
tvm._ffi._init_api("relay.backend", __name__)

abs 示例

以TVM中求绝对值的函数abs为例，这个函数实现在tir模块，函数的功能很简单，不会造成额外的理解负担，我们只关注从python调用是怎么映射到c++中的，先看在c++中abs函数的定义和注册：

// src/tir/op/op.cc
// 函数定义
PrimExpr abs(PrimExpr x, Span span) { … }
​
// 函数注册 注意： 注册名 以 tir. 开头
TVM_REGISTER_GLOBAL(“tir.abs”).set_body_typed(tvm::abs);
再看python端的调用：

# python/tvm/tir/_ffi_api.py
# 把c++ tir中注册的函数以python PackedFunc
# 对象的形式关联到了_ffi_api这个模块
tvm._ffi._init_api("tir", __name__)# 导入以 tir. 为开头注册的 函数
​# python/tvm/tir/op.py
# 定义了abs的python函数，其实内部调用了前面
# 关联到_ffi_api这个模块的python PackedFunc对象
def abs(x, span=None):
    return _ffi_api.abs(x, span) # 调用 c++端的 abs函数

c++ 后端调用 前端 传入的 前端函数指针 分析:
(自定义优化pass 实例)

TVM IR visitor 框架分析

visit函数注册宏 ->visit函数指针数组

tvm/include/tvm/node/functor.h

template <typename R, typename... Args>
class NodeFunctor<R(const ObjectRef& n, Args...)> {
 private:
  /*! \brief internal function pointer type */
  typedef R (*FPointer)(const ObjectRef& n, Args...);
  /*! \brief refer to itself. */
  using TSelf = NodeFunctor<R(const ObjectRef& n, Args...)>;
  /*! \brief internal function table */
  std::vector<FPointer> func_; // 函数指针数组
//...
}

// 函数指针注册函数
 template <typename TNode>
  TSelf& set_dispatch(FPointer f) {  // NOLINT(*)
    uint32_t tindex = TNode::RuntimeTypeIndex();
    if (func_.size() <= tindex) {
      func_.resize(tindex + 1, nullptr);
    }
    ICHECK(func_[tindex] == nullptr) << "Dispatch for " << TNode::_type_key << " is already set";
    func_[tindex] = f;
    return *this;
  }

tvm/include/tvm/tir/expr_functor.h

#define IR_EXPR_FUNCTOR_DISPATCH(OP)                                                       \
  vtable.template set_dispatch([](const ObjectRef& n, TSelf* self, Args... args) {     \
    return self->VisitExpr_(static_cast(n.get()), std::forward(args)...); \
  });


using TSelf = ExprFunctor<R(const PrimExpr& n, Args...)>;
using FType = NodeFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;
  
//visit函数注册
  // initialize the vtable.
  static FType InitVTable() {
    FType vtable;
    // Set dispatch
    IR_EXPR_FUNCTOR_DISPATCH(VarNode);
    IR_EXPR_FUNCTOR_DISPATCH(SizeVarNode);
    IR_EXPR_FUNCTOR_DISPATCH(LoadNode);
    IR_EXPR_FUNCTOR_DISPATCH(BufferLoadNode);
    IR_EXPR_FUNCTOR_DISPATCH(ProducerLoadNode);
    IR_EXPR_FUNCTOR_DISPATCH(LetNode);
    IR_EXPR_FUNCTOR_DISPATCH(CallNode);
    IR_EXPR_FUNCTOR_DISPATCH(AddNode);
    IR_EXPR_FUNCTOR_DISPATCH(SubNode);
    IR_EXPR_FUNCTOR_DISPATCH(MulNode);
    IR_EXPR_FUNCTOR_DISPATCH(DivNode);
    IR_EXPR_FUNCTOR_DISPATCH(ModNode);
    IR_EXPR_FUNCTOR_DISPATCH(FloorDivNode);
    IR_EXPR_FUNCTOR_DISPATCH(FloorModNode);
    IR_EXPR_FUNCTOR_DISPATCH(MinNode);
    IR_EXPR_FUNCTOR_DISPATCH(MaxNode);
    IR_EXPR_FUNCTOR_DISPATCH(EQNode);
    IR_EXPR_FUNCTOR_DISPATCH(NENode);
    IR_EXPR_FUNCTOR_DISPATCH(LTNode);
    IR_EXPR_FUNCTOR_DISPATCH(LENode);
    IR_EXPR_FUNCTOR_DISPATCH(GTNode);
    IR_EXPR_FUNCTOR_DISPATCH(GENode);
    IR_EXPR_FUNCTOR_DISPATCH(AndNode);
    IR_EXPR_FUNCTOR_DISPATCH(OrNode);
    IR_EXPR_FUNCTOR_DISPATCH(ReduceNode);
    IR_EXPR_FUNCTOR_DISPATCH(CastNode);
    IR_EXPR_FUNCTOR_DISPATCH(NotNode);
    IR_EXPR_FUNCTOR_DISPATCH(SelectNode);
    IR_EXPR_FUNCTOR_DISPATCH(RampNode);
    IR_EXPR_FUNCTOR_DISPATCH(ShuffleNode);
    IR_EXPR_FUNCTOR_DISPATCH(BroadcastNode);
    IR_EXPR_FUNCTOR_DISPATCH(IntImmNode);
    IR_EXPR_FUNCTOR_DISPATCH(FloatImmNode);
    IR_EXPR_FUNCTOR_DISPATCH(StringImmNode);
    IR_EXPR_FUNCTOR_DISPATCH(AnyNode);
    return vtable;
  }
};

//行为派发: 
  virtual R VisitExpr(const PrimExpr& n, Args... args) {
    static FType vtable = InitVTable();
    return vtable(n, this, std::forward<Args>(args)...);
  }

TVM调度原语（Schedule Primitives）

一个调度过程由多个阶段组成，一个阶段表示操作的一个调度。

1. 分裂split：split通过factor分裂指定轴为两个轴。

#分裂0轴为两个轴，先计算内循环再计算外循环，xo为外循环，xi为内循环
xo, xi = s[B].split(B.op.axis[0], factor=32)

2. 重排序reorder：reorder能按照指定顺序重新排列轴（类似于permute）。

xo, xi = s[B].split(B.op.axis[0], factor=32)
s[B].reorder(xi,xo)

3. 平铺tile：tile通过平铺两个轴执行计算图块

xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)

4. 融合fuse：fuse能融合一个计算的两个轴

#首先平铺成4轴（i.outer,j.outer,i.inner,j.inner）
xo,yo,xi,yi = s[B].tile(B.op.axis[0],B.op.axis[1], x_factor=10, y_factor=5)
#然后融合(i.inner,j.inner)进一个轴：（i.inner.j.inner.fused）
fused = s[B].fuse(xi,yj)

5. 从哪里开始计算compute_at
   对于包含多个算子的调度，TVM默认从root开始遍历计算张量。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]+1, name='B')
C = tvm.compute((m,), lambda i: B[i]*2, name='C')

s = tvm.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))
# 移动B循环到C循环的第一个轴
s[B].compute_at(S[C], C.op.axis[0])
print(tvm.lower(s, [A, B, C], simple_mode=True))

6. 计算内联compute_inline
   compute_inline可以将一个计算阶段标记为内联，然后将计算体扩展并插入需要张量的地址处。（和C中的内联函数一个意思）

7. compute_root
   compute_root可以将一个计算阶段的计算移动到root。（compute_at的逆过程）

TVM 涉及到的 Pass 优化操作

pass概念

一种 Pass 代表一种优化。比较多在 LLVM 中提到。

ModulePass是将整个程序视作一个单元处理的pass。

FunctionPass是以单个函数为作用域的pass, 每个函数间是相互独立的。

Pass 优化遍历遍

TVM提供了两个级别的优化，如下图所示。计算图优化，用于执行高级操作员融合，布局转换和内存管理等任务。然后是张量运算符优化和代码生成层，它优化张量运算符

high-level pass 高层 优化遍历遍：

header file: include/tvm/relay/transforms.h
src files: src/relay/pass
python: python/tvm/relay/transform.py

tensort level IR pass 低层优化遍历遍:

header file include/tvm/tir/ir_pass.h
src files: src/tir/pass python: python/tvm/tir/ir_pass.py

tensort level IR pass 更多的是底层代码生成时做的优化，如 VectorizeLoop，UnrollLoop，InjectPrefetch，StorageRewrite 等。

ModulePass

1. LambdaLift 将局部函数提升为全局函数。

在 src/relay/backend/vm/compiler.cc, src/relay/backend/vm/lambda_lift.cc。

2. RemoveUnusedFunctions 去掉 relay 的 IR 模块中未使用的函数。

Inline 将一个被 inline 标记的全局函数嵌入到 relay 的IR模块中。

3. EtaExpand 为构造函数添加抽象，或者给一个函数添加全局变量。
   如： square 被转化为 fn (%x: int32) -> int32 { square(x) } See https://en.wikipedia.org/wiki/Lambda_calculus#%CE%B7-conversion

4. PartialEval 在编译时评估静态的代码碎片，尽可能多的做常量传播，常量折叠，代码嵌入等优化，以减少运行时开销，获得更多的融合优化。作为代价生成的code量会增加。

5. PrintIR 打印IR。通常作为一系列 pass 的结果，用来debug。

6. PartitionGraph 将一个混合后端的relay程序分割到各自后端的计算区域。

7. ToANormalForm 将一个数据流图转化为行政范式 (Administrative Normal Form, A-Normal Form, ANF)。 将一个表达从隐式共享的图的形式转化为显式共享， 也就是 ANF。 The scope of the root expression is the global scope. The scope of any non root expression is the least common ancestor of all it’s scope. Values are ordered by post-DFS order in each scope.

FunctionPass

1. DeadCodeElimination 去掉不影响程序结果的表达。将没有被索引到的 let bindings 去掉，将被使用过一次的 let bindings 作为代码嵌入。

2. FoldConstant 折叠常量函数。

3. FuseOps 融合算子，可以指定融合的优化级别

4. RewriteAnnotatedOps 重写被注释的算子， 比如 on_device 会标记在哪一种设备上进行调度计算。帮助进行异构计算执行。

5. ToCPS 将一个表达转化为连续传递式样 (continuation passing style, CPS)。 CPS 意思是每一个函数将不会直接返回结果，而是传递一个另外的函数，作为参数，然后将结果传到下一个续集。这样每一个函数调用时将会多一个参数，表示其余的计算。每一个中间计算都会被传入一个续集。

6. ToGraphNormalForm 去除所有的 let binding, 并将所有的变量转化为直接的指针索引。返回的表达叫做 graph normal form。

7. SimplifyInference 在推理时简化一些特定的算子，比如 batchnorm 会被代替为一些更加简单的算子。这部分 Efficient Deep Learning Inference on Edge Devices 有写。

8. FastMath 将非线性激活函数替换成近似计算的算子以获得更快的计算速度。一些情景下有损失计算精度的风险。 InferType 推理一个表达式的类别。获得的结果是一个有显式类别信息的新表达式，以及它的返回类型。

9. EliminateCommonSubexpr 寻找并去掉共同的子表达式。比如有两个表达式被评估为等同，则新建一个变量来替代它们两个。

10. CombineParallelConv2D 将并行运算的二维卷积合成为一个卷积，如果合成的卷积的分支数大于某个最小值。

11. CombineParallelDense 将并行运算的稠密算子合并为批处理的矩阵乘法 (batch_matmul)，如果合成的算子的分支数小于某个最小值。

12. BackwardFoldScaleAxis / ForwardFoldScaleAxis 将轴的缩放折叠进卷积或者稠密算子的weights中 https://github.com/apache/incubator-tvm/pull/2020

13. CanonicalizeOps 将一些深度学习的算子规范化为一系列简化的算子，比如将 bias_add 规范化为先进行 expand_dims 然后做 broadcast_add

14. AlterOpLayout 转换算子的布局，或者将元算子替换为其他表达式。可以用来计算定制数据布局的卷积，或者更加通用的weights的前置转换。

15. ConvertLayout 给定一个目标的数据布局，将表达式大多数算子的数据布局转化为目的布局。在理想情况下，只有两个起始和结束的布局需要转换。一般在relay转化别的框架时需要。主要使用 AlterOpLayout 和 InferCorrectLayout。

16. RFC - https://discuss.tvm.ai/t/layout-conversion-pass/4009

17. Legalize 将一个表达式转化为另一个表达式 (expr)，可以基于算子的shape，dtype，或者 data layout 转化为另一个算子或者一系列算子。一般在target平台相关的优化中，两个语等价的算子性能不一定会一样。

18. CanonicalizeCast 规范化 cast 表达式，使得算子的融合更加方便。

19. MergeComposite 将多个模式相匹配的算子合并为一个复合算子。主要用在使用外部代码生成工具时多个算子map到同一个外部算子的情况。定义作为Wrapper使用。 在 src/relay/pass/annotate_target.cc

20. AnnotateTarget 将表达式中的算子标注编译器或者部署平台，使得算子被wrap为 compiler_begin／subgraph_start 和 compiler_end ／subgraph_end，从而用做之后其他编译器的代码生成。定义作为Wrapper使用。 在 src/relay/pass/annotate_target.cc

tensort level IR pass 低层优化遍历遍

自定义优化pass 编写

# 遍历IR树函数
loops = []
def find_width8(op):
    """ Find all the 'tir.For' nodes whose extent can be divided by 8. """
    if isinstance(op, tvm.tir.For): # for 节点
        if isinstance(op.extent, tvm.tir.IntImm):# 范围为常量
            if op.extent.value % 8 == 0:# 可被8整除
                loops.append(op) #添加找到的　节点

# IR修改转换函数
def vectorize8(op):
    """ Split can vectorize the loops found in `find_width8`. """
    if op in loops:
        extent = op.extent.value
        name = op.loop_var.name
        # 定义两个新的循环轴
        lo, li = te.var(name + ".outer"), te.var(name + ".inner")
        # IR 变量替换，原来的循环变量 op.loop_var 替换为 lo * 8 + li
        body = tvm.tir.stmt_functor.substitute(op.body, {op.loop_var: lo * 8 + li})
        # 构建新的 for节点 内层可 向量化的for循环
        body = tvm.tir.For(li, 0, 8, tvm.tir.ForKind.VECTORIZED, body)
        # 构建外层 不可向量化的 for节点
        body = tvm.tir.For(lo, 0, extent // 8, tvm.tir.ForKind.SERIAL, body)
        return body
    return None


@tvm.tir.transform.prim_func_pass(opt_level=0)
def vectorize(f, mod, ctx):
    global loops
    
    #  post_order_visit 搜索到 目标节点  分析IR
    tvm.tir.stmt_functor.post_order_visit(f.body, find_width8)

    if not loops:
        return sf

    # The last list arugment indicates what kinds of nodes will be transformed.
    # Thus, in this case only `For` nodes will call `vectorize8`
    
    # 转换IR   ir_transform
    return f.with_body(tvm.tir.stmt_functor.ir_transform(f.body, None, vectorize8, ["tir.For"]))