文章目录

元张量函数
TensorIR
端到端模型执行
自动程序优化
与机器学习框架的整合

课程主页：机器学习编译
具体内容查看vedio和notes，本文先当与一个精简笔记和脉络梳理，如有错误还请指出

元张量函数

张量（Tensor）：数组

张量函数（Tensor function）：对数组的运算，机器学习编译过程可以看作是张量函数的变换（元张量函数替换为更高效的实现）和映射

抽象（Abstraction）：用来表示张量函数

元张量函数（Primitive Tensor Function）：单个单元计算

张量函数抽象：数组，循环，计算，不同层次的抽象有助于实现自动化的张量函数的变换和优化

TensorIR

TensorIR：标准机器学习编译框架Apache TVM中使用的张量程序抽象，用TensorIR表示元张量函数

TensorIR和numpy存在一些直接对应的元素

开数组
for循环

TensorIR中特有的元素

block：一个block包括一组axis和相应计算
axis：包括绑定位置，范围，属性
- 添加的附加信息使得axis独立于外部循环
函数属性：
装饰器：
- @tvm.script.ir_module：表示修饰的类是一个IRModule，IRModule是保存张量函数集合的容器对象
- @T.prim_func：表示修饰的函数是一个元张量函数

元张量函数的变换

sch = tvm.tir.Schedule：对类进行修改的辅助类
- sample_perfect_tile(loop, n)：loop表示拆分的循环，n表示拆分的长度，随机分割循环，返回j_factors
- split：拆分循环
- reorder：循环重排
- reverse_compute_at：将一个循环放在另一个循环的内循环（某层循环同级）的后面，类似于算子融合
- decompose_reduction：将包含某个axis的循环语句拆分开，形成并列的另一个循环
- get_block：
- get_loops：
- mod.script()：每次变换后的结果（的代码）
- trace：历史变换轨迹

构建与运行：将变换后的代码构建和运行

构建：将IRModule转换为runtime.Module（过程类似于构建TensorRT中的context，runtime.Module表示可运行函数的集合，target指定部署环境的详细信息）
- tvm.build(Module, target)
运行：从runtime.Module集合中取出相应的函数

除了在类之前使用装饰器@tvm.script.ir_module（TVMScript）创建TensorIR和进行变换创建TensorIR，还可以使用张量表达式（Tensor Experssion，类似于numpy的apply函数）创建TensorIR

# placeholder类似于输入节点
A = tvm.placeholder((128, 128), "float32", name="A")
B = tvm.placeholder((128, 128), "float32", name="B")k = tvm.reduce_axis((0, 128), "k")# compute的参数是输出维度，计算表达式，compute 函数描述了我们要如何计算给定索引的每个元素 [i, j] 的值
Y = tvm.compute((128, 128), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = tvm.compute((128, 128), lambda i, j: tvm.max(Y[i, j], 0), name="C")# create_prim_func参数是TensorIR函数的输入和输出，可以创建TensorIR函数
tvm_func = tvm.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
# 可以进行算子融合，将矩乘计算Y和ReLu计算C进行融合# 通过tvm_func.show()可以打印出具体的代码
MyModuleFromTE = tvm.IRModule({"mm_relu": tvm_func})

端到端模型执行

计算图

矩形框：tensor function
节点：数据

destination passing：不会将输出作为返回值返回，而是在函数外部显式的将返回值的内存开好，将地址（或out）作为函数的参数

优点：底层实现不进行内存管理
缺点：在函数外部开辟数组就相当于创建输出节点，改变了计算图的形式，失去了计算图的一些特性

destination passing方式：主动创建输出节点

call_tir：希望隐藏可能的分配或对函数的显式写入，即函数是pure或side-effect free的（函数只获取输入返回输出，不会改变程序的其他部分或造成影响（例如递增全局计数器，无外部性）），对destination passing的元张量函数进行了一层封装

优点：保持计算图模式，便于在计算图层级进行优化

call_tir方式：保持计算图结构

dataflow block：dataflow block中不包含有side-effect的操作，显式标记这些操作可以更安全高效的实现变换

@tvm.register_func：注册函数，使得可以调用类外部的函数

参数绑定：将输入参数绑定为IPModule的常量，通过函数参数名和参数字典中的key相匹配来创建绑定，注意这也是一种端到端的变换（会改变原来的代码）

自动程序优化

随机程序采样（或随机程序变换）：将不确定的部分放在搜索空间中，对于专家知识使用确定性的方法写死

meta_schedule：智能搜索，多次运行，搜索随机变换空间找到最优的调度变换

meta_schedule.tune_tir：可以指定搜索空间（具体指定哪些部分要进行随机程序采样），如果不指定搜索空间则是进行自动调度（分析程序中潜在的全部搜索空间，进行搜索，搜索空间会大，因此用时会长）

使用优化好的变换替换原来的实现：

# MyModuleMixture是一个自己实现的神经网络（MLP）的class，nd_params是网络权重参数的dict，类似于Pytorch中的state_dict
# 进行参数绑定，构建一个新的模型
MyModuleWithParams2 = relax.transform.BindParams("main", nd_params)(MyModuleMixture)# sch_tuned_linear是自动程序优化meta_schedule.tune_tir的返回值，得到新的linear0
new_func = sch_tuned_linear.mod["main"].with_attr("global_symbol", "linear0")# 获取原来模型（新构建的模型）的linear0，get_global_var返回一个指向IRModule中函数的pointer引用
gv = MyModuleWithParams2.get_global_var("linear0")# 用调优后的函数new_func替换原来的函数gv
MyModuleWithParams2.update_func(gv, new_func)
# MyModuleWithParams2是经过替换的Module

与机器学习框架的整合

小结：

使用TVMScript创建元张量函数
使用TE（Tensor Expression）创建张量函数
端到端执行：
- 在TVMScript中创建元张量函数（@T.prim_func），在主函数（@R.function）main中通过R.call_tir调用元张量函数，实现端到端的执行
- 使用BlockBuilder构建IRModule：将使用TE创建的元张量函数构建为端到端的执行，使用BlockBuilder自动生成上面TVMScript的代码（优点是其他框架也使用类似的计算图抽象，方便进行优化）

BlockBuilder：

tvm.relax.Var对应tvm.placeholder，是输入节点
bb.function(“main”)对应@R.function修饰的main函数
bb.dataflow()对应R.dataflow()
中间结果是tvm.relax.expr.DataflowVar类型，是tvm.relax.Var的子类，因为中间结果不会被访问到
bb.emit_te：比如bb.emit_te(te_matmul,A,B)，首先为A和B创建placeholder，然后调用create_prim_func创建TensorIR函数（将使用TE创建的元张量函数构建TensorIR函数），然后通过调用call_tir函数生成对te_matmul函数的调用并执行
bb.emit_output：创建（或叫做标记）dataflow block的输出变量
bb.emit_func_output：标记函数输出，每个函数作用域内只能调用一次

TorchFX：PyTorch中计算图的操作工具，可以将PyTorch模型表示为计算图

目标：将TorchFX的计算图转换为IRModule的计算图

map_param：将nn.Parameter转换为relax.const

from_fx：将Torch中的计算图转换为IRModule的计算图

创建BlockBuilder

编译Torch中计算图节点的拓扑排序，根据不同的逻辑进行转换

如果是placeholder：fn_inputs(list)存放输入节点input_var(relax.Var)，node_map(dict)存放Torch中计算图的节点node到IRModule中计算图的节点input_var（注意input_shapes是有序的）
如果是Parameter：将Parameter转换为relax.const并存放到node_map的映射中

如果是call_function：传入自定义的dict，说明遇到某个操作时如何进行转换（比如matmul）：

def te_matmul(A: te.Tensor, B: te.Tensor) -> te.Tensor:assert A.shape[1] == B.shape[0]n = A.shape[0]m = B.shape[1]k = te.reduce_axis((0, A.shape[1]), name="k")return te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="matmul")def map_matmul(bb, node_map, node: fx.Node):A = node_map[node.args[0]]B = node_map[node.args[1]]return bb.emit_te(te_matmul, A, B)from_fx中传入参数：
call_function_map = {torch.matmul: map_matmul,
}

如果是call_module：传入自定义的dict，比如linear module：

def map_nn_linear(bb, node_map, node, nn_mod):x = node_map[node.args[0]]w = map_param(nn_mod.weight)if nn_mod.bias is not None:b = map_param(nn_mod.bias)//  使用TVM TOPI (TVM operator inventory) 的预定义TE库，不用自己手写TE张量表达式了y = bb.emit_te(topi.nn.dense, x, w) // 进行转置矩阵乘法：x @ w.Treturn bb.emit_te(topi.add, y, b) // add广播操作的加法：y+bfrom_fx中传入参数：
call_module_map = {torch.nn.Linear: map_nn_linear，
}

如果是output：标记输出节点

标记函数输出并返回IRModule

转换成高层算子：在call_function_map和call_module_map中，除了将Torch的计算图节点转换为IRModule的计算图节点，还可以转换为比IRModule高级一点的原始算子（tvm.relax节点对应于Torch节点的counterpart，由此可以进一步确定不同的tvm.relax节点使用何种具体实现），注意使用的是bb.emit函数而非是bb.emit_te函数

机器学习编译MLC 笔记 1-5章（上）相关推荐

《机器学习实战》笔记——第三章：决策树实战
1 说明该书主要以原理简介+项目实战为主,本人学习的主要目的是为了结合李航老师的<统计学习方法>以及周志华老师的西瓜书的理论进行学习,从而走上机器学习的"不归路".因 ...
《机器学习》读书笔记，第二章概念学习和一般到特殊序
本章展示了几种概念学习算法,并讨论了这些算法能收敛到正确假设的条件.这里还分析了归纳学习的本质以及任意程序能从训练数据中泛华的理由. 2.1概念学习: 从有关某个布尔函数的输入输出训练样例中推断出该布 ...
《机器学习》阅读笔记第三章
Contents 1. 不同学科中的线性模型[^1] 2. 线性模型:回归任务 2.1 估计方法 2.2 正则化 2.3 广义线性模型 3. 线性模型:分类任务 3.1 对数几率回归线性判别分析(L ...
吴恩达机器学习系列课程笔记——第五章：Octave教程(Octave Tutorial)
提示:这章选学,可以去学python,第六节可以看看. 5.1 基本操作 https://www.bilibili.com/video/BV164411b7dx?p=26 本章学习以种编程语言:Oct ...
javascript高级程序设计学习笔记第五章上
第五章引用类型的值(对象)是引用类型的一个实例.在 ECMAScript 中,引用类型是一种数据结构, 用于将数据和功能组织在一起.它也常被称为类,但这种称呼并不妥当.尽管 ECMAScri ...
吴恩达机器学习系列课程笔记——第十一章：机器学习系统的设计(Machine Learning System Design)
11.1 首先要做什么 https://www.bilibili.com/video/BV164411b7dx?p=65 在接下来的视频中,我将谈到机器学习系统的设计.这些视频将谈及在设计复杂的机器学 ...
Head First设计模式读书笔记八第九章上迭代器模式
之前的总结: https://blog.csdn.net/u011109881/article/details/59677544 个人觉得本章节,HeadFirst讲的没有之前看到的网站讲的清晰,至少 ...
《Web安全之机器学习入门》笔记：第九章 9.4 支持向量机算法SVM 检测DGA域名
DGA(Domain Generation Algorithm)域名生成算法是一种利用随机字符等算法来生成C&C域名,从而逃避安全设备域名黑名单检测的技术手段. 1.黑样本 def load_ ...
深度实践SPARK机器学习_学习笔记_第二章2.3加载数据
2.3加载数据 1.下载数据文件u.user head -3 u.user ##查看文件前几行 cat u.user |wc -l 或者 more u.user |wc -l ##数文件记录数 ...

机器学习编译MLC 笔记 1-5章（上）