背景

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

图1 Attention结构示例

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

J6平台高效支持Reshape和Transpose

架构介绍

图2 J6芯片架构图

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

下面介绍一下当前编译器为了高效支持Reshape和Transpose，所做的一些优化手段。一方面，将Reshape和Tranpose转换为硬件高效支持的操作，另一方面，对Reshape和Transpose实现高效的Tiling支持，从而充分的利用数据局部性，提高并行性。

Reshape和Transpose的变换

一般情况下，对于加速器来说，不同的硬件计算单元会设置不同的layout(同cutlass中的layout概念)，从而实现相应的并行度。为了更好的衔接不同的硬件计算单元，也会有一个部件来支持不同layout之间的高效转换。这里，我们会利用硬件对于不同layout之间的高效转换以及View操作来实现Reshape和Transpose操作。

下面通过例子分别介绍一下如何通过layout之间的转换和View操作来实现Reshape和Transpose。

Layout表示

通过block和order两组信息来表示一个层次化的layout，其中block和order都是层次化的。

图3 tensor的layout的图示

对于tensor来说，上图是实际数据存储的图示。其中block表示内层数据的排布方式，上述例子中block = "2W2H"表示对于内层数据来说，第一个维度在内，第二个维度在外。order表示block之间的排布方式，order = [1, 0]表示第一个维度的block排布在内，第二个维度的block排布在外。

Reshape变换

给定IR为

通过将输入layout设置为 {block = "128W", order = [0, 1]}，输出layout设置为 {block = "128W", order = [0, 1, 2]}，则可以将该reshape转换成没有开销的View操作，即变换后的IR为

Transpose变换

同理，对于Transpose来说，也可以通过这种方式将其转换为没有开销的View操作

给定IR为

变换后的IR为

Reshape和Transpose的Tiling

Tiling的收益主要在两个地方：提高局部性和并行性。相对于普通的算子来说，Reshape和Transpose的tiling处理相对复杂一些。例如，对于perm = [1, 0, 2]的transpose来说，输出在H方向tiling，对应的输入需要在W方向tiling。因此，对于包含reshape和tranpose操作的子图来说，需要记录每个op输出的tiling信息，而不是只用一个tiling方向来表示整个子图的tiling。

图4 Reshape和Tranpose子图的tiling

如上图4所示，对于子图中linear，reshape和transpose输出来说，tiling方向均不同。由于有reshape操作将一个维度映射成两个维度，因此需要两次tiling，才能将reshape的维度拆分完。

总结

编译器通过将reshape和transpose转换为layout转换操作和view操作，以及实现对reshape和transpose精细的tiling，高效的原生支持transformer模型。相比于优化前，编译出的模型性能提升接近40%。