MLIR 简介
MLIR(多级中间表示)是编译器实用工具的表示格式和库,它位于模型表示与生成硬件特定代码的低级编译器/执行器之间。
MLIR 本质上是用于现代优化编译器的灵活基础架构。这意味着它由一个中间表示 (IR) 规范和一个用于对该表示执行转换的代码工具包组成。(用编译器领域的话来说,当您从高级表示转换为低级表示时,此类转换可以被称为"降级"。)
MLIR 深受 LLVM 的影响,并且明显重用了后者的许多绝佳创意。它有一个灵活的类型系统,并允许在同一编译单元中结合多个级别的抽象来表示、分析和转换计算图。这些抽象包括 TensorFlow 运算、嵌套多面体循环区域,甚至包括 LLVM 指令和固定硬件运算与类型。
我们希望 MLIR 能引起许多群体的关注,其中包括:
希望优化机器学习模型的性能与内存消耗的编译器研究员和实现者
正在想办法将自己的硬件(例如 TPU、手机中可移植的神经网络硬件以及其他自定义专用集成电路 (ASIC))连接至 TensorFlow 的硬件制造商
在编写语言绑定时想利用优化编译器和硬件加速的人士
TensorFlow 生态系统包含许多在软件和硬件堆栈的多个级别上运行的编译器和优化器。我们希望逐步采用 MLIR 来简化此堆栈的各个方面。
这里引用一下官网的介绍:
MLIR is intended to be a hybrid IR which can support multiple different requirements in a unified infrastructure. For example, this includes:
- The ability to represent dataflow graphs (such as in TensorFlow), including dynamic shapes, the user-extensible op ecosystem, TensorFlow variables, etc.
- Optimizations and transformations typically done on such graphs (e.g. in Grappler).
Ability to host high-performance-computing-style loop optimizations across kernels (fusion, loop interchange, tiling, etc.), and to transform memory layouts of data. - Code generation "lowering" transformations such as DMA insertion, explicit cache management, memory tiling, and vectorization for 1D and 2D register architectures.
- Ability to represent target-specific operations, e.g. accelerator-specific high-level operations.
- Quantization and other graph transformations done on a Deep-Learning graph.
- Polyhedral primitives.
- Hardware Synthesis Tools / HLS.
架构分析中关注三个方面的表达,分别是计算架构(Computation Element),存储结构(Memory Hierarchy )和互联结构(Interconnect)。
数据流
根据定义,AI的数据流可以分为三类,输出静止(Output Stationary),权重静止(Weight Stationary)和行静止(Row Stationary)。
1. 输出静止(Output Stationary)
在输出静止中,输入数据在每一层神经元之间传输,而权重参数会被复制到每个神经元中,因此权重是随着输入数据一起移动的。
这意味着神经网络的每个节点(神经元)有自己的权重参数,并且这些权重参数在整个网络的训练和推理过程中不发生改变。
2. 权重静止(Weight Stationary)
在权重静止中,权重参数在所有神经元之间共享,在数据流过程中保持不变,而输入数据会根据不同的神经元进行移动。
这种方式下,权重参数对于整个神经网络都是固定的,而输入数据会根据不同的节点进行处理,权重参数不会随着数据移动而改变。
3. 行静止(Row Stationary)
在行静止中,输入数据在神经网络的每一层之间传输,而权重参数则按行(或列)进行共享,即每个神经元共享一行(或列)的权重。
这种方式允许在神经网络中共享权重参数,减少了参数数量,同时提高了计算效率。
这些不同的数据流模式影响了神经网络的训练速度、内存占用和计算效率。选择适合任务需求的数据流模式可以优化神经网络的性能。在实际应用中,根据具体的情况来选择合适的数据流模式是非常重要的。
GEMM 编译优化
1. 矩阵分块(Tile)
当前的处理器性能主要受限于内存墙,即计算速度要大于数据存储的速度。为了打破内存墙的约束,各类硬件包括CPU及其他专用处理器,会设置不同层次的存储单元,而这些不同层级的存储单元往往大小以及读写速度不同,一般越靠近计算单元的存储其存储容量也越小但访问的速度也越快。如果可以将计算过程的数据局部化分块,而这些分块的数据可以独立完成计算,那么分块的数据就可以放在层次化的存储中,然后通过不同存储间建立Ping-Pong的数据传输方式,将数据存储与计算解耦,从而可以有效得隐藏存储墙的问题,提高计算效率。矩阵运算就有这种特点,因而可以通过矩阵分块来加速运算。
2. 向量化(Vectorize)
向量化的操作,主要是利用硬件的向量化指令或者SIMD(单指令多数)指令的特性,实现一个指令周期对多个值操作的能力。如下图2所示,通过将4个数据组成向量,利用处理器可以处理4个元素的新向量的计算能力,可以将4个指令周期的处理时间,压缩成1个指令周期的处理时间,从而极大提高运算处理能力。
3. 循环展开(Unroll)
由于矩阵乘法有多层循环构成,如果底层硬件有一定的并行化能力,包括多线程多进程处理能力等,那么可以对循环进行适当展开,从而提高计算的并行度,增加并发执行的机会。比如将一个次数为1024的循环,展开成256次循环,新的循环内又包含4条可以并行执行的展开计算,如果处理器能够并行处理循环内部的展开计算,那么通过对原来的循环展开,可以获得接近4倍的性能提升。
MLIR 相关项目
- tensorflow:没有tf就没有MLIR
- mhlo:tensorflow组件,相当于支持动态规模的XLA
- tfrt:tensorflow组件,tf新的runtime
- torch-mlir:连接pytorch与mlir生态
- onnx-mlir:连接onnx与mlir生态
- iree:深度学习end2end编译器
- circt:硬件设计及软硬件协同开发
- flang:FORTRAN的编译器前端
- polygeist:C/C++ source code变成mlir Affine