Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs

文章标题：Hidet：面向深度学习张量程序的任务映射编程范式
作者/机构：
- Yaoyao Ding, Bojian Zheng, Gennady Pekhimenko (University of Toronto, Toronto, Canada)
- Cody Hao Yu, Yizhi Liu, Yida Wang (Amazon Web Services, Santa Clara, USA)

A1 主要贡献

随着深度学习模型在云服务和边缘设备上的广泛应用，降低模型推理延迟对于提供高效的模型服务至关重要。然而，由于现代加速器（如NVIDIA GPU和Google TPU）的高度复杂性以及算子数量的快速增长，为深度学习算子开发高效的张量程序变得极具挑战性。现有深度学习编译器（如Apache TVM）采用声明式调度原语来降低开发张量程序的门槛，但这种方法不足以覆盖最先进的张量程序优化（如双缓冲）。

本文旨在解决现有深度学习编译器在张量程序优化方面的局限性。核心问题是，基于循环操作的声明式调度原语表达能力有限，无法实现如双缓冲等需要细粒度程序操纵的关键优化，同时其搜索空间巨大导致调优时间过长。研究目标是提出一种新的编程范式，既能简化张量程序的开发，又不牺牲表达复杂优化的能力，并能显著缩短调优时间。

为此，本文提出了任务映射（task-mapping）编程范式。该方法将调度过程嵌入到张量程序中，使用称为“任务映射”的专用映射来直接定义计算的分配和排序。这种新方法允许开发者在更细的粒度上（例如，程序语句级别）操纵张量程序，从而极大地丰富了可表达的优化。此外，本文还提出了调度后融合（post-scheduling fusion）优化，允许开发者专注于调度单个算子，并在调度后自动进行融合，显著减少了算子融合的工程量。该范式还构建了一个高效的、与程序输入大小无关的硬件中心（hardware-centric）调度空间，从而大大减少了调优时间。

本文的主要贡献总结如下：

  • 识别并指出现有深度学习编译器局限性：本文指出，最先进的深度学习编译器所采用的面向循环的调度（loop-oriented scheduling）在高效编译复杂张量程序（如矩阵乘法）方面存在表达能力不足的根本性限制。
  • 提出任务映射编程范式：本文引入了任务映射编程范式，以简化张量程序的开发，同时在优化表达能力上不逊于手工实现。基于此范式，提出了调度后融合技术，用于将调度好的程序与周围算子融合。该范式还允许在硬件中心的调度空间中搜索，从而显著减少调优时间。
  • 实现并开源Hidet编译器：基于上述思想，本文实现了一个名为Hidet的新型深度学习编译器。大量实验表明，Hidet在现代卷积和Transformer模型上，性能优于最先进的深度学习推理框架ONNX Runtime和编译器TVM（及其调度器AutoTVM和Ansor），最高提升1.48倍（平均1.t22倍），同时将调优时间分别缩短了20倍和11倍。Hidet已在GitHub上开源。

A3 背景知识与设计动机

2.1 CUDA编程模型

CUDA编程平台【40，The GPU Computing Era，2010，IEEE Micro】是深度学习系统在NVIDIA GPU上广泛使用的平台。本节简要介绍现代GPU上的CUDA编程模型。

内核、线程块和线程。在GPU上运行工作负载时，会执行数千个线程。每个线程执行相同的代码，称为内核代码。启动内核时，一个线程块网格（grid）被分派到GPU上，如图1所示。每个网格通常包含数十到数千个线程块，而每个线程块则包含数十到数百个线程。在内核代码中，使用预定义的变量threadIdx和blockIdx及其后缀x、y、z来访问线程在线程块内的三维索引和线程块在网格中的三维索引。

硬件实现。每个现代GPU拥有数十到数百个流式多处理器（Streaming Multiprocessors, SMs）。每个SM支持调度多达数千个并发线程【40，The GPU Computing Era，2010，IEEE Micro】。一个线程块中的线程被划分为warp，每个warp包含32个连续执行相同指令的线程。有两种可编程的片上内存：共享内存和寄存器。寄存器是每个线程私有的，而共享内存分配给每个线程块，并且只有该线程块内的线程可以访问。启动内核时，线程块会一波一波地被分派到SM上【23，Demystifying the Placement Policies of the NVIDIA GPU Thread Block Scheduler for Concurrent Kernels，2021，SIGMETRICS Perform. Eval. Rev.】。每个线程块只会被分派到一个SM，而每个SM可能包含多个线程块。每个SM上驻留的最大线程块数量受到共享内存大小、寄存器文件和warp调度单元的限制。深度神经网络中的算子被实现为GPU内核。运行神经网络时，我们按照满足算子依赖关系的顺序启动这些内核。在这些算子中，矩阵乘法（也称为线性层或全连接层）是最重要的算子之一。

图1：CUDA编程模型概述。
图1：CUDA编程模型概述。

2.2 高效的矩阵乘法

高效矩阵乘法实现流程。本节通过一个在现代NVIDIA GPU上使用Tensor Cores【13，Nvidia a100 tensor core gpu: Performance and innovation，2021，IEEE Micro】实现的矩阵乘法$C = AB$（所有矩阵均为1024×1024）的例子，来说明其高效实现。图2展示了其工作流程。首先，在步骤1中，通过对M和N维度进行分块（tiling），将矩阵乘法分解为独立的子任务。分块后，会产生$M_{tileSize} \times N_{tileSize}$个独立的子任务，每个子任务是一个大小为$M_{tileSize} \times N_{tileSize} \times K$的矩阵乘法，并被分配给一个线程块。在每个线程块内部，K维度会进一步分块成$K_{tileSize}$大小的块，线程块会对每个K块执行步骤2-3。在步骤2中，线程块内的线程协同地将A和B矩阵的片段从全局内存加载到共享内存。线程块内的所有线程会进行同步，以确保数据加载完成后再进入下一步。在步骤3中，线程块中的4个warp共同处理$4 \times 4 = 16$个矩阵乘累加（MMA）操作，每个操作为$C_{16 \times 16} = A_{16 \times 8}B_{8 \times 16} + C_{16 \times 16}$。每个warp使用Tensor Core【13，Nvidia a100 tensor core gpu: Performance and innovation，2021，IEEE Micro】通过4次连续迭代执行4个MMA操作。当累加完每个K块的矩阵乘法结果后，在步骤4中，将累加寄存器中的结果存回全局内存。图3给出了这4个步骤的伪代码。实现该内核有两种方式：（1）直接编写CUDA C代码，如内核库【12，cuDNN: Efficient Primitives for Deep Learning，2014，ArXiv】，【26，Basic Linear Algebra on NVIDIA GPUs，2022，NVIDIA Inc.】，【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】；（2）使用声明式的面向循环的调度原语。

图2：CUDA平台上的高效矩阵乘法。
图3：矩阵乘法的伪代码。

2.3 声明式面向循环的调度

Halide与TVM的调度范式。为了简化张量程序的优化，Halide【43，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】提出了一种将计算定义与调度解耦的编程范式。这种范式被最先进的DNN编译器TVM【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】及其调度器（如AutoTVM【11，Learning to optimize tensor programs，2018，NeurIPS】和Ansor【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】）所采用。由于该范式提供了一系列声明式调度原语来操纵张量程序的循环结构，我们称之为声明式面向循环的调度。

面向循环的调度工作流程。如图4所示，开发者首先提供算子的数学计算定义，该定义描述了输出张量中每个元素是如何计算的。图中例子给出了矩阵乘法的定义，其中输出的(i, j)元素是一个规约求和。给定计算定义后，调度器首先在步骤1中，通过将compute和reduce原语翻译成嵌套循环，自动生成一个默认的张量程序。然后，在步骤2中，应用一系列声明式调度原语来变换默认张量程序的循环结构，以在特定硬件上获得更好的性能。表1展示了TVM【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】中的调度原语。在图4的示例中，我们仅列出了实现矩阵乘法的前几个调度原语，而TVM实际上使用了超过80个原语来调度矩阵乘法。从默认程序开始，我们首先将i和j循环以因子64进行拆分，得到(oi, ii)和(oj, ij)，然后将循环重排为(oi, oj, ii, ij)，最后将oi和oj分别绑定到blockIdx.x和blockIdx.y。通过这些原语，我们可以得到图4中的调度后程序。

调度范式的应用方式。在深度学习编译器中，有几种方式可以利用这种编程范式。最直接的方法是为每个工作负载（即在特定硬件上具有具体输入的算子）手动编写一个调度【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】，【43，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】。然而，这种方法需要大量的工程努力才能为所有常用算子及其典型输入尺寸实现最优性能。因此，引入了可调参数（如分块大小和循环顺序），供开发者在调度中指定。这样，一个手动调度就变成了调度模板，可以通过自动调优框架【11，Learning to optimize tensor programs，2018，NeurIPS】针对不同的输入形状和硬件进行优化。为了进一步节省编写调度模板的时间，还提出了自动调度方法，通过对计算定义应用预定义规则来自动生成调度【2，Learning to Optimize Halide with Tree Search and Random Programs，2019，ACM Trans. Graph.】，【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】。然而，如下一节所述，面向循环的调度范式所产生的调度空间仍然效率低下，导致难以达到内核库的高度优化性能，并且调度器需要数小时到数天才能在调度空间中找到最佳配置。

图4：面向循环的调度工作流程。
表1：TVM中的面向循环的调度原语。fuse、split、reorder和bind原语通过融合循环、将循环拆分为子循环、重排循环以及将循环绑定到硬件特定轴来变换程序。

3.1 有限的优化支持

双缓冲优化的局限性。声明式面向循环的调度原语在支持关键优化方面存在局限性。我们以双缓冲（double buffering）【6，CudaDMA: Optimizing GPU memory bandwidth via warp specialization，2011，SC】，【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】为例来说明这一根本性限制。双缓冲是一种重要的优化技术，已被多个供应商库（如cuBLAS【26，Basic Linear Algebra on NVIDIA GPUs，2022，NVIDIA Inc.】和CUTLASS【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】）采用，但TVM【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】并不支持。图3中的矩阵乘法实现是次优的，因为同一线程块中的所有线程很可能被单一类型的硬件资源（步骤2中的内存带宽或步骤3中的计算单元）阻塞，而使其他资源闲置。这是因为图3中的数据加载（L7）和计算（L10）使用相同的缓冲区，并且需要使用同步（L8）来满足数据依赖。

双缓冲工作原理。图5展示的双缓冲优化通过使用两个缓冲区来缓解上述问题：一个用于预加载下一次迭代的数据片段（L8和L10），另一个用于当前迭代的计算（L9）。我们首先将下一个A和B矩阵的分块预加载到寄存器（L8），并在当前分块计算后将其存入共享内存（L10）。这种方式更高效，因为L9的计算可以与L8的全局内存加载通过线程级并行同时进行。通过双缓冲，一个线程块中的线程可以同时利用内存访问单元和计算单元。

现有原语的表达能力不足。然而，这种优化无法使用表1中现有的声明式面向循环的调度原语来实现。这是因为没有任何一个调度原语能够以细粒度的方式操纵循环体。因此，尽管面向循环的调度简化了张量程序的编写，但其声明式的调度风格阻碍了开发者实现需要对张量程序进行细粒度操纵的优化。需要强调的是，双缓冲优化只是现有面向循环调度表达能力有限的一个例子。除双缓冲外，线程块洗牌（thread block swizzle）【7，Optimizing Compute Shaders for L2 Locality using ThreadGroup ID Swizzling，2020，NVIDIA Developer Blog】，【53，Locality-Aware CTA Scheduling for Gaming Applications，2021，ACM Trans. Archit. Code Optim.】、Tensor Core MMA PTX指令的高效使用【28，Parallel Thread Execution ISA，2022，NVIDIA Inc.】以及多阶段异步预取【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】等都是内核库【26，Basic Linear Algebra on NVIDIA GPUs，2022，NVIDIA Inc.】，【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】中广泛使用的优化，但用声明式面向循环的调度难以实现。要实现这些优化，我们需要一种更具表达力的方法来编写张量程序并调度其计算。

图5：双缓冲优化。
图5：双缓冲优化。

3.2 融合需要专用的调度模板

TVM中的子图融合机制。编译器相比内核库的一个重要优势是能够优化任意工作负载，特别是包含多个融合算子的工作负载（例如，卷积神经网络中的Conv2d-BN-ReLU【25，Deep residual learning for image recognition，2016，CVPR】和Transformer模型中的Reshape-Matmul-Transpose【16，BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding，2018，CoRR】）。图6展示了TVM【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】如何将Conv2d-BN-ReLU融合成一个单一的内核。具体来说，TVM将算子分组形成子图。每个子图只能包含一个锚点算子（anchor operator），这通常是计算最密集的算子（如卷积或矩阵乘法），并为其精心设计了调度模板。然后，锚点算子的调度模板将用于调度整个子图，这意味着该调度模板必须支持所有可能的融合场景，这极大地增加了编写调度模板的复杂性。

自动调度器的局限性。尽管提出了自动调度器（如Ansor【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】）来根据计算定义和预定义的自动调度规则自动生成调度模板，但用新规则扩展自动调度器是具有挑战性的。这是因为新规则必须与所有现有规则兼容，并且需要足够通用以支持所有算子。因此，在不增加编写专用调度模板复杂性的情况下支持融合，仍然是一个挑战。

图6：TVM子图融合的工作流程。
图6：TVM子图融合的工作流程。

3.3 漫长的调优时间

调度空间的低效性。除了表达能力和可扩展性，现有最先进的调度器【2，Learning to Optimize Halide with Tree Search and Random Programs，2019，ACM Trans. Graph.】，【11，Learning to optimize tensor programs，2018，NeurIPS】，【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】的调优时间通常从几小时到几天不等，原因在于其调度空间效率低下。它们的大部分调度空间由循环分块因子组成。为了限制调度空间大小并避免if-else条件分支，现有调度器只考虑完美分块大小（即，只用能整除循环长度$L$的因子来分块）。例如，长度为10的循环的潜在分块因子只包括1、2、5和10。因此，这些调度器用面向循环的调度构建的空间依赖于目标工作负载的输入形状。我们将这类调度空间称为以输入为中心的调度空间（input-centric schedule space）。

以输入为中心的调度空间的挑战。我们观察到以输入为中心的调度空间存在两个挑战。（1）调度空间大小随输入尺寸因子数量呈指数级增长。图7显示了ResNet-50【25，Deep residual learning for image recognition，2016，CVPR】中每个卷积层的调度数量。对于单个卷积层，需要搜索的调度多达$10^8$个。（2）调度空间可能不包含性能最优的调度，因为没有考虑非完美分块大小。一个极端的例子是，对于M=N=K=2039的矩阵乘法，Ansor和AutoTVM都无法找到有效的调度，因为2039是一个素数。

现有解决方案及其问题。为了应对第一个挑战，最先进的调度器【11，Learning to optimize tensor programs，2018，NeurIPS】，【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】采用成本模型来预测调度的性能，并使用遗传进化搜索来提高搜索效率。然而，搜索过程仍然需要大约半小时来调优单个算子，导致调优一个Inception V3模型【50，Rethinking the inception architecture for computer vision，2016，CVPR】需要8到15个小时。漫长的调优时间阻碍了现有调度器与图级优化【30，TASO: optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】，【37，Optimizing {CNN} Model Inference on {CPUs}，2019，USENIX ATC】和上层应用（如神经架构搜索【71，Neural architecture search with reinforcement learning，2016，arXiv】）的协同优化。这两者都需要内核的延迟信息来指导其优化和网络搜索，并要求调优时间很短。

图7：AutoTVM采用的调度空间大小。
图7：AutoTVM采用的调度空间大小。

A2 方法细节

4.1 任务映射编程范式

面向循环调度的局限性。面向循环的调度通过声明式的循环调度原语来操纵张量程序，这简化了张量编程，但同时也限制了细粒度的操纵和优化。

任务映射的核心思想。我们观察到，面向循环的调度原语的目标要么是（1）将计算分配给并行处理单元（如线程或warp），要么是（2）指定分配给每个处理单元的计算的执行顺序。以图8中矩阵乘法中矩阵A的协同加载为例（为简化，省略了块偏移，只显示矩阵A的加载），循环调度应用了三个原语（循环拆分、融合和绑定）来将512个（64x8）元素的加载任务分配给128个线程，每个线程按顺序加载4个元素。

任务映射的实现方式。我们提议不再通过应用声明式原语进行调度，而是将调度嵌入到张量程序中，并使用称为任务映射（task mappings）的专用映射来直接在程序中定义计算的分配和排序。我们用图8中的例子来展示如何使用任务映射来完成期望的调度。在步骤(1)中，首先定义一个任务映射，它将64x8个任务分配给128个线程。然后，在步骤(2)中，通过使用线程索引threadIdx.x调用任务映射来迭代分配给一个线程的每个任务(i, k)。最后，在步骤(3)中，使用任务的索引(i, k)来实现该任务。这三个步骤解耦了任务分配和每个任务的实现，极大地简化了张量程序的开发。与声明式面向循环的调度相比，它直接在张量程序中进行调度，并允许更细粒度的优化。此外，它还允许开发者在某些维度上回退到传统循环，以实现如双缓冲【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】等优化。由于任务映射是这三个步骤中使用的关键组件，我们将我们构建张量程序的新方法称为任务映射编程范式（task-mapping programming paradigm）。

任务映射的组合。步骤(1)中定义的任务映射是由两个基本任务映射（即repeat(4, 1)和spatial(16, 8)）的任务映射组合派生而来的。图8中的表格给出了所有出现的任务映射的详细信息。任务映射及其组合的正式定义在第5.1节中给出。

范式优势。所提出的范式简化了张量程序的开发，而没有牺牲优化的表达能力。除了单个算子的调度，调度一个融合的子图也很重要，因为算子融合可以极大地减少内存流量，从而加速端到端的DNN执行【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】，【18，Ios: Inter-operator scheduler for cnn acceleration，2021，MLSys】，【30，TASO: optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】。

图8：使用声明式面向循环调度和任务映射编程范式来调度协同加载。在声明式面向循环调度中，开发者对自动生成的程序应用一系列声明式调度原语，以将张量程序转换为更高效的程序。任务映射编程范式不采用声明式原语，而是允许开发者直接在张量程序中嵌入调度，并与面向循环的调度相比，能实现更广泛的优化。
图8：使用声明式面向循环调度和任务映射编程范式来调度协同加载。在声明式面向循环调度中，开发者对自动生成的程序应用一系列声明式调度原语，以将张量程序转换为更高效的程序。任务映射编程范式不采用声明式原语，而是允许开发者直接在张量程序中嵌入调度，并与面向循环的调度相比，能实现更广泛的优化。

4.2 调度后融合

调度后融合的两步法。我们提出将融合子图的调度分解为两个步骤，如图9所示。在步骤1中，我们像TVM【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】一样选择锚点算子，但只单独调度该锚点算子。在步骤2中，我们将周围的算子自动融合到锚点算子的调度后张量程序中。通过这种解耦，锚点算子的调度无需考虑整个子图，只需关注其自身的实现，这与AutoTVM【11，Learning to optimize tensor programs，2018，NeurIPS】相比，极大地减少了设计子图调度模板所需的工程量。

命名与机理。由于融合是在我们调度算子之后完成的，我们称这种方法为调度后融合（post-scheduling fusion）。在调度后融合中，锚点算子可以与它之前（作为prologues）和之后（作为epilogues）的算子融合。我们根据算子的特性来决定其是否可融合。如果一个算子没有规约计算，则它被定义为内射的（injective），并有资格成为prologue算子。如果一个算子是内射的，并且输入张量中的每个元素对输出张量中的单个元素有贡献，则它被定义为双射的（bijective），并有资格成为epilogue算子。例如，所有元素级算子（如加法、ReLU【3，Deep Learning using Rectified Linear Units (ReLU)，2018，ArXiv】）和变换算子（如reshape、transpose）都是双射算子，既可以作为prologue也可以作为epilogue算子。通过调度后融合，我们可以专注于单个算子的调度，同时支持灵活有效的融合。

图9：调度后融合的两个步骤。
图9：调度后融合的两个步骤。

4.3 硬件中心的调度空间

以输入为中心的调度空间的问题。现有最先进的调度器【11，Learning to optimize tensor programs，2018，NeurIPS】，【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】采用第3.3节讨论的以输入为中心的调度空间，其中调度选择循环范围的因子作为拆分或分块因子，这使得调度空间不可扩展，并且无法覆盖那些由非循环范围因子的分块大小所带来的最优性能。

硬件中心的调度空间。除了基于输入大小构建调度空间外，另一种方法是基于硬件设计调度空间，称为硬件中心的调度空间（hardware-centric schedule space）。硬件中心的调度空间通过采用谓词加载（predicated loading，即通过检查访问索引是否越界来保护数据加载）来将调度空间与输入大小解耦，这种方法被内核库【12，cuDNN: Efficient Primitives for Deep Learning，2014，ArXiv】，【26，Basic Linear Algebra on NVIDIA GPUs，2022，NVIDIA Inc.】，【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】广泛使用。

本范式的优势。通过我们提出的范式，我们可以提供一个小型但高效的硬件中心调度空间。由于分块因子是基于硬件资源的（例如，64x64、128x64、16x32等），硬件中心的调度空间比以输入为中心的调度空间小几个数量级。例如，我们为矩阵乘法采用的调度空间包含不到200个调度，这比AutoTVM【11，Learning to optimize tensor programs，2018，NeurIPS】中典型的调度空间平均小$10^5$倍。简单地枚举所有调度就足够了，并且可以在一分钟内完成。

5. HIDET：系统设计

Hidet整体设计。基于上述关键思想，我们设计并实现了一个名为Hidet的DNN编译器。图10展示了其整体设计。Hidet首先在步骤1中从广泛使用的框架（如PyTorch【41，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】）或ONNX【5，ONNX: Open Neural Network Exchange，2019】格式的模型文件中导入一个深度神经网络，然后在步骤2中执行图级优化，如常量折叠和可融合子图的划分。图级优化后，可融合子图中的每个锚点算子将被降级以进行调度。在Hidet中，我们在步骤3中使用任务映射编程范式（第5.1节）将算子调度成一个张量程序，并在硬件中心的调度空间中进行调优。然后，在步骤4中，应用调度后融合（第5.2节）将锚点算子的调度后张量程序与其周围的算子自动融合。最后，在步骤5中，Hidet中间表示（IR）中的融合张量程序将被优化和降级。一个代码生成器会将降级后的IR转换为CUDA内核。

图10：Hidet的整体设计。
图10：Hidet的整体设计。

5.1 任务映射编程范式

任务分配的挑战与目标。在为具有并行处理单元（如现代CPU、GPU和TPU）的特定硬件优化张量程序时，一个关键挑战是如何将独立的（子）任务分配给这些并行处理单元。以图8中的协同加载为例，当从全局内存加载形状为64x8的矩阵A片段到共享内存时，512个任务被分配给线程块中的128个线程，每个线程分配了4个加载任务。在这个例子中，任务被分配给称为工作单元（workers）的并行处理单元，分配给每个工作单元的任务将以特定顺序执行。本节将首先将任务分配和排序形式化为任务映射，然后介绍一个用于组合任务映射的二元运算符，最后讨论基于任务映射的调度。

5.1.1 任务映射

任务映射的正式定义。我们正式定义一个工作单元集$W_n$为一个包含$n$个工作单元的集合，其id从0到$n-1$：
$W_n = \{0, 1, \dots, n-1\}.$
我们还定义一个任务域T为：
$$T = \{(t_0, t_1, \ldots, t_{m-1}) \mid 0 \leq t_i < d_i, t_i \in \mathbb{Z}\}$$
以表示我们关注的所有任务，其中$m$是任务维度，$d = (d_0, d_1, ..., d_{m-1})$是任务形状。

映射函数。一个任务映射$f$被定义为一个函数，它将工作单元集中的每个工作单元映射到任务域中的一个任务列表，即：
图11：两种基本任务映射的示例：repeat(2, 2)和spatial(2, 2)。数字表示分配给同一任务的执行顺序，而颜色表示任务被分配到的工作单元。
$f(w) = [t^{(0)}, t^{(1)}, \dots, t^{(q-1)}]$
其中$w \in W$且$t^{(i)} \in T$。

基本任务映射。我们发现两种非常有用的基本任务映射。repeat(d1, ..., dm)任务映射将一个任务网格(d1, ..., dm)映射到单个工作单元，而spatial(d1, ..., dm)任务映射将一个任务网格(d1, ..., dm)映射到相同数量的工作单元，每个工作单元只处理一个任务。图11展示了这两种任务映射的两个例子。除此之外，Hidet还允许开发者通过指定任务形状、工作单元数量和映射函数来定义自定义任务映射。尽管所有示例都是二维的，但任务映射可以有任意数量的任务维度。

5.1.2 任务映射组合

组合的层次化结构。在协同加载的例子中，我们可以观察到一个层次化结构。64x8=512个任务可以被划分为4组，每组包含16x8=128个任务。每组中的128个任务由128个线程执行。如果我们将每个任务组看作一个宏任务，将128个线程看作一个宏工作单元，那么宏任务到宏工作单元的任务映射就是一个将4个任务映射到单个工作单元的任务映射，记为repeat(4, 1)。这个例子表明，一个任务映射中的所有任务可以被整体视为一个任务，所有工作单元可以被整体视为另一个任务映射中的一个工作单元，从而创建一个组合的任务映射。

组合的正式定义。我们将这个思想形式化如下。设$f_1, f_2$是两个具有相同任务维度的任务映射。设$n_1, n_2$是工作单元数量，$d_1, d_2$是两个任务映射的任务形状。我们定义$f_3$为$f_1$和$f_2$的组合任务映射，它有$n_1n_2$个工作单元，任务形状为$d_3 = d_1 \odot d_2$。其映射函数定义为：
$f_3(w) = [t_1 \odot d_2 + t_2 \mid t_1 \in f_1(\lfloor w/n_2 \rfloor), t_2 \in f_2(w \% n_2)]$
任务映射组合记为$f_3 = f_1 \circ f_2$。任务组合是满足结合律的，即对于任意任务映射$f_1, f_2, f_3$：
$(f_1 \circ f_2) \circ f_3 = f_1 \circ (f_2 \circ f_3),$

组合的强大功能。任务映射组合是构建新任务映射的强大工具。图12给出了一些任务映射组合的例子。除了这些例子，任务映射spatial(4, 2) * repeat(2, 2) * spatial(4, 8) * repeat(4, 4)被用于使用CUDA Core【40，The GPU Computing Era，2010，IEEE Micro】的矩阵乘法中。它们分别对应于一个块中的warp（4x2）、每个warp的重复次数（2x2）、一个warp中线程的布局（4x8），以及每个线程处理的C元素数量（4x4）。

任务与工作单元的抽象性。任务映射中的任务和工作单元是抽象概念，可以用来描述不同层次结构上的任务和工作单元。例如，除了单个线程，一个工作单元也可以代表一个warp、一个线程块或其它加速器中的一个处理单元。图13展示了一个以warp作为任务映射中工作单元的例子。它实现了前面提到的矩阵乘法中使用的block_mma函数（见图3和图5）。在这个例子中，我们使用一个任务映射将一个4×4的任务网格分配给4个warp，每个warp执行4个warp级别的矩阵乘累加（MMA）任务，其对应的分配如图2的步骤3所示。

范式总结。任务映射及其组合可以极大地简化张量程序的编写，因为它采用专用的映射来定义任务分配和排序，使开发者无需编写复杂的循环和索引计算来达到相同的目标。我们称这种基于任务映射的张量程序编写范式为任务映射编程范式。

图12：任务映射组合的示例。(a)和(b)显示任务映射组合不满足交换律。(c)显示了组合三个任务映射的示例。由于任务映射组合满足结合律，应用组合的顺序无关紧要。(d)显示了按列主序分配任务的示例。
图13：使用任务映射来调度warp级别的任务。

5.1.3 调度机制

Hidet中的两种调度机制。基于该范式，我们在Hidet中进一步实现了两种调度机制：基于模板的调度（template-based scheduling）和基于规则的调度（rule-based scheduling）。受Ansor【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】和Halide-AutoScheduler【2，Learning to Optimize Halide with Tree Search and Random Programs，2019，ACM Trans. Graph.】的启发，基于规则的调度直接从一个算子的计算定义生成张量程序，无需任何额外的工程努力，并用于Hidet中的大多数算子。另一方面，对于像矩阵乘法这样的关键算子，基于规则的调度可能无法生成足够高效的张量程序。受AutoTVM【11，Learning to optimize tensor programs，2018，NeurIPS】的启发，我们也允许开发者提供一个张量程序模板来支持这些算子的高效调度。图14阐释了这两种调度机制。

基于规则的调度。该机制根据给定的计算定义自动生成张量程序。它以有向无环图（DAG）的形式遍历计算定义，并应用预定义的规则将DAG中的每个节点转换为最终张量程序的一部分。因为这种机制不需要开发者编写专用的调度模板，它在Hidet中被广泛用于不包含规约的算子，如reshape、transpose、slice以及所有元素级算术算子。

基于模板的调度。该机制使用给定的模板来调度算子。一个调度模板是一个用参数化任务映射编写的张量程序。每个调度模板都配备了一个调度空间，其中包含一系列可用于参数化任务映射的参数，模板可以用调度空间中的任意选择来实例化。以图5和图13中的矩阵乘法为例，我们可以使用不同数量的warp并为每个warp重复不同次数来实现矩阵乘法。这些不同的选择构成了矩阵乘法的调度空间。在调度期间，Hidet首先从调度空间中枚举调度选择，然后用该选择来创建给定程序模板的任务映射。最后，Hidet将模板实例化为一个张量程序并测量其性能。性能最好的调度将被使用。我们将此过程称为调优（tuning）。

添加新算子的工程量。向Hidet添加新算子不需要很高的工程量。Hidet中的大多数算子都是通过基于规则的调度自动调度的，添加起来很容易。像卷积和矩阵乘法这样的计算密集型算子通常需要基于模板的调度来实现高性能。添加一个新的Hidet模板的复杂性与AutoTVM【11，Learning to optimize tensor programs，2018，NeurIPS】模板的复杂性相似。

图14：Hidet中的两种调度机制。
图14：Hidet中的两种调度机制。

5.2 调度后融合

调度后融合的优势。为了减轻像AutoTVM中调度子图的复杂性，我们提出将子图调度解耦为两个阶段：（1）调度锚点算子和（2）将调度后的张量程序与周围算子融合。这种解耦允许开发者专注于锚点算子的调度，而不是整个子图，并自动化了调度后的张量程序与子图中其他算子的融合过程。在调优期间，融合后的张量程序的性能将被用作最大化的目标，因此这种解耦不会损害最终性能。

融合示例。图15展示了一个调度后融合的例子。在步骤1中，图级优化阶段的一个优化遍会将计算图划分为子图。给定子图后，在步骤2中，一个选定的锚点算子将通过第5.1.3节中的一种调度机制被调度成一个张量程序。最后，在步骤3中，剩余的算子将被融合到调度后的程序中。这些算子被分为两类：每个输入张量的prologue算子和每个输出张量的epilogue算子。每个prologue算子定义了对输入张量的每次元素访问是如何计算的，而epilogue算子定义了输出张量元素如何被进一步计算并存储在融合算子的输出中。在这个例子中，对A[99 - i]的访问将被替换为C[99 - i] * 2.0，输出张量的第i个元素将被进一步计算（即乘以3.0）并以索引(i / 50, i % 50)存储到融合后的输出张量D中。

调度后融合的应用。调度后融合简化了算子融合。它还允许我们重用现有的高度优化的算子（如矩阵乘法）来支持新算子（如卷积）。在Hidet中，我们可以使用img2col算法【8，High performance convolutional neural networks for document processing，2006，IWFHR】将卷积算子实现为四个算子，其中一个是矩阵乘法，另外三个是简单的变换算子。通过调度后融合，我们将其他三个算子融合到矩阵乘法中，并为卷积重用了所有为矩阵乘法实现的优化（例如，k维度上的并行规约【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】）。

图15：调度后融合的示例。
图15：调度后融合的示例。

A4 实验环境

  • 实现：Hidet 从零开始实现，使用约2万行Python和C++代码。它采用两级IR：图级IR表示DNN模型的计算图，张量级IR表示带调度的张量程序。Hidet将使用任务映射编写的张量程序降级为CUDA C代码，并用CUDA编译器编译。值得注意的是，仅为矩阵乘法和规约算子实现了两个高效的调度模板，就覆盖了所有被评估模型中的算子。大多数算子要么由基于规则的调度机制调度，要么转换为矩阵乘法以重用现有模板（如卷积）。

  • 硬件配置：

    • CPU：16核24线程的Intel i9-12900K（启用超线程）。
    • 内存：64 GiB DRAM。
    • GPU：一块 NVIDIA RTX 3090。

  • 软件配置：

    • 操作系统：Ubuntu LTS 20.04。
    • 驱动与库：NVIDIA驱动 510.73.08，CUDA 11.6。
    • 依赖框架：PyTorch 1.11，torchvision，transformers，ONNX Runtime 1.11.1，TVM 0.9.dev。

  • 模型与数据集：

    • 模型架构：
      • ResNet-50【25，Deep residual learning for image recognition，2016，CVPR】：常用的图像分类CNN。
      • Inception-V3【50，Rethinking the inception architecture for computer vision，2016，CVPR】：采用多路径不同核尺寸卷积的CNN。
      • MobileNet-V2【45，Mobilenetv2: Inverted residuals and linear bottlenecks，2018，CVPR】：基于可分离卷积的轻量级CNN。
      • Bert【16，BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding，2018，CoRR】：广泛使用的基于Transformer的NLP模型。
      • GPT-2【42，Language models are unsupervised multitask learners，2019，OpenAI blog】：用于序列到序列任务（如翻译、问答）的NLP模型。
    • 模型参数：语言模型的序列长度统一设为128。
    • 来源：模型实现来自torchvision和transformers包，并导出为ONNX格式进行评估。

A4 实验结果

端到端评估

  • 实验内容：在ResNet50, Inception-V3, MobileNet-V2, Bert, GPT-2等模型上，比较了Hidet与PyTorch 1.11, Onnx Runtime 1.11.1, TVM 0.9.dev (AutoTVM和Ansor)的端到端推理延迟（batch size=1）。
  • 实验结果：如图16所示，Hidet在大多数模型上均优于所有基线，性能最高提升1.48倍，平均提升1.22倍。例外情况是Ansor在MobileNetV2上表现更好，因其为深度可分离卷积找到了更优的调度。AutoTVM在Bert和GPT-2上性能较差，因其调度模板缺乏优化。
  • 分析结论：Hidet的优势在于能够自动融合子图、针对给定输入尺寸进行调优，并能表达更多优化（如双缓冲），从而超越了依赖预制库的PyTorch/OnnxRuntime以及表达能力受限的AutoTVM/Ansor。
图16：Hidet与最先进的DNN推理框架和编译器的端到端比较。
图16：Hidet与最先进的DNN推理框架和编译器的端到端比较。
  • 实验内容：比较AutoTVM, Ansor, 和Hidet的调优成本（即调优过程所需时间）。
  • 实验结果：如图17所示，与Ansor和AutoTVM相比，Hidet的调优成本分别降低了11倍和20倍。
  • 分析结论：Hidet采用小而高效的硬件中心调度空间（如矩阵乘法只有180个调度），只需几分钟即可完成穷举搜索。相比之下，AutoTVM和Ansor的调度空间有$10^5$到$10^8$个候选项，即使有成本模型也需要很长时间才能找到最优调度。AutoTVM在Bert和GPT-2上调优时间短是因为其调度空间小（<20个调度），但这导致了性能不佳。
图17：AutoTVM、Ansor和Hidet的调优成本。
图17：AutoTVM、Ansor和Hidet的调优成本。

案例研究

  • 调度空间比较：
    • 实验内容：比较AutoTVM, Ansor, 和Hidet调度空间的效率，使用ResNet50中的一个卷积层作为基准工作负载，绘制了三个调度空间中调度延迟的分布。
    • 实验结果：如图18所示，Hidet的调度空间（180个调度）中的大多数调度性能优于从AutoTVM（1000个样本）和Ansor（800个样本）空间中采样的调度。
    • 分析结论：由于所提范式具有更好的表达能力，Hidet的调度空间整体质量更高，大部分调度方案的延迟都小于73μs。
图18：AutoTVM、Ansor和Hidet的调度空间延迟分布。X轴为对数尺度。
图18：AutoTVM、Ansor和Hidet的调度空间延迟分布。X轴为对数尺度。
  • 输入尺寸的性能敏感性：
    • 实验内容：在连续的输入尺寸上对矩阵乘法进行基准测试，比较AutoTVM, Ansor, 和Hidet的性能稳定性。
    • 实验结果：如图19所示，AutoTVM和Ansor的性能波动很大，在输入尺寸为素数（如2039）时甚至找不到有效调度。而Hidet在这些输入尺寸上实现了稳定的性能。
    • 分析结论：AutoTVM和Ansor的以输入为中心的调度空间导致其对输入尺寸敏感，而Hidet的硬件中心调度空间则避免了这个问题。
图19：AutoTVM、Ansor和Hidet在连续输入尺寸的矩阵乘法上的比较。
图19：AutoTVM、Ansor和Hidet在连续输入尺寸的矩阵乘法上的比较。
  • 不同批量大小的评估：
    • 实验内容：在ResNet50上使用不同批量大小（1、4、8）进行基准测试。
    • 实验结果：如图20所示，在小批量（1和4）时，AutoTVM和Ansor优于Onnx Runtime。在大批量（8）时，Onnx Runtime反超AutoTVM和Ansor。Hidet在所有批量大小下均表现最佳。
    • 分析结论：AutoTVM/Ansor能找到足够的线程块来饱和SMs，但在大批量时，由于缺乏双缓冲等优化，其每个线程块的延迟比Onnx Runtime更长。Hidet则能兼顾线程块数量和效率，因此在各种批量大小下都表现出色。
图20：在ResNet50上对批量大小为1、4和8的比较。
图20：在ResNet50上对批量大小为1、4和8的比较。
  • 调度后融合评估：
    • 实验内容：评估Hidet中通过img2col算法实现的卷积性能，这依赖于调度后融合将变换算子融合到矩阵乘法中。比较了ResNet50中Conv-Bn-ReLU子图的性能。
    • 实验结果：如图21所示，Hidet在大多数卷积上优于Onnx Runtime和Ansor。
    • 分析结论：通过调度后融合，Hidet能重用矩阵乘法的优化（如k维度并行规约），使其生成的卷积核能饱和GPU计算资源，从而获得更好的性能。
图21：在ResNet50的Conv2d-Bn-ReLU子图上比较Onnx Runtime、Ansor和Hidet。
图21：在ResNet50的Conv2d-Bn-ReLU子图上比较Onnx Runtime、Ansor和Hidet。
  • 与TensorRT的比较：
    • 实验内容：将Hidet与NVIDIA的高性能推理引擎TensorRT 8.4.1.5进行比较。
    • 实验结果：如图22所示，Hidet在三个CNN模型上优于TensorRT。TensorRT在Transformer模型（Bert和GPT-2）上优于Hidet。
    • 分析结论：Hidet的优势在于能够针对给定输入尺寸进行调优并自动融合算子。TensorRT可能对流行的Transformer模型中的自注意力层进行了专门的、未公开的优化。
图22：TensorRT与Hidet的比较。
图22：TensorRT与Hidet的比较。

A7 补充细节

7. 相关工作

现有DL编译器的普遍方法。许多现有的DL编译器采用面向循环的调度原语【9，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI】，【43，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，Acm Sigplan Notices】，并在此基础上建立了自动调优框架【2，Learning to Optimize Halide with Tree Search and Random Programs，2019，ACM Trans. Graph.】，【11，Learning to optimize tensor programs，2018，NeurIPS】，【46，Tensor Program Optimization with Probabilistic Programs，2022，ArXiv】，【54，Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions，2018，ArXiv】，【57，UNIT: Unifying Tensorized Instruction Compilation，2021，CGO】，【58，Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance，2022，MLSys】，【63，DietCode: Automatic Optimization for Dynamic Tensor Programs，2022，MLSys】，【65，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI】-【67，FlexTensor: An Automatic Schedule Exploration and Optimization Framework for Tensor Computation on Heterogeneous System，2020，ASPLOS】，这些框架都使用以输入为中心的调度空间。相比之下，Hidet利用任务映射编程范式和硬件中心的调度空间，从而能够以更短的调优时间获得更好的性能。

其他张量程序优化方法。除了面向循环的调度，还有其他优化张量程序的方法。深度学习框架如PyTorch【41，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】和TensorFlow【1，TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems，2015】利用现成的内核库（如cuDNN【12，cuDNN: Efficient Primitives for Deep Learning，2014，ArXiv】和cuBLAS【26，Basic Linear Algebra on NVIDIA GPUs，2022，NVIDIA Inc.】）以及手工编写的内核来覆盖常用算子。CUTLASS【32，CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES，2018】是一个开源的C++模板库，提供高效的CUDA矩阵乘法内核。Tiramisu【4，Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code，2019，CGO】和AKG【62，AKG: automatic kernel generation for neural processing units using polyhedral transformations，2021，PLDI】采用多面体模型来调度张量程序。Roller【69，ROLLER: Fast and Efficient Tensor Compilation for Deep Learning，2022，OSDI】采用自底向上的方法构建张量程序，并将分块大小与硬件规格对齐。AI-Template【59，AITemplate，2022，GitHub】采用源代码级别的模板来构建张量程序，支持更细粒度的优化，但牺牲了程序变换的灵活性。

近期相关进展。TVM社区也注意到了现有声明式面向循环调度机制的表达能力有限的问题。TensorIR【22，TensorIR: An Abstraction for Automatic Tensorized Program Optimization，2022，arXiv】是与Hidet同期的工作，它允许开发者直接编写张量程序，而不是对自动生成的张量程序应用一系列声明式原语。此外，XLA【44，XLA : Compiling Machine Learning for Peak Performance，2020】是针对线性代数的领域特定编译器。FreeTensor【51，FreeTensor: A Free-Form DSL with Holistic Optimizations for Irregular Tensor Programs，2022，PLDI】和CoRa【20，The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding，2022，MLSys】研究了不规则或参差张量程序的编译。AStitch【68，AStitch: Enabling a New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures，2022，ASPLOS】和Apollo【61，Apollo: Automatic Partition-based Operator Fusion through Layer by Layer Optimization，2022，MLSys】研究了内存密集型内核的融合以减少内存消耗。Fireiron【24，Fireiron: A Data-Movement-Aware Scheduling Language for GPUs，2020，PACT】提出了一种面向GPU的数据移动感知调度语言。Triton【52，Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations，2019，MAPL】提出以分块（tile）为基本数据类型、线程块为主要并行处理单元来编写张量程序。

正交的研究方向。Nimble【47，Nimble: Efficiently Compiling Dynamic Neural Networks for Model Inference，2021，MLSys】、DISC【70，DISC: A Dynamic Shape Compiler for Machine Learning Workloads，2021，EuroMLSys】、Cortex【21，Cortex: A Compiler for Recursive Deep Learning Models，2020，ArXiv】和DietCode【63，DietCode: Automatic Optimization for Dynamic Tensor Programs，2022，MLSys】研究动态模型的编译，这与Hidet是正交的。除了优化单个算子，Rammer【38，RAMMER: Enabling Holistic Deep Learning Compiler Optimizations with Rtasks，2020】和IOS【18，Ios: Inter-operator scheduler for cnn acceleration，2021，MLSys】提出并行化网络中的独立算子。TASO【30，TASO: optimizing deep learning computation with automatic generation of graph substitutions，2019，SOSP】、Fang等人【19，Optimizing DNN Computation Graph Using Graph Substitutions，2020，VLDB】、TENSAT【60，Equality Saturation for Tensor Graph Superoptimization，2021，MLSys】和PET【56，PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections，2021，OSDI】应用自动生成的重写规则在图级别优化DNN。Checkmate【29，Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization，2020，MLSys】、Chen等人【10，Training Deep Nets with Sublinear Memory Cost，2016，arXiv】、Echo【64，Echo: Compiler-Based GPU Memory Footprint Reduction for LSTM RNN Training，2020，ISCA】和DTR【33，Dynamic Tensor Rematerialization，2021，ICLR】旨在减少内存占用。这些工作与Hidet是正交的，可用于增强Hidet的不同方面（如图级优化、内存消耗和动态形状支持）。

8. 讨论

优化表达能力。加速器（如GPU和TPU）通常具有层次化的内存系统和基于向量或张量的计算引擎。这两者都需要专门的优化才能达到峰值性能，而这些优化通常很难通过一系列循环变换来表达。我们在本文中讨论的双缓冲例子就是这种挑战的一个很好的例证。Hidet不依赖于声明式的调度机制，而是提出在张量程序中直接用任务映射来表达任务分配和排序，这极大地提高了Hidet的表达能力，同时降低了张量程序编写的复杂性。

支持更多硬件。尽管本文我们只关注GPU，但任务映射的概念是通用的，可以用来描述其他处理器的任务分配和排序。任务映射中的工作单元可以是（1）单核CPU循环中的迭代，（2）多核CPU的CPU线程，（3）GPU的线程、warp或线程块，以及（4）其他加速器中的并行处理单元。而任务映射的任务可以是任意可索引的、同构的、可并行的操作。

未来工作。我们计划在未来支持CPU和其他加速器（如Amazon Inferentia和Trainium）。除此之外，我们还计划支持训练。由于TVM的调优时间很长，很难直接用于加速训练。得益于Hidet采用的硬件中心调度空间，Hidet的调优时间已大大减少，这使得基于Hidet构建一个训练系统成为可能。

A5 结论

本文观察到，基于面向循环调度的最先进DNN编译器无法表达需要对张量程序进行细粒度操纵的重要优化。为解决此局限性，我们提出了任务映射编程范式，这是一种编写和调度张量程序的新范式，它简化了张量程序的编写和调度，同时不像内核库那样牺牲表达优化的能力。基于此范式，我们实现了一个名为Hidet的新型DNN推理框架。实验表明，与最先进的DNN推理框架（如Onnx Runtime）和编译器（如配备AutoTVM和Ansor的TVM）相比，Hidet实现了高达1.48倍的加速（平均1.22倍）。与Ansor相比，Hidet还将调优成本降低了11倍。

A6 附录

A.1 摘要

附录目的。本附录帮助读者通过Hidet工件【17，yaoyaoding/hidet-artifacts: DOI Release，2022】复现评估章节中的所有实验。第6节中有6个实验（一个端到端实验和5个案例研究）。这些实验从执行延迟、优化时间、调度空间、输入敏感性和不同批量大小等角度，将Hidet与其他DNN框架和编译器在代表性DNN模型上进行了比较。在公开的工件中，我们提供了自动启动这6个实验的脚本。使用第A.3.2和A.3.3节中描述的硬件和软件，该工件应能复现评估章节中的所有实验结果。

A.2 工件清单

  • 编译：NVIDIA CUDA编译器 (nvcc)。
  • 模型：ResNet50, InceptionV3, MobileNetV2, Bert, 和 GPT-2。
  • 运行时环境：Linux Ubuntu 20.04+。
  • 硬件：配备Intel Core i9-12900K, NVIDIA RTX 3090, 和 64 GiB RAM的工作站。
  • 指标：端到端推理延迟和自动调优时间。
  • 所需磁盘空间（约）：2 GiB。
  • 准备工作流程所需时间（约）：2小时。
  • 完成实验所需时间（约）：60小时。大部分时间（约50小时）将用于基线AutoTVM和Ansor的模型调优。
  • 是否公开可用：是。
  • 代码许可证（如公开）：Apache 2.0。
  • 存档(提供DOI)：10.5281/zenodo.7429879。

A.3 描述

  • A.3.1 如何访问：源代码可从Zenodo存档(https://doi.org/10.5281/zenodo.7429879)或GitHub仓库(https://github.com/yaoyaoding/hidet-artifacts)下载 。
  • A.3.2 硬件依赖：为获得评估中的确切数字，需要确切的CPU和GPU：Intel Core i9-12900K CPU和NVIDIA RTX 3090 GPU。要功能性地运行实验，唯一的要求是支持CUDA 11.6+的现代NVIDIA GPU。
  • A.3.3 软件依赖：工件需要：
    • NVIDIA驱动 510.73.08
    • NVIDIA CUDA工具包 11.6 【39】
    • NVIDIA内核库 cuDNN 8.4 【12】
    • torch 1.11包 (PyTorch 【41】)
    • torchvision 0.12包 (CNN模型 【25, 45, 50】)
    • transformers 4.19.2包 (NLP模型 【16, 42】)
    • onnxruntime-gpu 1.11.1包 (ONNX Runtime 【15】)
    • nvidia-tensorrt 8.2.5.1包 (Tensor RT 【27】)
    • Apache TVM 0.9.dev，提交号c07a46327 【9】
  • A.3.4 模型：我们用五个DNN模型进行实验：ResNet50【25】，InceptionV3【50】，MobileNetV2【45】，Bert【16】，和GPT-2【42】（原文此处为GPT-3，应为笔误）。三个卷积网络来自torchvision模型库，两个transformer模型来自transformers包。所有模型都将自动下载。

A.4 安装

安装步骤。下载源代码或克隆A.3.1节中的git仓库。按照源代码目录根下的README.md文件中的安装部分的命令来构建和安装hidet及基线。

A.5 实验工作流

实验脚本。在hidet/artifacts目录下有6个以0, 1, 2, 3, 4, 5开头的子目录，分别对应第6.2, 6.3.1, 6.3.2, 6.3.3, 6.3.4, 和6.3.5节的6个实验。每个子目录包含一个python脚本main.py，可以直接启动以进行相应的实验。

A.6 评估与预期结果

结果格式与复现性。每个实验脚本都会有多个类似如下的输出：

BatchSize Model Executor Latency Std 1 resnet50 hidet 1.329 0.000

这表示一个执行器在特定批量大小下多次运行一个模型的平均延迟。这个例子显示，Hidet平均需要1.329毫秒（标准差为0.000毫秒）来运行单批次的ResNet50【25】模型。为简化，此处省略了一些列。当使用第A.3.2节和第A.3.3节中描述的硬件和软件进行实验时，工件应能复现每个评估章节中的所有实验结果。

方法细节中的引用汇总

  • 引用 [9, TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, 2018, OSDI]
    • 引用段落: 2.3节，4.1节，4.2节
    • 引用描述: 描述了TVM采用的将计算定义与调度解耦的编程范式，并列举了TVM中的调度原语。提到了TVM中的子图融合机制，即选择锚点算子并使用其调度模板来调度整个子图。
  • 引用 [11, Learning to optimize tensor programs, 2018, NeurIPS]
    • 引用段落: 2.3节，3.3节，4.2节，5.1.3节
    • 引用描述: 提到了AutoTVM作为TVM的自动调优框架，它为调度模板引入了可调参数。指出了AutoTVM的调优时间长、调度空间效率低的问题。Hidet的基于模板的调度机制的设计受到了AutoTVM的启发。
  • 引用 [32, CUTLASS: CUDA TEMPLATE LIBRARY FOR DENSE LINEAR ALGEBRA AT ALL LEVELS AND SCALES, 2018]
    • 引用段落: 2.2节，3.1节，4.1节，5.2节
    • 引用描述: 引用CUTLASS作为使用CUDA C++模板生成张量程序的例子。特别提到了双缓冲、多阶段异步预取等优化在CUTLASS中的应用，而这些优化在现有DL编译器中难以实现。还提到Hidet的卷积实现重用了矩阵乘法的并行规约优化。
  • 引用 [40, The GPU Computing Era, 2010, IEEE Micro]
    • 引用段落: 2.1节，5.1.2节
    • 引用描述: 引用了CUDA编程平台和模型的基础知识，包括SM支持并发线程数。在任务映射组合的例子中，提到了CUDA Core。
  • 引用 [65, Ansor: Generating High-Performance Tensor Programs for Deep Learning, 2020, OSDI]
    • 引用段落: 2.3节，3.2节，3.3节，4.3节，5.1.3节
    • 引用描述: 提到了Ansor作为TVM的自动调度器，它能自动生成调度模板。指出了Ansor的调优时间长、调度空间效率低的问题。Hidet的基于规则的调度机制的设计受到了Ansor的启发。
  • 引用 [43, Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, Acm Sigplan Notices]
    • 引用段落: 2.3节
    • 引用描述: 引用Halide作为首次提出计算定义与调度解耦编程范式的开创性工作。
  • 引用 [2, Learning to Optimize Halide with Tree Search and Random Programs, 2019, ACM Trans. Graph.]
    • 引用段落: 2.3节，5.1.3节
    • 引用描述: 提到了Halide-AutoScheduler作为自动调度方法的例子。Hidet的基于规则的调度机制的设计受到了其启发。
  • 引用 [8, High performance convolutional neural networks for document processing, 2006, IWFHR]
    • 引用段落: 5.2节
    • 引用描述: 提到了Hidet中使用img2col算法将卷积实现为矩阵乘法。