背景
相对于传统 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 转换为硬