[征程 6 编译优化] Transformer 模型在 征程 6 平台上的高效支持

背景

相对于传统 CNN 模型来说，Transformer 模型的最大的一个特点就是灵活性。这个灵活性主要体现在模型中穿插大量的数据重排操作，即 Reshape 和 Transpose。如下图 1 所示，对于一个典型的 Attention 结构来说，Reshape 和 Transpose 操作的数量大概可以占到所有操作的一半。

图 1 Attention 结构示例

由于在 Transformer 模型中，Reshape 和 Transpose 操作出现的概率很高，并且对于大部分加速器来说，都没有相应的硬件单元来直接支持。因此，高效的支持 Reshape 和 Transpose 就成为 Transoformer 模型优化的关键。

征程 6 平台高效支持 Reshape 和 Transpose

架构介绍

图 2 征程 6 芯片架构图

如图 2 所示，征程 6 芯片为异构架构设计，包含 Tensor 单元、Vector 单元以及 Scalar 单元。Tensor 单元中又包含多种计算单元，如 TAE、VAE、AAE 等，每个计算单元对数据的排布格式（即 layout）有不同的要求。因此，芯片中有强大的 layout 转换单元来支持数据在不同 layout 之间转换。Reshape 和 Transpose 即可利用该 layout 转换单元来高效实现。

下面介绍一下当前编译器为了高效支持 Reshape 和 Transpose，所做的一些优化手段。一方面，将 Reshape 和 Tranpose 转换为硬