派大星 · 2022年11月16日 · 北京市

谈谈AI编译软件栈、软硬件协同、TVM和MLIR的看法

作者是一位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编译器变得异常复杂,那解决方案自然需要更长的时间成熟和稳定。

image.png

图1 memory hierachy for mobile vs. laptop vs. server

另外,片上的内存层级越少越好,能做成cache就做成cache,当然最好是scratch pad和cache让软件可以自己选择模式。scratch pad模式下,软件可以做一些比较极致的调度,把片上缓存成分利用,计算部件的算力才能充分发挥。
image.png
图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的支持,做少量的优化,就差不多了

image.png

图3 RISC-V "V" Vector Extension Intrinsics

2)理解了1),算子库的状态就很好理解了。llvm编译器还不太理想的,只能堆更多人去写汇编算子。打磨几年,基于c/c++异构的算子团队才会起来,同时才有机会去做更高抽象的算子编程工具。但无论是写汇编还是写c/c++异构算子,对于最终用户自定义算子,门槛还是挺高的

3)如何兼容cuda,是很多公司想做的事情。很多cuda代码已经在那里了,如果新的架构可以很容易的跑起来,性能也不会损失很多,对于芯片的推广,就是一个巨大的生态优势。AMD已经在这方面做很好的一个范例。但实际去分析AMD做的硬件和软件的事情,你会发现,这本身的设计和投入不是一个小的事情

image.png

图4 移植CUDA到AMD的HIP,AMD提供复杂项目的移植服务,cuda-to-hip@amd.com

4)基于tensor的ai编译器,工作在这个领域已经5-6年了,感觉没什么特别的东西了,可能新的idea还在出来(例如基于micro kernel),但要解决的就是那些问题。
image.png
图5 常见ai compiler设计架构

image.png

图6 DL compiler比较:TVM, nGraph, TC, Glow, and XLA(2020年)

关键是基于芯片体系结构去解决问题,最终衡量标准是模型和程序的性能,如果基于达到手写或者硬件性能的90%以上了,应该也无所谓是基于tvm,还是mlir,或者是完全自研了。

image.png

图7 Google lasted work: MLIR Primer

当然,ai编译器的稳定性和泛化性还是很重要的,直接决定了芯片的推广范围有多大

5)自定义算子。用户都想要,最好是能兼容cuda,或者更好用的编程方式。但都不太好做提供c + intrinsic,用户基本很难用起来值得关注类似triton或者tensorir,或者太极的基于tensor抽象的编程方式,但距离成熟可用,还有距离。
image.png
图8 OpenCL for Low-level Parallel Programing

6)训练和推理。推理感觉没啥好说的。训练对于硬件来说,必须搞定内存和片上带宽问题,另外还有片间和跨机互联问题。没有片上带宽问题,就不太需要特别极致的算子融合了。图层和算子就可以做的比较简单,就能利用xla,或者mhlo这种小集合的算子去表达tf或者pytorch的2k多算子。

如果硬件的指标不太好看,就会给软件带来更复杂的工作,例如需要更强的算子融合能力,可能一些现成的开源不一定能用了。这些资源投入都是十分巨大的,动辄就是几十人团队,需要2年以上时间去构建和打磨。

3. TVM和MLIR

我还是那个观点,两个不同的数据结构而已,怎么解决你的问题才最重要
image.png
图9 AI模型在硬件上部署的一般流程

我们也看到了基于mlir的项目。总体来说,mlir是一个很符合传统编译器工程师的设计,dialect机制方便定制和扩展,从前端但最后生成机器指令,这个过程比较平缓

image.png

图10 Chris@ASPLOS 2021

relay到tir算子expand的过程太陡峭,算子的负担可以稍微迁移一部分给图层,图层会有机会做更全局的优化。这些,对于交付来说,其实不痛不痒,mlir并没有带来巨大的收益。同时目前mlir还缺自定义算子的能力,类似tvm的dsl还没有构建出来,以及算子tiling或者tuning的工具也还没有。这些估计都需要时间去构建,估计2-3年?

image.png

图11 AutoTVM自动优化架构

总体来说,tvm当前更能满足产品交付的需求。mlir想做到相同的能力(算子能力的差距,图层能力其实差不多),还需要更长时间。当然二者都在不停演进,相互吸收优点。

作者:dxq
文章来源:NeuralTalk

推荐阅读

更多嵌入式AI干货请关注 嵌入式AI 专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
16328
内容数
1218
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息