深度学习的编译优化方法师承程序语言,相比而言,程序语言的输入是语法树,深度学习的输入是计算图,有相似的编译结构,深度学习后端可以借鉴LLVM架构灵活支持新后端硬件。
不少社区都提出了自己的编译优化框架,定义了自己的一套中间表达(IR,Intermediate Represent),在发展过程中,这些框架和IR也在相互借鉴相互融合。
# XLA
XLA(Accelerated Linear Algebra)是TensorFLow给出的编译优化框架,嵌入在Tensorflow的图层面的Grappler优化中,提供了子图划分,IR转换,Codegen(代码生成)等功能,属于JIT(Just-In-Time)流程。
XLA定义一套Tensorflow算子集的子集Instruction,在IR转换阶段将Tensorflow的算子转换为等价表达,随后执行编译优化以及Codegen。
对于通用计算硬件,能在子图划分阶段进入候选集的算子基本都能完成整套编译优化流程,XLA后端将XLA Instruction转换为LLVM IR使用LLVM生成代码,能做到较好的模型覆盖。
但对于专用硬件,一是支持的算子较少,分割的子图大小和数量都较多,会引入较高的跨设备调度开销,通常会设定编译触发阈值,二是支持的子图pattern受限,纯粹使用算子集划分子图,后端硬件可能不支持,需要fallback回Tensorflow图,或者提前加入图匹配。
但XLA对子图输入Shape敏感,导致每种Shape都会触发一次编译,XLA的编译结果Cache很有可能溢出,尤其是对于输入极不固定的模型,频繁编译会大大影响模型性能,甚至导致负优化。
另一方面,XLA的代码生成仍然是启发式的,不能在指令层面优化,也没有动态搜索,导致XLA在大部分常见场景上不敌手工图优化。
因此,XLA的一个优化方向是使用第三方库负责计算密集型算子的计算,XLA编译机制作为fusion优化,能较好平衡优化性能。
# TensorRT
TensorRT是NVIDIA提供的一套独立的模型构建工具,提供了算子级别的API,用户可以用来直接构建计算图,并提供了转换工具将框架的模型文件转换为Tensor RT的执行参数。
同时,TensorRT也在图层面用和XLA类似的方法接入了TensorFlow,TensorRT内部负责优化子图和生成GPU kernel。
# OpenVINO
OpenVINO是Intel提供的一套独立的模型构建工具,并提供了转换工具执行全图转换,支持的算子的覆盖度有待提高。内部应该也是基于OneDNN(原MKL-DNN)提供的构建接口。
# TVM
TVM的主要思想是搜索计算密集型算子的所有可调的参数空间,比如循环展开的基本宽度(可能影响Cache),循环顺序(可能影响分块逻辑),数据Layout等。
TVM使用TVM IR定义了基本操作,可以直接通过TVM IR描述某个计算密集型算子的计算结构,使用Schedule定义在计算结构上搜索(Tuning)的空间。
在计算图层面,TVM定义了一套Relay IR作为计算图层面的表达,各个框架的算子可以转换为Relay IR,再统一优化,执行fusion等操作。
可以通过XLA接入Tensorflow。
# PolyHydral
多面体优化,TVM所使用的方法来源,但更侧重于使用数学方法启发式地寻找最优解,结合少量Profiling,定义每套硬件的启发参数。
# MLIR
MLIR旨在统一多种框架的多层IR,共享优化方法,并将图信息跨IR保留,解决Dynamic Shape的问题。