作者是一位AI编译器从业者，谈一谈工作以来对AI芯片软硬件协同、AI编译软件栈、TVM/MLIR的泛泛看法。
作者主页：https://www.zhihu.com/people/dxq-15-8，欢迎关注~共同交流AI编译相关的内容。
1. 软硬件协同
2个核心要点，都是钱砸出来的感悟。
1）对于SIMD为主的DSA，指令的图灵完备决定了芯片的可编程性。例如某些算子不能使用NPU编程，还需要在另外一个cpu上编程，这种异构就会把现有的事情变得复杂。尤其是AI编译器就很难做的健壮，客观需要一个过程去逐步识别哪些是软件本身的问题，还是硬件本身就无法做到。
2）片上的内存和片内带宽。从DDR到计算部件的内存层级设计，非常关键，也决定了芯片的可编程性。如果数据从DDR到计算部件的开销比较大，自然是希望中间结果可以驻留在片上（片上的缓存最好也足够大），不溢出到DDR上，这样会让ai编译器变得异常复杂，那解决方案自然需要更长的时间成熟和稳定。
图1 memory hierachy for mobile vs. laptop vs. server
另外，片上的内存层级越少越好，能做成cache就做成cache，当然最好是scratch pad和cache让软件可以自己选择模式。scratch pad模式下，软件可以做一些比较极致的调度，把片上缓存成分利用，计算部件的算力才能充分发挥。
图2 Memory hierarchy of the GeForce GTX780 (Kepler)
2. AI编译栈
1）底层llvm编译器做好还是蛮难的。很多AI芯片选择VLIW架构，或者自己设计一套标量+向量张量的ISA，需要从头搭建一个llvm/gcc的编译器；对gpu公司来说，就是一个类cuda的编译器。国内能把这件事情做好的公司真的不多。
我们往往可以发现，不少公司的算子库，第一个版本都是基于汇编。不是大家不想用c/c++的异构语言，是因为支持这个语言的编译器本身就不好做，写出来的算子性能能够比上汇编算子，往往需要好几年时间的打磨。例如寒武纪，最开始是汇编算子，后来逐步到bang c，然后才有机会去做更高抽象的算子编程工具，例如基于tvm。
其他不少ai芯片公司，或者gpu公司，感觉也面临这样的问题。而对于选择riscv的公司，就省事很多，基于开源的riscv加上一堆intrinsic的支持，做少量的优化，就差不多了。
图3 RISC-V "V" Vector Extension Intrinsics
2）理解了1），算子库的状态就很好理解了。llvm编译器还不太理想的，只能堆更多人去写汇编算子。打磨几年，基于c/c++异构的算子团队才会起来，同时才有机会去做更高抽象的算子编程工具。但无论是写汇编还是写c/c++异构算子，对于最终用户自定义算子，门槛还是挺高的。
3）如何兼容cuda，是很多公司想做的事情。很多cuda代码已经在那里了，如果新的架构可以很容易的跑起来，性能也不会损失很多，对于芯片的推广，就是一个巨大的生态优势。AMD已经在这方面做很好的一个范例。但实际去分析AMD做的硬件和软件的事情，你会发现，这本身的设计和投入不是一个小的事情。
图4 移植CUDA到AMD的HIP，AMD提供复杂项目的移植服务，cuda-to-hip@amd.com
4）基于tensor的ai编译器，工作在这个领域已经5-6年了，感觉没什么特别的东西了，可能新的idea还在出来（例如基于micro kernel），但要解决的就是那些问题。
图5 常见ai compiler设计架构
图6 DL compiler比较：TVM, nGraph, TC, Glow, and XLA（2020年）
关键是基于芯片体系结构去解决问题，最终衡量标准是模型和程序的性能，如果基于达到手写或者硬件性能的90%以上了，应该也无所谓是基于tvm，还是mlir，或者是完全自研了。
图7 Google lasted work: MLIR Primer
当然，ai编译器的稳定性和泛化性还是很重要的，直接决定了芯片的推广范围有多大。
5）自定义算子。用户都想要，最好是能兼容cuda，或者更好用的编程方式。但都不太好做。提供c + intrinsic，用户基本很难用起来。值得关注类似triton或者tensorir，或者太极的基于tensor抽象的编程方式，但距离成熟可用，还有距离。
图8 OpenCL for Low-level Parallel Programing
6）训练和推理。推理感觉没啥好说的。训练对于硬件来说，必须搞定内存和片上带宽问题，另外还有片间和跨机互联问题。没有片上带宽问题，就不太需要特别极致的算子融合了。图层和算子就可以做的比较简单，就能利用xla，或者mhlo这种小集合的算子去表达tf或者pytorch的2k多算子。
如果硬件的指标不太好看，就会给软件带来更复杂的工作，例如需要更强的算子融合能力，可能一些现成的开源不一定能用了。这些资源投入都是十分巨大的，动辄就是几十人团队，需要2年以上时间去构建和打磨。
3. TVM和MLIR
我还是那个观点，两个不同的数据结构而已，怎么解决你的问题才最重要。
图9 AI模型在硬件上部署的一般流程
我们也看到了基于mlir的项目。总体来说，mlir是一个很符合传统编译器工程师的设计，dialect机制方便定制和扩展，从前端但最后生成机器指令，这个过程比较平缓。
图10 Chris@ASPLOS 2021
relay到tir算子expand的过程太陡峭，算子的负担可以稍微迁移一部分给图层，图层会有机会做更全局的优化。这些，对于交付来说，其实不痛不痒，mlir并没有带来巨大的收益。同时目前mlir还缺自定义算子的能力，类似tvm的dsl还没有构建出来，以及算子tiling或者tuning的工具也还没有。这些估计都需要时间去构建，估计2-3年？
图11 AutoTVM自动优化架构
总体来说，tvm当前更能满足产品交付的需求。mlir想做到相同的能力（算子能力的差距，图层能力其实差不多），还需要更长时间。当然二者都在不停演进，相互吸收优点。
作者：dxq
文章来源：NeuralTalk
推荐阅读
更多嵌入式AI干货请关注 嵌入式AI 专栏。欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
