面向ASIC设备的编译器框架：TVM or MLIR？

最新推荐文章于 2024-08-15 12:08:12 发布

Linux基金会AI&Data基金会

最新推荐文章于 2024-08-15 12:08:12 发布

阅读量1.8k

点赞数 5

文章标签：编译器芯片 java 人工智能深度学习

本文链接：https://blog.csdn.net/lf_ai/article/details/124763065

版权

本文探讨了面向ASIC设备的编译器框架TVM和MLIR的优缺点。TVM是一个深度学习模型编译器，支持多种框架，提供完备的优化工具，而MLIR作为编译器基础设施，旨在简化自定义编译器的构建。TVM的VTA硬件加速器提供全栈解决方案，但硬件设计受限于其ISA。MLIR则允许更多硬件设计自由，但其生态尚未提供完整的DSA开发和编译解决方案。在实际工程开发中，两者都需要深度定制。

摘要由CSDN通过智能技术生成

2019~2021年，“摩尔定律失效”这一关键词频频出现于各大技术网站，在此背景下，市面上多如牛毛的AI芯片公司不约而同地给出了通用CPU+专用ASIC芯片的方案，以应对日益增长的AI边、端侧推理计算需求。在AI DSA芯片的开发实践中，棘手的问题除了底层硬件的设计，更多的还是AI模型在DSA芯片上优化、部署执行这一过程所需软件栈的实现，也即“AI编译器”技术栈，在这一领域最常常被大家提起并衡短论长的，莫过于TVM和MLIR。

严格来说，MLIR和TVM并不适合在一起对比：TVM是面向深度学习的模型编译器，用户可借此可直接获得编译/优化模型为推理blob的能力（可以看做机器学习时代的GCC、Clang），有兴趣了解TVM的同学可以参考这里；MLIR则是编译器基础设施类软件（可以看做机器学习时代的LLVM），面向的是需要构建自定义编译器的用户，它的基本设想是通过MLIR的Dialect共享生态减少用户开发编译器的工作量。

以上图为例，TensorFlow本身的优化推理过程如蓝色和绿色部分所示：TensorFlow Graph首先转换为XLA的HLO IR，应用XLA的优化Pass后再lower到目标设备的IR，例如对于x86/arm可lower为LLVM IR等，这样做的问题在于：

XLA的图优化Pass是封闭的，用户如果想在其他框架例如TVM下实现同样的优化，需要阅读XLA的源码和TVM的源码，并添加Pass代码到TVM中，开发成本很高
XLA HLO到LLVM IR的跨度太大，实现开销大，此处的开销包括各类针对微架构的带宽、缓存、指令集的优化

MLIR解决问题的方式如上图红色部分所示：先把外部IR转化为MLIR格式的Dialect IR（Translation），再把该IR lower到其他的MLIR方言IR（Conversion），最后从Target Dialect IR转出为Target IR（Anti-Translation），这么做的好处是，如果lower过程中所需的IR Dialect以及Target Dialect已经由MLIR生态中的其他用户实现了，那么你就可以直接从这些Dialect及其对应的优化Pass中受益。假设方言Me