TensorIR: An Abstraction for Automatic Tensorized Program Optimization

文章标题: TensorIR: 一种用于自动化张量化程序优化的抽象
作者/机构:
Siyuan Feng (上海交通大学), Bohan Hou (卡内基梅隆大学), Hongyi Jin (卡内基梅隆大学), Wuwei Lin (OctoML), Junru Shao (OctoML), Ruihang Lai (卡内基梅隆大学), Zihao Ye (华盛顿大学), Lianmin Zheng (加州大学伯克利分校), Cody Hao Yu (亚马逊网络服务), Yong Yu (上海交通大学), Tianqi Chen (卡内基梅隆大学, OctoML)

A1 主要贡献

核心问题: 在各种设备上部署深度学习模型已成为一个重要课题。硬件专业化的浪潮为多维张量计算带来了多样化的加速原语。这些新的加速原语，连同新兴的机器学习模型，带来了巨大的工程挑战。特别是，现代硬件（如Nvidia Tensor Core，Google TPU）引入了专门的指令（称为张量化内联函数, tensorized intrinsics）来加速张量计算，但利用这些原语来优化程序（即张量化程序优化）通常需要领域专家手动编写高度优化的库（如cuDNN, MKL-DNN），这耗费大量工程精力，并且难以适应快速变化的机器学习模型。

研究目标: 本文旨在通过一种自动化的编译方法来解决张量化程序优化问题。现有的机器学习编译器大多专注于循环嵌套变换，无法自动处理张量化程序。本文的目标是创建一个能够自动利用现代硬件中领域特定加速优势的编译器。为此，需要解决两大挑战：
1. 张量化程序的抽象: 需要一种能够务实地捕捉给定机器学习算子的等效张量化计算的抽象。该抽象必须能表示多维内存访问、线程层次结构以及来自不同硬件后端的张量化计算原语，并足够表达机器学习中的大多数算子。
2. 巨大的优化设计空间: 需要一种有效的方法来自动地在一个由循环分块、线程化、数据布局变换以及张量化计算组合而成的巨大搜索空间中，找到一个优化的张量化程序。

创新点: 为应对上述挑战，本文提出了TensorIR，一个用于自动化张量程序优化的抽象。其核心贡献如下：
* 新颖的张量化程序抽象: 提出了一种新颖的抽象，通过一个名为block的新结构，将张量化计算与循环变换分离开来。这个block允许将内部的张量化计算区域与外部的循环嵌套隔离开，从而实现“分而治之”的优化策略。同时，该抽象能够统一表示张量内联函数和硬件约束。
* 丰富的变换原语: 建立了多种程序变换原语，用于生成一个包含张量化程序优化的丰富搜索空间，并带有正确性验证机制。
* 张量化感知的自动调度器: 设计并实现了一种新的、能够感知张量化的自动调度算法。该算法能够通过分析计算模式生成张量化候选项，构建张量化程序草图，并利用基于学习的成本模型和进化搜索来找到最优程序。

本文将TensorIR与一个端到端的编译框架集成，实验表明，该框架能自动利用硬件后端的张量计算原语，其性能与业界最先进的手工优化系统相比具有竞争力。

图1：硬件专业化的趋势。经典的加速技术使用向量单元同时处理多个标量计算，这在CPU平台上仍被广泛使用。然而，为了满足日益增长的计算吞吐量需求，现代加速器通常包含专门的高维张量计算指令，从而产生了对张量化程序优化的需求。
图1：硬件专业化的趋势。经典的加速技术使用向量单元同时处理多个标量计算，这在CPU平台上仍被广泛使用。然而，为了满足日益增长的计算吞吐量需求，现代加速器通常包含专门的高维张量计算指令，从而产生了对张量化程序优化的需求。

A3 设计原则

本文方法的关键洞见：本文的关键洞见源于对领域专家如何优化张量化程序的观察。专家们通常采用“分而治之”的方法：首先，将原始程序问题分解为两个子问题——一个对应于硬件支持的张量化计算原语（如一个4x4的矩阵乘法），另一个是使用这些原语的外部循环嵌套；然后，分别对这两个子问题进行优化。这种方法允许开发者专注于一个子问题而不必担心其他部分，同时也便于针对多种张量化计算实现进行优化。

图2：一位专家开发者可以选择将问题分解为4x4矩阵乘法和使用该4x4矩阵乘法的循环，然后分别进行优化。这样我们可以有效地利用目标硬件中的专用张量指令。
图2：一位专家开发者可以选择将问题分解为4x4矩阵乘法和使用该4x4矩阵乘法的循环，然后分别进行优化。这样我们可以有效地利用目标硬件中的专用张量指令。

与现有方法的对比：大多数现有的机器学习编译器采用两种方法。第一种是“自下而上”的方法，如Halide【35，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】、TVM【9，{TVM}: An automated end-to-end optimizing compiler for deep learning，2018，{USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】、Tiramisu【5，Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code，2019，IEEE/ACM International Symposium on Code Generation and Optimization (CGO 2019)】等，它们将搜索空间建模为围绕标量操作体的循环嵌套迭代器，通过寻找最佳的循环嵌套变换来优化程序。第二种是“自上而下”的方法，如HTA【16，The Hierarchically Tiled Arrays Programming Approach，2004，Workshop on Languages, Compilers, and Run-Time Support for Scalable Systems (LCR ’04)】、Fireiron【17，Fireiron: A Data-Movement-Aware Scheduling Language for GPUs，2020，ACM International Conference on Parallel Architectures and Compilation Techniques】和Stripe【46，Stripe: Tensor compilation via the nested polyhedral model，2019，arXiv】等，它们通过嵌套的多面体结构逐步将问题分解为子问题。本文旨在将“分而治之”的洞见引入机器学习编译器设计中，并给出了肯定的答案。

图3：我们方法的概述。我们使用一个名为block的关键抽象来划分和隔离张量化计算，并支持对带有张量化操作的循环进行进一步变换。
图3：我们方法的概述。我们使用一个名为block的关键抽象来划分和隔离张量化计算，并支持对带有张量化操作的循环进行进一步变换。

引入Block抽象：为了实现“分而治之”，本文在循环嵌套中引入了一个名为block的新抽象。一个block包含了适量的签名信息，足以将内部问题空间与外部问题隔离开来。有了block，我们可以利用其签名作为接口，继续独立地对外部和内部问题进行变换。block的常见用途是表示硬件后端的张量化计算原语，但也可以用来隔离更大的子问题。在TensorIR中，张量计算是一等公民，带有block的循环嵌套可以被视为迭代空间的广义抽象。

自动化优化流程：为了自动化张量化程序优化，我们构建了一个搜索空间，该空间由硬件张量计算原语指导下的问题划分方式组成，然后进一步搜索解决子问题的可能程序变换。这种“分而治之”的方法覆盖了之前编译器方法（自下而上、自上而下）的搜索空间，并将高性能计算（HPC）和机器学习工程中典型的优化技术推广到一个允许自动张量化的抽象中，从而能够自动生成与供应商特定库相媲美的解决方案。

A2 方法细节

3 TENSORIR 抽象

本节介绍 TensorIR 抽象。图4给出了一个TensorIR程序的例子。我们引入了TensorIR的Python-AST（抽象语法树）方言，让开发者可以直接在Python中构建和转换程序。一个TensorIR程序包含三个主要元素：多维缓冲区、循环嵌套（在GPU设置中可能带有线程绑定）和块（block）。一个块可以包含一个或多个嵌套的循环嵌套，其叶子节点是子块或对应于计算内容的命令式语句序列。这种表示允许我们将计算划分为相应的子（块）区域，并利用存储在块签名中的依赖信息进行有效的程序转换。

图4：一个TensorIR程序示例，包含三个主要元素——多维缓冲区、循环嵌套和计算块。为简化起见，省略了块的细节。
图4：一个TensorIR程序示例，包含三个主要元素——多维缓冲区、循环嵌套和计算块。为简化起见，省略了块的细节。

3.1 Block

Block的基本定义：在TensorIR中，一个block代表对多维缓冲区子区域的张量化计算。图5展示了一个用于矩阵乘法（matmul）计算的block示例。block的主体由一组块迭代器变量（$vy, vx, vk$）参数化，这些变量代表一个抽象的张量化计算。通过将这些块迭代器变量实例化为不同的值组合，block会映射到不同的具体运行块实例。这些迭代器变量可以绑定到包含外部循环迭代器的表达式上，从而隐含了块实例的执行顺序。

图5：块包含用于依赖性分析的完整签名，我们将其作为主体计算和外部循环嵌套之间的隔离层。
图5：块包含用于依赖性分析的完整签名，我们将其作为主体计算和外部循环嵌套之间的隔离层。

设计原理：block的主要设计原理是隔离张量化计算——我们希望能够在不查看其主体的情况下转换block外部的循环嵌套。然而，与标量计算不同，我们可能无法从一个不透明的张量计算主体中提取转换所需的依赖信息。因此，我们引入了一个block签名，其中包含足够的依赖信息以支持转换。此外，该签名可用于在转换过程中独立验证迭代器绑定的正确性。

块迭代器域：虽然理论上可以通过将块迭代器绑定到任何循环嵌套来实例化块的主体计算，但大多数实例化并不对应于相同的计算。为了确保转换过程中计算的一致性，我们在块签名中存储了迭代器域信息和迭代器的约束。在图5的例子中，我们知道 $vy, vx$ 和 $vk$ 必须绑定到域 [0, 16) 内的迭代器。此外，由于 $vk$ 是一个归约轴，我们知道不能将其绑定到并行循环上，除非该归约是原子的。域约束仍然为外部循环转换留下了巨大的空间，因为有多种方法可以构造满足约束的循环。我们的域签名可以被看作是表示迭代器的整数域集和关系的一种特定方式，选择这种表示是因其实现效率高和推理简单，但同样的设计理念也适用于其他形式化的整数集和关系域表示【42，Polyhedral code generation in the real world，2006，International Conference on Compiler Construction】。

访问区域和依赖关系：为了提供充分的依赖信息，一个块签名包含了该块相对于多维缓冲区的访问区域以及读/写依赖关系。在图5中，该块通过读取 $A[vy*4:vy*4+4, vk*4:vk*4+4]$ 和 $B[vk*4:vk*4+4, vx*4:vx*4+4]$ 来写入区域 $C[vy*4:vy*4+4, vx*4:vx*4+4]$。这些依赖信息在转换过程中使用。我们只标记每个块相对于多维缓冲区的依赖关系，而不是其他语句（块）。这种间接性支持了更广泛的转换，例如数据布局转换和重计算，这些在张量化程序优化中至关重要。

归约块和初始化：一个归约计算通常包含一个初始化步骤和一个更新步骤。我们可以很自然地将归约计算映射为两个块。但另一方面，联合地对这两个步骤做出调度决策（如分块和计算位置）通常是有帮助的。我们为执行归约的块引入了一个可选的初始化语句。初始化语句在归约的第一次迭代期间执行。这种归约块的表示主要在转换过程中有用。我们提供了在基于两个块的表示和基于初始化块的表示之间进行转换的原语，以便我们可以选择最佳的表示方式进行底层代码生成。

3.2 调度变换

调度变换的目标：对于一个给定的输入程序，我们需要生成一个具有等价语义的丰富搜索空间。我们引入了将TensorIR程序转换为等价优化程序的原语。遵循张量程序优化的现有惯例【5，Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code，2019，IEEE/ACM International Symposium on Code Generation and Optimization (CGO 2019)；9，{TVM}: An automated end-to-end optimizing compiler for deep learning，2018，{USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)；35，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】，我们称此过程为调度（scheduling）。一个块如果只包含以子块为叶节点的循环嵌套，那么它就是可调度的。我们可以通过分析子块签名及其依赖信息来转换可调度块内的循环嵌套和子块计算位置。值得注意的是，一个可调度块可以包含不可调度的子块（例如，不透明的Tensor Core计算）。一个不透明的块也可以包含一个可调度的子块。基于块的隔离性，我们仍然可以有效地独立探索可调度部分的搜索空间，同时保持不透明块不变。

图6：循环变换会改变外部循环嵌套，但块内部没有任何变化。
图6：循环变换会改变外部循环嵌套，但块内部没有任何变化。

循环变换：循环变换，如循环分块（拆分、重排）和计算位置变更，是优化程序以获得更好内存局部性的重要方式。我们也提供了这些循环变换原语（如图6所示）。与现有的张量编译器直接提取每个叶子标量计算语句的依赖关系不同，我们仅通过检查块签名来计算依赖关系。除了循环变换，我们还支持将循环绑定到GPU线程的原语，并提供如向量化和展开等注释提示。请注意，块的隔离性并不妨碍许多重要的跨块协同优化（例如内联、协作抓取）。我们的循环变换涵盖了以前工作提供的循环变换，这使得TensorIR能够重现它们在第2节中提到的搜索空间。

图7：块化（Blockization）创建一个新的块，以隔离内部计算和外部循环嵌套。
图7：块化（Blockization）创建一个新的块，以隔离内部计算和外部循环嵌套。

块化 (Blockization)：循环变换原语保留了块的整体层次结构。正如我们在第2节中提到的，有时通过将子区域计算隔离到一个新的子块中来划分问题是有帮助的。我们称这种变换为块化（blockization），如图7所示。一个经过块化的程序不再是基于标量的，因为新的子块对应于一个张量化计算。我们可以使用块化来隔离可能的张量化候选项。除了块化，我们还引入了可以改变块层次结构的原语。例如，我们提供了缓存原语，它引入子块将输入数据缓存到共享内存中。我们还提供了在单个归约块和相应的初始化-更新块之间来回转换的原语。

调度与TensorIR的分离：许多以前的张量编译器【9，{TVM}: An automated end-to-end optimizing compiler for deep learning，2018，{USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)；35，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】依赖于一种声明式的调度语言来构建一个调度树。在这些编译器中添加新的调度原语需要同时修改调度树数据结构和相应的底层代码生成规则。我们采取了不同的方法，将每个调度原语实现为从一个TensorIR程序到另一个的独立转换。这种设计更容易扩展，因为不同的开发者可以基于一个稳定的TensorIR抽象并行开发新的原语。此外，开发者可以在任何转换阶段打印出程序以进行调试，并将自动重写与手动调度转换混合使用。

3.3 验证

验证机制的核心：block及其缓冲区读/写关系捕捉了原始计算的完整图景，并被用于验证循环嵌套和线程分配的正确性。

循环嵌套验证：循环嵌套验证检查由循环嵌套提供的迭代器绑定是否匹配迭代器域的约束，包括域大小和迭代器独立性信息。例如，如果两个数据并行的块迭代器被绑定为 $v_1 = i; v_2 = i * 2$，那么相应的程序是无效的，因为 $v_1$ 和 $v_2$ 不是独立的。但是 $v_1 = i/4; v_2 = i\%4$ 可以是一个合法的绑定。我们构建了模式匹配器来找到从循环迭代器到块迭代器变量的准仿射映射，并使用该模式来验证绑定的独立性和域。除了迭代器域验证，检查生产者-消费者关系也很重要，以确保写入缓冲区区域的生产者块总是覆盖下游消费者的读取区域。

线程验证：在为GPU和其他支持线程的加速器构建程序时，我们还需要对线程和内存层次结构进行额外的验证。我们进行三种验证：
* 线程绑定：确保绑定到同一线程的不同迭代器是一致的，并满足后端的启动约束。
* 协作内存访问：对于跨线程协作产生存储在共享内存中缓冲区的块，我们需要确保该块能覆盖同一线程组中所有线程的下游需求。同时，为该块提供输入的上游块需要覆盖该块来自该组中所有线程的读取需求。
* 执行范围：验证张量内联函数在正确的执行范围内运行（例如，TensorCore需要在warp级别运行）。

调度原语的正确性：我们为每个调度原语添加检查，以确保转换的正确性。当调度原语只改变循环嵌套时，我们也可以使用验证程序来确保正确性，因为在这些情况下，块迭代域和依赖关系保持不变。对于改变块的调度原语（例如，块化），我们找到了原语特定的必要条件。

验证机制的收益：循环嵌套验证和线程验证被用作过滤器，以筛除无效的TensorIR程序；而调度原语检查则用于确保TensorIR程序在转换前后的等价性。当用户错误地手动构建、导入和调度TensorIR程序时，会收到警告或错误信息。当用户使用编译器自动生成程序时（将在第4节讨论），验证可以帮助在搜索空间探索期间过滤掉假阳性案例。因此，用户程序和编译生成的程序都将从验证中受益。

3.4 编程投入

编程接口与自动化：如图4所示，我们提供了一个TensorIR的Python-AST方言，允许开发者直接在Python中构建、转储、检查、修改和转换TensorIR程序。如果用户需要手动指定所有的计算和优化，编程投入将会很高。我们的框架允许用户从TensorFlow/PyTorch导入模型，并从高级算子自动生成TensorIR程序。此外，系统通过自动调度（第4节）为给定的硬件平台自动提供优化，如分块和缓存。我们仍然允许用户在需要自定义算子时使用Python方言编写TensorIR。在这些情况下，系统会自动提供优化转换。因此，编程投入通常被最小化。

4 自动化调度张量化程序

上一节介绍了TensorIR抽象和一系列转换原语。为了充分利用这些改进，我们需要一个自动化解决方案来在一系列转换中进行优化，并将计算映射到本地张量内联函数。本节描述了一个能够感知张量化的自动调度器来解决这个问题。

图8展示了我们方法的概览。我们的系统将用户的工作负载描述和关于硬件平台的张量内联函数描述作为输入。自动调度器首先通过检查计算模式来生成张量化的候选项。然后，它生成使用张量化计算的程序草图候选项，并根据计算模式决定数据移动。对于由张量化程序草图构成的给定搜索空间，我们执行由基于学习的成本模型指导的进化搜索。整个过程围绕张量化展开，并利用新的block抽象来隔离张量化计算。

图8：带有硬件内联函数的张量化程序自动优化流程。我们以64x64x64矩阵乘法后接一个RELU算子作为输入工作负载，并以一个由点积指令实现的4x4x4矩阵乘法作为合成的张量内联函数。张量化候选项生成步骤将64x64x64的GEMM分块为4x4x4的子块并隔离子计算。然后张量化程序草图生成步骤调度计算并插入由此产生的数据移动（AutoCopy）块，这些块被独立调度。最后，我们使用进化搜索来填充草图中的随机决策，并使用验证机制过滤掉不正确的程序。
图8：带有硬件内联函数的张量化程序自动优化流程。我们以64x64x64矩阵乘法后接一个RELU算子作为输入工作负载，并以一个由点积指令实现的4x4x4矩阵乘法作为合成的张量内联函数。张量化候选项生成步骤将64x64x64的GEMM分块为4x4x4的子块并隔离子计算。然后张量化程序草图生成步骤调度计算并插入由此产生的数据移动（AutoCopy）块，这些块被独立调度。最后，我们使用进化搜索来填充草图中的随机决策，并使用验证机制过滤掉不正确的程序。

4.1 张量内联函数的抽象

描述张量内联函数：为了在我们的优化中利用张量内联函数，我们需要一种方法向系统提供其语义和后端实现。我们利用相同的TensorIR抽象来描述给定硬件后端的张量内联函数。对于每个张量化指令，我们引入一个由两个块组成的TensorIntrin结构。一个块描述计算语义，另一个提供张量化计算的底层实现。在图8的例子中，我们使用一个带有标量主体的普通循环嵌套 $C[i, j] += A[i, k] * B[k, j]$ 来表示计算语义，并使用内部点积指令 __nv_dot_product... 来实现该内联函数。我们还通过TensorIntrin中的多维缓冲区规范包含了数据类型、存储范围、内存布局和连续性约束。这些约束在验证步骤中使用。

与平台特性的结合：值得注意的是，在常见平台上，张量内联函数通常与特殊的内存范围、数据布局以及相应的加载/存储指令一起应用。例如，在NVIDIA GPU上，如果我们决定使用nvcuda::wmma::mma_sync API来执行密集计算，那么我们需要应用nvcuda::wmma::load_matrix_sync和nvcuda::wmma::store_matrix_sync来分别准备输入操作数和检索输出结果。在ARM CPU上，像a64_gemm_u8_8x12这样的微内核要求操作数以交错布局存储。开发者可以通过为张量计算的每个输入和输出操作数指定特殊的存储范围来告知系统这些约束。

4.2 张量化候选项生成

匹配过程：给定一个后端目标和一个输入程序，我们首先将程序主体与可能的TensorIntrin进行匹配，以生成张量化候选项。匹配是逐步进行的。我们首先匹配表达式模式 $C[.] += A[.] \times B[.]$，而不考虑索引。然后，我们通过提出索引之间可能的映射来细化匹配。图9给出了一个遍历匹配过程的例子。在这个例子中，我们以一个矩阵乘法内联函数作为后端描述。该张量内联函数的计算可以用以下公式描述：

C[x, y] += A[x, k] * B[k, y]
C[x, y] += A[x, k] * B[k, y]

对于像批量矩阵乘法这样的工作负载，很容易将张量内联函数与之匹配，其计算可以描述为：

C[b, i, j] += A[b, i, r] * B[b, r, j],
C[b, i, j] += A[b, i, r] * B[b, r, j],

通过将 $x, y, k$ 映射到 $i, j, r$。但对于像二维卷积（Conv2D）这样具有更复杂索引表达式模式的工作负载：

C[n, h, w, co] += A[n, h * s_h + r_h * d_h, w * s_w + r_w * d_w, r_c] * B[r_c, r_h, r_w, co].
C[n, h, w, co] += A[n, h * s_h + r_h * d_h, w * s_w + r_w * d_w, r_c] * B[r_c, r_h, r_w, co].

$x, y, k$ 和 $n, h, w, rc, rh, rw, co$ 之间的映射并不直接。

图9：张量化候选项生成流程示例。我们以标准的NHWC 2D卷积作为输入工作负载，以16x16x16矩阵乘法作为硬件后端内联函数。首先，系统将缓冲区访问表达式转换为中间迭代器。基于缓冲区访问模式，我们为每个迭代器计算特征函数，并建立共享相同特征向量的迭代器之间的映射。该映射进一步指导块实例空间的转换和ReIndex缓冲区。注意，虽然B和C的ReIndex阶段是冗余的，但它们将在草图生成阶段内联到消费者中，因此不影响性能。
图9：张量化候选项生成流程示例。我们以标准的NHWC 2D卷积作为输入工作负载，以16x16x16矩阵乘法作为硬件后端内联函数。首先，系统将缓冲区访问表达式转换为中间迭代器。基于缓冲区访问模式，我们为每个迭代器计算特征函数，并建立共享相同特征向量的迭代器之间的映射。该映射进一步指导块实例空间的转换和ReIndex缓冲区。注意，虽然B和C的ReIndex阶段是冗余的，但它们将在草图生成阶段内联到消费者中，因此不影响性能。

使用ReIndex进行转换：在这些更一般的情况下，我们将计算表达式重写为等价形式：
$C[n, h, w, co] += A_R[n, h, w, rh, rw, rc] \times B[rc, rh, rw, co]$，其中 $A_R[n, h, w, rh, rw, rc] = A[n, h*s_h + rh*d_h, w*s_w + rw*d_w, rc]$。
我们称这种转换为ReIndex，它使用出现在缓冲区访问索引中的中间迭代器来重写缓冲区访问表达式。为了将新的计算与张量内联函数匹配，我们检查每个迭代器出现的缓冲区访问。例如，我们注意到 $n, h, w$ 和 $rc$ 出现在 $A_R$ 的索引中， $co$ 出现在 $C$ 的索引中， $rh, rw, co$ 和 $rc$ 出现在 $B$ 的索引中。然后，我们可以通过检查它们的出现模式来将计算中的迭代器与张量内联函数中的迭代器进行匹配。具体来说，我们将 fuse(n, h, w) 映射到 $x$，co 映射到 $y$，并将 fuse(rc, rh, rw) 映射到 $k$。这里 fuse() 是将多个迭代器融合在一起，可以递归定义为：
$$ \begin{array}{l} \text{fuse}(i_1) = i_1 \\ \text{fuse}(i_1, i_2, \ldots, i_r) = \text{fuse}(i_1, i_2, \ldots, i_{r-1}) * \text{extent}(i_r) + i_r, \end{array} $$
其中 extent() 是迭代器 $i_r$ 的范围。然后我们可以将计算转换为：
C_t[fuse(n, h, w), co] += A_t[fuse(n, h, w), fuse(r_h, r_w, r_c)] * B_t[fuse(r_h, r_w, r_c), co],
其中
C_t[fuse(n, h, w), co] = C[n, h, w, co], A_t[fuse(n, h, w), fuse(r_h, r_w, r_c)] = A_R[n, h, w, r_h, r_w, r_c], B_t[fuse(r_h, r_w, r_c), co] = B[r_h, r_w, r_c, co].
我们使用这个映射来重塑块实例空间和外部循环，并转换ReIndex缓冲区的布局。我们插入布局重写块来分别将 $A, B, C$ 重写为 $A_t, B_t, C_t$，并使用 $A_t, B_t, C_t$ 来重写计算主体。经过这些步骤，计算主体与张量内联函数兼容。

循环重组和早期块化：除了计算主体，我们还需要确保张量计算区域与TensorIntrin提供的描述相匹配。重组后的块实例空间的形状可能不能被张量内联函数的子问题大小整除。对于上一步的每个计算主体，我们对计算块和输入/输出操作数进行必要的填充，使其达到最接近的可整除形状。然后我们执行分块以创建内部循环，以匹配张量内联函数的循环嵌套，并进一步块化内部循环以隔离相应的张量计算。值得注意的是，此步骤中生成的候选项并不总是能成功进行张量化。这是因为其他约束，如内存布局和线程，取决于后续的转换。这些约束嵌入在张量化候选项中，并在验证期间进行检查。

过程的正式描述：假设内联函数的标量表达式可以形式化为：
O[v_0] = f(O[v_0], I_1[v_1], I_2[v_2], ..., I_k[v_k]).
其中 $O$ 是输出操作数， $I_{[1:k]}$ 是输入操作数，$v$ 是参数化此计算的迭代器集合，$v_{[0:k]}$ 都是属于 $v$ 的迭代器列表，$f$ 是检测到的表达式模式。这可以容纳常见的点积和矩阵乘法内联函数。此外，假设工作负载的标量表达式可以形式化为：
Õ[g_0(ṽ_0)] = f(Õ[g_0(ṽ_0)], Ĩ_1[g_1(ṽ_1)], ..., Ĩ_k[g_k(ṽ_k)]).
其中 $Õ, Ĩ_{[1:k]}, ṽ_{[0:k]}$ 对应于它们的对等部分，$f$ 完全相同。$g_{[0:k]}$ 是将迭代器列表映射到实际缓冲区访问位置的映射。为了将问题简化为更简单的规范形式，我们应用ReIndex调度转换。如果我们对所有操作数应用ReIndex，工作负载标量表达式被简化为：
\hat{O}[ṽ_0] = f(\hat{O}[ṽ_0], \hat{I}_1[ṽ_1], ..., \hat{I}_k[ṽ_k]),
其中缓冲区访问索引直接对应于迭代器。为了匹配这两个公式，我们定义迭代器 $v \in \mathbf{v}$ 的特征向量 $\chi(v) \in \{0, 1\}^{k+1}$，通过检查 $v_0, v_1, ..., v_k$ 中的每一个是否包含 $v$。形式上：
χ(v)_i = [v ∈ v_i] i ∈ [0, k],
其中 [] 是Iverson括号。只要 $\forall v \in \mathbf{v}, \exists \tilde{v} \in \tilde{\mathbf{v}}, \chi(v) = \tilde{\chi}(\tilde{v})$，我们就可以成功提出映射。在当前实现中，我们可以进一步安全地假设 $\mathbf{v}$ 中的迭代器都具有不同的特征向量。然后对于所有 $v \in \mathbf{v}$，我们融合所有满足 $\chi(v) = \tilde{\chi}(\tilde{v})$ 的 $\tilde{v}$，并将融合后的迭代器映射到 $v$。

4.3 张量化程序草图生成

生成搜索空间：对于一组给定的张量化候选项，我们需要构建一个包含该张量化的庞大程序搜索空间。我们将现有的层次化搜索空间生成【48，Ansor: generating high-performance tensor programs for deep learning，2020，{USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 20)】推广到张量计算。我们通过生成包含张量化计算的程序草图来构建搜索空间，然后枚举由生成的草图引发的选择。如图8右侧所示，程序草图固定了部分程序结构，同时为循环分块大小和计算缓存决策等参数留下了选择空间。我们通过迭代应用预定义的草图生成规则来生成草图。重要的是，我们需要构建能够在张量化计算上工作的草图生成规则，这需要查看块签名并利用我们分析中的访问区域信息。

数据移动作为一等公民：现有的张量程序自动调度器将其设计重点放在计算的调度上，而将不同内存范围之间的数据移动视为次要优先级。然而，由于张量内联函数极大地提高了计算吞吐量，数据移动成为张量程序的瓶颈。此外，数据移动决策通常依赖于计算调度决策，如分块、线程绑定、执行范围和生产者-消费者数据流粒度。我们采纳这些见解，将数据移动作为我们自动调度器中的一等公民，并将其与计算调度解耦。具体来说，我们在草图生成规则决定执行数据移动的地方插入AutoCopy块（图8）。复制块隐藏了内存调度细节，仅在块签名级别暴露必要的缓冲区访问信息。隔离的复制块允许草图生成独立地做出计算调度决策，而无需考虑如何进行数据移动。AutoCopy块的主体描述了数据移动任务的细节，包括缓冲区位置映射、线程和存储范围要求。数据移动调度器将这些信息作为输入，并执行与内存相关的调度转换，例如插入中间缓存阶段、利用数据移动张量内联函数、向量化、协作抓取或步幅填充以避免bank冲突。

4.4 进化搜索

搜索过程：在张量化程序草图生成阶段之后，我们可能得到数十亿个可能的导出程序。我们使用进化搜索来探索这个空间并找到一个优化的张量化程序。我们的搜索从对给定程序草图的随机选择初始化开始。然后我们对当前程序集进行变异。接着我们从变异的候选中选择有前途的程序，并在我们感兴趣的硬件后端上对其进行基准测试。我们从评估阶段收集数据以更新学习到的成本模型。

张量化计算的成本模型：我们构建了一个基于提升树集成【7，XGBoost: A Scalable Tree Boosting System，2016，KDD ’16】的成本模型，该模型使用从程序中提取的特征。特征向量包含与内存访问模式、重用和循环注释相关的信息。重要的是，我们以隔离的方式从块签名以及块的主体中提取特征（例如，标记Tensor Core的使用）。我们的成本模型可以看作是先前方法对张量化程序的推广。我们相信，一个有效的张量化程序成本模型是未来研究的一个有前途的领域。

验证：在搜索过程中随机变异程序可能会因为未满足张量内联函数的约束或无效的循环嵌套候选项而生成无效程序。假阳性的可能性使得在搜索过程中需要一个验证步骤。我们应用第3.3节中的技术来验证进化搜索中的程序候选项，以识别和拒绝无效程序。验证步骤减轻了进化搜索算法的负担，并允许我们在搜索过程中生成少量假阳性。

A4 实验

我们基于 Apache TVM 【9，{TVM}: An automated end-to-end optimizing compiler for deep learning，2018，{USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】 实现了 TensorIR。本节评估旨在回答以下问题：
* TensorIR 能否优化常见的机器学习算子？
* TensorIR 能否为端到端网络执行带来性能提升？
* TensorIR 能否支持不同硬件平台上的张量内联函数？

实验环境

  • GPU平台:
    • 硬件: NVIDIA RTX 3080 (带Tensor Cores)。
    • 软件与基线:
      • TVM (commit: 27b0aad5, with auto-scheduler)
      • AMOS (commit: 6aee6fe2)
      • CUTLASS (v2.9)
      • TensorRT (PyTorch-TensorRT container Release 22.06)
      • PyTorch (v1.13.0.dev20220612)
    • 数据类型: 所有算子使用 float16 作为输入和累加器数据类型。
  • CPU平台:
    • 硬件: AWS Graviton2 CPU。
    • 软件与基线:
    • 内联函数: 使用8位整型点积 (sdot) 指令。

实验结果

5.1 单算子评估

GPU平台: 我们在NVIDIA RTX 3080上评估了一系列常见深度学习算子，包括一维/二维/三维卷积（C1D/C2D/C3D）、深度卷积（DEP）、空洞卷积（DIL）、通用矩阵乘法（GMM）、分组卷积（GRP）和转置二维卷积（T2D）。

  • 与机器学习编译器的比较 (图10): TensorIR 相较于现有的机器学习编译解决方案（AMOS 和 TVM）取得了高达7.5倍的性能提升。TVM在计算密集度较低的工作负载（如DEP）上表现尚可，但在重度计算负载（如C2D、C3D、GMM）上因Tensor Core支持有限而性能不佳。AMOS虽然能为每个工作负载使用Tensor Core，但效果不如TensorIR。这些提升来自于更好的抽象和自动调度，它充分利用了张量计算内联函数和相应的数据移动。
图10：在Nvidia GPU上与现有机器学习编译器的单算子比较。TensorIR在各种工作负载中带来了高达7.5倍的速度提升。
图10：在Nvidia GPU上与现有机器学习编译器的单算子比较。TensorIR在各种工作负载中带来了高达7.5倍的速度提升。
  • 与平台特定库的比较 (图11): TensorIR 与两种平台特定解决方案（CUTLASS 和 TensorRT）进行了比较。在C1D、C2D、DEP、T2D和DIL上，TensorIR的性能最高可达基线的13.9倍。在C3D、GRP和GMM这些由专门工程团队深度优化的工作负载上，TensorIR也达到了超过75%的性能。这表明，即使在这些被高度优化的算子上，TensorIR也能接近或匹配现有的供应商特定解决方案。
图11：与平台特定库的单算子比较。我们没有展示CUTLASS在DEP、GRP和T2D上的数据，因为该库不支持它们。TensorIR在C1D、C2D、DEP、T2D和DIL上比基线快了高达13.9倍，并在C3D、GMM和GRP上达到了超过75%的吞吐量。
图11：与平台特定库的单算子比较。我们没有展示CUTLASS在DEP、GRP和T2D上的数据，因为该库不支持它们。TensorIR在C1D、C2D、DEP、T2D和DIL上比基线快了高达13.9倍，并在C3D、GMM和GRP上达到了超过75%的吞吐量。

5.2 端到端模型评估

GPU平台: 我们在NVIDIA RTX 3080上评估了四个广泛使用的模型：ResNet-50、MobileNetV2、BERT_large和Vision Transformer (ViT)。

  • 性能结果 (图12): TensorIR的性能比PyTorch、TVM和AMOS高出1.2至8.8倍。与TensorRT相比，TensorIR在MobileNet V2上的性能提高了30%，并在ResNet-50和BERT_large上达到了88%-100%的吞吐量。此外，TensorIR能够自动支持像ViT这样的新兴模型，而当时的TensorRT尚不支持。
  • 调优时间 (表1): TensorIR的调优速度比TVM快2倍。这主要有两个原因：首先，TensorIR的自动张量化利用了Tensor Core，生成的程序运行更快，从而减少了硬件评测时间；其次，“分而治之”的方法将问题空间划分为外部循环嵌套和内部张量化主体，减小了搜索空间。
图12：在NVIDIA GPU上的端到端模型评估。TensorIR显著优于现有的机器学习编译解决方案，并在流行网络上实现了与GPU推理库相当或更好的吞吐量。TensorIR在新兴模型ViT上表现更优，而TensorRT尚不支持该模型。
图12：在NVIDIA GPU上的端到端模型评估。TensorIR显著优于现有的机器学习编译解决方案，并在流行网络上实现了与GPU推理库相当或更好的吞吐量。TensorIR在新兴模型ViT上表现更优，而TensorRT尚不支持该模型。
表1：NVIDIA GPU上端到端模型的调优时间比较。TensorIR的调优速度快2倍。
表1：NVIDIA GPU上端到端模型的调优时间比较。TensorIR的调优速度快2倍。

5.3 ARM CPU 评估

ARM平台: 我们在AWS Graviton2 CPU上评估了TensorIR的泛化能力，为其提供了8位整型点积（sdot）内联函数的描述。

  • 单算子结果 (图13): 在C2D和GMM上，得益于利用原生硬件加速，TensorIR比TVM快了12.5倍。同时，TensorIR达到了ARMComputeLib性能的85%-105%，显示了其在该平台上达到与供应商特定解决方案同等水平性能的能力。
图13：在ARM CPU上的单算子评估。由于使用了原生张量指令加速，TensorIR比TVM快了高达12.6倍。它也达到了与高度优化的平台特定库（ArmComputeLib）相同水平的性能。
图13：在ARM CPU上的单算子评估。由于使用了原生张量指令加速，TensorIR比TVM快了高达12.6倍。它也达到了与高度优化的平台特定库（ArmComputeLib）相同水平的性能。
  • 端到端结果 (图14): 在该平台上，TensorIR的端到端神经网络执行性能比PyTorch和TVM高出2.3倍。值得注意的是，尽管PyTorch通过QNNPACK后端支持专门的量化模型，但QNNPACK尚未添加sdot支持。这个结果暗示了这些框架为了跟上硬件变化所面临的维护成本，而TensorIR可以通过自动化来减轻这种维护负担，同时仍然提供与手工优化系统相媲美的性能。
图14：在ARM CPU上的端到端评估结果。TensorIR比PyTorch和TVM快1.2倍–2.5倍。
图14：在ARM CPU上的端到端评估结果。TensorIR比PyTorch和TVM快1.2倍–2.5倍。

A7 补充细节

深度学习框架: 现有的深度学习框架，如TensorFlow【1】、MXNet【8】、PyTorch【33】，通常通过调用供应商优化的库（如cuDNN【11】、MKL-DNN【20】、TensorRT【32】、ArmComputeLibrary【3】）来优化深度神经网络。这些库需要高昂的工程开发成本，且特定于某一硬件。TensorIR通过自动提供与供应商库相当的性能，补充了库开发，从而能够实现更好的覆盖范围并降低开发成本。

高性能计算（HPC）: 像矩阵乘法和点积这类计算密集的线性代数算子，由于其在科学计算中的重要性，一直是HPC社区的优化目标（例如CUTLASS【23】）。“分而治之”是HPC和机器学习工程中一种典型的优化技术。TensorIR借鉴了这些思想，并将它们泛化成一个允许自动进行张量化的抽象。

机器学习与张量编译器: 现有的编译器引入了不同的张量程序抽象。Halide【35】和TVM【9】使用一种调度语言来描述带有标量主体的循环嵌套的循环优化原语。Tensor Comprehensions【43】、Tiramisu【5】和MLIR/Affine【26】使用多面体模型来分析循环嵌套依赖。这些工作以自下而上的方式优化带有标量计算的循环嵌套。Fireiron【17】和Stripe【46】使用嵌套的多面体结构以自上而下的方式对张量程序建模。TensorIR结合了两种方法的见解，并将表示推广到张量化程序。

自动化: 自动化是机器学习编译和张量程序优化的一个重要课题。AutoTVM【10】引入了一种基于学习的方法，通过学习的成本模型和模板引导的搜索来优化张量程序。Ansor【48】通过层次化的搜索空间改进了自动调度。我们的自动调度算法借鉴了这些方法，并将它们推广到最适合领域特定硬件加速的张量化计算。

自动向量化/张量化: 自动向量化是一个长期的编译器研究课题。张量化可以看作是向量化问题的推广，以支持现代加速器中的张量内联函数。已有的一些工作，如AKG【47】使用多面体方法探索张量化搜索空间，UNIT【45】引入了张量化的通用流程，而AMOS【49】通过张量表达式实现了到张量化内联函数的自动映射。我们的方法通过提出一种新颖的张量化计算抽象，并与其他优化联合进行张量化，从而推广了这些先前的方法。

A5 结论

本文提出了TensorIR，一个用于自动化张量化程序优化的抽象。我们设计了一个名为block的关键抽象，它可以隔离张量化计算，并为程序优化提供有效的转换原语。我们构建了一个自动调度算法，该算法将张量化与其他优化联合执行，并生成高性能的程序。我们希望这项工作能鼓励对张量化程序优化进行更多的研究，并为硬件和软件的专业化发展提供新的机遇。