Page 99 - 《软件学报》2024年第6期
P. 99
张洪滨 等: AutoConfig: 面向深度学习编译优化的自动配置机制 2675
调用 AutoConfig 的各个模块对关键的深度学习操作进行分析和配置, 随后选择适当的代码生成策略进行编译优
化, 并将代码编译下降后翻译至 LLVM IR, 最终生成针对目标硬件的汇编代码. Buddy Compiler 支持面向 CPU SIMD/
Vector、GPU、加速器的代码生成, 通过集成 AutoConfig 的优化与调优方法, 它能够实现面向多种硬件后端的性
能迁移和优化复用, 从而打造出高效通用的编译通路.
除了在 Buddy Compiler 中的集成之外, AutoConfig 更重要的贡献是为 MLIR 生态提供了新的调优方法和思
路. MLIR 是深度学习编译器的主流生态之一, MLIR 上游提供 PDL 和 Transform 方言作为基础对高层中间表示进
行调度和优化, 例如循环分块、顺序调整、向量化等. 基于 MLIR 的深度学习编译器 IREE 除了采用上游提供的
基础方言之外还通过自定义方言进行调度, 并且使用强大的运行时环境实现多种平台的部署. 不同于使用特定方
言进行调优和调度, AutoConfig 的创新点在于使用多层级的方言 (例如 Vector 方言, Affine 方言等) 编写可配置的
优化算法重写模式, 并为优化算法建立分析模型, 搭配硬件信息收集策略提供的开销函数, 可以在编译时确定优化
重写模式的配置参数, 从而实现面向特定计算负载和硬件的代码生成. 同时, AutoConfig 与 MLIR 生态中的其他调
优方式并不冲突, 使用上游调度方言或者自定义方言的方式也可以集成 AutoConfig 使得其调优过程具备可解释
性和可预测性.
DM , 所包含的特殊指令种数为
4 优化分析模型
AutoConfig 的可解释性源于其优化分析模型. 该分析模型可以对优化算法的计算特征进行建模, 从而近似预
测其执行开销. 本节提出一种向量尺寸可配置的优化分析模型建模方法, 可以在编译时根据高层中间表示中计算
负载的输入尺寸和 SIMD 硬件平台特征来预测计算负载优化算法的期望开销, 从而有效指导 AutoConfig 进行代
码生成. 值得注意的是, AutoConfig 并不耦合与某个特定的优化分析模型, 而是开放出可注册的接口, 用户可以按
需注册自己的优化分析模型.
4.1 模型总体描述
在深度学习编译优化中, 对优化算法的选择以及算法内的参数配置都会对性能产生影响. 这些影响主要体现
在计算开销、访存开销和特殊指令执行开销 3 个方面. 一般来说, 优化算法与朴素算法相比具有更低的计算开销,
但是由于对算法的优化通常会引入额外的内存排布操作和特殊指令, 这会带来额外的访存开销和特殊指令执行开
销. 设 Cost All 为一个优化算法的理论总开销, 则有:
(1)
Cost All = Cost Arith +Cost Memory +Cost SpecIns
Cost Arith 表示计算开销, 在 SIMD 浮点计算场景中主要指融合乘加操作 Cost Memory 表示
其中, (FMA) 的执行开销.
Cost SpecIns 则涵盖了除了计算指令和访存指
访存开销, 主要指访问内存、缓存和寄存器组等存储模块所需的开销.
令以外其他特殊指令的执行开销. 不同的优化算法通常会引入不同的特殊指令, 比如 Broadcast 向量化算法会引入
广播操作指令, 针对卷积优化的 Im2Col 算法会引入维度变换指令等. 朴素算法一般不引入特殊指令, 它的理论总
开销只由计算开销和访存开销构成.
为了衡量这些开销, 设优化算法中包含的浮点运算指令个数 (floating point operations) 为 FPO , 内存访问指令
个数 (data movement) 为 M , 且第 条特殊指令的个数为 Num SpecIns i , 则上述公式可
i
进一步描述为:
M ∑
Cost All = λ Arith FPO+λ Memory DM + λ SpecIns i Num SpecIns i (2)
i=1
其中, λ Arith 、 λ Memory 和 λ SpecIns 分别表示计算开销、访存开销和特殊指令开销的重要性系数. 以计算开销为例, 由于
它的大小与浮点运算指令个数成正比, 因此 λ Arith 的含义就是单位运算指令开销. 重要性系数与优化算法的参数配
置和执行平台紧密相关, 优化算法的参数调优越有效, 越能充分发挥执行平台的硬件特点, 则重要性系数越小. 在
实际优化中, 该系数将通过第 5.2 节的动态测量流程确定.