TLP: A Deep Learning-based Cost Model for Tensor Program Tuning

TLP：一种用于张量程序调优的基于深度学习的成本模型
作者/机构：
Yi Zhai, Yu Zhang, Shuo Liu, Xiaomeng Chu, Jie Peng, Jianmin Ji, Yanyong Zhang
(University of Science and Technology of China)

A1 主要贡献

本文旨在解决张量程序调优中的两大核心问题。首先，现有的基于深度学习的成本模型在特征提取上严重依赖于硬件架构的专家级领域知识，即便如此，提取的特征仍不理想，且需要为CPU和GPU分别设计，通用性差。其次，在一个硬件平台上训练的成本模型在另一平台上性能会急剧下降，这一问题被称为“跨硬件不可用性”。

为了应对这些挑战，本文提出了两种方法：TLP和MTL-TLP。
* TLP (Tensor Language Processing)：这是一种基于深度学习的成本模型，它创新地从调度原语（schedule primitives）而非张量程序本身提取特征。作者将调度原语序列视为一种“张量语言”，从而将预测张量程序延迟的任务转化为一个自然语言处理（NLP）的回归任务。这种方法避免了复杂的、需要专家知识的特征工程，并且简化了特征提取过程。
* MTL-TLP (Multi-Task Learning TLP)：该方法结合了多任务学习（Multi-Task Learning）和TLP，专门用于解决成本模型的跨硬件不可用性问题。通过将不同硬件平台设置为不同的学习任务，利用共享参数学习硬件无关的通用特征，利用非共享参数学习特定于硬件的特征，从而使得模型能够在仅有少量目标硬件数据的情况下，实现高效的跨平台迁移。

本文的主要贡献如下：
1. 提出了一种简单、高效且通用的张量程序特征提取机制：通过直接从调度原语中提取特征，将问题转化为NLP任务，避免了对张量程序源代码（如AST）进行复杂的分析，减少了对硬件先验知识的依赖。
2. 提出了一种解决成本模型跨硬件平台不可用性的多任务学习方法：MTL-TLP能够有效利用来自多个硬件平台的数据，显著提升模型在目标硬件上的性能，即使目标硬件数据量很少。
3. 实现了TLP和MTL-TLP并进行了全面的评估：将所提方法集成到先进的搜索框架Ansor中，并在多种深度学习模型和硬件平台上进行了验证。实验结果表明，与现有最先进技术相比，TLP和MTL-TLP在搜索效率和性能上均取得了显著的提升。据作者所知，TLP是首个直接从调度原语提取特征的张量程序成本模型，而MTL-TLP是首个有效解决跨平台不可用性问题的开源工作。

A3 背景知识

基于搜索的编译器
- 通用流程：深度学习编译器接收以高级语言编写的计算图作为输入，并为特定硬件平台输出可执行文件。流程包括：计算图优化、子图划分、搜索空间定义、搜索策略（如遗传算法、蒙特卡洛树搜索等）引导下的调度原语组合应用、生成数以万计的张量程序、使用成本模型筛选候选程序、在目标硬件上实测以找到最优程序，最后通过LLVM/CUDA等后端生成最终的可执行文件。TLP的工作空间聚焦于从高级图IR通过调度原语降低到低级张量IR的阶段。
- 示例说明：以一个融合了dense和ReLU激活的计算子图为例（图2），不同的调度原语组合会生成计算上等价但延迟不同的多个张量程序。

TenSet
- 数据集：TenSet【39, Tenset: A large-scale program performance dataset for learned tensor compilers, Lianmin Zheng et al., NeurIPS Datasets and Benchmarks Track, 2021】提出了一个大规模的张量程序数据集，并在其上预训练了一个多层感知机（MLP）成本模型（后文分别称为TenSet数据集和TenSet MLP）。
- 数据收集：该数据集在6个硬件平台（包括Intel CPU、AMD CPU、ARM CPU和NVIDIA GPU）上收集。每个平台上都从120个典型的深度学习工作负载（如ResNet, MobileNet, BERT）中划分出2308个子图。对每个子图，使用TVM中的Ansor自动搜索框架生成最多4000个张量程序。
- 数据集规模：总共包含约5157万个张量程序及其延迟数据，每个硬件平台约有859万条数据。
- 基准：TenSet是Ansor的延续，其TenSet MLP相比Ansor可将搜索时间最多加速10倍。本文使用Ansor和TenSet MLP作为TLP和MTL-TLP的基准。作者还额外收集了一个在Intel i7-10510U CPU上的TenSet数据集，称为TenSet-TLP数据集，包含约865万条数据。

自然语言处理
- 核心思想：自然语言处理（NLP）是人工智能的一个分支，利用深度学习处理和解释文本数据，已在情感分析、机器翻译等领域取得巨大成功。
- 处理流程：NLP任务通常处理由词语组成的有序序列（句子）。处理过程首先将每个词映射为一个数字（称为token），然后将每个token转换为高维向量（称为嵌入向量），这一操作称为嵌入（embedding）。一个好的嵌入方法能使同义词之间的欧氏距离或余弦距离更近。

A3 设计原则

图3展示了集成了TLP的自动搜索框架（auto-tuner）从计算子图生成张量程序的过程。右侧黑框是TLP的流水线，该图描绘了TLP的训练和推理两种状态。

• 训练阶段：自动调优器根据计算子图生成多个调度原语序列。接着，代码生成器结合计算子图和调度原语序列生成张量程序，并在目标硬件上测量其延迟。TLP以调度原语序列为输入，经过预处理、特征提取和后处理后，将最终提取的特征前向传播通过成本模型。同时，将对应张量程序的归一化延迟作为标签来计算损失，并通过反向传播更新模型权重。
• 推理阶段：自动调优器生成调度原语序列后，通过成本模型获得预测分数。根据分数筛选出top-k个潜在的候选者（图中黄色块）。最终，自动调优器只为这些候选者生成张量程序，并在目标硬件上测量它们的延迟。

在TLP的流水线中，“具体原语”指自动搜索框架中的调度原语，“抽象原语”指经过预处理后符合TLP特征提取器输入规范的原语。预处理器、提取器和后处理器的细节将在第4节中详细描述。

A2 方法细节

4 TLP

现有成本模型特征提取的局限性
- 特征提取是核心：特征提取是构建成本模型的关键，它决定了成本模型性能的上限。
- Ansor：Ansor【38, Ansor: Generating {High-Performance} tensor programs for deep learning, Lianmin Zheng et al., OSDI 20, 2020】为最内层的赋值语句从计算、内存访问和算术强度等五个方面手动提取了164个特征。
- TenSet MLP：TenSet MLP沿用了Ansor的特征提取方法，并增加了10个高级计算图IR特征。
- TIRAMISU：TIRAMISU成本模型【5, A deep learning based cost model for automatic code optimization, Riyadh Baghdadi et al., MLSys, 2021】同样从最内层赋值语句提取特征，共提取了2534个特征，但大部分为零，导致特征稀疏。与Ansor不同，它将抽象语法树（AST）结构作为计算流来前向传播输入特征。然而，将AST结构视为计算流是不可行的，因为每个张量程序的AST都不同，导致成本模型只能逐个前向传播，无法并行化，并且只能在CPU上运行。这些方法的其他缺陷在前面章节已讨论过。

4.1 TLP的特征提取

从调度原语提取特征
- 创新性：据我们所知，TLP是第一个直接从调度原语提取特征的张量程序成本模型。从调度原语中高效提取特征并保留所有有价值信息是一项挑战。
- 三种方法对比：
1. 将调度原语序列视为字符文本，使用NLP任务处理。这种简单粗暴的端到端方法难以取得好性能。
2. 将每个调度原语视为一个NLP单词，整个序列视为一个句子，然后将每个单词转换为token。这种方法会丢失一些有用信息，例如，类型相同但参数不同的原语会被编码为两个毫无关联的token，只能依赖嵌入算法和深度学习网络去挖掘它们之间的同义关系。
3. TLP采用的方法，将调度原语视为三个基本元素的组合：原语类型、数值参数和字符参数。分别对这三个元素提取特征，然后拼接在一起。这种方法强有力地保留了同义关系。

TLP特征提取流程
- 预处理：TLP特征提取首先对输入的调度原语序列进行预处理。对于序列中的每个原语，只保留三个基本元素：原语类型、数值参数和字符参数，并剥离无关字符。这三个元素基本包含了调度原语中的所有语义信息。预处理算法的实现与具体的自动搜索框架相关。在大多数框架中，此预处理算法是可逆的，即可以从这三个基本元素恢复调度原语序列，因为每个原语都有其语义和输入规范。
- 特征提取规范：图4是TLP特征提取的规范。图4a描绘了调度原语的抽象表示：一个调度原语序列（PrimitiveSequence）由多个调度原语（Primitive）组成，每个原语以原语类型（PrimitiveType）开始，后跟多个数值参数（Number）和字符参数（NameParam）。图4b是TLP特征提取的提取器实现：
- 原语类型：转换为one-hot向量。
- 数值参数：保持其值不变。
- 字符参数：转换为token，处理方式与NLP任务处理单词相同。我们将不同的字符参数映射到不同的token。
- 特征拼接与后处理：最后，所有特征按照元素的原始位置拼接起来。之后，提取的特征会经过裁剪、填充和归一化等后处理。

特征提取示例
- 示例说明：我们以图5为例来说明TLP特征提取的过程。图中的张量程序3和调度原语3取自图2。TLP的任务是估计张量程序3的延迟。我们将张量程序3和调度原语3视为等价物。图5自上而下展示了TLP特征提取的示例，即调度原语3、提取的特征和最终提取的特征。此示例中没有后处理过程。
- 与NLP的类比：我们可以将一个原语看作一个NLP单词，编码后的结果看作经过嵌入层后的NLP任务的嵌入向量。TLP特征提取是一种词嵌入算法，它极大地保留了同义关系。由于类型相同但参数不同的两个原语的编码结果有很多相同的值，因此归一化后这两个原语之间的欧氏距离相对较短。此外，TLP直接从调度原语中提取特征，无需生成张量程序，从而减少了调优时间。

4.2 TLP在TenSet数据集上的特征提取

术语定义
- 序列长度 (Sequence length)：一个调度原语序列的长度。
- 嵌入大小 (Embedding size)：从一个调度原语中提取的特征数量。
- 特征大小 (Feature size)：为一个张量程序从其调度原语序列中提取的所有特征数量，即序列长度 × 嵌入大小。

TenSet数据集上的特征统计
- 原语类型：TenSet使用Ansor框架来转储张量程序。Ansor在CPU和GPU上分别使用了11种调度原语，因此所有调度原语的类型被替换为一个11维的one-hot向量。
- 序列长度和嵌入大小：根据我们的统计，在TenSet CPU数据集中，张量程序的调度原语序列最多包含54个原语，每个原语类型加上其参数最多为40个（类型变为11位one-hot向量后）。图6显示了TenSet CPU数据集中张量程序序列长度的分布。其中，最多有1,807,960个张量程序的调度原语序列长度相同，为21。
- 嵌入大小的不确定性：一个原语的嵌入大小是不确定的。例如，一个split原语可以将一个循环变量拆分为3个或4个循环变量。表1显示了TenSet CPU数据集中每个调度原语的最大嵌入大小。这些数据表明，即使将一个张量程序的特征大小从54 × 40 = 2,160限制到25 × 22 = 550，也不会损害绝大多数特征。后续的分析和实验也证明了这一点。TenSet GPU数据集的统计数据类似，不再赘述。

4.3 TLP特征提取的可行性分析

调度原语序列与张量程序的对应关系
- 唯一性问题：不同的张量程序是否一定对应不同的调度原语序列？答案是否定的。但是，我们认为不同张量程序拥有相同调度原语序列的概率很低。相同的计算子图结合相同的调度原语序列必然生成相同的张量程序。对于不同的计算子图，要么其计算流不同，要么其计算参数不同。我们认为，一个高性能的调度原语序列包含了计算子图的大部分甚至全部计算参数，并且是根据计算子图的计算流量身定制的。这导致了不同计算子图但调度原语序列相同的碰撞概率很低。
- 碰撞处理：退一步讲，即使这种碰撞概率上升，TLP的特征提取方案仍然可用。一个调度原语序列对应的多个张量程序的标签的最优值可以作为该调度原语序列的标签。这种方法的缺点是可能导致成本模型将一个低性能的张量程序误判为高性能的张量程序，从而导致更长的调优时间。尽管如此，它保证了高性能的调度原语序列不会被错过。

数据验证
- 重复率分析：我们分析了TenSet CPU数据库，发现在总共865万个张量程序中，有856万个不同的调度原语序列，重复率仅为1.0430%。即使将特征限制在25×22，在856万个张量程序中也有853万个是不同的，重复率为1.4034%。
- 等价性结论：因此，我们可以近似地认为一个调度原语序列唯一地表征了一个张量程序。也就是说，我们将调度原语序列等同于一个张量程序。

调度原语序列的优势
- 特征稠密且信息完整：如前所述，在TenSet CPU数据集中，一个张量程序几乎可以完全由550个特征来表征。这些特征是稠密的，并且极大地保留了所有信息。
- 结构规整：调度原语序列是一个规则的序列，可以利用一些现有的NLP技术。
- 种类少且通用：调度原语的种类很少，例如在Ansor中只有14种（包括CPU和GPU），并且大部分在CPU和GPU上是通用的。
- 无需复杂工程：对于调度原语，我们不需要复杂的特征工程，只需记录其类型和参数，避免了对AST中复杂的嵌套循环进行操作。

结论：张量语言处理
- 一一对应：综上所述，调度原语序列和张量程序几乎是一一对应的，包含相同数量的语义信息。
- 结构对比：张量程序的源代码是嵌套的循环树结构，而调度原语序列是规则的序列结构。从语言学的角度看，调度原语序列的结构与自然语言相同，都是序列结构。
- 新范式：我们将调度原语称为张量语言。因此，TLP是一个张量语言处理任务，它将用成本模型预测张量程序延迟的任务转化为一个NLP回归任务。

4.4 模型架构

TLP模型结构
- 模型流程：图7是TLP的模型架构。模型首先通过多个线性层将维度上采样到256或更高。调度原语具有很强的上下文关系。自注意力机制可以有效捕捉上下文特征，避免记忆遗忘的缺陷，并且可以并行运行。此外，LSTM也能有效捕捉上下文特征但无法并行执行。因此，我们使用自注意力或LSTM模块（我们称之为主干基础模块）来捕捉上下文特征。之后是两个残差块。最后，使用多个线性层和一个求和操作来获得预测分数。
- 标签与损失函数：目标标签是归一化的延迟，计算公式为：$label = min\_latency/latency$，其中$min\_latency$指的是一个子图所有张量程序中的最小值。在TenSet数据集中，一个子图的张量程序数量从几十到4000不等。标签的取值范围是(0, 1]。我们使用均方误差（MSE）损失函数或排名损失【9, Learning to rank: from pairwise approach to listwise approach, Zhe Cao et al., ICML, 2007】函数【35, The lambdaloss framework for ranking metric optimization, Xuanhui Wang et al., CIKM, 2018】来计算损失。

5 MTL-TLP

成本模型发展历程
- 发展阶段：表2展示了张量程序成本模型的发展历程。早期，研究人员直接使用经验公式来评估张量程序的性能，例如Halide16【25, Automatically scheduling halide image processing pipelines, Ravi Teja Mullapudi et al., TOG, 2016】使用数据复用率和浮点计算量来评估不同循环分块和阶段折叠策略的优劣。随后，研究人员开始使用简单的深度学习模型（如MLP）、机器学习模型（如XGBoost）、强化学习模型等。这些模型结构简单，计算成本低。近年来，研究人员开始尝试复杂的深度学习模型，如TIRAMISU使用的动态计算流LSTM模型，以及Benoit Steiner等人【30, Value learning for throughput optimization of deep learning workloads, Benoit Steiner et al., MLSys, 2021】使用的复杂双向LSTM模型。
- 演化规律：这种演化的规律是，随着性能的显著增长，所需的数据量也显著增长。在线成本模型只使用调优过程中生成的测量数据，不使用任何离线数据集。离线成本模型依赖于海量的离线数据集，其性能与数据集的规模高度相关。

5.1 跨硬件不可用性

问题描述
- 领域差异：由于硬件架构和硬件性能的领域差异，离线收集的数据在不同硬件间是无效的。这导致在离线训练的成本模型在不同硬件平台上的性能显著下降。经验公式成本模型和在线学习成本模型不存在这个问题，因为它们不使用离线数据集。
- 挑战：这就提出了一个挑战：如何使用少量目标硬件数据来训练一个高性能的成本模型。

潜在解决方案分析
- 自监督学习：当使用少量标记数据获得高性能模型时，首先想到的是自监督学习方法，如BERT【16, Bert: Pre-training of deep bidirectional transformers for language understanding, Jacob Devlin et al., arXiv, 2018; 19, Tinybert: Distilling bert for natural language understanding, Xiaoqi Jiao et al., arXiv, 2019】和GPT【8, Language models are few-shot learners, Tom Brown et al., NeurIPS, 2020; 27, Improving language understanding by generative pre-training, Alec Radford et al., 2018; 28, Language models are unsupervised multitask learners, Alec Radford et al., OpenAI blog, 2019】。然而，BERT通常使用6-12个transformer层，输入数据维度为512x768等。但我们的特征尺寸远小于此，与BERT的权重数量不在一个数量级。例如，前面提到，TenSet中提取的特征仅为25x22。如此大的权重数量不仅容易使模型过拟合，还会严重拖慢调优速度，这在编译期间是不可接受的。
- 迁移学习：另一种用少量数据获得高性能模型的方法是迁移学习，其类型多样，如微调、训练一个本地模型来预测两个域之间的差距，以及多任务学习。

选择多任务学习的原因
- 原因分析：我们认为多任务学习技术最适合解决离线模型跨硬件不可用的问题，主要有以下三个原因：
1. 跨硬件不可用性问题非常符合多任务学习的范畴，可以减轻由领域差异引起的模型性能下降。
2. 多任务学习与张量编译器的多层IR共享共同的设计理念。张量编译器使用图IR和张量IR来分别抽象硬件无关和硬件相关的特征。多任务学习在这里非常契合：它可以使用共享参数来拟合硬件无关的特征，使用非共享参数来拟合硬件相关的特征。
3. 我们的验证结果也表明，在微调、多任务学习、训练本地模型和自监督学习中，多任务学习表现最好。我们将在实验部分展示这些比较数据。

5.2 MTL-TLP

模型结构
- 多头设计：MTL-TLP的模型结构与TLP相似。不同之处在于，MTL-TLP设置了多个头（head），即多个任务，每个任务对应一个硬件平台，如图8所示。
- 数据表示：如果TLP的一个带标签数据表示为元组(feature of primitives, latency)，那么MTL-TLP对应的元组是(feature of primitives, [latency1, latency2, ..., latencyn])。我们将任务1设为目标平台的任务。当数据在目标平台上没有标签时，元组为(feature of primitives, [no label, latency2, ..., latencyn])，在计算损失时将忽略这个目标任务，并且在反向传播期间不会更新这个头的权重。

损失函数
- 设计：损失函数设计如下：
$Loss_{total} = \sum_{i=1}^{n} Loss\_i(pred_i, label_i)$, 其中 $label_i$ is not no label.
其中$Loss\_i$是每个任务的损失函数（通常相同），例如MSE损失函数、排名损失函数等。
- CPU/GPU隔离：张量程序在CPU和GPU之间不通用，因此我们不讨论跨CPU和GPU的多任务学习。

5.3 MTL-TLP的可行性分析

MTL-TLP之所以能发挥巨大作用，主要归功于以下几点：
* 有效增加样本量：MTL有效地增加了我们用于训练模型的样本量。尽管目标平台上的标记数据量很少，但其他平台上确实有大量数据。目标平台成本模型可以利用其他平台的数据作为辅助任务进行训练。
* 避免局部最优：在单任务学习中，梯度的反向传播容易陷入局部最小值。在多任务学习中，不同任务的局部最小值位于不同位置。不同任务的相互作用可以帮助隐藏层逃离局部最小值。
* 正则化效应：MTL通过引入归纳偏置（inductive bias）起到正则化的作用。多个任务在浅层共享权重，这可能会削弱网络的能力，减少网络过拟合，并提高泛化能力。
* 特征解耦：MTL-TLP可以使用共享参数来拟合硬件无关的特征，使用非共享参数来拟合硬件相关的特征。

A4 实验环境

• 数据集

  • TenSet: 在6个硬件平台上收集的大规模张量程序性能数据集，包括Intel CPU、AMD CPU、ARM CPU和NVIDIA GPU。每个平台包含从120个深度学习模型中划分的2308个子图，总计约5157万个张量程序及其延迟。
  • TenSet-TLP: 作者在Intel i7-10510U CPU上额外收集的数据集，包含约865万条数据，用于补充和验证。
  • 测试集: 从TenSet数据集中，为每个硬件平台划分出由ResNet-50, MobileNet-V2, ResNext-50, BERT-tiny, 和 BERT-base五个网络组成的测试集（batch size为1，图像尺寸224或序列长度128）。其余数据按9:1划分为训练集和验证集。

• 模型架构

  • TLP/MTL-TLP: 基于自注意力机制（或LSTM）和残差块构建。特征输入维度为25x22，通过线性层上采样至256。自注意力模块设置为8个头，后接两个残差块。
  • 基准模型: Ansor的默认成本模型，以及在TenSet上预训练的TenSet MLP。

• 硬件配置

  • CPU实验平台:
    • 笔记本：8核 Intel i7-10510U, 16GB内存, 2GB NVIDIA MX350。
    • TenSet数据集平台：Intel Xeon Platinum 8272 (16核), Intel E5-2673 v4 (8核), AMD EPYC 7452 (64核), ARM Graviton2 (16核)。
  • GPU实验平台:
    • 服务器：8核 Intel Platinum 8255C, 32GB内存, 16GB NVIDIA Tesla T4。
    • TenSet数据集平台：NVIDIA Tesla K80, NVIDIA Tesla T4。

• 软件配置

  • 框架: Ansor (TVM编译器中的自动搜索框架)。
  • 实现: 将TLP和MTL-TLP技术集成到Ansor中进行端到端评估。
  • 调优设置: 所有端到端实验均进行200轮调优，每轮选择10个张量程序进行实测，总计2000次测量。

A5 实验结果

评估指标分为两类：基于数据集的指标（模型在静态数据集上的准确率）和基于搜索的指标（模型集成到搜索算法后的端到端搜索效率或质量）。

6.1 基于数据集的TLP性能评估

使用top-k分数作为评估标准，其表达式如下：
$top-k = \frac{\sum_m \sum_s min\_latency_{m,s} \times weight_{m,s}}{\sum_m \sum_s \min(latency_{m,s,i}) \times weight_{m,s}}, 1 \le i \le k$
$$top-k = \frac{\sum_M \sum_S min\_latency_{m,s} \times weight_{m,s}}{\sum_M \sum_S min (latency_{m,s,i}) \times weight_{m,s}}, 1 \le i \le k$$
其中$min\_latency_{m,s}$是模型$m$的子图$s$所有张量程序中的最小延迟，$weight_{m,s}$是子图$s$在模型$m$中出现的次数，$latency_{m,s,i}$是模型$m$的子图$s$所有张量程序中，成本模型输出分数第$i$大值对应的延迟。

• 损失函数与主干模块选择

  • 实验内容：在Intel Xeon Platinum 8272的TenSet CPU数据集上，比较了MSE损失和Lambda Rank损失，以及LSTM和自注意力（Attention）作为主干基础模块的组合。
  • 实验结果：如表3所示，"Attention + Rank"的组合在Top-1和Top-5分数上均表现最佳。因此，后续实验默认采用此组合。
    

• 特征大小裁剪

  • 实验内容：验证将特征大小从最大值（54x40）裁剪到25x22是否会提升性能。
  • 实验结果：如表4所示，将序列长度和嵌入大小分别缩减至25和22时，模型准确率得到提升，同时减少了计算和存储开销。这可能是因为裁剪后的特征更稠密，避免了大量零值的影响。
    

• 与TenSet MLP对比

  • 实验内容：在所有平台的TenSet和TenSet-TLP数据集上，将优化后的TLP模型与SOTA的TenSet MLP进行比较。
  • 实验结果：如表5所示，在所有CPU平台上，TLP的Top-1和Top-5分数均大幅超过TenSet MLP。在GPU平台上，两者各有优劣。但TLP性能在不同CPU和GPU上表现稳定，且与需要为CPU/GPU分别设计特征的TenSet MLP不同，TLP使用统一的特征提取机制。这证明了TLP特征提取是一种简单、有效且通用的方法。
    

6.2 基于数据集的MTL-TLP性能评估

• MTL-TLP有效性验证

  • 实验内容：在CPU（目标平台Intel E5-2673）和GPU（目标平台NVIDIA Tesla T4）上进行多任务学习实验，比较仅使用500K目标平台数据与结合其他平台数据进行多任务学习的效果。
  • 实验结果：如表6和表7所示，多任务学习显著提升了成本模型的准确率。在CPU上，仅用500K目标数据结合其他平台数据，性能就超过了使用全部数据的TenSet MLP，并与使用全部数据的单任务TLP相当。同时，适当增加任务数量（如增加到3个）能提升准确率，但任务过多（4个）则会因任务间干扰而降低性能。
    
    

• 迁移学习与自监督学习方法比较

  • 实验内容：比较了四种方法：微调（Fine-tuning）、多任务学习（MTL）、GPT和BERT（自监督学习）。目标平台为Intel i7-10510U，使用500K数据。
  • 实验结果：如表8所示，MTL在解决跨硬件不可用性问题上优于微调。GPT和BERT由于参数量巨大，对于小尺寸输入特征容易过拟合，表现不佳。
    

• 跨架构多任务学习

  • 实验内容：以Intel i7-10510U（X86架构）为目标平台，分别与Intel Platinum 8272（X86）、Intel E5-2673（X86）、AMD EPYC 7452（X86）和ARM Graviton2（ARM）进行双任务学习。
  • 实验结果：如表9所示，与同为X86架构的平台进行多任务学习时，对目标平台的精度提升最大。这表明架构相似性对多任务学习效果有积极影响。
    

• 数据量与性能关系

  • 实验内容：探究MTL-TLP准确率与目标平台数据量的关系，数据量从50K增加到2.0M。
  • 实验结果：如图9所示，当数据量增加到500K时，性能提升最显著，此时MTL-TLP的准确率已超过TenSet MLP。
    

6.3 基于搜索的端到端评估

• 搜索时间对比

  • 执行速度：如图10所示，由于TLP从调度原语提取特征而无需生成张量程序，其调优执行速度比TenSet MLP在CPU和GPU上平均快1.7倍和1.8倍。
    

• 调优曲线与性能加速比

  • 收敛速度：如图11所示，TLP和MTL-TLP的调优曲线能更快地收敛到更低的延迟值，尤其在CPU上更为明显。在某些情况下，Ansor调优2000次的性能仍远不如TLP/MTL-TLP的初始性能。
    
  • 与TenSet MLP比较：如图12所示，为达到TenSet MLP调优2000次的性能，TLP在CPU和GPU上平均可将搜索时间加速9.1倍和3.0倍；MTL-TLP（仅使用7%目标数据）可加速4.7倍和2.9倍。
    
  • 与Ansor比较：如图13所示，为达到Ansor调优2000次的性能，TLP在CPU和GPU上平均可加速16.7倍和16.0倍；MTL-TLP可加速10.0倍和15.8倍。实际加速比可能更大，因为TLP/MTL-TLP的起点性能已经优于Ansor的终点性能。
    

• 对计算资源的影响

  • 资源使用：与TenSet MLP相比，TLP在CPU上的资源使用减少，而在GPU上的内存使用从882MB增加到1634MB（batch_size=2048），但仍在普通笔记本GPU可接受范围内。执行遗传算法的时间从约20秒减少到约6秒。

A7 补充细节

局限性 (Limitation)
- 数据收集成本：尽管MTL-TLP仅用7%的目标平台数据就超越了使用全部数据的TenSet MLP，但收集50万条数据仍然需要数十小时。跨硬件不可用性问题仍需更多研究来进一步降低数据收集成本。

未来工作 (Future Work)
- 方法推广：本文已证明将调度原语作为张量程序等价物的可行性。在Ansor中实现的TLP和MTL-TLP可以轻松迁移到其他张量编译器自动搜索框架中。
- 借鉴NLP技术：本文将张量程序调优转化为NLP回归任务。基于此，未来可以尝试使用更多成熟的NLP技术来解决张量程序相关问题。

A6 结论

本文提出了TLP和MTL-TLP两种方法，以改进张量程序调优中的成本模型。TLP设计了一种新颖、简单且通用的特征提取机制，直接从调度原语中提取特征，将调优问题转化为张量语言处理任务。MTL-TLP则利用多任务学习技术有效解决了离线成本模型在不同硬件平台间的不可用性问题。本文从理论上分析了这两种方法的可行性和优势，并通过基于数据集和基于搜索的详尽实验验证了它们的有效性。实验结果表明，TLP和MTL-TLP在搜索效率和最终性能上均显著优于当前最先进的实现。