背景

相对于传统 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 转换为硬件高效支持的操作，另一方面，对 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<4x8xsi8>的 layout 的图示

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

Reshape 变换

给定 IR 为

%1 = "hbir.reshape"(%0) <{shape = [4, 32, 128]}> : (tensor<128x128xsi8) -> tensor<4x32x128xsi8>

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

%1 = "hbir.view"(%0) <{}> : (tensor<128x128xsi8, layout = {block = 128W, order = [0, 1]}>) -> tensor<4x32x128xsi8, layout = {block = 128W, order = [0, 1, 2]}>

Transpose 变换

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

给定 IR 为

%1 = "hbir.transpose"(%0) <{perm = [1, 0, 2]}> : (tensor<4x32x128xsi8>) -> tensor<32x4x128xsi8>

变换后的 IR 为

%1 = "hbir.view"(%0) <{perm = [1, 0, 2]}> : (tensor<4x32x128xsi8, {block = 128C, order = [0, 1, 2]}>) -> tensor<32x4x128xsi8, {block = 128C, order = [1, 0, 2]}>

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%。

专栏文章内容及配图由作者撰写发布，仅供工程师学习之用，如有侵权或者其他违规问题，请联系本站处理。 联系我们