The Tensor Algebra Compiler

The Tensor Algebra Compiler

0. Abstract

Tensor代数是一种功能强大的工具，可用于机器学习，数据分析，工程和物理科学中。张量通常是稀疏的，并且复合运算必须经常在单个内核中计算，以提高性能并节省内存。程序员要为每个操作编写内核，并使用不同格式的稠密和稀疏张量的不同混合。组合是无限的，这使得无法手动实现和优化它们。本文介绍了第一种编译器技术，该技术可自动为稠密和稀疏张量上的任何复合张量代数运算生成内核。该技术在名为taco的C++库中实现。它的性能与流行库中同类最佳的手动优化内核相比具有竞争力，同时支持更多的张量操作。

1. Introduction

Tensor代数是用于多维数据计算的强大工具，在机器学习，数据分析，工程和科学中具有许多应用。Tensor将向量和矩阵泛化为任意数量的维度并允许进行多线性计算。感兴趣的张量通常是稀疏的，这意味着它们大多包含可以压缩掉的零。例如，在数据分析中使用的大型现实世界数据集（例如Netflix Ratings，Facebook Activities和Amazon Reviews）可以方便地转换为稀疏张量。 Amazon Reviews张量包含1.5×$10^{19}$个组件，对应于107EB（1EB=1024PB, 1PB=1024TB, 1TB=1024GB​）的数据（假设每个组件使用8个字节），但是只有1.7×$10^9$个组件（13 GB）是非零的。

计算复合张量运算的内核优化可能非常复杂。首先，将稀疏的输入张量压缩为内核必须在其上操作的索引数据结构。其次，仅应生成非零输出值，并且计算得出的分量之间的计算会有所不同，具体取决于贡献非零值的输入张量。最后，稀疏张量数据结构通常不支持恒定时间随机访问，因此内核必须仔细协调多个张量的协同迭代。

当前的方法是手动编写用于张量操作的高性能代码。库提供了一组有限的手动优化操作，程序员使用临时张量通过一系列支持的操作来计算复合操作。这降低了局部性和效率，对于某些复合运算，临时变量比内核的输入和输出大得多。另一方面，由于所有可能的复合运算，张量阶数和张量格式的组合爆炸，手工为每个复合运算编写优化的代码是不可行的。因此，需要一种编译器方法来从高级表示法（例如用于张量表达式的张量索引表示法）生成快速内核。

本文的贡献是：

张量存储格式，分别将每个维指定为稠密或稀疏，并指定维的存储顺序，它可以描述几种广泛使用的格式，但可以概括为更多格式（第3节）；

描述如何在复合张量表达式中稀疏操作数的多级索引数据结构进行迭代的迭代图（第4节）；

合并格，描述如何合并在相同张量表达式中使用的稀疏张量的索引数据结构（第5节）；

以及使用上述概念生成有效代码的代码生成算法，该代码计算具有稠密，稀疏和混合操作数的张量表达式（第6节）。

为了演示我们的方法，我们开发了一个名为taco的C++库，它是Tensor Algebra Compiler的缩写（第7节）。我们将taco生成的代码与广泛使用的最新技术的手写实现进行了比较。比较表明，taco可以为较简单的内核（如稀疏矩阵向量乘法）以及复杂的内核（如矩阵张量乘以Khatri-Rao乘积）生成有效的代码（第8节）。

2 Example

张量将矩阵推广到任意数量的维度（尽管在本文的其余部分中我们将使用“维度”，但通常也称为模式）。张量具有的维数是其阶数。因此，标量是0阶张量，向量是1阶张量，矩阵是2阶张量。考虑一个简单的三阶张量内核，即张量矢量乘法（TTV）：
$$A_{ij}=\sum _k{B}_{ijk}{c}_k$$
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-8jKpXfCd-1578556932426)(.\typora-user-images\image-20200101125738652.png)]

我们在演示文稿中使用了这种表示法，并作为我们技术的输入，它是Ricci-Curbastro和Levi-Civita 开发的张量指数表示法的变体。该表示法使用索引变量将结果中的每个组件与操作数中的组件相关联，这些操作数被组合以产生其值。在此示例中，输出的每个分量$A_{ij}$是从B的最后一个维度到c的向量的内积。

计算此表达式的内核完全取决于三个张量的格式。最简单的情况是所有张量都稠密时，其内核如图1所示。循环在张量维度的m×n×p矩形迭代空间上迭代，并且由于张量很密，所以可以计算每个组件的物理位置。

如果B中的大多数分量为零，则仅存储非零会更有效。这可以通过稀疏张量数据结构（例如压缩的稀疏维（CSF））来实现，其中每个维度都被压缩。但是，必须编写内核以遍历数据结构以找到非零值。图2显示了将B存储在CSF中时用于计算表达式的代码。循环遍历包含非零的B的每个维度的子集（第1-13行），执行乘法运算，并将结果存储在A的正确位置（第18行）。

如果c中的大多数分量也为零，则也可以将其压缩。但是，内核现在必须同时迭代并合并由B和c中的k索引的维。这意味着内核仅在B和c都在某个位置具有非零值时才计算值。合并代码如图3所示。 在B和c都有值时While循环迭代，并且仅在两个位置都具有值时才进行计算。

这些示例表明，根据张量的格式，同一表达式的不同内核看起来完全不同。而且，对于我们的技术仅支持此表达式的所有格式组合，这仅是384种可能内核中的三种。手工实现每个都是不现实的。此外，可能的复合表达式的数量不受限制。可以将多个更简单的内核并置在一起，从而高效地进行计算，而另一些则最好作为单个复合内核进行评估。本文介绍的技术由用户决定，但无需手工编写内核。这使得混合和匹配格式并为任何张量代数运算自动生成内核成为可能。

稀疏张量代数内核复杂度有两个主要原因。首先，稀疏数据结构只能在一个方向上有效地访问。例如，B的索引遵循维度的顺序。有时最好选择其他顺序，例如先存储B的最后一个维度。这在稀疏线性代数中很常见，其中压缩稀疏行（CSR）格式从行到列，而压缩稀疏列（CSC）格式从列到行。但是，任何维度的排序都必须遵循内核循环嵌套的排序。在第4节中，我们使用称为迭代图的概念解决循环排序问题。

稀疏内核复杂性的第二个原因是稀疏张量索引的合并。我们在B和c都稀疏时看到了这一点，但考虑了表达式$a_i=b_i{c}_i+d_i$，其中两个矢量的按分量乘积被加到第三个矢量上。与d的合并与bc合并不同，因为该操作是加法运算，如果两个操作数中的任何一个非零，该运算都会产生一个值。 （如果两个操作数都不为零，则乘法将产生一个值。）因此，如果b和c都具有非零或d都具有非零，则b，c和d的三元合并将仅产生一个值。图4示出了一个例子。当只有b或c具有非零值而没有d时，将不产生输出。当d为非零值时，总是产生输出；但是，如果b和c的值均非零，则将其加到b和c的乘积中。最后，如果b和c具有非零值但不具有d，则也会产生输出。为了提高效率，仅访问输入中的非零值，并且仅生成输出中的非零值。如图所示，每个非零输出都是使用仅包含非零输入的简化表达式来计算的，从而避免了不必要的操作。我们使用第5节中称为合并格的新概念解决合并的复杂性。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-0Tnkh6WV-1578556932428)(.\typora-user-images\image-20200101131013756.png)]

3 张量存储

张量可以具有任何顺序（维度），因此枚举所有张量的格式并非易事。为了支持任何顺序的张量，有必要从有限数量的原语构造格式。本节介绍了一种方法，该方法通过指定每个维度是稠密维度还是稀疏维度并声明维度的存储顺序来定义任何顺序的张量的存储格式。对于稀疏级别，我们使用与CSR格式相同的压缩技术。但是，我们允许将这些级别组合起来以构造任何张量的稀疏存储。由于每个级别可以独立地稠密或稀疏，因此用户可以为特定张量构建自定义格式。分别描述每个维度的第二个好处是，它导致可组合的代码生成算法，该算法支持任何阶数的张量。

考虑图5a所示的二阶张量（矩阵）A。最简单的存储格式是稠密的多维数组，但是如果大多数分量为零，则很浪费。直观地，我们发现将张量视为一个树，树的每一层代表张量的一个维度是方便的，如图5b–c中图1和2中矩阵A所示。在此公式中，从根到叶的每个树路径代表一个具有非零值的张量坐标。沿着路径的非根节点是坐标，并且非零值附加到叶节点。最后，根据A的维的存储顺序，树的层可以以不同的顺序产生，例如（d1，d2）或（d2，d1），其中di表示矩阵的维数。

在我们的方法中，张量值始终存储在单独的数组中，并且张量格式索引数组用于确定它们在数组中的位置。对于每种级别的存储，我们都存储与相应维相关联的索引元数据：

稠密仅需要存储维的大小，因为它存储维中每个坐标的值（零和非零）。

稀疏仅存储具有非零值的相应维度的子集。这需要两个索引数组pos和idx，它们一起形成分段向量，在前一维（树中的父节点）中每个条目有一个分段。 idx数组存储维度中的所有非零索引，而pos数组存储idx数组中每个段开始的位置。因此，段i存储在idx数组中的位置pos [i]：pos [i + 1]中（我们在pos数组的末尾添加了idx数组大小的标记）。我们将每个段按排序顺序存储在idx中。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-6PcMfuc5-1578556932430)(.\typora-user-images\image-20200101131601324.png)]

请注意，稀疏维度中的索引数组与CSR矩阵格式中的索引数组相同。除了高阶张量，我们的公式还可以表示几种常见的稀疏矩阵格式。图5的d–k显示了使用我们的技术存储二阶张量的所有八种方法。第一列显示稠密的行和列主存储。第二列显示（$spars{e}_{d_1}$，$dens{e}_{d_2}$）格式及其第一列等效格式，这两种格式都不常用，但在某些情况下可能有用（请参阅8.5节）。第三列显示CSR和CSC格式。两者都表示为（稠密，稀疏），但是维度顺序已切换。最后，第四列显示（稀疏，稀疏）格式，等效于双压缩稀疏行（DCSR）和相应的列主要格式DCSC 。格式的数量随着张量维数的增加而呈指数增长（实际上是$d!2^d$，其中d是张量的阶数），使得手工编码变得难以处理。此技术支持的其他重要稀疏格式包括稀疏矢量和高阶张量的CSF格式，它们对应于对每个张量维使用稀疏级别的格式。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-fzIagQBW-1578556932433)(.\typora-user-images\image-20200101195652690.png)]

稀疏的存储维不允许有效地随机访问索引和值，但已针对特定顺序的迭代进行了优化。图6a显示了如何在稀疏级别上进行迭代。 pos数组给出idx数组中每个段的边界，并以步幅对段中的索引进行代码迭代。图6b显示了通过（$spars{e}_{d_1}$，$spars{e}_{d_2}$）二阶张量进行迭代的代码与张量的存储之间的对应关系。

我们支持的两种级别的存储表示张量格式和操作的空间很大，包括块稀疏矩阵矢量乘法，我们将其转换为4阶张量的乘法（一个块矩阵，其中两个内部稠密维存储小密度非零的块）和二阶张量（一个块向量）。以这种方式描述张量存储格式的空间使我们能够支持无数格式，为我们免费提供了对块线性代数的支持，并使代码生成针对每个存储级别模块化（第6节）。

4 迭代图

迭代图是一个编译器中间表示形式（IR)，它描述如何迭代张量表达式的非零值。这需要生成代码以同时遍历表达式操作数的张量存储树。迭代图由张量索引表达式构造而成，并用于生成此类代码。它们是通用的，可以表示任何张量表达式，从简单的稀疏矩阵向量乘法（SpMV）到复杂的复合表达式，例如矩阵张量与Khatri-Rao乘积（MTTKRP），我们接下来将展示。

定义4.1 迭代图是有向图$G=(V，P)$，它具有一组索引变量 V = { v 1 ， … ， v n } V = \{v_1，…，v_n\} V={v1​，…，vn​}和一组张量路径 P = { p 1 ， … ， p m } P = \{p_1，…，p_m\} P={p1​，…，pm​}。张量路径是索引变量的元组。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ToJFc0rU-1578556932434)(.\typora-user-images\image-20200101200105396.png)]

如前一部分所述，张量可以用树表示，其中每个级别对应于可以是稠密或稀疏的维。图7a示出了两个示例，三阶张量B和向量c，而图7b示出了张量访问表达式$B_{ijk}$和$c_k$的迭代图。迭代图中的每个节点对应于循环嵌套中的循环。这些迭代图的每一个都包含通过访问表达式中的索引变量的路径。我们把他们叫做张量路径，它们象征性地总结了B和c中从根到叶的所有遍历。迭代图的目的是为了推理这些遍历。图7c示出了张量向量乘法的迭代图。每个张量访问子表达式有一个张量路径，并表示B和c的张量存储树上的同时迭代。它还显示了输出A的张量路径。通常，迭代图可以具有任意数量的张量路径。

通往同一节点的张量路径边缘表示需要在此维度上合并的张量。也就是说，要同时迭代多个张量存储树，必须将它们合并。在张量-向量乘法示例中，这意味着在遍历k个索引变量时将B的最后维度与c合并。由于它是一个乘法，因此合并是一个“与”（and），因为两个值都必须非零才能使结果非零。如果是加法，则合并将是"或"（or）。如果要合并两个以上的张量维，则合并将是“与”和“或”的组合，如第5节将进一步讨论。

图8显示了八个更多的内核迭代图示例，范围从图8a中的简单SpMV到图8h中的MTTKRP。如果所有操作数都是稀疏的，则所有这些示例都需要合并。注意，图8f中的块稀疏矩阵矢量乘法被转换为4阶张量乘以2阶张量乘法。最后，请注意图8d中的采样稠密稠密矩阵乘积（SDDMM）内核，其中B是稀疏的。这是机器学习的一个内核，并且由于B稀疏，因此可以通过评估单个循环嵌套中的整个表达式将CD中要评估的内积数量从$O(n^2)$减少到$O(nnz(B))$。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-wLw2ldiX-1578556932436)(.\typora-user-images\image-20200101204051042.png)]

通过为每个张量访问表达式创建一个张量路径，从索引表达式构建迭代图。张量路径基于访问表达式中索引变量的顺序和张量存储树中维的顺序进行排序。因此，如果第一个维度是存储树的第一级（例如CSR），则$B_{ij}$表示路径i→j；否则，则是j→i（例如CSC）。最后，将迭代图的索引变量排序到林中，以便每个张量路径边缘从较高的索引变量向上移动至较低的索引变量。这是确保张量树从根到叶遍历的必要条件。

可以为任何张量索引表达式构造迭代图，但是第6节中描述的代码生成技术不支持带有循环的迭代图。原因是后边缘需要沿非自然方向遍历某种格式，例如从列到行遍历CSR矩阵。这将需要用于扫描索引以查找组件的代码，而这超出了本文的范围。当表达式对具有不兼容格式的张量进行运算时，就会发生循环，例如将CSR矩阵添加到CSC矩阵或转置矩阵。解决方案是提供一种更改张量格式的功能，该功能可用于中断循环和转置张量。

5 合并格

访问多个张量维的索引变量必须遍历那些维的合并索引。合并的类型取决于张量索引表达式。如果张量相乘，则合并是一个conjunction（∧），因为两个因子都必须非零才能使结果非零（a*0 = 0）。这意味着合并的迭代空间是合并维的迭代空间的交集。相反，如果将两个张量相加，则合并是disjunction（∨），因为只有一个项需要为非零才能使结果为非零（a + 0 = a）。这对应于合并维的迭代空间的并集。如果合并了两个以上的维度，则合并是反映索引表达式的conjunction和disjunction的组合。例如，$（b_i+c_i）d_i\Rightarrow （b_i\vee c_i）\wedge d_i$。

合并格是由disjunction的成本推动的，在disjunction中，每次循环迭代必须检查每个合并索引是否还有更多值，以避免越界访问。双向合并算法用于避免这些昂贵的检查。它具有三个循环：一个循环直到两个索引中的任何一个耗尽（值用完），另一个循环处理其余的未耗尽索引。我们用合并格来概括这种方法，合并格可用于为任何合并表达式生成有效的循环。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-nMEBOp5g-1578556932437)(.\typora-user-images\image-20200101204435842.png)]

考虑一个具体的例子，稀疏向量加法，它需要disjunction。图9（左）显示了一个示例，结果为a，两个操作数b和c，其中c具有更多的值。双向合并对两个操作数进行迭代，直到用尽任何一个。在每个步骤中，将b + c，b或c添加到a上，具体取决于b和c，仅b还是仅c都在该坐标处具有值。向量之一耗尽后，双向合并将迭代另一个向量中的其余条目。在此示例中，b首先耗尽，因此将c的剩余值添加到a。

图9（中心）显示了用于稀疏矢量相加的合并格，其中三种情况中的每一种都有一个格点：其中b和c具有值，仅b具有值，或仅c具有值。每个格点包含这种情况所需的值的条件（bi∧ci，bi或ci）以及要计算的子表达式。格子箭头对应于何时用完一个维度。

合并格可以合并任意数量的维度，因此可以具有任意数量的格点。格点有两个目的。首先，它们描述了合并维的必要循环。顶部的格点描述了每个维度仍有剩余值的循环。当用完任何维度时，我们沿着箭头移动到另一个格点，该格点描述了其余维度上的迭代。重复此过程，直到达到最低点为止，在此不再有任何值可以合并。第二个目的是描述每个循环中必须考虑的情况。由任何格点描述的迭代可能会遇到不同的情况，具体取决于合并维的哪个子集在给定位置具有值。用给定格点（包括其自身）控制的格点来描述必要的情况，并且每个格点都包含要为其情况计算的表达式。

**定义5.1 **合并格L是包含n个格点 { L 1 ， … ， L n } \{L_1，…，L_n\} {L1​，…，Ln​}和一个运算符的格。每个格点$L_p$都有一组张量维度 T p = { t p 1 ， … ， t p k } Tp = \{t_{p1}，…，t_{pk}\} Tp={tp1​，…，tpk​}要conjunction合并（比如$t_{p1}\dots \dots t_{pk}$）和表达式$exp{r}_p$要计算。分别具有相关张量维度$T_1$和$T_2$的两个格点$L_1$和$L_2$的结合是具有张量维度$T_1\bigcup T_2$的格点。我们说$L_1\le L_2$当且仅当$T_1\subseteq T_2$，换句话说，如果$L_2$具有在$L_1$中耗尽的张量维度，反之则不然。

图9（右）显示了从合并格生成的代码。每个格点都会产生一个while循环：格点$bi\wedge ci$在6-21行的循环中，$b_i$在22-22行的循环中，$c_i$ 25-27行的循环中。此外，$b_i\wedge c_i$支配了三个非底部格点，因此，最终的循环考虑了第12-17行的三种情况。每个案例将案例中每个张量维度的索引变量与循环迭代的合并索引变量进行比较，后者是维度索引变量中最小的变量，并在第9行计算得出。最后，维度位置变量递增为需要19-20行。

5.1 构建

在本节中，我们描述一种为索引表达式中的索引变量构造合并格的算法。该算法从下至上遍历索引表达式树，在叶子处构造合并格，并使用以下运算符在内部节点处连续合并合并格：

定义5.2 令两个格点$L_p$和$L_q$与算子$op$的结合为具有相关张量维度$Tp\bigcup Tq$和表达式$exp{r}_{p}opexp{r}_q$的新格点。

定义5.3 令两个格$L^1$和$L^2$与算子$op$的合并表示为$L^1{\wedge }_{op}{L}^2$，这是一个新的格，其格点是通过对$（L_0^1，...，...，L_n^1）\times （L_0^2，...，L_m^2）$中的每对格点进行笛卡尔乘积。

定义5.4 假设两个格$L1$或$L2$与算子$op$的disjunction表示为$L^1{\vee }_{op}{L}^2$，是一个新的格，其中包含$L^1{\wedge }_{op}{L}^2$，$L^1$和$L^2$中的所有格点。

以下算法根据索引表达式和索引变量$i$构造合并格。该算法从下至上对索引表达式进行操作，并对每个子表达式应用以下规则之一：

• 张量访问：构造一个合并格，其中一个格点包含在访问表达式中由i索引的维。格点表达式是张量访问。如果i没有访问任何维，则构造一个空格点。
• conjunction运算符（例如乘法）：计算并返回操作数合并格与运算符的conjunction。
• disjunction运算符（例如加）：计算并返回操作数合并格与运算符的disjunction。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-ueOK4Thx-1578556932439)(.\typora-user-images\image-20200101212150514.png)]

要建立直观的认识，请考虑表达式$a_i=b_i+c_i{d}_i$，这是组合加法和逐分量乘法的组合。该算法首先使用张量访问规则从叶中创建格，然后通过计算操作数合并格的conjunction，为$c_i{d}_i$创建合并格。此合并格具有一个格点（$c_i\wedge d_i$）。在内联符号中，圆括号和格点包围着格点，格点按顺序排列，以便每个点都在其祖先之后出现。最后，由于顶级表达式是加法运算，因此该算法计算disjunction，$(b_i){\vee }_{+}(c_i\wedge d_i)=(b_i{c}_i\wedge d_i，b_i，c_i\wedge d_i)$。最终的合并格如图10b所示。顶部导致所有三个维度上的循环。如果$c_i$或$d_i$耗尽，则格循环退出，$b_i$循环接管。否则，如果先耗尽$b_i$，则$c_i\vee d_i$的循环将接管。

5.2 优化

使用第5.1节中的算法构造的合并格会合并所有访问的维。如果它们是稀疏的，这是必需的，但是如果它们是稠密的，则可以通过三种方式来优化格：

稠密的迭代空间是相同的，因此，如果多个合并维之一是稠密的，那么我们只需要迭代其中一个即可他们。图10c和图10d分别示出了向量乘积和向量加法中的合并的稠密维度。

稠密迭代空间是稀疏迭代空间的超集，因此当耗尽稠密维时，我们就完成了。这意味着我们可以删除不合并每个稠密维度的顶部点下方的格点，因为如果在顶部格点中耗尽了稠密维度，则将没有更多的值可合并。图10d显示了稠密的矢量加法，其中除一个格点外的每个格点均被删除。

稠密的迭代空间支持随机访问，因此当与稀疏的迭代空间合并使用时，我们可以在稀疏的迭代空间上进行迭代并从稠密的空间中选取值。图10e示出了具有稀疏b和稠密c的矢量乘积格。格在b上迭代，因为c是具有随机访问权限的超集。此优化用于SpMV内核。

6 代码生成

在本节中，我们描述了如何在递归算法中使用张量存储描述符，迭代图和合并格来从复合索引表达式生成内核。该算法会产生循环，循环遍历表达式的合并迭代空间（第6.1节）。它还用计算结果值（第6.2节）和组合结果索引（第6.3节）的语句填充这些循环。这些语句可以合并为一个内核，该内核同时组装索引和计算值，或者分离为一个仅执行汇编的内核和另一个仅计算的内核。当在非零结构保持不变的情况下重新计算值时，后者很有用。最后，我们将讨论并行代码生成（第6.4节）。

6.1代码生成算法

递归代码生成算法的伪代码在图11a中给出。编译器代码显示为蓝色，发出的代码显示为带引号的黑色和红色。引号内的蓝色文本发出编译器变量的值，例如索引变量的名称。图11b显示了生成的用于256×256稀疏矩阵加法的代码，其中A是稠密的，输入B和C是CSR矩阵。该代码假定稠密输出矩阵的值已预先初始化为零。带编号的标签将代码生成算法与发出的代码相关联。代码生成算法是递归函数。通过索引表达式和该表达式的迭代图中的第一个索引变量进行调用，然后以森林顺序在迭代图中的索引变量上递归。图11c示出了矩阵相加的迭代图。在每个递归级别，代码生成算法首先根据索引变量和索引表达式构造一个合并格。无花果图11d和11e示出了在稀疏矩阵相加示例中为i和j创建的合并晶格。接下来，它初始化稀疏idx变量（1），然后是一个循环，该循环按级别顺序（2）在每个晶格点发出一个while循环。每个while循环都从加载（3）和合并（4）稀疏idx变量开始。合并后的idx变量就是当前索引变量维度中的坐标。然后，通过将合并的idx变量添加到变量的相应张量切片的开头，来计算稠密的pos变量（例如pB1）（5）。接下来，该算法在由循环晶格点（7）支配的子晶格中每个晶格点发出一个if-else if case。例如，顶部晶格点的子晶格是整个晶格（一个子晶格包括其自身）。其他晶格点的子晶格包括从它们可到达的所有点。在格点的if-else if情况下，code-gen函数递归调用自身，以迭代图森林顺序为当前索引变量的每个子代生成代码。接下来，它发出代码以在此循环级别插入索引条目并根据需要计算值；这将在6.2和6.3节中更详细地描述。最后，如果稀疏pos变量参与了当前的合并坐标（8），该算法将发出代码以有条件地增加它们。在我们描述的算法中，pos变量是基于它们访问的张量来命名的。但是，在存在相同张量的重复访问的情况下，例如在表达式Aij =“ k BikBkj中，存在一个单独的机制，可以确保唯一的名称与与每个张量访问相关的pos变量相关联。

6.2计算代码

我们已经展示了如何生成循环遍历输入张量的循环，但到目前为止，我们没有将哪些语句插入到这些循环嵌套中以计算结果值。图11中的代码生成算法调用了三个函数来发出计算代码，即emit-available-表达式（6），emission-reduction-compute（7）和emit-compute（7）。对于任何给定的索引变量，这些函数中只有一个函数会根据上面是否是递归中的最后一个自由变量来发出代码。最后一个自由变量，或低于它。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-gWDAmtRT-1578556932447)(.\typora-user-images\image-20200101215743827.png)]

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-zoDPTEPy-1578556932450)(.\typora-user-images\image-20200101215803068.png)]

最后一个free变量是特殊的，因为它的循环嵌套是代码写入输出张量的地方。嵌套在其上方的循环计算可用的表达式（emit-available-expressions），这有助于避免多余的计算。嵌套在它下面的循环是归约循环，该归约循环将子计算添加到归约变量（emit-reduction-compute）。最后，最后一个自由变量的循环结合了其他循环准备的临时变量，以计算最终表达式（emit-compute）。

6.3 索引汇编代码

emit-index-assembly函数发出的代码将输出张量中的稀疏级别的索引结构进行汇编，该结构由一个pos数组和一个idx数组组成，该pos数组存储每个索引段的起点，而idx数组包含该段中条目的索引。对于每个输出的稀疏级别，都会发出代码，将与该级别相对应的非零坐标插入到idx中。例如，

A2_idx [pA2 ++] = j;

对于输出中高于稀疏级别的级别，emit-index-assembly还发出代码以更新下一个级别的pos数组，以匹配插入到idx中的位置。但是，由于子循环不一定会产生任何值，因此在最后一个自由变量之上可能会出现复杂情况。例如，在逐分量矩阵乘法中，两行可能都具有非零值，但交集可能没有。为了防止结果索引结构中的空位置（合法但在压缩方面次优），发出的代码检查子循环是否产生非零值。例如，

A2_pos [pA1 + 1] = pA2;
if(A2_pos [pA1 + 1]> A2_pos [pA1]){
    A1_idx [pA1 ++] = i; 
}

条件测试子循环的当前位置和先前位置是否相同。如果不是，则子循环产生非零值，因此汇编代码将新坐标插入索引。

最后，有必要为结果张量分配内存。通过发出代码来检查idx，pos和vals数组中是否剩余空间，然后再写入这些代码来处理它们。如果没有更多空间，则发出的代码会使内存分配加倍。未来的工作包括设计启发式方法以设置数组的初始大小。

6.4 并行代码生成

代码生成算法使用OpenMP并行pragma批注循环。它仅使外部循环并行，从而提供足够的并行度和良好的并行粒度。但是，这种方法应由用户控制。循环并行化受三个限制：

• 循环不得合并张量维，即，它必须是for循环。可以从合并格确定此条件。
• 自由变量不能由迭代图中的约简变量控制，因为这会导致分散行为，并且我们还没有发出并行同步构造。
• 输出张量必须在所有维度上稠密。

未来的工作包括放宽这些限制。尽管如此，在不放松的情况下，代码生成算法仍然能够为许多实际的实际内核发出并行代码。

7 TACO C++库和命令行工具

我们已经在名为taco（Tensor Algebra Compiler的缩写）的编译器中实现了本文中描述的技术。 taco可以用作C ++库，使程序员可以在其应用程序中的张量上进行计算。它也可以用作命令行工具，使用户可以生成C内核以包含在应用程序中或用作进一步开发的起点。

图12演示了如何使用C++库计算第2节中的张量矢量乘法。可以通过指定张量的维度，其条目的类型和存储格式来创建张量对象。张量的存储格式可以依次通过创建Format对象来声明，该对象描述每个张量级别的存储种类以及级别的存储顺序，遵循第3节中的公式。在第1-8行中，A声明为一个CSR矩阵，B声明为CSF张量，c声明为稀疏向量。通过指定所有非零分量的坐标和值，然后调用pack方法以所需的存储格式存储张量，可以使用用户定义的数据来初始化输入到计算中的张量，如第10-17行所示。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-wzbU0GJY-1578556932454)(.\typora-user-images\image-20200101221442281.png)]

张量代数计算在taco中用张量索引符号表示，如第19-20行所示。注意第20行与第2节中提出的张量矢量乘法的数学定义非常相似；这可以通过操作符重载来实现。 Var对象对应于张量索引表示法中的索引变量。求和减少是指仅出现在表达式右侧的索引变量。

在张量代数计算的目标（示例中为A）上调用compile会提示taco生成内核以组装索引结构并进行计算。在当前实现中，taco生成C代码，调用系统编译器将其编译为动态库，然后将其动态链接到dlopen。这使得这些函数可用于以后的调用以进行组装和计算。系统编译器处理低级性能优化。这些步骤是自动发生的，对用户不可见。这种方法很适合开发，但是我们计划也实现不需要系统编译器的LLVM JIT后端。

接下来，assemble方法组装输出张量的稀疏索引结构并为其预分配内存。最后，通过调用计算方法以执行编译生成的代码来执行实际计算。通过将输出程序集与实际计算分开，我们使用户可以一次组装输出张量，然后多次重新计算其值。这是有用的，因为在许多实际应用中，重复计算会更改输出张量的值，但不会更改其非零结构。

图13显示了如何使用taco命令行工具生成计算相同张量向量乘法操作的C代码。和以前一样，张量索引表示法用于指定生成的内核计算的操作。再次按照第3节中的描述，-f开关可用于指定输入和输出的存储格式。在后台，命令行工具使用与C ++库相同的代码生成机制。我们使用taco命令行工具生成了本文中描述的所有内核，并将它们作为补充材料包括在内。

8 评估

为了评估本文中描述的技术，我们使用taco以及一些现有的流行的稀疏线性和张量代数库来计算许多实际表达式，并以实际矩阵和张量作为输入。在第8.2节和第8.3节中，我们证明了我们的技术生成了一系列稀疏的线性代数内核，这些内核与手动优化的内核的性能具有竞争力，同时消除了现有库在性能和完整性之间的取舍。我们在8.4节中进一步证明，这些观察结果对于涉及高阶张量的稀疏张量代数也成立。最后，我们在8.5节中表明，支持各种稠密和稀疏张量存储格式对于通过实际的线性和张量代数计算实现高性能确实是必不可少的。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-sTjGCfV0-1578556932456)(.\typora-user-images\image-20200101221811534.png)]

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-Y4VYEp2s-1578556932458)(.\typora-user-images\image-20200101221828361.png)]