GPU|BladeDISC 0.2.0更新发布 CPU|深度学习

文章图片

在BladeDISC正式开源三个月后，我们发布了0.2.0版本，该更新包含了大量的性能优化与功能增强。
BladeDISC是目前业界领先的支持动态shape的深度学习优化编译器。深度学习优化编译器负责将上层的神经网络计算图转换为底层硬件可执行的程序，当前流行的深度学习优化编译器（TVM[1
、XLA[2
、TensorRT[3
等）对静态shape的支持力度较大，对动态shape的支持则有所欠缺。其中， XLA目前只支持静态shape ， TensorRT可以支持ranged shape（即指定尺寸范围内的动态shape）。 BladeDISC编译器提供了对动态shape的完整支持，可以将包含动态shape语义的上层模型描述转换为高效的底层可执行程序。更多关于BladeDISC的介绍内容，请查看此处[5
。
本文描述BladeDISC v0.2.0版本相对于最初开源版本（即v0.1.0版本）的主要更新内容。
性能优化当前的神经网络计算图主要由访存密集型算子（主要包括element-wise/elemental算子及reduce算子）和计算密集型算子（主要包括GEMM和Convolution算子）组成。 BladeDISC v0.2.0版本相对v0.1.0版本在这两方面都做了大量优化。
访存密集算子的优化本章重点描述v0.2.0新增的GPU stitch优化（以及附带的shape constraint功能增强）。
GPU上的stitch优化方法来源于我们发表在ASPLOS 2022上的论文AStitch[6
，我们正在逐步地将AStitch中的技术迁移到BladeDISC中，目前已完成部分迁移工作。 Stitch的方法旨在解决当前深度学习模型在GPU上执行时遇到的大量的kernel调度（scheduling）和发射（launch）开销以及片外访存开销，以及不同shape之下memory-intensive算子的高效schedule问题，其基本方法是通过GPU上多层次的中间结果访存管理（以寄存器， shared memory及global memory作为中间数据的buffer）、层次化的数据同步控制以及自动适配的parallelism-aware的codegen schedule生成，将原先多个不同的fusion kernel给stitch在一起，并生成高效的计算代码。
本次BladeDISC release实现了通过shared memory将原先被reduce算子分割开的多个fusion进行stitch融合的功能，达到计算融合加速效果（跨CUDA thread block的global memory的stitch以及自动适配的parallelism-aware的codegen在本次release中暂未实现）。我们提供了一组示例来展示如何开启stitch优化。

动态shape给stitch优化带来了新的挑战，其最重要的挑战在于，在不知道确切shape的情况下，正确地判断被stitch在一起的op之间的数据局部性，从而正确得为其中的数据传输分配片上存储（即GPU shared memory）。本次更新包含了一系列的shape constraint功能增强，作为GPU stitch功能的基础支撑。其基本思想是挖掘与传播producer op和consumer op之间以及sibling op之间的shape constraint ，通过symbolic的方式构建全局的shape等价关系；在此基础上，进一步构建shape dimension之间的乘法关系，用以解析reshape等算子的shape等价性关系。
值得说明的是，在前一个版本中（即BladeDISC开源版本v0.1.0），我们针对CPU上的模型，对部分访存密集型计算子图实现了CPU端的stitch优化，从而更好地增强数据局部性，提升访存效率。
计算密集算子的优化本章主要介绍本次更新中包含的GEMM合并优化及CPU上计算密集型算子的pre-packing和layout优化。
GEMM 合并优化本次release实现了以下两种GEMM合并优化：