大家好！感谢大家在百忙之中出席今天的分享会，我非常荣幸有机会向大家汇报过去半年在MLIR小队的工作成果。接下来，我将更详细地分享我的工作，涵盖从PyTorch到RISC-V V扩展验证流程的全面完成，以及与chipsalliance/t1项目的深度合作等方面。
首先我要聊的是PyTorch与buddy-mlir的联系。PyTorch作为目前流行的深度学习框架，API 设计的简单直观易懂，学习曲线较为平缓，而且拥有着丰富的图像语音多模态的扩展库支持，赢得了大量用户和强大的社区支持。PyTorch 还有着动态计算图的特性，计算图是在运行时构建而不是编译时构建。这使得在开发时 AI 模型的构建更加灵活，能更加轻松的动态调试和调整模型。
然而尽管PyTorch框架在AI生态中非常活跃，但其在高性能计算方面的局限性不可忽视。由于 Python 其本身的语言特性，比如其解释性质和动态类型系统，同时也没有 JIT，十分不适应高性能计算场景。其次，正是由于上述所提的灵活模型构建性质，使得在PyTorch中进行静态分析，以及针对架构进行优化变得相对困难。高性能计算往往需要对代码进行静态分析以进行更有效的优化，但是PyTorch的动态图特性使得这一过程变得复杂。
而单为 PyTorch 实现编译器并不可行。首先 PyTorch 本身带有各种不同的图 runtime，每一个都会需要做一个前端实现。而接下来，由于没有社区统一的约定，我们将会看到各个runtime或者各个实现，都有着自己定义的一层 IR。这些 IR 又需要各自对接硬件，可能在相同的硬件上出现相同反复的造车轮。而即使有了完整的编译器实现，由于高级语言在编译时会一步到位最终 IR，由于最后底层 IR 表达能力较弱，像对循环操作的优化，消除多余的数组复制操作，特殊指令调用等优化做起来成本很高，针对这些比较需要高抽象层级的操作，来对特定硬件做优化会消耗很大的人力。而 MLIR 为解决这些问题应运而生。
MLIR 是一套编译器基础设施，它的一个核心设计思路是：用一个统一的 IR 来表达不同层级的编译。在 MLIR 里有着可自定义的 Operation，每一个 Operation 都有着自己的层级，他们用着相同的 IR 表达，在编译期间相互转换。从前端词法解析，从高抽象级别的算子，到稍微低层级的循环控制，再到更低层级的逻辑运算，最后到最低层级的机器语言表达等。MLIR 通过把一个端到端的编译流程，分割成语义上的，可自定义的各个层级，开发者们能选择性的在某一层级实现自己需要的优化，同时由于 IR 表达相同，还可以直接复用其他原有的基建，完成常见的编译优化。除此之外，他还能允许用户部分的降低抽象层级。比如说对于一个异构设计的编译器，其可以在保持其他操作还在高抽象层级的基础上，将加速器特定的常量或向量指令降至底层机器指令级别。这可以使得开发者们在设计和使用 IR 时，可以更加专注于自己的特定领域，不需要花特别大的精力在一些基础设施上。而从硬件厂商的角度来讲，由于这种可定制性和可复用性，他们可以很轻松的拓展 MLIR 的后端支持。
buddy-mlir 利用了 MLIR 的这一套编译器基础设施，打造了一套不仅限于 PyTorch 的 DSL 到 DSA 的编译生态。分为前端中端后端三个部分，在前端，buddy-mlir 使用 Dynamo Compiler 做静态图的捕捉，将 Aten IR 转换到 MLIR 的 Linalg/TOSA Dialect 。在中端，对模型的循环操作做并行优化，针对后端实现向量化等编译优化。再到后端对接 X86/ RVV 等架构。基于 MLIR 的模块化和可拓展基建，完成了一套同样模块化，可拓展的，但具有领域特定，针对性优化的编译器生态。下面我来简单分享一个在 RVV 上跑通端到端 PyTorch 算子的用例。在这个例子中，我将会使用 buddy-mlir 的工具链，将 PyTorch 的 addmm 算子在 chipsalliance/t1 这个项目上跑通。t1项目是一个使用Chisel生成、支持RISC-V V拓展的长向量IP核心，是一颗完整支持RISC-V V 1.0 spec，多通道，全流水线 chaining，高内存带宽利用率的核心。在 t1 上跑通是我们用来对 RVV 向量化后端的验证。
这里是个简单的使用 addmm 算子做计算的 Python 代码。在 18 到 23 行，我们使用 buddy-mlir 提供的 Python Binding，通过 DynamoCompiler, 遍历从 TorchDynamo 传入的 FX Graph，并且针对每个 Aten IR 节点进行转换，再生成出对应的 MLIR 操作。最终生成的 MLIR module 将会被打印到 stdout 上。经过 DynamoCompiler 生成的 MLIR 代码如图所示，可以看到原来的 addmm 操作，已经变成了 MLIR 里 tosa Dialect 对应的操作了。
而正如上述所提，tosa Dialect 是一个高层级的抽象，我们还需要经过一系列的操作，将其做对应优化和层级下降。如图所示，在传递 pipeline 信息给 buddy-opt 之后，tosa Dialect 已经被降低到相对低层级的 linalg Dialect 上，我们也可以观测到，已经出现了具体的对 tensor 的读取和加法操作。而再经过一系列优化和下降操作后，这些原本高层级的 Dialect 最终会下降到最低层级的，对 llvm Dialect 的引用，而我们也可以使用 buddy-translate 工具，将 mlir 转换到 llvm ir 上。如图所示，上图是最终优化和下降完毕的，只剩下对 llvm dialect 引用的 MLIR 代码，而下图是经过 buddy-translate 工具转换后，生成的 llvm IR 代码。有了 LLVM IR 之后，我们就可以通过 buddy-llc 将 LLVM IR 编译到指定平台的二进制 object 文件上了。
值得一提的是，t1 的仿真环境是 baremetal 的，既没有 kernel，没有 syscall，也没有 filesystem。而 bert 依赖往 stdout 打输出，也需要 malloc 申请内存，buddy-mlir 是如何为这种嵌入式的后端做实现的呢？实际上并不需要。在完成编译优化和下降之后，模型的 Forward 函数将会被构建为共享库 。再加入 request llvm wrapper 这个参数后，最终会暴露出一个 C API mlir_ciface_forward ，不需要在编译器内为某平台做支持，有了 C API，平台方就可以根据自己的需求来安排输入和输出。不管是如何管理和存放权重数据，还是获取模型的输出放到指定地址上，把模型当作一个计算库，一个大的算子来使用之后，平台方就有了足够的自由去自行驱动这个模型，在这里，将模型的编译与平台的执行解耦。在 t1 上，我帮助实现了 sbrk，write 等函数，用上 t1 自己的 UART设备和内存模型来做输入输出以及内存分配，借此成功的在 t1 上跑通 PyTorch 的算子和模型。
接下来还有我们成功在 t1 上跑通的小语言模型 bert，同样是通过 buddy-mlir 的工具链，从 Python 到 RVV 端到端的编译跑通的。
由这两个示例可以看出 buddy-mlir 的优势：对于用户而言，他只需要调 PyTorch 的库，经过 buddy-mlir 工具链之后，就能几乎无感的将这些代码放在 rvv 硬件上运行，同时还能获得 rvv 的计算优化。对于编译器前中端开发者而言，他们完全不需要担心并发，异步，数据操作等底层的基础设施实现，所有的底层实现都是可复用的，他们可以再高抽象层级对应用做优化。对于硬件厂商而言，他们不需要非常操心如何给 Python 做自己的平台支持，不需要担心如何为特定 AI 库做计算优化，他们可以为 MLIR 实现自己的后端支持，就能复用前端中端的编译基建。这即是 MLIR 的核心思想：拆分领域特定的编译问题到不同的小的编译器上，在不同的层级对这些编译问题做实现和优化，这是一个合理且可持续发展的开发思路。