曾经的CPU芯片和OS操作系统可能对我们普通程序员或者研发遥不可及,特别是国内的情况。但是随着技术的逐步积累和开放,可以看到国内外很多大学里面动不动就攒一个CPU或者OS,推动了计算机体系结构的发展,也打破了这些高精尖技术的壁垒,让我们普通小菜鸟也可以动手攒一个,或者至少可以了解到其运行的原理了。
这一切还是要源于计算机在工业中的发展迅速,就需要出现各种各样的定制化软硬件,定制多了,差异多了,就冒出来很多五花八门的CPU和OS,来适应不同的场景。我们身处在时代之中,尽管比国外落后一些,但是依然可以参与其中,并奋力追赶。
上图可以看到传统CPU的增长曲线已经停滞,这也正是Intel的没落和Invidia的崛起,以及国内外众多芯片公司以及AI公司崛起的原因。甚至被誉为“第四次工业革命”。
20 世纪 80 年代计算机迎来了黄金年代,形成了目前的冯诺依曼架构,但是随着新的AI硬件的出现,
- 图灵奖获得者 David Patterson (计算机系统结构学科「圣经」《计算机体系结构:量化研究方法》作者之一)在 2019 年 5 月发表了一个名为“计算机架构新的黄金年代”的演讲,David Patterson 介绍了从最初诞生的复杂指令计算机到精简指令计算机,再从单核处理器到多核处理器的发展,再到如今随着 AI 的飞速发展,DSA(Domain-specific architectures)迅速崛起,比如 AI 芯片、GPU 芯片以及 NPU 芯片,这些都是 DSA。他认为高级、特定于领域的语言和体系结构,将为计算机架构师带来一个新的黄金时代。扩展阅读:zhuanlan.zhihu.com/p/477346109
- LLVM 之父 Chris Lattner 在 2021 年 ASPLOS 会议的主题演讲中发表了名为“编译器的黄金时代”的演讲,主要分享了关于编译器的发展现状和未来、编程语言、加速器和摩尔定律失效论,并且讨论业内人士如何去协同创新,推动行业发展,实现处理器运行速度的大幅提升。扩展阅读:zhuanlan.zhihu.com/p/502730940
随着各种不同 AI 硬件的爆发式增长,会逐渐导致基础软件的碎片化,这种碎片化的发展带来了巨大成本,也会反噬 AI 硬件行业。当前,AI 编译器的发展阶段似乎回到了 GCC 出现之前的时代。每家 AI 芯片公司都在推出自己的 AI 编译器、框架甚至软件栈,市场上出现了极度碎片化的现象。一个对软硬件兼容性好的类似LLVM这个AI编译器对目前的技术发展十分的迫切需要。
之前的文章XXX介绍了编译器基础技术,那么AI编译器有又有什么样的不同,本文进行下说明。
参考:AI系统:/infrasys-ai.github.io/aisystem-docs/03Compiler02AICompiler/README.html
1. AI编译器介绍
1.1 为什么需要AI编译器
硬件平台变得越来越复杂和多层次:
- 计算加速平台采用了多层架构,包括标量、向量、多核、多包、多机架等不同层次的并行处理能力。
- 高端 SoCs 和 FPGAs 中,集成了特定领域的加速器,这些加速器针对特定类型的计算任务进行了优化
- 许多加速器的 IP(知识产权)现在是可配置的,这意味着可以根据不同的应用需求进行定制。可选的扩展、瓦片计数、内存层次结构等都可以根据需要进行调整,以满足特定的性能和功能要求。
多种多样的 AI 算法与不同硬件厂商出品的不同类型的硬件加速器之间有着巨大的鸿沟,AI 算法无法高效的在不同的底层硬件中运行。这中间缺少了 AI 编译器。
- 下一代编译器需要提供硬件抽象,以跨越多样化的加速器
- 编译器必须支持异构计算平台,因为不同的异构核擅长处理不同类型的任务。
- 领域特定语言(DSL)和编程模型为特定类型的应用提供了定制化的编程环境,使得开发者可以更高效地表达和优化算法。
- 确保基础设施的质量、可靠性和可扩展性。随着软件系统变得越来越复杂,编译器生成的代码必须经过严格的测试和验证,以确保其质量和性能。同时,编译器本身也需要具备高度的可扩展性,以适应不断变化的硬件环境和应用需求。
上图是一些比较大的芯片公司及其主要研究方向,但是随着应用的深入都有融合AI的趋势。对于AI编译器的开发主要就是这些芯片公司来开发,所以有些跟他们的芯片一样是不开源的(核心技术),对于后来者就需要买一些可以买到的软硬件+开源来做。而且以后有生命力的必将是开源的才能统一各家的使用,用起来方便免费不支持的硬件芯片就会被淘汰,这就是开源生态的魅力。
AI 领域的快速发展,导致每天都有各种各样的不同的 AI 算法被提出,这就需要:
- 对算子基于硬件进行优化,充分发挥硬件性能。
- 算子针对不同的硬件可移植性需要加强
不同芯片厂商提供 XPU 的 ISA(Instruction Set Architectrue)千奇百怪,各不相同。而当前阶段又缺乏如 GCC、LLVM 等编译工具链,这使得针对 CPU 和 GPU 已有的优化算子库和针对语言的优化 Pass 很难移植到 NPU 上。这样的现状可能对于某个较为领先的硬件厂商来说是技术壁垒,是一种优势,但随着技术的不断发展,很多 idea 会被不断地相互借鉴,不断整合类似的 Pass,在 AI 编译器领域推出一个类似于 GCC 或者 LLVM 的编译器是非常有必要的。这会极大的促进 AI 领域的发展!
上面为AI 编译器框架图
1.2 AI编译器跟传统编译器区别
1.2.1 编译目标:
- 相同点:AI 编译器与传统编译器都是通过自动化的方式进行程序优化和代码生成,从而节省大量的人力对不同底层硬件的手动优化。
- 不同点:传统编译器的起点是高级编程语言,终点则是硬件能够执行的机器码。相对地,AI 编译器的输入是神经网络模型的计算图,而输出同样是机器码。这在输入层面构成了传统编译器与 AI 编译器最根本的区别。
对于传统编译器而言,其核心使命是简化编程过程,将人类可读的高级语言代码转化为机器可执行的代码,并进行编译优化。
对于 AI 编译器,其主要目标则是优化整个程序的性能,确保神经网络模型在硬件上高效运行。降低编程难度虽然也是 AI 编译器的目标之一,但相较于性能优化,它退居次要位置。
1.2.2 优化方式:
- 相同点:在编译优化层,AI 编译器与传统编译器都是通过统一 IR 执行不同的 Pass 进行优化,从而提高程序执行时的性能。
不同点:如上图左侧是传统编译器
传统编译器的前端专注于对高级编程语言进行深入的语义分析、语法分析和词法分析,将源代码转化为中间表示(IR)。在中间阶段,编译器执行一系列优化 Pass,专门针对高级语言代码进行性能提升。而在后端,编译器负责处理代码的具体布局、寄存器分配等任务。
相比之下,AI 编译器的架构则有显著的不同。
- 它的前端主要负责将神经网络的 API 表达为计算图,这一过程涉及到模型的构建和转换。
- 在中间优化阶段,AI 编译器专注于图算融合、算子融合、自动微分和并行切分等特定优化技术。
- 后端则根据目标硬件平台,对 Kernel 进行定制化优化,确保代码在不同硬件上都能高效运行。在某些情况下,如 CPU 或 TPU 等芯片,AI 编译器甚至可能利用类似 LLVM 这样的传统编译器技术。
1.2.3 中间表达 IR 的差异
软件栈结构类似,相同点:
- 它们都分成前端、优化、后端三段式,通过 IR 解耦前端和后端使得可以进行模块化表示。
- AI 编译器依赖传统编译器: AI 编译器对 Graph IR 进行优化后,将优化后的 IR 转化成传统编译器 IR,最后依赖传统编译器进行机器码生成。因为传统编译器经过几十年的发展已经趋于稳定,所以 AI 编译器的角色更像是对传统编译器的一种补充。
AI 编译器的 IR 和传统编译器的 IR 所抽象出来的概念和意义并不相同。这些区别主要体现在它们处理的中间表示(IR)的抽象层次上。
- AI 编译器的 IR
- AI 编译器通常针对的是神经网络模型,它们处理的是高度抽象的 IR,这种 IR 通常是为了描述复杂的数学运算而设计的。这些运算包括但不限于卷积层(Convolutional Layers)、矩阵乘法(Matrix Multiplications)、激活函数(Activation Functions)等。
- 高层次的 IR 使得 AI 编译器能够更直接地表达和优化神经网络模型中的操作,因为它们与模型的高级结构和数学概念紧密相关。
- 传统编译器的 IR
- 传统编译器处理的是更为底层的 IR,这种 IR 更接近于机器指令,用于描述基本的计算操作。
- 这些操作包括加载(Load)、存储(Store)、算术运算(Arithmetic Operations)、控制流(Control Flow)等。
- 低层次的 IR 允许传统编译器进行更细致的优化,比如指令调度、寄存器分配、内存访问优化等。
- 抽象层次的意义
- 高层次 IR(如 AI 编译器使用的)使得编译器能够更专注于算法和模型级别的优化,例如通过融合操作、量化、剪枝等技术来提高模型的效率和性能。
- 低层次 IR(如传统编译器使用的)则让编译器能够进行更底层的优化,比如利用特定硬件的特性来提升性能。
- DSL(领域特定语言)
- 当提到 AI 编译器在描述神经网络模型类 DSL 时更加方便,这指的是高层次 IR 能够更自然地映射到神经网络模型的结构和操作。
- DSL 通常是为了特定领域或应用而设计的,它们提供了一种更直观和高效的方式来表达特定类型的程序或模型。
总结来说,AI 编译器和传统编译器的 IR 在抽象层次上的差异,反映了它们服务的不同领域和优化目标。AI 编译器的高层次 IR 使得它们在处理神经网络模型时更为高效,而传统编译器的低层次 IR 则让它们能够进行更细致的底层优化。
- 优化策略的差异
AI 编译器与传统编译器在优化策略上有所不同,主要体现在它们针对的领域和优化目标上。
- 领域特定优化:AI 编译器面向 AI 领域,优化时会引入更多领域特定知识,如算子融合,即将多个连续的运算操作合并,以减少内存访问和提高计算效率。
- 计算精度:AI 编译器可以降低计算精度,使用如 int8、fp16、bf16 等较低精度的数据类型,因为神经网络模型对精度的容忍度较高。而传统编译器通常不会改变变量类型和精度,以保证程序正确性。
- 硬件适配:AI 编译器在优化时会考虑特定硬件的特性,如 GPU、TPU 等,这些硬件针对深度学习任务进行了优化。传统编译器则需要考虑更广泛的硬件平台。
2. AI编译器的发展
AI 编译器的发展应该分为三个阶段:朴素 AI 编译器(阶段一)、专用 AI 编译器(阶段二)、通用 AI 编译器(阶段三)
借助AI编译器的框图来理解下:
- 最开始各家加速器PU的硬件不同,针对自己的硬件研究了算子和算法,编译器也是自己改的。基本就是闭源的,这样只有对算子非常熟悉的人才能用,使用者需要直接调用使用算子进行编程。PyTorch的算法需要重写调用算子才能实现(算法移植)
- PyTorch写的算法可以直接调用硬件对应的算子,使用者可以不用关心算子的具体实现,进行直接调用,对使用者比较友好,大大拓展了应用范围(PyTorch有大量的应用和算法)
- 对计算图和算子可以进行自动优化,使用者可以不用关心优化问题,写出统一的调用代码即可。实现完全相同的代码在不同的硬件上跑。
总结:1.各家自己的软硬件2.支持Pytorch平台调用自己的接口3.算法代码通用,自动根据硬件调用优化
2.1 朴素 AI 编译器(阶段一)
在 AI 框架中,算子是执行具体计算的操作单元,例如矩阵乘法、卷积、激活函数等。Kernel 在深度学习中通常指的是在底层硬件上执行特定计算任务的函数或代码块。手写 Kernel 意味着开发者需要手动编写这些函数,以确保它们能够在特定的硬件上以最高效率运行。这通常需要深入理解目标硬件的架构和编程模型,例如 GPU 的内存层次结构、线程组织等。
在 TensorFlow 中,为了实现高性能的算子,开发者可能会直接编写 CUDA 代码来创建算子的实现,这些实现直接编译为 GPU 上的机器指令,并在 GPU 上执行。或者对于一些较为通用的算子,TensorFlow 可以直接使用 CuDNN 库中的算子实现,而无需开发者手动编写 CUDA 代码。这种方式简化了开发过程,同时确保了计算的高性能。
总结来说,TensorFlow 早期版本中算子层的手写 Kernel 和对 CuDNN 的依赖是为了在英伟达 GPU 上实现高性能计算。这种方式虽然提供了优化性能的可能性,但也带来了开发和维护上的挑战。
TensorFlow 早期的静态图设计,有其天然的劣势。
- 在 TensorFlow 早期版本中,静态图的概念不是 Python 语言原生支持的,因此,开发者需要使用 TensorFlow 框架提供的 API 来构建计算图。这意味着开发者不能直接使用 Python 的控制流语句来动态地构建图,这限制了表达的灵活性和直观性。
- 由于静态图的构建方式与 Python 的动态特性不完全兼容,新开发者可能会觉得难以理解和使用。静态图要求开发者在执行任何计算之前就定义好完整的计算流程,这与 Python 中动态构建和修改对象的习惯不同,因此可能会降低易用性。
- 随着专门为深度学习设计的 DSA 芯片(如谷歌的 TPU 等)的出现,编译器和算子实现需要更好地适应这些硬件的特性,以充分发挥它们的性能。这些芯片通常具有与传统 CPU 不同的架构特性,如并行处理单元、高带宽内存等,这对编译器的优化策略提出了新的要求。
- 在静态图中,算子的粒度(即算子的大小和复杂性)和边界(即算子之间的界限)通常在图构建时就已经确定。这种预先确定的粒度和边界可能限制了编译器在运行时根据具体硬件特性进行更细粒度优化的能力,从而无法完全发挥硬件的性能。
- 硬件厂商通常会提供针对其硬件优化的算子库(如英伟达的 CuDNN),但这些库可能不是最优的,因为:模型和 shape 确定情况下的优化:在某些情况下,可能存在比库中提供的实现更优的算子版本,特别是当模型结构和数据形状(shape)已经确定时。
- SIMT 和 SIMD 架构中的优化空间:在单指令多线程(SIMT)和单指令多数据(SIMD)架构中,编译器可以通过调度(Scheduling)和分块(Tiling)等技术进一步优化性能,这些是硬件厂商的库可能没有充分利用的。
2.2 专用 AI 编译器(阶段二)
PyTorch 框架以其动态图(也称为即时执行模式)而受到欢迎,它允许开发者以 Python 原生的方式编写和修改神经网络模型。这种灵活性甚至成为了当前所有 AI 框架设计的参考标准,促使编译器开发者考虑如何将类似 PyTorch 的表达方式转换为优化的中间表示(IR)。
尽管动态图提供了灵活性,但在性能关键的应用中,静态图的优化潜力更大。因此,编译器需要能够将动态图转换为静态图,以便于进行进一步的优化。同时,PyTorch 引入了 AI 专用编译器架构,打开开图算边界进行融合优化,够更好地处理 AI 工作负载的特点,自动管理图和算子的边界。
在阶段二的 AI 编译器中,性能优化是一个关键的焦点,特别是在如何充分利用硬件资源方面。目前工业界已经有了很多产品在这一方面进行了尝试,如:TVM、Meta 推出的 TC、谷歌推出的 XLA 等。
在之前的编译器中,算子的边界是固定的,这意味着每个算子作为一个独立的单元执行,与其他算子的交互有限。而在当前阶段编译器开始打破这些边界,允许算子之间的更深层次的交互和优化。通过重新组合优化,编译器能够更有效地利用硬件资源,如 CPU、GPU 或专用 AI 加速器。编译器可能将大的计算子图拆解为更小的算子,这样可以更细致地进行优化(子图展开);也可能将多个小算子合并为一个更大的算子,减少数据传输和内存访问,提高执行效率(算子融合)。
总的来说,阶段二的 AI 编译器在表达和性能上都进行了显著的改进。在表达上,编译器能够处理类似 PyTorch 的灵活表达方式,并通过转换为计算图 IR 来进行优化。在性能上,通过打开计算图和算子的边界,并运用先进的编译优化技术,编译器能够更有效地利用硬件资源,提高神经网络模型的性能。这些改进使得 AI 编译器更加强大和灵活,能够满足日益增长的 AI 应用需求。
我们目前处在阶段二,就是各个厂商不光自己用,也推出让用户使用,用户使用的PyTorch,那就去支持PyTorch,提供IR。
2.3 通用 AI 编译器(阶段三)
前端:
编译器前端的组成集中展示在上图中间靠左部分。输入的神经网络模型格式可以来自多种框架(如 TensorFlow、PyTorch 等);这些模型通过符号表示的转换(如 TVM、nGraph 等)生成计算图;高级 IR/图 IR(设备无关)包含了 DAG(有向无环图)和基于 let-binding 的表示形式以及张量计算;计算图经过多种优化,如代数简化、操作融合、操作下沉、公共子表达式消除(CSE)、死代码消除(DCE)和静态内存规划等,得到初步优化的计算图;随后通过模式匹配和图重写等方法进一步优化,最终生成优化后的计算图;同时,编译器前端也提供了调试工具(如 IR dumping)可以以文本或 DAG 形式呈现。
后端:
编译器后端(Compiler Backend)负责将优化后的计算图转换为特定硬件平台的低层次表示,并进行硬件特定优化和代码生成。
编译器后端的组成集中展示再上图中间靠右部分。首先进行硬件特定的优化,包括内在映射、内存分配、内存延迟隐藏、循环优化、并行化等;这些优化通过自动调度(如多面体模型)和手动调度(如 Halide)进行;自动调优模块包含参数化、成本模型和参数搜索,旨在进一步优化性能;后端还利用内核库(如 Intel DNNL、NV cuDNN/TensorRT、AMD MIOpen 等)来提高效率;低级 IR/操作符 IR(设备相关)采用 Halide、多面体模型等独特的 IR 实现;编译方案支持即时(Just-In-Time)和提前(Ahead-Of-Time)编译;最终,代码生成模块将生成 LLVM、CUDA、OpenCL、OpenGL 等代码,以支持各种目标平台,包括 CPU(X86、ARM、RISC-V)、GPU(英伟达、AMD)、ASIC(TPU、Inferentia、NNP 等)和 DSP 等,从而适配越来越多的加速器。
阶段三的通用 AI 编译器将带来一系列创新特点,成为发展目标:
- 统一表达:在阶段三,通用 AI 编译器将实现计算图(Graph)和算子(Operators)的统一表达。这一创新意味着编译器能够在统一的框架下,同时进行图级别的优化和算子级别的优化,提升了编译过程的效率和效果。
- 自动化优化:通用 AI 编译器将在算子层实现自动调度(Schedule)、自动分块(Tiling)和自动代码生成(CodeGen),从而大幅降低开发难度,提高开发效率。
- 泛化优化能力:通用 AI 编译器将具备更广泛的优化能力,包括动静统一、动态形状(Shape)处理、稀疏性优化、高阶微分以及自动并行化等高级特性。
- 模块化设计:通用 AI 编译器将编译器本身、运行时系统、异构计算支持以及从边缘设备到数据中心的部署需求划分为独立的、可重用和可组合的模块。这种设计不仅增强了系统的灵活性和可扩展性,使其能够适应多样化的硬件架构和计算环境,而且通过提供简单直观的 API、详尽的文档和强大的工具支持,显著降低了开发难度,加快了开发速度,使得 AI 应用的部署和维护变得更加容易和高效。
后记:AI SoC的文章最近会有一个加速结束过程,然后开启新的知识分享。