全文，Design of a High-Performance GEMM-like Tensor-Tensor Multiplication

Eloudy

已于 2025-02-20 14:28:52 修改

阅读量549

点赞数 15

文章标签：量子力学多线性代数

于 2025-02-18 17:40:55 首次发布

本文链接：https://blog.csdn.net/eloudy/article/details/145691377

版权

我们提出了一种“类GEMM张量-张量乘法”（GETT），这是一种用于密集张量收缩的新方法，其设计类似于高性能的通用矩阵-矩阵乘法（GEMM）。GETT 背后的关键见解是识别出三个参与张量收缩的索引集，使我们能够系统地将任意张量收缩简化为围绕高度优化的“宏内核”的循环。这个宏内核在指定的缓存层级中操作准备好了（“打包”）的子张量。与之前的张量收缩方法相比，GETT 展示了理想的特性，如单位步长内存访问、缓存感知以及完全向量化，而不需要辅助内存。我们将GETT与所谓的转置-转置-GEMM-转置和循环-GEMM 方法集成到一个开源的“张量收缩代码生成器”（TCCG）中。对于广泛的张量收缩，性能结果表明 GETT 有可能成为首选方法：虽然GETT在各方面表现出色，但其在带宽受限的张量收缩中的效果尤其令人印象深刻，性能优于现有方法高达12.4倍。更准确地说，GETT 在带宽受限的张量收缩中实现了高达 1.41 倍的加速，相对于等效大小的GEMM，同时在计算受限的张量收缩中达到了高达 91.3% 的峰值浮点性能。

1. 引言

密集张量收缩（TC）是科学计算中常见且关键的组成部分，涉及的领域包括机器学习、谱元方法、量子化学计算、多维傅里叶变换和气候模拟。尽管矩阵-矩阵乘积（GEMM）与 TC 之间有密切的联系，但后者的性能通常远低于优化的 GEMM。为了缩小这种性能差距，我们提出了一种新方法，类 GEMM 张量-张量乘法（GETT），它捕捉了高性能 GEMM 的基本特征。与高度优化的 GEMM 一致，并与之前的 TC 方法形成对比，我们的 GETT 方法的高性能实现是充分向量化的，利用了 CPU 的缓存层次结构，避免了显式的预处理步骤，并在保持步长为一的内存访问的同时操作任意维度的子张量。

从计算的角度来看，张量可以被解释为高维矩阵或简单的多维数组；同样，张量收缩是矩阵-矩阵乘法向高维的推广。鉴于这种联系，优化的 GEMM 与 TC 之间的性能差异是显著的。虽然 GEMM 实现通常能达到系统峰值浮点性能的显著部分，但这通常不是 TC 的情况。

TC效率较低的原因主要在于张量维度增加导致的复杂性增加（例如，6D 或 8D ），这通常导致次优的内存访问；因此，许多张量收缩的性能可能受限于系统的内存带宽（即带宽受限），而不是其浮点单元（即计算受限）。因此，开发一种可靠且系统的方法来利用系统的缓存对于提高 TC 的性能并将其从带宽受限状态推向计算受限状态至关重要。

之前关于张量收缩的研究主要可以分为三种方法：

1）基于向量化的嵌套循环代码，

2）转置-转置-GEMM-转置（TTGT），

3）以及最近的循环-GEMM（LoG）；

接下来我们简单讨论每种方法。

嵌套循环。基于嵌套循环的实现通过应用循环变换（例如，循环重排、循环融合）提高了性能。这些方法通常受到步长内存访问的影响，从而导致次优的内存访问模式和低内存子系统利用率。Stock等人提出了一种复杂的向量化方案，用于适合缓存的小张量收缩。虽然纯嵌套循环实现的缺点可以通过向量化来缓解，但这种策略对于不适合缓存的较大张量收缩是不够的，因此需要某种内存优化来提高缓存利用率。Ma等人提出了一种代码生成器，基于循环方法生成高性能的 CUDA 代码，用于所谓的正则化 CCSD(T) 方法中出现的张量收缩。Nelson等人后来也为 GPU 开发了另一种基于循环的代码生成器。

转置-转置-GEMM-转置。TTGT 背后的关键思想是利用经过调优的BLAS库（例如，ATLAS、OpenBLAS、BLIS、MKL）提供的高效GEMM实现。此方法需要一个准备步骤，通过显式的张量转置将任意维度的张量“展平”（或“展开”）为矩阵，以便将收缩转换为单个 GEMM；最后，将结果矩阵折叠为所需的张量布局——再次产生额外的开销。虽然这种方法有可能为计算受限的 TC 提供高性能，但它有两个固有的缺点：首先，转置的张量需要额外的内存，其次，转置过程纯粹是开销。在许多情况下，这种开销主导了运行时间，使得该方法在带宽受限状态下不可行。TTGT 的著名采用者包括 Cyclops Tensor Framework（CTF）、Tensor Toolbox、Tensorlab 和 libtensor。张量收缩引擎（TCE）是由计算化学软件包 NWChem 采用的代码生成器，用于进行耦合簇计算；其生成的代码在分布式内存级别使用 LoG，在共享内存张量收缩中使用 TTGT。在第4节中，我们讨论了 TTGT 的不同实现候选，并概述了通过其选择最有前途的候选的性能指标。与其他现有的 TTGT 实现相比，我们依赖于张量转置编译器（ TTC ）生成高效的张量转置，使得前后处理开销变得不那么明显。

循环-GEMM。最近关于LoG的工作建议将张量切片为一系列2D子张量（矩阵），并通过GEMM收缩它们。当涉及的子张量大小较大时，LoG特别有效。然而，根据 TC 的不同，2D切片可能较小和/或导致步长内存访问，导致性能不佳。Shi 等人在 “Tensor Contractions with Extended BLAS Kernels on CPU and GPU” 中，独立于我们开发了一个 CPU 和 GPU 实现，遵循 LoG 方法进行单模式收缩。他们为cuBLAS引入了stridedBatchedGemm实现，非常适合具有恒定步长的多个小GEMM批次。此外，他们提供了选择有前途的LoG候选的指南。

类GEMM 张量-张量乘法。在本文中，我们介绍了 GETT，这是一种旨在捕捉上述方法的优点，同时避免其缺点的方法。GETT 的灵感来自 Chetlur 等人在机器学习背景下对卷积的研究。卷积与 TC 类似，如果操作数被展平为矩阵，可以用矩阵-矩阵乘法来表示——这与 TTGT 方法对张量收缩的处理方式非常相似。然而，与 TTGT 不同，Chetlur 等人通过在将数据加载到缓存时隐式重组数据，避免了在调用GEMM之前和之后的昂贵准备步骤。同样，GETT 方法将任意张量收缩简化为围绕高度调优的“宏内核”的嵌套循环，其中准备好的操作数——多维子张量——驻留在缓存层次结构的指定级别。

GETT 的设计受到之前关于高性能矩阵-矩阵乘法研究的启发，其中一个大的 GEMM 被简化为一系列调用专用（较小）宏内核的操作；从概念上讲，GETT 的宏内核与这些高性能 GEMM 中使用的类似。此外，GETT 类似于 TTGT 方法，关键区别在于避免了在调用 GEMM 之前和之后准备张量的开销；相反，GETT 适当地准备（“打包”）张量的子块（不一定是二维的）到 CPU 的缓存中，因为它们是专用宏内核所需的，从而减少了到较慢主内存的数据流量。因此，可以将GETT视为“融合的TTGT”，其中转置被“融合”到GEMM中，使得GETT不需要额外的内存，并且不受在收缩之前和之后显式转置的开销影响。

最近，Devin Matthews独立开发了 TBLIS，一个用于张量收缩的C++库，也紧密遵循高性能矩阵-矩阵乘法的设计——就像GETT一样。我们的方法共享相同的关键见解，即：需要避免张量转置，并应将其移入类似GEMM的内核。除了这种相似性，GETT 和 TBLIS 在不同的方面有所侧重：首先，GETT的打包例程基于张量转置，从而利用了这些操作固有的空间局部性。其次，GETT使用一个自动微调框架——由性能模型指导——在多个实现中进行选择。另一方面，TBLIS 支持所有形式的张量收缩，包括那些不能映射到类似 GEMM 内核的收缩，并提供不需要额外编译步骤的库解决方案。

类GEMM张量-张量乘法面临的挑战是多方面的：需要识别适当大小的任意维度子张量，并开发一种系统的方法将它们打包成2D或3D连续张量，以增加空间局部性，使它们驻留在缓存层次结构的指定级别；此外，我们在确保这些子张量的步长为一的索引得以保留的同时解决这些挑战，以便打包尽可能高效。GETT背后的一个最关键的见解是，传递给专用宏内核的任意维度子张量可以逻辑地解释为更高（或更低）维度的张量，以便通过张量转置准备（打包和转置）它们。鉴于这一见解，我们的 GETT 实现依赖于张量转置编译器，保证使用单位步长内存访问，无论实际转置如何。因此，GETT 避免了非单位步长内存访问，无论考虑的实际TC如何。

与之前的方法形成鲜明对比，GETT 在给定张量收缩的情况下保持与等效大小 GEMM 相同的算术强度；这一特性对于高性能至关重要。GETT 的优势——其优越的内存访问模式、为各种缓存层次打包数据的能力以及其高度调优的向量化宏内核——转化为出色的性能特征，尤其是对于带宽受限的 TC。

张量收缩代码生成器。我们的 GETT、LoG 和 TTGT 实现被结合到一个统一的工具中，即张量收缩代码生成器（TCCG）。TCCG 生成的单精度和双精度代码在一系列张量收缩中相对于其他现有的基于 TTGT 的实现的加速如图1所示。在所有情况下，TCCG至少与其他方法一样快，单精度的加速在1.0×到12.4×之间（见图1a），双精度的加速在1.0×到3.7×之间（见图1b）。正如将在第7节中显而易见的，低加速，例如测试用例21-24，对应于效率非常接近峰值浮点性能的收缩。

组织。本文的其余部分结构如下。第2节提供了本文自包含所需的背景信息。第3节详细介绍了核心贡献并描述了GETT方法。第4节和第5节分别介绍了转置-转置-GEMM-转置和循环-GEMM方法。第6节介绍了一个统一的代码生成器（张量收缩代码生成器），第7节给出了广泛的性能评估。最后，第8节总结了我们的工作并概述了可能的未来方向。

2. 背景

为了使本文自包含，我们在此提供对通用高性能矩阵-矩阵乘法、张量收缩和张量转置的介绍。

2.1. 矩阵-矩阵乘法

设 $\mathbf{A} \in \mathbb{R}^(M \times K),\ \mathbf{B}\in \mathbb{R}^{K\times N}, \ and \ \mathbf{C} \in \mathbb{R}^{N\times N}$ 为输入矩阵；根据BLAS接口，通用矩阵-矩阵乘积表示为

$\forall i \forall j\ \ C_{i,j} \leftarrow \sum_k \alpha \times A_{i,k} \times B_{k,j}+\beta \times C_{i,j}$ (1)

列表1包含了将方程1直接翻译为代码的形式，采用了三个嵌套循环。由于对缓存的利用不佳（即，很多数据是从主内存中冗余获取的），这种简单的实现导致了极差的性能。关于高质量GEMM实现的详细讨论有很多；在此我们仅简要介绍高性能GEMM的基本思想，作为后续章节的预备知识。

任何高性能 GEMM 的必要成分是将计算组织成块（而不是标量），并将这些块（即A、B和C的子矩阵）打包成适合特定缓存级别的连续数组。这种技术提高了局部性，从而减少了缓存未命中和转换旁路缓冲区（TLB）未命中。

在列表2中，块被表示为 $\widehat{A} \in \mathbb{R}^{m_C \times k_C},\ \widehat{B} \in \mathbb{R}^{k_C \times n_C}, \ and \ \ \widehat{C} \in \mathbb{R}^{m_C \times n_C}$ ，而它们的打包对应物是辅助数组 $\widetilde{A}\ and \ \ \widetilde{B}$ ；参数 $m_C, \ n_C \ \ and \ \ k_C$ 根据给定 CPU 的缓存大小选择。每个块一旦加载到缓存中，在最终被逐出缓存之前会被多次重用；这减少了从慢速主内存中冗余获取数据的需要。

一旦子矩阵 $\widetilde{A} \ \ and \ \ \widetilde{B}$ 被准备好（打包），它们就通过一个宏内核进行乘法运算，宏内核本质上是一个高度调优的 GEMM，专为已知驻留在缓存中的操作数定制。这样的宏内核是早期 GEMM 实现的基本构建块；Van Zee等人的最新工作建议将这个构建块进一步分解为更小的微内核，以最大限度地控制 CPU 的资源。

Gunnels等人[Gunnels et al. 2001] 识别了三种将矩阵-矩阵乘法分解为一系列 1）矩阵-面板、2）面板-矩阵或 3）面板-面板乘法的变体(见Table1)。每种变体对应于列表2中循环的不同顺序。

2.2. 张量收缩

设 $\mathcal{A} \in \mathbb{R}^{S_1^{\mathcal{A}} \times S^{\mathcal{A}}_2 \times \dots S_d^{\mathcal{A}}},\ \mathcal{B} \in \mathbb{R}^{S_1^{\mathcal{B}}\times S_2^{\mathcal{B}}\times\dots S_{d_\mathcal{B}}^\mathcal{B}}$ ，和 $\mathcal{C} \in \mathbb{R}^{S_1^{\mathcal{C}} \times S_2^{\mathcal{C}} \dots S_{d_\mathcal{C}}^\mathcal{C} }$ ，是 $d_{\mathcal{A}}-, \ d_\mathcal{B}- ,\ d_\mathcal{C}-$ 维张量。扩展 Bader 等人 [Bader and Kolda 2006] 的“缩并张量积”，我们将张量收缩表示为

$\mathcal{C}_{\Pi^{\mathcal{C}}(I_m \cup I_n)} \leftarrow \sum_{k_1}\dots \sum_{k_\xi }\alpha \times \mathcal{A}_{\Pi^{\mathcal{A}}(I_m \cup I_k)} \times \mathcal{B}_{ \Pi^{\mathcal{B}}(I_k \cup I_n)}+\beta \times \mathcal{C}_{ \Pi^{\mathcal{C}}(I_m \cup I_n) }$ (2)

其中 $\Pi^{\mathcal{A}},\ \Pi^{\mathcal{B}}, \ \Pi^{\mathcal{C}}$ 符号索引， $I_m := \{ m_1, \ m_2, \dots ,\ m_{\gamma}\}$ ， $I_n := \{ n_1,\ n_2, \dots,\ n_{\varsigma }\}$ , 和 $I_k := \{ k_1,\ k_2,\dots ,\ k_{\xi} \}$ 的排列构成的集合。这些索引集分别表示A和B的自由索引（即，那些同时出现在 $\mathcal{C}$ 与 $\mathcal{A}$ 中, 或 $\mathcal{C}$ 与 $\mathcal{B}$ 中的索引），以及收缩索引（即，那些出现在 $\mathcal{A}$ 和 $\mathcal{B}$ 中但不在 $\mathcal{C}$ 中的索引）。注意以下关系成立：

$d_{\mathcal{A}} = \gamma + \xi$ ， $d_{\mathcal{B}} = \varsigma + \xi$ ，和 $d_{\mathcal{C}} = \gamma + \varsigma$

一个关键的观察是，通过采用索引集 $I_m$ ， $I_n$ 和 $I_k$ ，可以以类似 GEMM 的方式表示任意收缩。此外，为了简化符号，在本文中我们采用“爱因斯坦符号”，其中对收缩索引的求和是隐含的。方程2变为：

$\mathcal{C}_{\Pi^{\mathcal{C}}(I_m \cup I_n)} \leftarrow \alpha \times \mathcal{A}_{\Pi^{\mathcal{A}}(I_m \cup I_k)} \times \mathcal{B}_{ \Pi^{\mathcal{B}}(I_k \cup I_n)}+\beta \times \mathcal{C}_{ \Pi^{\mathcal{C}}(I_m \cup I_n) }$ (3)

示例。使用这种形式主义，通用矩阵-矩阵乘法可以表示为：

$I_m = \{ m_1\},\ I_n=\{ n_1\}, \ I_k=\{ k_1 \}$ ，

其中

$\Pi^{\mathcal{A}}(I_m \cup I_k) =(m_1,\ k_1)$ ，

$\Pi^{\mathcal{B}}(I_k \cup I_n) = (k_1, \ n_1)$ ，

$\Pi^{\mathcal{C}}(I_m \cup I_n) = (m_1,\ n_1)$ ，

得到 $\mathcal{C}_(m_1,n_1) = \mathcal{A}_{m_1,k_1}\mathcal{B}_{k_1,n_1}$ 。

在下文中，我们假设 $I_m$ ， $I_n$ 和 $I_k$ 不为空，以便用矩阵-矩阵乘法表示张量收缩；相反，当这一假设被违反时，收缩可能映射到较低级别的BLAS内核（例如，GEMV、DOT） [Di Napoli et al. 2014]。

此外，我们遵循Fortran内存布局：张量 $\mathcal{T}_{i_1,i_2,\dots , i_N}$ 的索引从左到右存储（即， $i_1$ 是步长为1的索引）。类似于任何索引 $i$ 的大小 $S_i \in \mathbb{N}$ ，我们使用相同的符号来表示索引集 $I$ 的大小（即， $S_I := \Pi_{i\in I} S_i$ ）。

【注：索引是索引，如 i，j，k 等等符号；而索引的大小则是索引的size，比如 i = 1，2，...，9，则 i的 size 就是9， $S_i = 9$ 】

2.3. 张量转置

设 $\mathcal{A}\in \mathbb{R}^{S_1^{\mathcal{A}}\times S_2^{\mathcal{B} }\times \dots S_d^{\mathcal{A}}}$ 是一个 d-维张量， $\alpha, \beta \in \mathbb{R}$ ， $\Pi^*(i_1, i_2, \dots , i_d)$ 为索引构成的任意排列，而且 $\widetilde{\mathcal{A}} \in \mathbb{R}^{\Pi^* (S_1^{\mathcal{A}},S_2^{\mathcal{A}},\dots ,S_d^{\mathcal{A}})}$ 。通用张量转置，允许输入和输出缩放，表示为

$\widetilde{\mathcal{A}}_{\Pi^* (i_1, i_2,\dots , i_d)} \leftarrow \alpha \times \mathcal{A}_{i_1,i_2,\dots,i_d} + \beta \times \widetilde{A}_{\Pi^* (i_1,i_2,\dots , i_d)}$ (4)

我们使用相同的符号（ $\widetilde{\mathcal{A}}$ ）来表示转置和操作数（ $\mathcal{A}$ ）的打包并非巧合。张量转置可以用来将任意维度的张量展平成矩阵形式，因此是 TTGT 方法的关键组成部分。图2 展示了输入张量 $\mathcal{A}_{m_1,k,m_2}$ （图2左上）被展平成矩阵 $\widetilde{A}_{m,k} \equiv \widetilde{A}_{(m_1,m_2),k}$ （图2右上）的情况，以便其大小不变（即， $S_m=S_{m_1}S_{m_2}$ ）。这个过程可以通过张量转置来完成，通过将矩阵 $\widetilde{A}_{(m_1,m_2),k}$ 重新解释为3D张量 $\widetilde{A}_{m_1,m_2,k}$ （图2底部）。

在之前的工作中，作者介绍了 TTC，一个为任何给定张量转置生成显式向量化和并行化C++代码的编译器。通过为每个索引接受一个步长，TTC 可以在子张量上操作；这一特性使得 TTC 成为 GETT 的理想构建块。正如稍后讨论的，TTC 用于生成 $\mathcal{A}$ ， $\mathcal{B}$ 和 $\mathcal{C}$ 的子张量的高性能打包例程。

3 类GEMM张量-张量乘法（GETT）

GETT背后的关键思想是将任意张量收缩简化为一系列小矩阵-矩阵乘法，其中操作数（即子张量）适合缓存。这种方法类似于实现大GEMM的技术，通过较小的矩阵-矩阵乘法来计算，由“宏内核”计算。观察到，GEMM（方程1）翻译为列表1中的三重循环，任意张量收缩（方程3）可以通过多个嵌套循环计算，如列表3所示——不同之处在于 GETT 可能需要多个M-、N-和K-循环。注意，辅助变量 tmp 的更新（第14行）可能需要对 $\mathcal{A}$ 和 $\mathcal{B}$ 进行分散的内存访问。

以索引集 $I_m$ ， $I_n$ 和 $I_k$ 表示的张量收缩类似于矩阵-矩阵乘法的数学表示；这些集合还使得可以以类似于高性能 GEMM 的方式表示任何张量收缩。因此，GETT 将任意维度的张量收缩简化为围绕宏内核的三个循环，其中循环索引 $m, \ n$ 和 $k$ 分别影响 $\mathcal{A}$ 和 $\mathcal{B}$ 的自由索引以及缩并索引。为此，必须沿 $M-$ ， $N-$ 和 $K-$ 维度对输入操作数进行分块，以创建将由宏内核处理的打包辅助数组 $\widetilde{\mathcal{A}} \in \mathbb{R}^{m_C \times k_C}$ 和 $\widetilde{\mathcal{B}} \in \mathbb{R}^{n_C \times k_C}$ 。虽然分块提高了时间局部性，但打包增加了子张量的空间局部性——与简单的张量-张量收缩相比，避免了在更新 $\mathcal{C}$ 时对 $\mathcal{A},\ \mathcal{B}$ 的分散内存访问。按照这种方法，高性能 GEMM（见列表2）和高性能 GETT（见列表4）之间的显著相似性变得显而易见。

与简单的张量-张量乘法实现相比，其中每个索引 $i \in (I_m \cup I_n \cup I_k)$ 作为单独的循环出现（见列表3），高性能 GETT（见列表4）用单个循环替换了多个M-、N-和K-循环。因此，每个循环计数器（即，m、n和k）可能影响张量 $\mathcal{A}$ 、 $\mathcal{B}$ 和 $\mathcal{C}$ 的多个索引。确切的机制由函数 identify_subtensor() 处理，该函数根据当前循环迭代，识别适当的子张量 $\widetilde{\mathcal{A}}, \widetilde{\mathcal{B}} \ or \ \widetilde{\mathcal{C}}$ ，它们来自 $\mathcal{A,B} \ or \ \mathcal{C}$ ；注意，这个 “识别处理” 完全发生在编译时，并不导致任何数据移动。现在我们假设这些子张量是给定的，并且它们的大小和维度与其打包对应物 $\widetilde{\mathcal{A}}, \widetilde{\mathcal{B}} \ and \ \widetilde{\mathcal{C}}$ 匹配；函数identify_subtensor() 将在 第3.2节 中详细讨论。

一个典型的张量收缩的程序流程， $\mathcal{C}_{m_1,n_1,m_2} \leftarrow \mathcal{A}_{m_1,m_2,k_1}\mathcal{B}_{n_1,k_1}$ ，如图3所示。即使在这个例子中子张量只有二维，我们强调在一般情况下 $\widetilde{\mathcal{A}}, \widetilde{\mathcal{B}} \ or \ \widetilde{\mathcal{C}}$ 可以是任意维度的；因此，这些子张量可以跨多个维度收集——不仅仅是两个维度。

首先， $\textcircled{1}\ n-$ 和 $\textcircled{2}\ k-$ 循环选择 $\mathcal{B}$ 的子张量 $\widehat{\mathcal{B}}$ ；这限制了非打包子张量 $\widehat{\mathcal{B}}$ 的大小。然后 $\textcircled{3}$ 将 $\widehat{\mathcal{B}}$ 打包到连续的辅助数组 $\widetilde{\mathcal{B}}$ 中。接下来是 $\textcircled{4}$ m-循环；在示例的上下文中，这影响了多个索引（即， $m_1$ 和 $m_2$ ），而 n-和k-循环只影响每个索引（即， $n_1$ 和 $k_1$ ，分别对应）。接下来， $\textcircled{5}$ 将 $\mathcal{A}$ 的非打包子张量 $\widehat{\mathcal{A}}$ 打包到连续张量 $\widetilde{\mathcal{A}}$ 中。最后，调用宏内核 $\textcircled{6}$ 计算 $\widetilde{\mathcal{A}}$ 与 $\widetilde{\mathcal{B}}$ 的矩阵-矩阵乘积以更新子张量 $\widehat{\mathcal{C}}$ ；注意 $\widehat{\mathcal{C}}$ 没有打包，仅表示对 $\mathcal{C}$ 的相应元素的引用。

3.1. 微内核和宏内核

图4（底部）展示了辅助数组 $\widetilde{\mathcal{A}}$ 和 $\widetilde{\mathcal{B}}$ 的存储方案。参数 $m_C,\ n_C, \ k_C, \ m_R$ 和 $n_R$ 的适当选择将在第3.3节中讨论；目前只需知道这些参数的选择使得 $\widetilde{\mathcal{B}}$ 保留在L3缓存中， $\widetilde{\mathcal{A}}$ 保留在L2缓存中，并且 $\widetilde{\mathcal{A}}$ 的 $m_R \times k_C$ 大小的子矩阵与 $\widetilde{\mathcal{B}}$ 的 $n_R \times k_C$ 大小的子矩阵同时适合存进 L1缓存。因此，所有打包的子张量将保留在某个缓存层次中，而不需要从主内存中额外加载。

宏内核通过一系列对微内核的调用计算 $\widetilde{\mathcal{A}}$ 和 $\widetilde{\mathcal{B}}$ 的矩阵-矩阵乘法（见图4，顶部）。微内核表示 $\widetilde{\mathcal{A}}$ 的 $m_R \times k_C$ 子矩阵与 $\widetilde{\mathcal{B}}$ 的 $n_R \times k_C$ 子矩阵的宽内积为一系列大小为 $m_R \times n_R$ 外积，结果辅助数组 Ce可以完全保留在寄存器中（例如，在Intel Haswell微架构上,对于单精度计算， $m_R=24,\ n_R=4$ ，对于双精度计算， $m_R=12,\ n_R=4$ ）。

虽然 $\widetilde{\mathcal{A}}$ 和 $\widetilde{\mathcal{B}}$ 是打包的，但 $\widehat{\mathcal{C}}$ 不是；与高性能GEMM（见第3.2节）相比而言， $\widehat{\mathcal{C}}$ 可以是任意维度的。因此，我们通过 $\widehat{\mathcal{C}}(\widetilde{m}_1,\widetilde{n}_1, \widetilde{m}_2, \widetilde{n}_2)$ 表示这个子张量，使得其维度不固定，并且其对 $\widetilde{\mathcal{A}}$ 和 $\widetilde{\mathcal{B}}$ 的自由索引的依赖是显式的。箭头标记为“更新”表明打包的子张量 $\widetilde{\mathcal{C}}_{\widetilde{m}_1, \widetilde{n}_1}$ 更新子张量 $\widehat{\mathcal{C}}(\widetilde{m}_1,\widetilde{n}_1, \widetilde{m}_2, \widetilde{n}_2)$ 的适当部分；我们将此过程称为“解包”。正如我们将在下一节中讨论的，这种更新——在一般情况下——表示张量转置（见第3.2节）。

3.2. 打包

关键的观察是，打包数组 $\widetilde{\mathcal{A}},\ \widetilde{\mathcal{B}}$ 和 $\widetilde{\mathcal{C}}$ 可以逻辑地描述为分别为三维、三维和二维的张量（即， $\widetilde{\mathcal{A}}_{\widetilde{m}_1,\widetilde{k}_1,\widetilde{m}_2}$ ， $\widetilde{\mathcal{B}}_{\widetilde{n}_1,\widetilde{k}_1,\widetilde{n}_2}$ ，和 $\widetilde{\mathcal{C}}_{\widetilde{m}_1,\widetilde{n}_1}$ ）——而不是二维数组。将这些辅助数组表示为打包子张量使我们能够以张量转置的形式表达（解）打包例程。这种“转换”是有利的，因为张量转置可以完全向量化，并且通常达到接近系统内存带宽的峰值。

当分块使我们能够有效地将张量收缩简化为围绕宏内核的三个嵌套循环，同时将数据打包到连续数组中避免了由于现代缓存的有限制的关联性而导致的严重冲突引发未命中；打包格式还增加了子张量的空间局部性，减少了宏内核内的TLB未命中。与仅处理2D子矩阵的矩阵-矩阵乘法形成鲜明对比，GETT 的打包例程可能需要从任意维度的子张量中获取数据并将其打包到连续的辅助数组中。

要打包 $\mathcal{A}$ ， $\mathcal{B}$ 和 $\mathcal{C}$ 的部分，必须：（1）从原始张量 $\mathcal{A}$ 、 $\mathcal{B}$ 、 $\mathcal{C}$ 中识别适当大小的子张量 $\widehat{\mathcal{A}}$ ， $\widehat{\mathcal{B}}$ 和 $\widehat{\mathcal{C}}$ ，然后（2）将这些非连续子张量 $\widehat{\mathcal{A}}$ ， $\widehat{\mathcal{B}}$ ， $\widehat{\mathcal{C}}$ 转置为 $\widetilde{\mathcal{A}}$ ， $\widetilde{\mathcal{B}}$ 和 $\widetilde{\mathcal{C}}$ 。此外，为了将这些打包例程表示为张量转置，必须确保元素总数以及子张量及其打包对应物的维度是相同的。作为一个例子，我们现在讨论 $\widehat{\mathcal{A}}$ Ab子张量的识别和打包； $\widehat{\mathcal{B}}$ 的打包过程和 $\widehat{\mathcal{C}}$ 的解包过程类似。

3.2.1. 识别子张量

给定所需的分块大小 $m_C$ 和 $k_C$ ，适合底层处理器的缓存容量，我们寻求识别两个子集 $I_{\widehat{m} }\subseteq I_m$ 和 $I_{\widehat{k}} \subseteq I_k$ ，使得

$S_{I_{\widehat{m}}} = m_C$

$S_{I_{\widehat{k}}}=k_C$

以形成子张量 $\widehat{\mathcal{A}}_{\Pi^{\widehat{\mathcal{A}}}(I_{\widehat{m}} \cup I_{\widehat{k}})}$ 。 $\Pi^{\widehat{A}}$ 表示与原始排列 $\Pi^{\mathcal{A}}$ 几乎相同的排列，但移除了下述所有索引 $i \in (I_m \cup I_k) \setminus (I_{\widehat{m}} \cup I_{\widehat{k}})$ ；更准确地说，索引的顺序保持不变。例如，给定原始索引集 $I_m=\{ m_1,\ m_2\},\ I_k=\{ k_1,\ k_2\}$ 和适当大小的打包张量的索引集 $I_{\widehat{m}} = \{ m_1\}$ 和 $I_{\widehat{k}}=\{ k_2\}$ 以及排列 $\Pi^{\mathcal{A}}(I_m \cup I_k) = (m_1, k_1,m_2,k_2)$ ，则 $\Pi^{\mathcal{A}}$ 表示为： $\Pi^{\widehat{A}}(I_{\widehat{m}} \cup I_{\widehat{k}}) =(m_1,\ k_2)$ 。

除了约束（5）和（6）式，我们还要求A的步长为一的索引是 $I_{\widehat{m}}$ 或 $I_{\widehat{k}}$ 的一部分。即使从正确性的角度来看，这一条件不是必需的，但对于高性能是必要的。步长为一的索引是 $\widehat{\mathcal{A}}$ 的一部分，使得张量转置中的向量化加载和存储成为可能；因此，产生更高效的打包例程。图5a展示了两个“合适”子张量的例子，这些子张量保留了步长为一的索引；另一方面，图5b描绘了一个“不合适”的子张量，其步长为一的索引未被保留。

根据 $\mathcal{A}$ 的大小和形状，可能无法立即识别出符合上述约束的合适子集 $I_{\widehat{m}}$ 和 $I_{\widehat{k}}$ 。因此，为了找到这样的子集，我们可能需要将 $\mathcal{A}_{\Pi^{\mathcal{A}}(I_m \cup I_k) }$ 重新解释为具有相同大小（即， $S_{I_m} = S_{I_{m'}}$ 和 $S_{I_k} = S_{I_{k'}}$ 但不同维度（即， $\left | I_m \right | \neq \left | I_{m'}\right |$ 或 $\left | I_k \right | \neq \left | I_{k'}\right |$ ）的不同张量 $\mathcal{A'}_{\Pi^{\mathcal{A'}}(I'_m \cup I'_k) }$ 。重要的是要提到，所讨论的张量的内存布局不会改变；张量的形状仅仅是重新解释。因此，这种视角的改变不需要任何内存访问，因此完全发生在编译时。为了给读者更好的直观理解这个过程，我们来看两个需要重新解释 $\mathcal{A}$ 的例子。

例1：设 $\mathcal{A}_{m_1,k_1}$ 为一个二维张量， $S_{m_1}=4m_C$ 和 $S_{k_1}=k_C$ (见图6a），我们的目标是识别一个大小为 $m_C\times k_C$ 的子张量 $\widehat{\mathcal{A}}_{\Pi^{\widehat{\mathcal{A}}}(I_{\widehat{m}} \cup I_{\widehat{k}})}$ ，其中 $S_{I_{\widehat{m}}}=m_C$ 并且 $S_{I_{\widehat{k}}}=k_C$ 。这可以通过将索引 $m_1$ 分割为 $m_1^{(1)}$ 和 $m_1^{(2)}$ 来实现，使得 $m_1^{(1)}$ 和 $m_1^{(2)}$ 的组合大小与原始索引 $m_1$ 的大小相同（即， $S_{m_1}=S_{m_1^{(1)}} \times S_{m_1^{(2)}}$ ——有效地将2D张量 $\mathcal{A}_{m_1,k_1}$ 重新解释为 3D 张量 ${A'}_{m_1^{(1)},m_1^{(2)},k_1}$

（见图6b）。一旦 $\mathcal{A}$ 被重新解释为 3D 张量，就可以定义子集 $I_{\widehat{m}}:={m_1^{(1)}}$ 和 $I_{\widehat{k}}:={k_1}$ 。因此，我们识别出符合方程 (5)-(6) 的所需子张量 $\widehat{A}_{m_1^{(1)},k_1}$ 。

例2：我们现在关注一个更复杂的例子。设 $\mathcal{A}_{m_1,k_1,m_2}$ 为一个 3D 张量，其中 $S_{m_1}$ 和 $S_{m_2}$ 都不是 $m_C$ 的倍数，而它们的乘积 $S_{m_1} \times S_{m_2}$ 是 $m_C$ 的倍数。此外， $S_{k_1} = k_C$ 与前一个例子相同。为了找到一个总大小为 $m_C$ 元素的子集 $I_{\widehat{m}}$ ，可以将 $\mathcal{A}$ 重新解释为5D张量 $\mathcal{A}_{m_1^{(1)},m_1^{(2)},k_1,m_2^{(1)},m_2^{(2)}}$ 并分配 $I_{\widehat{m}}:= \{ m_1^{(1)}, m_2^{(1)}\}$ ，其中 $S_{I_{\widehat{m}}} = m_C$ 。与前一个例子类似，我们强制分割索引的总大小保持不变（即， $S_{m_1}= S_{m_1^{(1)}}S_{m_1^{(2)}}$ 和 $S_{m_2} = S_{m_2^{(1)}}S_{m_2^{(2)}}$ 。因此，所需的子张量 $\widehat{\mathcal{A}}_{m_1^{(1)},k_1,m_2^{(1)}}$ 可以从 5D 张量 $\mathcal{A'}_{m_1^{(1)},m_1^{(2)},k_1,m_2^{(1)},m_2^{(2)}}$ 中识别出来；同样，索引的顺序不允许改变。

敏锐的读者可能已经注意到，任何索引的分割必须基于其大小的素因数。虽然对这一机制的详细讨论超出了本文的范围，但我们指出，我们的 GETT 实现基于其素因数分割所有索引，并选择合适的子集 $I_{\widehat{m}}$ ， $I_{\widehat{n}}$ 和 $I_{\widehat{k}}$ （适当大小），以便在子张量 $\widehat{\mathcal{A}}$ ， $\widehat{\mathcal{B}}$ 和 $\widehat{\mathcal{C}}$ 中保留 $\mathcal{A}$ ， $\mathcal{B}$ ， $\mathcal{C}$ 的步长为一的索引。

由于存在多个满足上述条件的索引集，我们目前探索所有候选（见第3.3节）并使用性能模型自动对它们进行排名（见第3.4节）。

3.2.2. 通过张量转置进行打包

虽然打包子张量 $\widetilde{\mathcal{A}}$ ， $\widetilde{\mathcal{B}}$ 和 $\widetilde{\mathcal{C}}$ 的内存布局（见图4）以及其索引的大小始终固定（即， $S_{\widetilde{m}_1} = m_R$ ， $S_{\widetilde{n}_1}=n_R$ ， $S_{\widetilde{m}_2}=\frac{m_C}{S_{\widetilde{m}_1}}$ ， $S_{\widetilde{n}_2}=\frac{m_C}{S_{\widetilde{m}_1}}$ 和 $S_{\widetilde{k}_1}=k_C$ ），但非打包子张量 $\widehat{\mathcal{A}}$ ， $\widehat{\mathcal{B}}$ 和 $\widehat{\mathcal{C}}$ 的维度可以从一个TC到另一个TC变化。因此，我们需要为任意 TC 定义这些张量之间的可靠映射。

回想一下，非打包子张量 $\widehat{\mathcal{A}}$ 的自由索引 $I_{\widehat{m}}$ 和收缩索引 $I_{\widehat{k}}$ 的索引集已经具有适当且固定的大小（即， $S_{I_{\widehat{m}}} = m_C$ 和 $S_{I_{\widehat{k}}} = k_C$ ）。此外， $m_C$ 的值选择为 $m_R$ 的倍数。与前一节类似，我们必须识别两个不重叠的子集 $I_{\widehat{m}}^0,\ I_{\widehat{m}}^1 \ \subset I_{\widehat{m}}$ 使得

$S_{I_{\widehat{m}}^0}=m_R$

$S_{I_{\widehat{m}}^1}=\frac{S_{I_{\widehat{m}}}}{S_{I_{\widehat{m}}^0}}$

由于这些子集可能不存在， $\widehat{\mathcal{A}}_{\Pi^{\widehat{\mathcal{A}}}(I_{\widehat{m}} \cup I_{\widehat{k}})}$ 的形状必须再次重新解释为 $\breve{A}_{\Pi^{\breve{A}}(I_{\widehat{m}}^0 \cup I_{\widehat{k}} \cup I_{\widehat{m}}^1)}$ ，通过前一节讨论的相同机制。

以类似的方式重新解释 $\widehat{\mathcal{B}}$ 和 $\widehat{\mathcal{C}}$ 的形状产生合适的子张量 $\breve{B}$ 和 $\breve{C}$ 。我们终于能够将非打包子张量 $\breve{A}$ ， $\breve{B}$ ， $\breve{C}$ 的打包编码为其打包对应物 $\widetilde{\mathcal{A}}$ ， $\widetilde{\mathcal{B}}$ ， $\widetilde{\mathcal{C}}$ 的张量转置形式：

$\widetilde{\mathcal{A}}_{I_{\widehat{m}}^0 , I_{\widehat{k}} , I_{\widehat{m}}^1} \leftarrow \breve{A}_{\Pi^{\breve{A}}(I_{\widehat{m}}^0 \cup I_{\widehat{k}} \cup I_{\widehat{m}}^1)}$ (9)

$\widetilde{\mathcal{B}}_{I_{\widehat{n}}^0 , I_{\widehat{k}} , I_{\widehat{n}}^1} \leftarrow \breve{B}_{\Pi^{\breve{B}}(I_{\widehat{n}}^0 \cup I_{\widehat{k}} \cup I_{\widehat{n}}^1)}$ (10)

$\breve{C}_{\Pi^{\breve{C}}(I^0_{\widehat{m}} \cup I^0_{\widehat{n}})} \leftarrow \alpha \times \widetilde{\mathcal{C}}_{I^0_{\widehat{m}} \cup I^0_{\widehat{n}}} \ + \ \beta \times \breve{C}_{\Pi^{\breve{C}}(I^0_{\widehat{m}} \cup I^0_{\widehat{n}})}$ (11)

示例。给定子张量 $\widehat{A}_{k_1,m_1,k_2,m_2}$ ，其中 $k_C=S_{k_1} \times S_{k_2}$ 和 $S_{m_1}=a \times b$ ， $S_{m_2}=c \times d$ ， $a,b,c,d \in \mathbb{N}$ 和 $m_R = a \times c$ 。然后我们可以将 $\widehat{A}$ 重新解释为 $\breve{A}_{m_1^{(1)},m_1^{(2)},k_1,m_2^{(1)},m_2^{(2)}}$ ，使得 $S_{m_1}^{(1)}=a,\ S_{m_1}^{(2)}=b,\ S_{m_2}^{(1)}=c,\ S_{m_2}^{(2)}=d$ ；结果是张量转置：

$\widetilde{\mathcal{A}}_{m_1^{(1)},m_2^{(1)},k_1,k_2,m_1^{(2)},m_2^{(2)}} \leftarrow \breve{A}_{m_1^{(1)},m_1^{(2)},k_1,m_2^{(1)},m_2^{(2)}}$ 。注意，其他排列也是可能的（例如， $\widetilde{\mathcal{A}}_{m_1^{(1)},m_2^{(1)},k_2,k_1,m_1^{(2)},m_2^{(2)}} \leftarrow \breve{A}_{m_1^{(1)},m_1^{(2)},k_1,m_2^{(1)},m_2^{(2)}}$ ）。

方便查看，再复制一遍 Table I:

3.3. 搜索空间

本节概述了任何以 GETT 表示的 TC 的可行实现（以下简称候选）的搜索空间。不同 GEMM 变体的组合（见表I）、分块大小的选择（即， $m_C, n_C, k_C, m_R, n_R$ ）以及自由和缩并索引集的排列（即， $I_m,\ I_n, \ I_k$ ）构成了一个大的可行候选搜索空间。我们没有穷举测试所有候选（即，生成、编译和计时），而是根据一个架构感知的测度对它们进行排名（见下一节）；感谢这个测度指标，它使得搜索空间被显著修剪。

根据索引集 $I_m,\ I_n, \ I_k$ 的大小，可能有必要偏爱三种 GEMM-变体之一（见第2.1节的表I）而不是另一个。因此，GETT 会为每个 TC 考察所有三种变体。

我们当前的 GETT 实现为每个 $m_C,\ n_C, \ k_C$ 探索多达四个不同的参数，同时保持 $m_R\ and \ n_R$ 固定。更准确地说， $m_R\ and \ n_R$ 的选择使得微内核内的 fused-multiply-add 融合乘加（FMA）指令的延迟完全隐藏，以实现峰值浮点性能[Low et al. 2015]。未来的 GETT 实现可能还会探索不同的 $m_R\ and \ n_R$ 值；这在带宽受限状态下可能是有用的，在这种情况下，人们愿意以较低效的微内核换取更高效的子张量的打包例程。

最后，正如我们在前一节中已经提到的，有几种方法可以从 $\mathcal{A}$ 、 $\mathcal{B}$ 和 $\mathcal{C}$ C 中识别合适的子张量。我们将搜索（可行子张量）限制在那些在相应张量 $\mathcal{A}$ 、 $\mathcal{B}$ 或 $\mathcal{C}$ 中保留步长为一(stride-one)的索引的子张量上。

3.4. 性能模型

我们根据所选的 GEMM 变体、打包子张量的大小以及所需打包例程的排列，来为为任意 GETT 的候选，做运行时间建模。这个过程分为两个阶段：（1）估计打包例程的时间，（2）估计宏内核的时间。当前候选的总时间估计为（1）和（2）的总和。

为了简单起见，让我们再次关注矩阵-面板乘法的例子（即，列表4中概述的算法）。在整个张量收缩过程中，将 $\widehat{\mathcal{B}}$ 打包到 $\widetilde{\mathcal{B}}$ 中读取和写入总共数据量为：

$\mathbf{Data}_{\mathcal{B}} = 2 \times S_{I_k}\times S_{I_n} \times \mathbf{sizeof(floatType)}$ (12)

字节。这个操作显然是带宽受限的；因此，我们可以通过

$\mathbf{Time}_{\mathcal{B}} = \frac{\mathbf{Data}_{\mathcal{B}}}{BW \times p_D}$ (13)

估计打包Bb所需的时间，其中 $BW$ 表示系统的 $\mathbf{SAXPY-Bandwideth}$ 带宽， $p_D \in (0, \ 1]$ 是一个惩罚，偏爱某些排列而不是其他排列。例如，性能模型略微偏爱小维度转置而不是高维度转置，因为前者通常表现出更规则的内存访问模式（即， $p_D$ 随着维度 d 的增加略微减少， $p_D=1.0-(d-1)\times0.015$ 。此外，之前的研究表明，步长为一的索引不变的转置（例如， $\widetilde{A}_{i,j,k} \leftarrow \widehat{A}_{i,k,j},\ \ \widetilde{A}_{i,j,k,l} \leftarrow \widehat{A}_{i,l,k,j}$ ）比其他转置更高效。因此，我们通过将 $p_D$ 额外减少 30% 来惩罚那些步长为一的索引改变的转置（即， $p_d \leftarrow p_d \times 0.7$ ）。对 $\mathbf{Time}_{\mathcal{A}}$ 和 $\mathbf{Time}_{\mathcal{C}}$ 进行类似的分析，唯一的区别是 $A$ 和 $C$ 将分别被（解）打包 $\frac{S_{I_n}}{n_C}$ 和 $\frac{S_{I_k}}{k_C}$ 多次。

最后，我们根据

$\mathbf{Time_{macro}} =\frac{2\times S_{I_m}\times S_{I_n} \times S_{I_k}}{FP^{\mathbf{peak}}\times p_F}$

估计宏内核的时间，其中. $FP^{\mathbf{peak}}$ 表示系统的理论峰值浮点性能， $p_F \in (0,\ 1]$ pF ∈ (0, 1] 还是表示一个惩罚。我们通过 30% 惩罚微内核的性能（即， $p_F \leftarrow p_F \times 0.7$ ），每当以下条件之一被违反时：1） $\widetilde{\mathcal{B}} \in \mathbb{R}^{n_C \times k_C}$ 适合L3缓存，2） $\widetilde{\mathcal{A}} \in \mathbb{R}^{m_C \times k_C}$ 以及 $\widetilde{\mathcal{B}}$ 的两个 $n_R \times k_C$ 子块适合L2缓存（ $\mathcal{B}$ 通过 L2 流动），以及3）一个 $\widetilde{\mathcal{B}}$ 的 $n_R \times k_C$ 子块以及 $\widetilde{\mathcal{A}}$ 的两个 $m_R \times k_C$ 子块同时适合 L1 缓存（ $\mathcal{A}$ 通过 L1 流动）。

候选的总估计时间为 $\mathbf{Time}_{\mathbf{macro}}+\mathbf{Time}_{\mathcal{A}}+\mathbf{Time}_{\mathcal{B}}+\mathbf{Time}_{\mathcal{C}}$ ；虽然这仍然是一个相当粗略的模型，但在实践中效果很好（见第7节）。

4. 转置-转置-GEMM-转置

在本节中，我们讨论了我们的 TTGT代码生成器的原理。TTGT 背后的想法是重新排列（转置）张量的元素，使其可以逻辑地解释为矩阵。这个初始步骤，也称为展平，使得可以通过普通的矩阵-矩阵乘法计算收缩，从而利用GEMM的高效率。根据实际的收缩，结果矩阵必须折叠（转置）回来。

列表5 概述了 TTGT 代码生成器的伪代码。在第2-4行，从 A 和 B 中提取了自由索引（即， ${I_m ,\ I_n}$ ）以及收缩索引（即， $I_k$ ）。接下来，遍历这些索引集的所有排列（第5-7行）。一旦索引集的排列固定，八个不同的 TTGT 候选被添加到候选集（第9-20行）。一个候选包括最多三个张量转置（第10、11、13行）以及一个矩阵-矩阵乘法（第12行）；注意，根据实际的张量收缩，这些转置中的一些（甚至全部）可能是多余的。然后根据一个最小化 “（解）折叠开销” 的指标选择最有前途的候选，即，它最小化了第10、11和13行中涉及的张量的组合大小。因此，我们的TTGT实现不需要任何搜索，而是完全依赖于启发式自动选择最有前途的候选。类似于第3.4节中讨论的 GETT 方法的性能模型，这个指标也偏爱那些步长为一的索引保持不变的转置。

示例。让我们考虑收缩 $\mathbf{C}_{a,b,c}=\mathbf{A}_{a,c,k}\mathbf{B}_{b,k}$ ；在这种情况下，A 和 B 都不需要转置。一个遵循 TTGT 方法的这个 TC 的示例实现如列表6所示。

虽然 TTGT 在计算瓶颈受限状态下表现非常好，但在带宽瓶颈受限的 TC 中表现不佳。这种次优行为是由于其两个主要缺点：

首先，（解）折叠的张量需要额外的内存；

其次，每个（需要转置的张量的）元素至少需要加载两次，实际上更加倍了带宽需求。

5. 循环-GEMM

本节概述了基于LoG方法的代码生成过程。列表7中采用的策略是将步长为一的索引作为循环-GEMM的一部分；这些索引被称为requiredIndices（第2行）。如果C的步长为一的索引属于B，则LoG交换A和B（第3-4行）。在第5-7行中识别A和B的自由索引以及收缩索引。只有在第8-11行的条件都不被违反时，才考虑LoG实现；这些条件确保2D切片可以通过普通GEMM处理，只有一个维度的步长访问，而不是两个维度的步长访问，这会导致次优的内存访问。

在第12-25行中，生成所有可能的候选，然后存储到候选列表中。第13-21行确定GEMM调用的m-、n-和k-索引；注意，如果mIdx、nIdx或kIdx尚未被requiredIndices覆盖，则可能有多个候选。剩余的索引构成了循环索引列表（第22行）。考虑循环索引的所有排列（第23-25行），并将结果候选附加到候选中。最后，在第27行，根据一个基于GEMM浮点运算数的指标选择最有前途的候选（即，2×SmIdx×SnIdx×SkIdx）；其基本原理是较大的GEMM通常表现出更高的性能。这个模型的一个可能扩展是还考虑相应GEMM调用的估计效率。

列表8显示了为一个示例张量收缩Ca,b,c,i,j,k = Ai,j,m,c × Bm,k,a,b生成的LoG代码，其中每个索引的大小为24。这个候选循环遍历k和c索引（第10、11行），并通过GEMM收缩m索引（第13-16行）。重要的是要注意，索引a、b和i、j被合并/展平为超级索引（i, j）和（a, b），分别；因此，可以将GEMM解释为C(a,b),(i,j) = Bm, T (a,b) × AT (i,j),m，其中T表示矩阵转置。

6. 张量收缩代码生成器

在介绍了三种不同的方法（GETT、TTGT和LoG）后，我们现在描述张量收缩代码生成器（TCCG），一个用Python编写的统一代码生成器。TCCG的输入是每个索引的大小的收缩（见列表9）；输出是高性能C++代码。TCCG在

www.github.com/springer13/tccg

上公开可用。

TCCG 的示意图如图7所示。在开始代码生成过程之前，TCCG 将张量中的连续索引合并为超级索引（阶段1）。仍然作为预处理阶段，查询已知实现的本地 SQL 数据库，以检查是否已经存在特定收缩（和大小）的解决方案；如果是，则不进行生成，并返回先前的实现。否则，进行阶段2： TCCG 在整个代码生成过程中维护一个有前途的候选列表；候选基于 GETT、LoG 和 TTGT 方法生成。在 阶段3 中，根据性能模型 “评估” 候选；如果当前候选的成本小于列表中最差候选的成本κ，则将候选存储到内部列表中。这个过程继续，直到考虑了所有候选。然后进行第4步编译和第5部计时最有前途的候选。最后，在第6步中，将最快的候选存储到 SQL 数据库中，并生成相应的 C++ 代码。

TCCG 为用户提供了通过命令行选项影响 TCCG 代码生成过程的可能性。例如，用户可以限制TCCG内部候选列表的容量（通过--maxImplementations）或限制 TTGT 允许使用的辅助工作空间的数量（通过--maxWorkspace）。TCCG 的当前实现支持单精度和双精度计算。此外，TTGT和 LoG 也可用于支持CUDA 的 GPU。

7. 性能结果

在本节中，我们分别评估GETT、LoG和TTGT的性能，然后通过TCCG在基于Intel Haswell微架构的Intel Xeon E5-2680 v3 CPU的单核上进行集体评估。启用了ECC，并且在所有测量中禁用了Intel Speedstep和Intel TurboBoost。选择的C++编译器是Intel的icpc 16.0.1 20151021，标志为-O3 -xhost -mkl。

我们报告三次运行中的最小运行时间，同时在每次运行之间清除缓存（冷数据）。生成代码的正确性通过与类似于列表3中概述的简单张量-张量收缩的简单循环实现进行检查；这个实现作为性能的下限（我们称之为“参考”）。此外，我们还报告了与给定TC（即，A SIk，B SIn，C SIn）大小相同的GEMM的性能；GEMM模仿TTGT方法，但省略了显式转置，因此产生不正确的结果——可以将其视为性能的上限。

7.1. 基准

为了便于比较不同张量收缩方法的性能，我们编制了一个张量收缩基准，包含从之前的出版物中收集的广泛用例。基准在www.github.com/springer13/tccg上公开可用，由48个不同的收缩组成。

[Apra ]：18个在CCSD(T)方法中遇到的带宽受限收缩。相应的张量具有高维度（即，两个4D输入张量形成一个6D输出张量）。

[Stock]：19个在耦合簇方法中遇到的收缩，使用2D到4D张量。

[Baumgartner]：3个“常用于量子化学计算中将一组双电子积分从原子轨道基变换为MO基”的收缩。

[Li et al]：8个收缩，将3D、4D或5D张量与2D张量收缩（即，张量-矩阵乘法）。

索引的大小选择为再现相应出版物中的大小（如果披露）。此外，基准确保每个TC的总内存消耗至少为200 MiB，因此，显著大于我们测试系统的最后一级缓存（实际大小可以在附录A中找到）。

在接下来的章节中，为了将读者的注意力集中在那些表现出不同性能特征的收缩上，我们只报告48个收缩子集的结果。完整基准的结果在附录A中提供。

为了简化表示，我们将收缩CΠC(Im∪In) ← Ak)编码为C(IIk)；例如，Cijkl ← AimjnBnlmk表示为i。

7.2. 性能评估

我们从对TTGT方法的更深入评估开始，通过将我们的TTGT实现（记为TTGT）与Cyclops Tensor Framework（CTF）、Tensor Toolbox（TT）和libtensor实现的性能进行比较。CTF、TT和libtensor都实现了TTGT方法。

图8报告了基准中单精度和双精度TTGT实现的性能；水平的黑线表示等效大小的GEMM实现的性能。图的顶部表示给定CPU和所选精度的理论峰值浮点性能。列根据其单精度GEMM性能排序。因此，图左侧的TC是带宽受限的，而右侧的是计算受限的。

我们观察到我们的TTGT实现相对于CTF、TT和libtensor在整个基准中的显著加速。正如这些结果所显示的，TTGT在计算受限状态下表现非常好——在这种情况下，矩阵-矩阵乘法主导了执行时间。然而，显然TTGT实现在带宽受限状态下并没有达到最佳性能。次优性能可以归因于显式转置，这纯粹是开销。更准确地说，这些转置表现出的开销由实线和报告的TTGT性能之间的差异显示；例如，带宽受限的TC在（解）折叠张量上花费了多达50%的运行时间，而计算受限的TC在这些转置上只花费了可忽略不计的运行时间。

图9继续我们对不同张量收缩方法的调查，并展示了GETT和LoG与TTGT的性能。观察GETT的结果（绿色条），显然GETT在带宽受限状态下（图左侧）相对于TTGT方法（红色条）表现出显著的加速，而在计算受限状态下（图右侧）并没有达到TTGT的性能；我们指出，GETT——与TTGT形成对比——不需要任何辅助内存。对于一些测试用例，GETT超过了报告的GEMM性能，高达1.41倍。更准确地说，GETT在基准中平均达到GEMM性能的98.1%（最小值：72.4%；最大值：141.4%）和97.0%（最小值：60.8%；最大值：132.9%），分别为单精度（见图9a）和双精度（见图9b）。

另一方面，LoG方法（图9中的橙色条）在整个基准中表现出可变的性能。LoG在某些TC中提高了相对于TTGT的性能（例如，abcdef），但在其他情况下表现较低（例如，abcd-ebad-ce，abcd-aebf-fdec）。在某些情况下（例如，abcde-ef bad-cf，abcdef-dega-gfbc），该方法在没有显式转置或允许多维度步长内存访问的GEMM实现的情况下不可用。

在图10中，我们将图8和图9的数据结合起来，并以性能曲线的形式展示，从而将GETT与CTF、LoG和TTGT进行比较。对于给定的方法M和x轴上的给定点α，相应的y轴上的值p表示M最多比四种方法中最快的方法慢α倍的概率。示例：对于α = 1.2，TTGT的p值约为0.4；这意味着在40%的测试中，TTGT是最快的方法，或者在最快的方法的1.2倍以内。这个图使得得出一些结论变得容易。（1）在单精度（双精度）测试用例中，GETT是最快的方法约70%（60%）。（2）在GETT不是最快的方法的情况下，它从未比最佳解决方案差1.22（1.24）倍。TTGT和CTF始终在最佳方法的2.2（2.0）和20.2（19.7）倍以内。（3）LoG的线在约p = 0.6处达到平台；这表明在约40%的测试用例中，该方法不可用。

图11以两个相似的张量收缩为例，突出显示GETT、LoG和TTGT的不同性能特征。对于两个TC，索引的大小选择为SIm = SIn = 1152，同时让SIk从8到1024变化。因此，我们有效地将TC从带宽受限状态（图左侧）推向计算受限状态（图右侧）。与我们之前的发现类似，可以观察到GETT在计算受限状态下达到高达91.3%的峰值性能，并在带宽受限状态下表现出色。此外，虽然LoG经历了显著的性能损失（比较实线和虚线），但GETT和TTGT在更复杂的张量收缩i1j1i2j2-i1ki2-j1kj2中仅受到轻微影响。LoG的性能损失可以通过GEMM中涉及的子矩阵大小变小来解释（即，n = Sj对于i1ji2-i1ki2-jk，但n = max(Sj1, Sj2) < Sj）；较小的子矩阵导致较低的算术强度，从而导致较低的性能。GETT和TTGT，另一方面，保持相同的算术强度，因此为两个TC提供类似的性能。

图12总结了我们的性能讨论，并结合GETT、TTGT和LoG方法（以绿色显示）以反映TCCG表现出的性能。TCCG生成的代码始终与参考实现（以红色显示）相当，且通常显著更快；注意，由于Tensor Toolbox和libtensor缺乏单精度支持，单精度的参考性能低于双精度（见图8）。与GEMM的性能相比，TCCG在单精度和双精度下分别实现了74.0%/101.1%/141.4%和64.6%/101.8%/132.9%的最小/平均/最大值。

7.3. GETT性能模型的评估

图13描绘了基准中最佳1、4、8、16和32个GETT候选的性能；每个条形图中心的数字表示每个TC的可行候选总数。我们观察到GETT的性能模型工作到一个点，在大多数情况下，经验搜索（在多个候选中）几乎变得多余。定量地说，即使在将搜索空间限制为单个候选（即，消除搜索）的极端情况下，GETT在单精度（双精度）上平均仍然达到最快候选性能的90.7%（92.3%）。在四个候选中进行主动搜索将单精度（双精度）的平均性能提高到98.3%（97.2%）。此外，我们观察到在超过16个候选中进行搜索不会为任何测试的TC带来任何性能收益。

为了给用户提供在GETT性能和搜索时间之间进行权衡的灵活性，TCCG提供了通过--maxImplementations命令行参数限制搜索的可能性（默认：16个候选）。

8. 结论和未来工作

我们提出了类GEMM张量-张量乘法，这是一种计算张量收缩的新型高性能方法。GETT与之前的方法有几个关键区别：

任何张量收缩都被简化为一个高度调优且完全向量化的宏内核，其中宏内核的操作数被打包到缓存层次结构的指定级别。

在整个打包过程中保留步长为一的索引，导致有利的内存访问模式。

相对于等效大小的GEMM保持算术强度。

不需要额外的辅助工作空间。

我们在跨越广泛测试用例的张量收缩基准上评估了GETT的性能；平均而言，GETT在单精度和双精度下分别达到MKL的GEMM性能的98.1%（最小值：72.4%；最大值：141.4%）和97.0%（最小值：60.8%；最大值：132.9%）。在计算受限状态下，GETT实现了高达91.3%的峰值浮点性能的效率；在带宽受限的TC中，GETT有利的内存访问模式及其为各种缓存层次进行分块的能力的积极效果尤为明显，性能优于现有的TC方法高达12.4倍。

为了进一步评估GETT的性能，我们进行了包括两种替代方法的全面调查：转置-转置-GEMM-转置（TTGT）和循环-GEMM（LoG）。调查揭示了TTGT在带宽受限的TC中的缺点和LoG在某些情况下的任意差的性能。虽然TTGT在计算受限的TC中略微优于GETT（以额外内存为代价），但我们不预期在GETT的未来实现中消除这种差异会有任何概念上的障碍。鉴于这些结果，我们认为专门的、类GEMM的方法——如GETT和TBLIS——具有普遍吸引力，因为它避免了之前方法的缺点——即，额外的内存需求和次优的算术强度——并且不依赖于高性能GEMM的存在。

通过将GETT、TTGT和LoG结合到一个统一的张量收缩代码生成器中，我们获得了带宽受限和计算受限张量收缩的高性能。虽然TCCG提供了自动搜索多个候选的可能性，但我们发现搜索空间可以通过架构感知的指标（例如，假设对缓存大小的了解）有效修剪。例如，即使在GETT的搜索限制为单个候选（即，没有搜索）的极端情况下，GETT仍然在单精度和双精度下分别达到最快实现的89.5%和87.7%。

GETT的多线程版本可在GitHub上获得。并行化紧密遵循BLIS库的方法，并将在第一作者即将发表的博士论文中详细讨论。

在不久的将来，我们希望研究GETT在Intel最新的Xeon Phi架构（Knights Landing）上的性能，并探索我们的方法如何转移到形式为DΠn∪Il)的三元张量收缩，其中收缩索引同时出现在三个张量中。