机器学习编译（Machine Learning Compiler，MLC）

机器学习编译

目的：

总的来说，MLC的目的是实现日渐丰富的前端特性（pytorch/tensorflow）和丰富的后端（GPU/NPU/TPU等硬件设备）之间的自动转换，减少人工适配工作，将整个过程尽可能交给系统完成，这也和传统编译器有相似的作用。

核心：

在MLC中贯穿始终的两个概念是：

1. 抽象表现形式：将前端表示转为中间表示形式（如计算图、IR），便于进行分析和优化；
2. 转换：各种表现形式之间的转换，能够有机涵盖各种可能的情况，同时尽可能模块化，解耦各模块，转换是为了进一步优化或映射到底层硬件实现；
   整个MLC的过程其实就是构造抽象表现形式，然后在不同层级的表现形式上进行优化，转换到下一层级表现形式进行进一步的优化，直到最终映射到底层硬件提供的API。

MLC的基本流程：

前端代码—》计算图—》计算图优化（例如融合、调度等）—》IR优化—》映射到库函数调用或代码生成
在底层IR到实际可执行文件之间有两种情况：一种是直接将某个IR映射到对应的实现，另一种是自动代码生成，即自动生成对应的实现。

未来可能的研究方向：

1. 跨层级优化：如上层能够和底层硬件优化联动，统一调度优化
2. 不同加速器的自动化解决方案：实现一种自动变换以用上各种硬件的特有加速属性（如tensor core等）

MLC的相关概念：

1. Pure function（纯函数）和side effect：

Pure function：指对输入只读而不写的函数实现，这种函数可以进行水平合并，因为彼此没有依赖关系。
Side effect：指函数会对输入进行修改，不仅读还会写，这样会导致水平节点之间不一定可以融合。
所以在MLC中为了避免side effect以实现pure function，很多编译器会出现高层IR，如TVM后续加的relax，是为了能够更好的和计算图对应，方便做一些操作和转换，但在底层实现的时候不一定会遵从计算图，即依旧会有输出空间分配并作为参数传入并往输出空间写的操作，这是因为计算图只是方便上层进行优化，而不是底层优化的工具。

2. Destination pass（目标传递）：

调用函数时会将输出作为参数传入，这样做的目的是在上层分配好内存方便内存管理，而底层单元算子只管计算即可，不需要自己单独创建内存，这有利于上层做内存共享或内存池。

3. 计算图一般不支持的情况：

操作全局变量、访问文件、访问网络、生成随机数、别名（alias）等，可能会访问全局变量和改写全局变量，存在side effect，不输出pure function，这在计算图中是表示不好的，一般这种情况不会表示成计算图而是单独进行操作。

4. Symbolic variable：

是为了记录后续可能会进行优化修改的中间结果

5. 底层硬件代码通常都会包含host侧代码和device侧代码：

1. Host侧代码：即cpu执行的代码，包括环境的准备，整体kernel执行顺序，device侧api的准备等。
2. Device侧代码：即GPU执行的代码，对应硬件执行的操作。
   常见的nvcc编译好的可执行代码里面其实也包含了host和device代码而不是单纯的device代码，会在cpu调用这个可执行代码里的host代码，再转入调用device侧代码。因此nvcc编译完的可执行文件并不是严格意义上的kernel

6. 矩阵乘法常见优化方法：

1. local blocking，即利用缓存的特性，一次性读取一块数据尽量完成所有计算，因为对于C方框中的计算需要多次用到A和B的同一元素，如果每次只算C的一个元素则会带来多次读取操作，因此转换成一次性计算C的一个小矩阵实现A、B一次性读取多次计算。—》那块的大小如何选择，需要根据底层硬件的结构进行选择，块越大则占用缓存越多，且线程块能划分的越少，需要做一个权衡。
   

2. 进阶的优化是share memory blocking，进一步考虑share memory的复用，即在share memory上创建一个区域读入一部分数据（这里的读取会涉及到cooperation fetch，即线程分工读取数据），然后thread block可以复用这块share memory上的数据

7. Tensor core：

NVIDIA在硬件上实现的一个矩阵乘法加速器，对应上层调用的时候会几个要求：1）将数据搬移到指定的buffer上；
2）调用指定的api在tensor core上进行计算；
3）将指定输出buffer上的输出写出；

8. 大模型方向的整体链条：

算法—》模型训练—》编译—》部署

TVM

TVM的整体流程：

关键的数据：

1. Tensor
2. Tensor function
   核心思想：抽象（更高级的表示） & 实现（底层详细实现），实现是抽象的细化。
   机器学习编译可以视为是对Tensor function不同层级抽象之间的转换，在转换过程中可以实现优化。

两个关键的实现：

Primitive function transformation：

直接对单个函数进行优化，替换成更高效的等价实现，换算子库映射或底层硬件实现（类似codegen或单个op的实现）。

一种提倡方式是手工开发和自动转换相结合，即融合专家知识和自动化生成，例如对于核心频繁调用的部分或算子会手写实现一个调优的方案。
但自动化过程涉及到如何选择最优的方案：

1. Stochastic transform：对编译过程中可能比较重要但不确定的部分进行随机生成（类似于概率式编程），如矩阵分块的大小等，采用随机数生成，从而可以生成多种可能的transform方案，组成space of possible program，即后续的搜索空间。而人为确定的部分可以直接写成固定的，从而实现不确定的调优部分交给系统自动完成寻优（这里不用遍历用随机的原因是因为搜索空间较大的时候可以用随机找到一个次优结果，同时如何随机也有相关研究，如遗传算法等），在实际的TVM中是通过meta_schedule分析代码生成随机策略进而生成搜索空间进行搜索。
   a) 如何分析代码：会分析计算图，如两个point-wise会进行内联、矩阵乘法尝试进行块切分等，是一种固定的策略，但策略生成的是随机transform方案
   b) 如何在实际的端到端过程中引入随机transform：对单元算子的具体实现可以用随机transform生成搜索空间然后进行择优从而实现较优的单元算子实现，这对于关键算子优化有用
   

c) 在转换过程中，有两个问题值得关注：用于表示的数据结构 & 如何进行转换

Tensor program abstraction：（IR和计算图构建）

Tensor program一般包含三个组成部分：buffers（存储数据）、loop（循环处理每个元素）、computation（元素计算），对Tensor program进行抽象以转换成其他表现形式，所有等价的表现形式组成搜索空间，至于怎么找到其他表现形式是个开放话题

TVM相关知识点：

1. relax.call_tir的出现是为了更好的和计算图贴合，能够将前端表示更丝滑的转成计算图，隐藏side effect，表示成pure function
2. with relax.data_flow是为了表示pure function区域，和存在side effect的区域区分开，类似torch.compile里面的子图概念
3. 在计算图中区分output和中间结果是为了方便管理计算图，例如对于output结构在后续会用到不能随意删除，而中间结果在结束函数后可删除方便做内存复用
4. dlpack：在一个语言中将数据结果封装成dlpack，就可以在其他任意语言中使用，相当于不同语言的数据结果指向同一块内存空间进行访问，如torch.from_dlpack可以将其余语言的数据封装成pytorch的tensor进行使用
5. packed_func：是tvm环境下的函数，将其他语言的函数封装成packed_func就可以为tvm使用（大概是这样的？确认一下）
6. tensor expression：为了用python等前端语言可以编写ir module等表现形式
7. tvm可以接收torch.symbolic_trace捕获的fx graph进行优化，只是中间会经过从fx graph到ir module的映射转换过程，只要写好每个节点的映射规则，然后按照拓扑排序遍历fx graph的节点进行映射构建即可。需要注意的是，中间转换过程其实也有很多不同的实现，例如在tvm中直接映射转换到low level实现的tir（会转换成带循环的操作），或者映射到high level的实现relax（隐藏循环操作）。转换high level实现的好处是可以隐藏底层细节，在优化过程中可以先不用考虑底层的实现，专注于当前层的优化
8. 在tvm中一个ir module中允许同时存在relax、tir.prim_func和其他库调用，这样保证了每次转换只是在ir module到ir module之间的转换
9. 在tvm中会有blocking操作，即对一个tensor的操作过程中会尝试进行分块处理，这是为了方便对应底层硬件加速（类似tensor core之类的），这一类的操作其实都是为了逐步将上层计算图进行转换，以便和底层硬件逻辑对应上
10. Tvm还有tensorization概念，即考虑提供底层硬件支持的tensor操作特性，将上层的计算转换成和底层硬件相对应的操作，提供一种自动转换的方式。在tvm中，核心通过tensor intrinsic（TensorIntrin）实现，包含两部分：
    1) Description：描述底层算子能做什么，描述了输入的大小，用scope描述了需要哪些特殊存储
    2) Implement：当某个操作符合前面的Description的时候应该执行的操作，一般对应调用一个实现（可以自定义算子实现，也可以是硬件提供的api）
11. TVM的转换流程：relax对应high level计算图，tensor IR对应low level（和底层硬件对应），FFI对应直接库函数调用，Tensorization对应自动代码生成（转为张量化操作能和底层硬件更好的映射）。
12. Loop reorder用于对循环进行重排，方便行/列优先访问，或和底层处理映射（如GPU的SM）。线程绑定在GPU中会出现，用于某个线程组绑定某个执行
13. 一个改善矩阵乘法的加速例子：
    为了尽量使用行优先的缓存特性，在矩阵的乘法的时候，会对B矩阵一次性计算多列，以增加读取B矩阵的缓存命中