Enabling Tensor Language Model to Assist in Generating High-Performance Tensor Programs for Deep Learning

  • 文章标题：赋能张量语言模型以辅助生成深度学习高性能张量程序
  • 作者/机构：Yi Zhai (中国科学技术大学); Sijia Yang (华为技术有限公司); Keyu Pan (字节跳动); Renwei Zhang (华为技术有限公司); Shuo Liu (中国科学技术大学); Chao Liu, Zichun Ye (华为技术有限公司); Jianmin Ji (中国科学技术大学); Jie Zhao (湖南大学); Yu Zhang, Yanyong Zhang (中国科学技术大学)

A1 主要贡献

本文提出了一种用于深度学习应用的张量程序生成框架，其核心思想是维持一个广阔的探索空间以确保高性能，同时借助语言模型进行强大的探索，从而高效地生成张量程序。

  • 核心问题: 现有获取高性能张量程序的方法存在两难：追求效率的方法通常通过启发式约束限制探索空间，但这缺乏通用性；而追求高性能的方法则构建了广阔的探索空间，但采用了效率低下的探索策略（如随机采样）。

  • 研究目标与方法: 本文旨在解决上述效率与性能的矛盾。作者将张量程序探索任务转化为语言模型生成任务。

    1. 设计张量语言（Tensor Language）: 为了方便语言模型处理，作者设计了一种对语言模型友好的张量语言。该语言通过记录决策信息来表示张量程序，一个张量语言句子（tensor sentence）唯一对应一个张量程序。与动辄上万词元的张量程序源码相比，张量语言句子更为简洁（不超过1024个词元）。
    2. 开发张量语言模型（Tensor Language Model, TLM）: 作者借鉴了ChatGPT的训练方法，利用数百万个张量句子和少量高性能的演示句子（demonstration sentences）来预训练和监督微调TLM。
    3. 高效探索: 在编译目标工作负载时，TLM结合离线学习的知识和先前做出的决策，对当前决策空间中的最佳决策进行概率性采样。这种方法比现有方法中常用的随机采样更具指导性。

  • 主要创新点:

    1. 设计了对语言模型友好的张量语言：该语言作为桥梁，连接了张量程序和语言模型。
    2. 开发了张量语言模型（TLM）：该模型能够结合离线学习的知识和上下文决策，对当前决策空间进行概率采样，从而实现更有效的空间探索。
    3. 实验证明了TLM兼具高效率和高性能：
      • 与完全调优的Ansor/MetaSchedule相比，TLM在达到同等性能的同时，编译速度提升了61倍。
      • 在相同的编译时间内，TLM的性能比Roller提升了2.25倍。
      • 在充足的探索时间内，TLM的编译时间与Ansor/MetaSchedule相当，性能分别提升了1.08倍和1.04倍。

下图展示了通用张量编译器与集成了本文提出的张量程序生成框架的架构对比。

图 1: 左图显示了通用张量编译器的架构，右图展示了整合我们张量程序生成框架的结构。

A3 背景知识

2.1 张量程序探索框架

张量程序探索框架的演进与挑战。在张量程序探索框架的发展中，先前研究主要采用搜索算法来自动定位最优的张量程序。因此，基于搜索的张量程序探索框架通常被称为张量程序搜索框架，其探索空间也用作搜索空间。早期的Halide自动调度器【4，Learning to optimize halide with tree search and random programs，ACM Transactions on Graphics (TOG)，2019】通过评估不完整的程序来积极剪枝搜索空间；AutoTVM和FlexTensor【6，Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system，ASPLOS，2020】则使用预定义的、手动编写的模板来定义其搜索空间。这些方法引入了限制搜索空间的约束，从而错失了许多有效的决策组合，导致性能次优。

Ansor/MetaSchedule的困境。随后，Ansor【13，Ansor: Generating {High-Performance} tensor programs for deep learning，OSDI，2020】和MetaSchedule【14，Tensor program optimization with probabilistic programs，NeurIPS，2022】利用派生规则构建了广阔的搜索空间，将性能推向了业界领先水平。然而，这也带来了一个非常棘手的问题——过长的编译时间。根据作者经验，在NVIDIA V100 GPU上，使用Ansor使BERT-base工作负载收敛需要21.6小时；在Intel i7-10510U CPU上，则需要13.1小时。这背后的原因是Ansor/MetaSchedule采用了一种随机采样后使用可学习成本模型进行评估的策略。在完成所有决策空间的随机采样之前，Ansor/MetaSchedule将子图降级为张量程序，然后提取统计特征，包括计算、内存访问和算术强度，以用于可学习成本模型的性能评估。然而，该策略存在三个主要问题：
- 随机采样是一种低效的探索形式，采样到最优或次优决策的概率相等。
- 评估存在滞后性。在初始采样期间做出的决策可能导致性能不佳，但这只能在评估时才能确定。
- 训练数据的不足和成本模型的小参数量削弱了模型的学习能力，限制了其在指导搜索算法方面的有效性。Ansor/MetaSchedule完全依赖于极少的在线数据来训练成本模型。此外，它们的模型主要基于低参数的机器学习或深度学习模型（例如XGBoost【22，Xgboost: A scalable tree boosting system，KDD，2016】、MLP、LSTM）。

近期工作的局限性。最近的研究，如Roller【23，{ROLLER}: Fast and efficient tensor compilation for deep learning，OSDI，2022】、TenSet【20，Tenset: A large-scale program performance dataset for learned tensor compilers，NeurIPS Datasets and Benchmarks Track，2021】和TLP【21，Tlp: A deep learningbased cost model for tensor program tuning，ASPLOS，2023】，旨在解决编译速度慢的问题。Roller通过将张量形状与硬件属性对齐来加快编译速度。然而，这样做导致Roller的性能下降（§6.4.2），这仍然是由于启发式约束的通用性问题。TenSet和TLP在编译目标工作负载之前收集离线数据集，以构建更强大的成本模型。虽然这些方法加快了编译速度，但它们存在两个关键缺点。首先，它们采用了Ansor/MetaSchedule的策略，即先进行随机采样，然后用成本模型进行评估。此外，它们依赖于860万条带标签的数据。

不同方法的概率视角对比。作者通过一个例子，从概率角度分析了以Roller为代表的启发式编译器、Ansor/MetaSchedule等基于搜索的框架以及本文的生成式TLM框架之间的差异。假设一个决策空间Di有四个有效决策，其中最优决策是d3，如图2所示。启发式编译器使用启发式约束来排除d1、d2和d4，使得d3有100%的概率被采样，从而加速编译但可能牺牲性能。基于搜索的框架保留所有选项，d3的采样概率为25%，通过充分搜索逼近最优解，保证性能但效率低下。而TLM则结合离线学习的知识和已做出的决策，对当前决策空间中的最佳决策进行概率性预测，从而实现更有效的探索。

图 2: 几种不同方法在决策空间中的概率分布示例。

2.2 语言模型

语言模型基础。深度学习语言模型主要分为两类：掩码语言模型（MLM）和因果语言模型（CLM），后者也称为自回归语言模型。CLM的著名例子包括GPT系列，如GPT-2【24，Language models are unsupervised multitask learners，OpenAI blog，2019】和GPT-3【25，Language models are few-shot learners，NeurIPS，2020】。CLM因其自然的文本生成方式而闻名，通常在需要连贯文本创作的任务中表现出色，如写作或聊天机器人对话。在本文中，语言模型特指CLM。

语言模型的工作流程。在训练语言模型之前，通过分词（tokenization）创建一个词汇表。分词通常指将输入句子分割成其组成部分（词元），所有这些词元共同构成一个词汇表。图3中的步骤①②和⑦⑧分别表示将句子分词和将词元转换回句子的过程。语言模型所做的唯一事情就是结合学到的知识和给定的输入，预测下一个词元是词汇表中特定词元的概率分布。初始输入被称为提示（prompt）。在语言模型训练期间，使用自然语言句子作为输入。通过反向传播算法【26，Learning representations by back-propagating errors，nature，1986】，语言模型的参数被迭代更新，以最大化生成的下一个词元与自然语言中相应词元对齐的概率。

语言模型的推理过程。在使用语言模型进行推理时，以一个现实生活中的例子为例：我们问一个语言模型，“什么是语言模型？”，它回答说，“语言模型是一种经过训练以理解、生成和回应人类语言的人工智能……”。生成过程如图3所示。在步骤③中，语言模型使用提示作为输入，为词汇表中的每个词元预测一个概率。这个概率表示每个词元在逻辑上跟随输入的可能性。然后根据这些概率对下一个词元进行采样。步骤④涉及使用提示和在步骤③中生成的词元作为输入，以概率方式采样下一个词元，这个过程持续进行，直到生成“</s>”词元。“</s>”表示句子的结束。

图 3: 使用语言模型生成自然语言句子的过程。

2.3 ChatGPT/InstructGPT

ChatGPT的训练方法。ChatGPT【19，ChatGPT，URL: https://openai.com/blog/ chatgpt】和InstructGPT【27，Training language models to follow instructions with human feedback，NeurIPS，2022】采用了类似的训练方法，包括：预训练、监督微调（SFT）和基于人类反馈的强化学习（RLHF，包含奖励模型RM + 强化学习RL）。TLM部分采用了这些训练方法。InstructGPT的训练过程包括：
- 预训练GPT-3：InstructGPT基于拥有1750亿参数的GPT-3，利用约45TB的互联网文本进行训练，旨在让GPT-3学习语言的基本结构和语义。
- 监督微调（SFT）：通过收集演示数据（约1.4万条）来训练一个监督策略，通过监督学习对GPT-3进行微调。这些由人类提供的演示数据代表了与提示相对应的期望输出行为。
- 奖励建模（RM）：通过收集比较数据（约5.1万条）来训练一个奖励模型。这些同样由人类提供的比较数据，为模型输出提供了从最好到最差的排名。
- 强化学习（RL）：使用PPO【28，Proximal policy optimization algorithms，arXiv，2017】强化学习算法，根据奖励模型优化一个策略。奖励由前述的奖励模型确定，此步骤会迭代多次。

A2 方法细节

3 系统概述

TLM框架的设计理念。TLM框架是一个旨在高效生成高性能张量程序的张量程序生成框架。维持一个大的探索空间可以确保高性能的张量程序不会被启发式约束所排除。然而，这需要更强的探索能力。利用深度学习模型从离线数据中汲取知识来辅助在线张量程序探索是一种有前景的策略。考虑到当前的学习能力，语言模型是可用的最强大的方法。因此，本文提出利用语言模型来辅助生成高性能张量程序，将张量程序探索任务转化为语言模型生成任务。

TLM框架的系统流程。图4展示了系统概览，并标记了详细介绍每个系统组件的相应小节。第4节重点介绍为TLM预训练收集大规模离线数据集。第4.1节详细说明了用于数据采样的探索空间及其构建过程，通过对探索空间的公式化为张量语言设计提供了理论基础。第4.2节讨论了张量语言的设计，强调其在以更适合语言模型的格式保存从探索空间中采样的张量程序方面的作用。在讨论了采样空间和保存方法之后，第4.3节介绍了大规模采样，其中广泛的随机采样实现了对探索空间的无偏估计。第5节专注于张量语言模型的开发。首先，第5.1节阐述了TLM的模型架构、训练方法和训练过程，包括预训练和监督微调。完成预训练后，TLM应该能够生成探索空间内的任何有效张量程序。接着，第5.2节详细说明了TLM框架如何在TLM的决策支持下生成张量程序。为了生成高性能的张量程序，我们采用演示数据（对应于高性能的张量程序）通过监督学习对TLM进行微调，使其生成的张量句子与我们期望的演示句子对齐。第5.3节讨论了通过迭代算法获取这些演示数据的方法。

图 4: TLM框架的系统概览。

4 空间构建器

空间构建器的功能。空间构建器的功能是创建一个大的探索空间以确保高性能。这个探索空间由所有可能的张量程序组成。关于大探索空间能保证高性能的观点，需要从两个方面来讨论。首先，更大的空间意味着与较小的空间相比，它构建时受到的约束更少，因此能够生成更多的张量程序。在某种程度上，可以说大空间包含了小空间。其次，确保高性能意味着保持高性能张量程序存在于探索空间中的潜力。空间越大，被约束剪枝的可能性就越小，但这也要求更强的探索能力。

4.1 探索空间

探索空间的构建。本节介绍两个主要议题。首先是关于构建探索空间。其次是形式化张量程序的生成过程，为设计张量语言提供理论基础。张量编译器通常为优化目的提取两个中间表示（IR），即图层（graph layer）和张量层（tensor layer），它们分别针对硬件无关和硬件相关的增强。如图5所示，图层接收一个工作负载，通过多个优化遍和算子融合与划分算法进行优化。然后，它产生子图，每个子图由一个或多个算子组成。子图描述了预期的计算结果，而张量程序探索框架则将此子图转换为张量程序，提供详细的计算实现。所有可能的张量程序的集合构成了张量层的探索空间S。先前的探索框架利用专家知识，包括模板（AutoTVM）、多面体模型（AKG）和派生规则（Ansor, MetaSchedule），将子图映射到张量程序。对于专家知识不足的领域，这些框架使用可调参数（AutoTVM, AKG）、注解（Ansor）或随机变量（MetaSchedule）来创建一个决策空间Di，然后采用模拟退火【29，Optimization by simulated annealing，science，1983】和遗传算法【30，Adaptation in natural and artificial systems: an introductory analysis with applications to biology, control, and artificial intelligence，MIT press，1992】等搜索算法来定位最优解。

图 5: 通用张量编译器的两个中间表示。

决策空间的组成。决策空间主要涉及确定循环轴的切分大小、设置展开步长、为算子选择计算位置、制定并行化和向量化策略，以及在GPU上特别确定线程绑定。

探索空间的形式化。根据以上分析，张量程序是通过一系列决策优化输入子图的结果。我们将张量层的探索空间形式化如下：
$$S = \left\{ s^{(n)} \left| \begin{array}{l} s^{(i)} = \text{apply}\left(s^{(i-1)}, d_i\right), \\ \forall d_i \in D_i, 1 \leq i \leq n \end{array} \right. \right\}$$
其中$s^{(0)}$表示输入子图的初始程序，而$d_i$是从集合$D_i$中随机抽样的一个样本。因此，探索空间的大小与决策组合的数量一致。我们有：
$|S| = |D_1| \times |D_2| \times \dots \times |D_n|.$
本文仅讨论张量层的探索空间。在CPU上，对应于一个子图的探索空间大小约为$10^6$；在GPU上，约为$10^9$。

探索空间的复用与扩展。我们复用了先前工作（如Ansor, MetaSchedule）的决策空间，因为我们认为这个空间已经足够大，可以产生高性能的张量程序。将张量层的探索空间扩展到更大是一个相当大的挑战。在未来的工作中，我们主张主要通过整合图层的决策空间来探索更大的空间。当然，TLM框架支持不同大小的探索空间，因为它们都以相同的方式与TLM集成。

4.2 张量语言

张量语言的设计动机。本节我们关注如何存储从探索空间中采样的张量程序。我们的目标是让语言模型学习这些张量程序，使其能够在目标工作负载编译期间辅助生成高性能的张量程序。张量程序源代码通常超过一万个词元，这给语言模型生成如此冗长、连贯且有效的源代码带来了挑战。因此，我们不追求端到端生成张量程序源代码。相反，我们利用语言模型来辅助决策制定。为此，我们明确设计了对语言模型友好的张量语言，该语言通过记录决策信息来表示张量程序。

张量语言的构成。从4.1节可知，一个张量程序$s^{(n)}$只与输入子图的初始程序$s^{(0)}$和决策$d_1, d_2, \ldots, d_n$有关。每个决策$d_i$都是从决策空间$D_i$中随机抽样的，其设计取决于硬件平台。因此，为了确保一个张量语言句子唯一地对应一个张量程序，一个张量句子必须封装输入子图、硬件规格和决策。我们利用算法1从探索空间中采样数据。每个句子从输入子图中提取信息，包括子图中每个算子的类型和形状信息。此外，它还检索硬件细节，包括处理核心数和支持的向量指令。算法1展示了如何从tile size决策空间中提取词元，保留了相应算子、轴、决策空间类型、从空间中采样的决策以及其他关键信息的细节。类似的方法也应用于其他决策空间。

张量语言的特性与示例。张量语言是一种自然语言形式，而非编程语言，因此不严格遵循巴科斯-诺尔范式（Backus-Naur Form）【31，Report on the algorithmic language algol 60，Communications of the ACM，1960】。其主要意图是使用单个句子来表示一个张量程序，强调其记录而非编程的角色。这个记录过程提供了很大的灵活性，唯一的约束是离线收集的张量句子和在线生成的张量句子之间的一致性。这种一致性对于深度学习模型确保训练和测试数据独立同分布至关重要。鉴于张量语言的灵活性，关于需要记录输入子图、硬件规格和决策信息的具体细节没有严格的指导方针。在保持一致性的前提下，这些细节可以根据具体的工程项目动态调整。图6展示了为Ansor（上）和MetaSchedule（下）量身定制的张量句子样本。

图 6: 为Ansor（上）和MetaSchedule（下）量身定制的张量句子样本，包含输入子图、硬件规格和决策信息。例如，上方的 "SP 2 0 1024 32 1 4 1" 代表 "split operator_index axis_index axis_extent tile_size_0 tile_size_1 tile_size_2 save_manner"。

4.3 大规模采样张量句子

大规模采样的目的。本节介绍采用采样算法来创建一个大规模的离线数据集。大规模采样有两个主要目的。首先，通过广泛的随机采样创建探索空间S的无偏估计，这使得TLM能够学习张量语言的基本结构和语义，从而能够生成空间S内的任何张量句子。其次，建立尽可能大的词汇表（§2.2），其中所有决策选项如 "i.0=16" 和 "i.0=32" 都被分词为离散的词元。如果一个决策，如 "i.0=17"，不在词汇表中，它将永远不会被生成。

数据集的构建与收集。为TLM配置的工作负载数据集借鉴了TenSet。这些工作负载来源于PyTorch的视觉模型库和Huggingface的Transformer模型库，涵盖了计算机视觉（CV）和自然语言处理（NLP）的代表性任务。我们调整输入形状以生成各种子图。请注意，我们在这个数据集中关注小批量大小，因为张量编译器主要用于优化已训练模型以进行推理。总共，该数据集包含138个工作负载，其中12个作为测试集保留；从剩余的工作负载中提取了大约3000个子图。

数据收集的效率。我们为3000个子图收集了200万个张量句子来预训练TLM。值得注意的是，这200万条数据与TenSet/TLP中使用的860万条数据不同。这里的数据是未标记的，即不需要测量执行延迟；测量通常是最耗时的部分。在一台96核服务器上，收集200万个CPU数据条目大约需要2小时，GPU数据大约需要10小时。GPU数据收集时间较长是由于额外的检查，例如确保线程绑定满足硬件约束。

5 张量语言模型

5.1 模型细节

模型架构。考虑到学习能力和资源开销，TLM采用了GPT-2 Small的架构，该架构包含约1亿个参数。TLM由12个Transformer层组成，每层具有12个注意力头和768个隐藏单元。

训练方法。TLM的训练部分受到ChatGPT所使用的训练程序的启发，如§2.3所述。简而言之，我们将公式从1次预训练 + 1次SFT + n次(RM + RL)修改为1次预训练 + n次(测量 + SFT)。我们将奖励模型（RM）替换为测量（measurement）。RM为模型输出分配一个奖励值并评估其质量。在这里，张量语言具有天然优势。一个张量句子可以被转换成一个张量程序（§5.2），使其执行延迟可以直接在硬件上测量，延迟越低表明张量句子越好。虽然ChatGPT只进行一次监督微调（SFT），但我们进行了多次。SFT的输入是演示数据，由人类提供，象征着与提示相对应的期望输出行为。在自然语言处理领域，很难说哪个演示数据对于一个提示是“最好”的，但这在张量语言中是可行的。也就是说，执行延迟最低的句子对其提示来说是“最好”的。多次执行测量和SFT就是为了持续寻找最好的张量句子（§5.3）。我们没有使用强化学习（RL），主要是因为它需要调整各种超参数，并且在训练收敛方面存在挑战，而且我们发现在多次执行SFT后性能已经足够好。

训练细节。通过预训练得到的模型称为TLM-base，而TLM是通过对TLM-base进行SFT得到的。我们使用§4.3中收集的离线数据集，以与其他语言模型预训练相同的方式预训练TLM-base。值得注意的是，TLM-base仅在2个epoch内就收敛了。这种更快的收敛是由于张量语言相比于自然语言的广泛多样性具有更明显的规律性。在200万个数据条目中，输入子图类型、硬件类型和决策类型都是有限的。使用4个NVIDIA V100预训练TLM-base 2个epoch大约需要10小时。预训练后，TLM-base具备生成有效张量句子的能力。此外，对于输入子图，我们期望TLM-base能辅助生成高性能的张量程序。为此，我们采用演示数据通过监督学习对TLM-base进行微调。演示数据指的是对应于给定输入子图的一小部分执行延迟最低的张量程序的张量句子。用演示数据对语言模型进行SFT的目的是使模型对提示的响应与演示数据中反映的人类意图对齐。类似地，我们将演示数据应用于TLM，旨在使其能够响应提示生成高性能的输出。对TLM-base进行SFT需要一小批（约3K，为每个子图选择最佳的一个）演示数据。我们采用迭代优化（§5.3）来持续优化这些演示条目。使用3K演示数据，在4个NVIDIA V100上对TLM-base进行一次SFT大约需要10分钟。

5.2 张量程序生成

TLM辅助的张量程序生成流程。TLM框架在张量程序生成过程中利用TLM进行决策，如算法2所详述。算法2中的步骤与算法1中的步骤一致，遵循生成张量句子所需的一致性。TLM生成一个决策后，框架会检查其在决策空间内的有效性。如果决策无效，框架会丢弃该张量程序并启动重新生成。由于决策是通过采样获得的，下一次生成可能会采样到不同的决策，从而避免了连续的重新生成失败。幸运的是，经过预训练和微调，我们观察到无效决策极为罕见。平均而言，生成一个有效的张量程序需要调用算法2的次数不超过1.1次。

图 7: 为一个维度m=1024, n=512, k=1024的矩阵乘法算子生成张量句子的过程，图中描绘了切分大小（tiling size）的部分。

决策生成的具体过程。图7展示了算法2中response_tokens = TLM(tokens)的过程，该过程针对一个维度为m=1024, n=512, k=1024的矩阵乘法算子，并描绘了其中的切分大小部分。这个张量句子与图6顶部的句子相匹配，并在此处以更易读的格式呈现（省略了算子索引）。在步骤①中，已知信息——包括输入子图、硬件规格和“split i=1024 to”（对应于m=1024）——作为TLM的输入提示。TLM随后根据此提示预测下一个词元的概率分布，然后通过从该分布中进行概率采样选择一个词元，假设为“i.0=32”。在步骤②中，将步骤①中的提示和预测的下一个词元组合起来，形成一个新的输入，供TLM预测下一个词元。步骤③重复步骤②。完成这三个步骤后，TLM最终确定了i轴的切分大小，并返回到算法2。当需要为j轴生成切分大小时，将再次调用TLM。在步骤④中，将步骤③的输入、其预测的下一个词元和“split j=512 to”（对应于n=512）合并成一个新的输入，然后将其输入到TLM中以预测下一个词元。TLM可以但不采用将输入图和硬件规格作为提示一次性生成所有决策信息的方式。相反，框架在算法2中重复调用TLM，每次只生成一部分决策。这种精细化的方法能有效过滤掉无效数据，增强了TLM的可用性和稳定性。

5.3 迭代优化

演示数据的获取。使用演示句子对TLM-base进行微调对其运行至关重要。本节重点介绍获取一批演示数据的方法。在获取演示数据的整个过程中，测量是最耗时的步骤，并且位于关键路径上，是整个过程的瓶颈。为了解决这个问题，我们开发了一个流水线系统，执行迭代优化算法来收集演示数据。该系统通过使用两个独立的进程来最大化测量硬件的利用率：一个用于执行SFT，另一个用于测量，前者的执行时间被控制得比后者短，以确保测量的连续性。在GPU上操作时，这些进程在不同的GPU卡上运行。

迭代优化流程。图8左侧显示了迭代优化的流程图，右侧是结合了流水线的流程图。我们将工作负载数据集中的所有子图（每个子图对应一个提示）分成kb（例如4）个批次，然后循环选择每个批次。在批次i的测量完成后（称为记录i），从数据库中提取最佳的一批数据，记为演示数据i，并用于微调TLM-base，得到TLMi+2。随后，使用TLMi+2生成句子i+2（一个提示生成kp个，例如16或32个句子）。TLM0和TLM1被TLM-base替代。

图 8: 迭代优化过程的流程图。

迭代优化的关键细节。有几个细节值得注意：
- 子图划分策略：在划分 子图时，我们按子图类型（例如，fused_nn_dense_add, fused_nn_conv2d_add_nn_relu, 和 fused_nn_adaptive_avg_pool2d）进行排序，每kb个选择一个，其基本原理是，一个优越的决策通常对相同类型的子图也具有一定的优化效果。
- TLM生成更优数据的原理：只有当TLM能够生成更优的张量句子时，这个迭代优化算法才能正常运行。从遗传学的角度来看，当为某个子图生成句子时，TLM已经学习了当前子图对应的演示数据，以及其他子图的演示数据，从而继承了自身和其他子图的优点。在为提示生成下一个词元时，使用的是概率采样，由于其随机性，提供了额外的探索，因此可能引入突变。继承和突变的结合可能潜在地生成更高质量的数据，这与遗传算法的设计理念相符。
- 收敛性保证：演示数据的好坏取决于执行延迟。所有子图的总执行延迟可以表示为 $latency_{total} = \sum_{s_i \in \text{subgraphs}} \min(\text{all record latencies of } s_i)$。演示数据的延迟不会上升，因为它总是从所有测量中选择最好的，所以它是单调递减的。此外，演示数据的最佳延迟不能低于其物理极限，所以它是有界的。从数学上讲，单调有界的极限会导致收敛。评估的第6.2节讨论了需要测量多少数据才能收敛。
- TLM的训练稳定性：TLM旨在辅助在决策空间内做出决策，通过梯度下降策略进行训练。$TLM_i$ 可能比 $TLM_{i-1}$ 表现更差，因为梯度下降不保证TLM总能被训练到其最优状态。然而，由于$TLM_i$总是从TLM-base开始训练，它不受$TLM_{i-1}$的影响，也不影响$TLM_{i+1}$。因此，$TLM_i$不太可能出现累积误差，最终的TLM只与TLM-base和最终的演示数据有关。
- 演示数据的灵活性：图8中的演示数据是通过从头开始的迭代优化获得的。同样，也有其他方法可以获得一批演示数据。例如，可以使用其他搜索算法获得；如果在硬件A上有一批演示数据，可以通过迁移学习转移到硬件B；数据也可以由专家知识直接编写。好消息是，这些演示数据仍然可以继续使用迭代优化算法直到收敛。
- 迭代优化的适用场景：迭代优化算法也可以被看作一个调优系统，特别适用于需要同时调优多个工作负载的情况。例如，在动态形状场景中，需要同时调优具有不同形状的目标工作负载。

A4 实验环境与结果

6.1 实验环境

  • 软件配置
    • TLM支持的决策空间：TLM支持多种决策空间，包括从Ansor (V0.12), MetaSchedule (V0.12), AKG (V2.1) 和 AKG-MLIR (V0.1) 适配的空间。后续评估中，这些TLM分别称为TLM-Ansor, TLM-Meta, TLM-AKG和TLM-AKG-MLIR。
    • 实现：评估中重点关注的TLM-Ansor和TLM-Meta，是使用Python和C++实现的，代码量约1万行。
  • 数据集
    • 来源：TLM配置的数据集包含138个工作负载，其中12个被选为测试集（如表1所示）。
    • 测试集工作负载：包括ResNet-50, MobileNet-V2, ResNeXt-50, BERT-base, BERT-tiny, GPT-2, DenseNet-121, BERT-large, Wide-ResNet-50, ResNet3D-18, DCGAN, LLAMA等模型，输入形状各异。
      表 1: TLM测试集中的工作负载。
  • 硬件配置
    • CPU平台：一台配备4核Intel(R) Core(TM) i7-10510U CPU（支持AVX2指令集）、16GB内存、运行Ubuntu 20.04的笔记本电脑。
    • GPU平台：一台配备48核Intel(R) Xeon(R) Gold 6226 CPU、376GB内存和四块32GB NVIDIA Tesla V100 GPU的服务器。运行环境为Ubuntu 20.04，CUDA 11.6和cuDNN 8.4.0。

6.2 演示数据的收敛行为

本节讨论通过迭代算法获取演示数据的收敛行为，这是整个流程中最慢的阶段。图9展示了TLM-Ansor和TLM-Meta在GPU和CPU上的四条收敛曲线。
- 实验设置：这四种场景分别使用了2169, 3120, 2169, 2657个子图（从126个工作负载中提取）。横轴表示已测量的数据条目数（所有子图总和），纵轴表示归一化总延迟。
- 收敛定义：当2万次测量（所有子图）带来的性能提升小于1%时，定义为“收敛”。
- 结果分析：
- 达到收敛所需测量次数：要达到收敛状态，每个子图平均需要104、69、145和130次测量，对应所有子图的总测量次数分别为225K、213K、314K和344K。这意味着在GPU上，TLM需要约20万次测量，在CPU上需要约30万次测量才能使126个工作负载收敛。
- 与TenSet/TLP的对比：TLM所需的数据量比使用约860万次测量的TenSet和TLP小一个数量级。
- 更高精度收敛：若要求性能提升小于千分之一，每个子图分别需要272、150、375和363次测量。
- 后续实验配置：后续章节实验中使用的TLM是使用约30万个标记数据进行微调的。

图 9: TLM-Ansor和TLM-Meta在GPU和CPU上的演示数据收敛曲线。

6.3 子图基准测试

在SFT之后，我们在NVIDIA V100上进行了子图实验。
- 实验设置：
- TLM测试集中的12个工作负载为TLM-Ansor和TLM-Meta分别产生了232和364个子图，这些子图分属23和40个类别。我们从每个类别中选择一个代表性子图进行实验。
- 对比组为：TLM-Ansor vs. Ansor，以及TLM-Meta vs. MetaSchedule。
- 各框架的测量次数分别为10K、10K、1K和10K。
- 结果分析（见表2和表3）：
- TLM的快速收敛：TLM-Ansor在1K次测量时就达到了其10K次测量性能的99%。TLM-Meta在64次测量时就达到了其1K次测量性能的99%。这表明增加测量次数带来的性能增益非常有限，进一步提升性能需要探索更大的空间。
- TLM的高效性：仅需10次测量，TLM就能分别达到Ansor和MetaSchedule经过10K次测量后性能的103%和95%。在32次测量时，TLM的性能已经超过了Ansor和MetaSchedule。
- TLM的性能优势：在相同的测量次数下，TLM始终比Ansor和MetaSchedule表现更好。例如，TLM-Ansor-1K比Ansor-1K快1.13倍，TLM-Ansor-10K比Ansor-10K快1.08倍。

表 2: 23个TLM-Ansor子图的总体加速比。数值越高越好。表中“Times”代表每个子图的测量次数。

表 3: 40个TLM-Meta子图的总体加速比。

6.4 端到端工作负载基准测试

6.4.1 与Ansor/MetaSchedule的比较

  • 实验设置：
    • GPU与CPU：在GPU和CPU上分别进行四组实验，对比Ansor-20K、MetaSchedule-20K与TLM-Ansor/TLM-Meta在不同测量次数（1×g, 10×g, 32×g, 20K，其中g为子图数量）下的端到端性能。
  • GPU结果 (图10)：
    • 充足测量下：TLM-Ansor-20K相比Ansor-20K平均加速1.08倍。TLM-Meta-20K相比MetaSchedule-20K平均加速1.04倍。
    • 有限测量下 (高效性)：
      • TLM-Ansor在10×g次测量下性能与Ansor-20K持平（平均0.99倍），在32×g次测量下性能超过Ansor-20K（平均1.03倍）。
      • TLM-Meta在10×g次测量下性能与MetaSchedule-20K持平（平均1.00倍），在32×g次测量下性能超过MetaSchedule-20K（平均1.02倍）。
    • 零测量潜力：仅需1×g次测量，TLM就能达到Ansor/MetaSchedule性能的80%以上，这意味着TLM无需测量即可达到基准83%的性能，展示了其在深度学习训练框架中应用的潜力。
  • CPU结果 (图11)：
    • 结果与GPU一致。TLM-Ansor-10×g比Ansor-20K快1.01倍。TLM-Meta-32×g与MetaSchedule-20K性能持平。
  • 编译时间：
    • 本文使用测量次数来计算编译时间加速比。
    • TLM-Ansor-10×g相比Ansor-20K，在达到同等性能的同时，编译速度提升了95倍。
    • TLM-Meta-10×g相比MetaSchedule-20K，编译速度提升了61倍。

      图 10: 在V100上与Ansor/MetaSchedule的工作负载推理性能比较。

      图 11: 在Intel CPU上与Ansor/MetaSchedule的工作负载推理性能比较。

6.4.2 与Roller在V100上的比较

  • 实验设置：
    • Roller基于旧版TVM(V0.8)和CUDA(10.2)。为了公平比较，引入了同样环境下运行的Ansor-20K-V0.8作为桥梁。
  • 结果分析 (图12)：
    • 在相同环境下，Ansor-20K-V0.8优于Roller-10×g。
    • TLM-Ansor-10×g和新版Ansor-20K-V0.12性能相当，且略优于Ansor-20K-V0.8。
    • 排除软件版本影响后，TLM-10×g明显优于Roller-10×g。
    • 直接比较下，TLM-Ansor-10×g相比Roller-10×g平均加速2.25倍。原因是Roller为追求高效率，通过对齐张量形状和硬件特征，过度剪枝了探索空间。

      图 12: 在V100上与Roller的工作负载推理性能比较。

6.4.3 与TensorRT/PyTorch在V100上的比较

  • 实验设置：
    • 与基于静态核函数库的TensorRT (V8.6) 和 PyTorch (V1.13.1) 进行性能比较。
  • 结果分析 (图13)：
    • vs. TensorRT：TLM-Ansor-10×g平均加速1.04倍。但在BERT、GPT-2、LLAMA等模型上，TensorRT表现更优。
    • vs. PyTorch：TLM-Ansor-10×g平均加速3.42倍。但在BERT-large、LLAMA等模型上，PyTorch表现更优。
  • 结论：TensorRT/PyTorch的优势在于对常用算子（如batch matmul, matmul, transposed 2D convolution）在核函数库中进行了深度优化并利用了专用硬件计算单元。总的来说，张量编译器在支持广泛算子方面表现出色，而静态核函数库更擅长深度优化常用算子。

    图 13: 在V100上与TensorRT/PyTorch的工作负载推理性能比较。

A7 补充细节

7 相关工作

Halide与TVM的演进。Halide【44，Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，Acm Sigplan Notices，2013】引入了计算与调度分离的概念，使用领域特定语言（DSL）定义计算，并用调度原语抽象硬件特性，通过自动调度器【4，Learning to optimize halide with tree search and random programs，ACM Transactions on Graphics (TOG)，2019】【45，Differentiable programming for image processing and deep learning in halide，ACM Transactions on Graphics (ToG)，2018】【46，Automatically scheduling halide image processing pipelines，ACM Transactions on Graphics (TOG)，2016】来寻找最佳原语组合。TVM【3，{TVM}: An automated {End-to-End} optimizing compiler for deep learning，OSDI，2018】继承了Halide的理念，使用调度原语实现算子。它目前拥有三代张量程序搜索框架：第一代使用模板将子图映射到张量程序，并通过AutoTVM【12，Learning to optimize tensor programs，NeurIPS，2018】进行优化。第二代Ansor【13，Ansor: Generating {High-Performance} tensor programs for deep learning，OSDI，2020】通过派生规则构建张量程序并使用遗传算法搜索高效程序，解决了基于模板的探索空间的局限性。第三代包括TensorIR【47，Tensorir: An abstraction for automatic tensorized program optimization，ASPLOS，2023】和MetaSchedule【14，Tensor program optimization with probabilistic programs，NeurIPS，2022】，通过引入块抽象来解决支持TensorCore的挑战，该抽象隔离了张量化计算以便映射到张量计算单元。TenSet【20，Tenset: A large-scale program performance dataset for learned tensor compilers，NeurIPS Datasets and Benchmarks Track，2021】和TLP【21，Tlp: A deep learningbased cost model for tensor program tuning，ASPLOS，2023】提出使用离线数据集来解决Ansor/MetaSchedule带来的长时间搜索问题。Roller【23，{ROLLER}: Fast and efficient tensor compilation for deep learning，OSDI，2022】引入了一个封装张量形状的tile抽象，将其与底层加速器的关键特性对齐以限制形状选择。FlexTensor【6，Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system，ASPLOS，2020】是一个调度探索和优化框架，提出了自动通用模板，用于将张量算法映射到不同硬件平台的低级实现。

其他相关编译器技术。Tiramisu【48，Tiramisu: A polyhedral compiler for expressing fast and portable code，CGO，2019】、AKG【15，Akg: automatic kernel generation for neural processing units using polyhedral transformations，PLDI，2021】和Tensor Comprehensions【5，Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions，arXiv，2018】应用了基于多面体的技术，将程序优化问题表述为整数线性规划（ILP）问题。Triton【49，Triton: an intermediate language and compiler for tiled neural network computations，MLPL，2019】引入了一种基于tile的模板表示，程序员可以指定块大小并管理其调度以实现有效的程序优化。CUTLASS【50，CUTLASS，URL: https://github.com/NVIDIA/ cutlass】是用于在CUDA中实现各级别和规模的高性能矩阵-矩阵乘法及相关计算的模板抽象集合。MLIR【51，Multi-Level Intermediate Representation Overview，URL: https://mlir.llvm.org】构建了可重用和可扩展的编译器基础设施，以解决软件碎片化问题并改进异构硬件的编译 。

8 讨论

局限性。1) TLM本质上是一个深度学习模型，它利用离线数据来指导探索。目标场景和训练场景的数据分布需要保持一致。为了在差异显著的场景中获得最佳性能，可能需要加入额外的训练数据。2) TLM的设计理念是用预编译时间换取编译时间。训练TLM的开销可能长达数十小时，这一点不应被忽视。如果目的仅仅是编译一个或少数几个模型，TLM可能不是最经济的选择。相反，TLM更适合用于编译大量模型。

未来工作。1) TLM拥有1亿参数，其训练和推理的时间和硬件开销应予以关注。探索在保证TLM性能稳健的同时大幅减少参数规模是值得的。2) 相反，仅有1亿参数，仍有机会大幅扩展TLM的参数规模和训练数据。这些改进可能带来卓越的泛化能力。例如，直接从TLM生成的结果进行编译而无需测量，或实现跨多种硬件平台的强大泛化。3) 深入探索更广阔的探索空间。凭借TLM强大的探索能力，可以导航广阔的领域而不受探索成本的限制。

A5 结论

我们介绍了一种新颖的张量程序生成框架——TLM框架，它由两个解耦的组件组成：一个空间构建器和一个生成器。TLM框架开创性地将语言模型整合到张量程序领域。由于TLM在熟练学习离线数据和有效捕获上下文方面的双重优势，它展示了强大的生成能力。实验结果表明，TLM始终能够提供高效率和高性能。

A6 附录

A 编译加速比指标

使用测量次数计算编译加速比的合理性。本节我们讨论使用测量时间来计算编译时间加速比的合理性。编译一个工作负载的总时间($T_{\text{total}}$)如以下公式所示，它包括两个主要部分：用于识别高性能张量程序的探索时间($T_{\text{exploration}}$)，以及最终编译工作负载的时间($T_{\text{post exploration compilation}}$)。探索过程涉及采样c个张量程序($c \times T_{\text{sample}}$)和测量其中一个程序的执行延迟($T_{\text{measurement}}$)。在Ansor/MetaSchedule中，采样多个张量程序的目的是使用成本模型选择最有希望的张量程序；如果TLM采样的程序被发现无效，将重新采样。采样系数$c_{\text{Ansor}}$和$c_{\text{MetaSchedule}}$约为128，而$c_{TLM}$不超过1.1。测量一个张量程序包括其编译($T_{\text{compilation}}$)和执行($T_{\text{execution}}$)。
$$T_{total}=T_{exploration}+T_{post~exploration~compilation} \\ =\sum_k (c \times T_{sample}+T_{measurement}) \\ +T_{post~exploration~compilation} \\ =\sum_k (c \times T_{sample}+(T_{compilation}+T_{execution})) \\ +T_{post~exploration~compilation}$$
当测量次数k超过某个阈值$k_0$（例如100）时，$T_{\text{exploration}}$成为$T_{\text{total}}$的主导因素，使得$T_{\text{post exploration compilation}}$可以忽略不计。采样和测量过程可以通过流水线方法进行并行优化。总的来说，测量是整个编译过程的关键部分，减少测量次数需要方法论上的创新，而不仅仅是工程上的努力。
$$ T_{total} = T_{exploration} + T_{post \ exploration \ compilation} \\ \approx T_{exploration} \quad \text{if } k \geq k_0 \\ = \sum_k (c \times T_{sample} + T_{measurement}) \\ = \sum_k T_{measurement} \quad \text{if } T_{measurement} \text{ hides } c \times T_{sample} \\ = \sum_k (T_{compilation} + T_{execution}) $$
TLP使用$T_{\text{total}}$的加速比来计算编译时间的加速，而Roller使用$\sum_k T_{\text{compilation}}$来计算编译时间的加速。虽然它们在一定程度上确实反映了编译时间的加速，但这些指标由于受系统负载的影响而不稳定。在本文中，我们使用测量次数k来计算编译时间的加速比。

B 子图延迟比较曲线

附录B提供了两张详细的子图延迟比较曲线图，用于支持§6.3节的结论。
- 图14：展示了TLM-Ansor与Ansor在23类代表性子图上的延迟比较。曲线显示，在不同测量次数下，TLM-Ansor通常能以更少的测量次数找到与Ansor相当或更优的解。

图 14: TLM-Ansor vs. Ansor的子图延迟比较曲线。为了增强细微细节的比较，曲线仅包括不超过TLM-Ansor或Ansor最低延迟1.1倍的延迟。在图中，如果某个曲线在整个过程中都不可见，则表示其最低延迟超过了另一个的10%。

  • 图15：展示了TLM-Meta与MetaSchedule在40类代表性子图上的延迟比较，结论与图14类似。

    图 15: TLM-Meta vs. MetaSchedule的子图延迟比较曲线。为了增强细微细节的比较，曲线仅包括不超过TLM-Meta或MetaSchedule最低延迟1.1倍的延迟。在图中，如果某个曲线在整个过程中都不可见，则表示其最低延迟超过了另一个的10%。

参考文献引用说明

  • 【3, {TVM}: An automated {End-to-End} optimizing compiler for deep learning, 2018, OSDI】: 在相关工作中被引用，作为继承Halide理念并发展了三代张量程序搜索框架的代表性张量编译器。
  • 【4, Learning to optimize halide with tree search and random programs, 2019, ACM Transactions on Graphics (TOG)】: 在背景知识和相关工作中被引用，作为早期通过评估不完整程序来积极剪枝搜索空间的自动调度器代表。
  • 【5, Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions, 2018, arXiv】: 在相关工作中被引用，作为应用多面体技术的编译器之一。
  • 【6, Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system, 2020, ASPLOS】: 在背景知识和相关工作中被引用，作为使用预定义模板来定义搜索空间的早期方法，并提出自动生成通用模板。
  • 【12, Learning to optimize tensor programs, 2018, NeurIPS】: 在相关工作中被引用，作为TVM的第一代优化框架AutoTVM。
  • 【13, Ansor: Generating {High-Performance} tensor programs for deep learning, 2020, OSDI】: 在背景知识和相关工作中被多次引用，作为构建广阔搜索空间但编译时间过长的代表性框架，是本文方法的重要对比基线。
  • 【14, Tensor program optimization with probabilistic programs, 2022, NeurIPS】: 在背景知识和相关工作中被多次引用，作为Ansor的后续者，同样面临编译时间长的问题，是本文方法的另一重要对比基线。
  • 【15, Akg: automatic kernel generation for neural processing units using polyhedral transformations, 2021, PLDI】: 在相关工作中被引用，作为应用多面体技术的编译器之一。
  • 【19, ChatGPT, URL: https://openai.com/blog/ chatgpt】: 在背景知识中被引用，作为TLM训练方法借鉴的范例。
  • 【20, Tenset: A large-scale program performance dataset for learned tensor compilers, 2021, NeurIPS Datasets and Benchmarks Track】: 在背景知识和相关工作中被引用，作为通过收集离线数据集来加速编译的代表，但其数据集规模远大于本文方法。
  • 【21, Tlp: A deep learningbased cost model for tensor program tuning, 2023, ASPLOS】: 在背景知识和相关工作中被引用，与TenSet类似，也是通过离线数据集构建成本模型，但数据量需求巨大。
  • 【22, Xgboost: A scalable tree boosting system, 2016, KDD】: 在背景知识中被引用，作为Ansor/MetaSchedule中成本模型可能使用的一种低参数机器学习模型。
  • 【23, {ROLLER}: Fast and efficient tensor compilation for deep learning, 2022, OSDI】: 在背景知识和相关工作中被引用，作为通过启发式约束来加速编译但可能牺牲性能的代表性工作。
  • 【24, Language models are unsupervised multitask learners, 2019, OpenAI blog】: 在背景知识中被引用，作为因果语言模型（CLM）的代表（GPT-2）。
  • 【25, Language models are few-shot learners, 2020, NeurIPS】: 在背景知识中被引用，作为因果语言模型（CLM）的代表（GPT-3）。
  • 【26, Learning representations by back-propagating errors, 1986, nature】: 在背景知识中被引用，作为语言模型训练中使用的基础算法（反向传播）。
  • 【27, Training language models to follow instructions with human feedback, 2022, NeurIPS】: 在背景知识中被引用，作为TLM训练方法借鉴的范例（InstructGPT）。
  • 【28, Proximal policy optimization algorithms, 2017, arXiv】: 在背景知识中被引用，作为InstructGPT中用于强化学习的PPO算法。
  • 【29, Optimization by simulated annealing, 1983, science】: 在方法细节中被引用，作为早期张量编译器中用于搜索最优解的算法之一。
  • 【30, Adaptation in natural and artificial systems: an introductory analysis with applications to biology, control, and artificial intelligence, 1992, MIT press】: 在方法细节中被引用，作为早期张量编译器中用于搜索最优解的算法之一（遗传算法）。
  • 【31, Report on the algorithmic language algol 60, 1960, Communications of the ACM】: 在方法细节中被引用，来说明本文的张量语言不遵循严格的巴科斯-诺尔范式。
  • 【44, Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, Acm Sigplan Notices】: 在相关工作中被引用，作为引入计算与调度分离理念的开创性工作。
  • 【45, Differentiable programming for image processing and deep learning in halide, 2018, ACM Transactions on Graphics (ToG)】: 在相关工作中被引用，作为Halide自动调度器的相关研究。
  • 【46, Automatically scheduling halide image processing pipelines, 2016, ACM Transactions on Graphics (TOG)】: 在相关工作中被引用，作为Halide自动调度器的相关研究。
  • 【47, Tensorir: An abstraction for automatic tensorized program optimization, 2023, ASPLOS】: 在相关工作中被引用，作为TVM第三代搜索框架的核心组件，用于支持TensorCore。
  • 【48, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】: 在相关工作中被引用，作为应用多面体技术的编译器之一。
  • 【49, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MLPL】: 在相关工作中被引用，作为引入基于tile的模板表示的编译器。
  • 【50, CUTLASS, URL: https://github.com/NVIDIA/ cutlass】: 在相关工作中被引用，作为CUDA中高性能矩阵乘法实现的模板库。
  • 【51, Multi-Level Intermediate Representation Overview, URL: https://mlir.llvm.org 】: 在相关工作中被引用，作为可扩展的编译器基础设施。