MegEngine 使用小技巧：如何解读 MegCC 编译模型几个阶段 Pass 的作用

本文详细介绍了MegCC，一个轻量级的深度学习模型编译器，其主要Pass包括MGBToKernelPass、MGBFuseKernelPass、MemoryForwardingPass、KernelMaterializationPass和StaticMemoryPlanningPass，分别涉及图优化、算子融合、内存规划等。通过这些Pass，MegCC实现了模型的高效编译和优化，生成高性能的Kernel代码。

MegCC 是一个真真实实的深度学习模型编译器，具备极其轻量的 Runtime 二进制体积，高性能，方便移植，极低内存使用以及快启动等核心特点。用户可在 MLIR 上进行计算图优化，内存规划，最后通过预先写好的 code 模版进行代码生成。

MegCC 中主要的 Pass

MGBToKernelPass：这个 Pass 主要将 MGB IR 转换为 Abstract Kernel IR，转换过程中主要完成几件事情：
- 将 MGB IR 中的所有输入输出 Tensor 类型转换为 Buffer 类型。
- 将 MGB IR 中的所有枚举参数转换为对应的字符，这样 Abstract Kernel IR 就可以完全和 MegEngine 解耦。
- 将一些内存搬运相关的 Opr 全部转换为 Relayout，如：Concat，SetSubtensor 等 Opr（node-level optimizations）。
- 将判断 Opr 是静态 shape 还是动态 shape，动态 shape 就是输入 tensor 的 shape 需要依赖输入的值才能计算出来的，如：输出一个 tensor 中所有大于 1 的数。如果是静态 shape 直接转换到 Abstract Kernel IR，如果是动态 shape 直接转换到 Kernel IR 的 Instruction 中。
MGBFuseKernelPass：应用在 MGB IR 上，基于 mlir 的模板匹配的方法尽可能的完成 kernel 的融合，比如连续两个 typecvt 合并成为一个 typecvt 等（block-level optimizations，算子融合）。
MemoryForwardingPass：将遍历 Abstract Kernel IR 所有可能不用计算，直接 share 输入内存的 Opr，如果这些 Opr 确实不用计算，则直接 forward memory，如果这些 Opr 需要进行内存搬运，则会用 Relayout Opr 替换原来的 Opr（node-level optimizations）。KernelMaterializationPass：将所有 Abstract Kernel IR 都装载上真正 Kernel code 并转化为 KernelCall，然后添加对应的 KernelDef。KernelCall 和 KernelDef 之间通过 symbol 进行匹配。
StaticMemoryPlanningPass：将所有静态 shape 的 memref 进行内存规划，内存规划算法使用改进的 MegEngine 的内存规划算法–PushDown 算法，能够极大程度的压缩运行时内存使用量。同时将 mlir 的 memref.Alloc 替换为 Kernel IR 的 MemPlan，MemPlan 中主要记录了内存规划的一整块 memref 以及该 Tensor 在规划的内存中的偏移量（dataflow-level optimizations，静态内存规划）。

上面的 Pass 就完成模型的图优化、内存规划以及 Kernel 生成，上文提到的后端优化即在 Kernel 生成阶段体现，目前 MegCC 主要使用人工优化的 Kernel 模版。最终可以根据 Runtime 中定义的模型格式 dump 编译之后的模型，以及生成计算模型所需的 Kernel 文件。下面以一个简单的模型为例，使用 MegCC 的辅助工具(下载 Release 包) mgb-imp