TCL框架：基于持续学习的跨硬件张量程序优化编译器

1. 项目概述：为什么我们需要一个更聪明的编译器？

如果你在GPU上跑过深度学习模型，大概率用过像TVM、TensorRT或者XLA这样的编译器。它们干的事儿，本质上就是把我们写的那些高级计算图（比如PyTorch的nn.Module），变成能在特定硬件（比如NVIDIA A100或者Intel CPU）上高效执行的机器码。这个过程的核心，就是“张量程序优化”——编译器会尝试成千上万种不同的“调度方案”（比如怎么切分循环、怎么安排数据在内存中的布局），然后找出跑得最快的那一个。

听起来很美好，对吧？但现实很骨感。这个“找最优”的过程，传统上严重依赖“在线调优”（Auto-Tuning）。简单说，就是编译器生成一堆候选程序，一个个扔到真实硬件上跑，测出实际耗时，再根据结果指导下一轮搜索。这相当于每来一个新模型、新硬件，就得重新做一遍耗时巨大的“暴力实验”。在论文里的数据，为一块Intel i7-12700F CPU收集足够的调优数据要40天，为一块RTX 3080Ti GPU则要超过60天。这成本，无论是时间还是算力，对大多数团队来说都是难以承受的。

因此，近几年的研究趋势转向了“离线训练成本模型”。思路是：我先花大力气，为某几款硬件收集一个庞大的“程序-性能”数据集（比如Tenset数据集），然后训练一个神经网络模型，让它学会根据程序的“特征”（比如用了哪些算子、循环怎么嵌套）来预测其运行时间。以后来了新程序，让模型预测一下性能就行，不用再真机跑测了。

但这又引出了三个新问题：

1. 数据饥渴：构建那个初始数据集成本太高，且每增加一个新硬件平台，这个痛苦的过程就要重来一遍。
2. 模型笨重：为了预测准确，成本模型本身不能太简单。像Transformer这类强大模型，虽然预测准，但训练和推理的计算开销（O(n²)复杂度）又成了新瓶颈。
3. 知识孤岛：为一个硬件（比如V100）训练好的模型，很难直接用到另一个硬件（比如A100）上。传统迁移学习通常是“一对一”的，而同时从多个硬件（比如CPU、GPU、各种AI芯片）学习并积累通用优化知识的“多对一”能力很弱。

TCL框架，就是为了系统性地解决这三个痛点而生的。它的全称是“基于持续学习的跨硬件张量程序优化编译器框架”。它不是某个单一算法的改进，而是一个整合了数据策略、模型架构和学习范式的完整系统。其核心目标是：用尽可能少的数据、一个轻量但强大的模型，实现知识在多个硬件平台间的有效积累和迁移，最终让编译器能快速、精准地为任意新硬件找到最优程序。

接下来，我会带你深入TCL的三个核心组件，看看它们是如何协同工作，把理想变成现实的。

2. 核心设计解析：TCL的三驾马车

TCL的威力来自于三个紧密耦合的核心组件：RDU采样器、Mamba成本模型和持续知识蒸馏框架。它们分别对应了数据、模型和知识迁移这三个层面。

2.1 RDU采样器：用10%的数据，办100%的事

主动学习（Active Learning）的核心思想是“好钢用在刀刃上”——与其盲目地标注所有数据，不如让模型自己挑选最有价值、最能帮助自己进步的样本去标注。在编译器场景下，“标注”就是去真实硬件上运行一个张量程序并记录其耗时，这是最昂贵的操作。

RDU采样器的聪明之处在于，它用一个综合评分公式，为海量未标注的程序候选池中的每一个程序打分，然后只挑选分数最高的一小批（比如10%）去实际运行。这个评分综合了三个维度：

1. 代表性（Representativeness）：选出来的样本分布，应该和全体数据的分布一致。比如，你的数据集中卷积（Conv）和矩阵乘法（MatMul）这类计算密集型算子占了80%，池化（Pooling）占了5%。那么你采样时，也应该大致按这个比例来选，确保模型对所有类型的算子都有学习机会，而不是偏科。这通过统计算子类型分布并依此分配采样预算来实现。

2. 多样性（Diversity）：选出来的样本彼此之间应该差异大，覆盖尽可能广的“程序形态空间”。假设你已经标注了一些程序，RDU会计算一个新候选程序与所有已标注程序在模型预测性能上的最小差异。差异越大，说明这个新程序越“与众不同”，越可能带来新的信息。公式如下，其中f(x)是模型对程序x的性能预测值：dist(x_i, x_j) = min |f(x_i) - f(x_j)|（x_i属于未标注集，x_j属于已标注集） 我们倾向于选择dist值大的样本。

3. 不确定性（Uncertainty）：模型对自己预测越没把握的样本，往往越有价值。因为纠正模型的这些“困惑点”，能最有效地提升模型能力。对于回归任务，一个常用的不确定性度量是预测值的方差。RDU会计算将候选程序加入已标注集后，整个集合预测值的方差，方差越大，不确定性越高。

最终的采样分数是这三者的加权组合：总分数 = f(x_i) * 多样性分数 + 不确定性分数。这里有个精妙的设计：给多样性分数乘上了模型对该程序的预测值f(x_i)。这意味着，当两个程序多样性得分相同时，我们会优先选择模型预测其性能更好的那个。这完美契合了编译器的终极目标——找到最快的程序，引导采样过程向高性能区域倾斜。

实操心得：在实际实现RDU时，关键是要平衡好三个指标的权重。初期，数据少，模型不准，应更侧重不确定性，快速提升模型基础能力。中期，当模型有一定判别力后，加强多样性，探索未知区域。后期，当搜索收敛时，代表性和向高性能区域倾斜的机制能帮助精细调优。这个动态权衡的过程，往往需要根据具体硬件和算子集的特性进行微调。

2.2 基于Mamba的成本模型：快、准、狠的预测引擎

成本模型是TCL的大脑，负责吃进一个张量程序的“特征描述”，吐出对其运行时间的预测。它的设计必须在“预测精度”和“计算效率”之间取得平衡。

特征工程：从硬件无关到硬件感知早期的成本模型（如Ansor）主要提取程序本身的特征（循环次数、内存访问模式等），这些是硬件无关的。TLP等工作在此基础上，进一步从“调度原语序列”（即编译器对程序做的各种变换步骤）中提取特征，取得了更好效果。TCL继承了这一思路，但做了关键增强：引入了硬件感知特征。

• 对于CPU：我们关心核心数、线程数、各级缓存大小（L1/L2/L3）、支持的SIMD指令集宽度等。
• 对于GPU：我们关心CUDA核心数、显存带宽、显存大小、计算单元峰值算力（FP16/FP32/INT8）等。 这样，模型的输入就变成了一个“调度序列特征”和“硬件平台特征”的拼接向量。例如，在GPU上，一个调度序列可能对应一个41x20的特征矩阵。这让模型能同时理解“程序做了什么”和“它在什么硬件上跑”。

模型架构：用Mamba取代TransformerTransformer在序列建模上很强，但其注意力机制的计算和内存复杂度是序列长度的平方（O(n²)）。当调度序列很长时，这将成为训练和推理的沉重负担。TCL选择了Mamba这一基于结构化状态空间模型（SSM）的新架构。 Mamba的核心优势在于：

1. 线性复杂度：它对序列长度的复杂度是线性的（O(n)），在处理长序列时，速度和内存占用远优于Transformer。
2. 选择性状态机制：Mamba有一个“选择机制”，能根据输入动态地决定关注或忽略哪些历史信息。这在调度序列优化中非常有用，因为不是每一个过去的调度步骤都对当前性能有同等影响，模型需要学会聚焦于那些关键决策点。
3. 硬件友好：Mamba的设计包含专门优化的并行扫描算法，在现代硬件（尤其是GPU）上能实现高效的并行计算。

TCL的成本模型结构很清晰：输入特征先经过一个由三个线性层组成的编码器（Encoder）进行升维和融合；然后送入Mamba Block捕捉序列依赖；最后再通过一个解码器（Decoder）输出最终的预测值。

损失函数：我们不在乎绝对时间，只在乎谁更快这里有一个非常重要的认知：对于编译器优化，我们不需要模型精确预测出一个程序要跑1.23毫秒还是1.24毫秒。我们只需要模型能正确排序，即判断出程序A比程序B快即可。因此，TCL没有使用均方误差（MSE）这类回归损失，而是使用了LambdaRank损失。 LambdaRank是一种列表排序（Listwise）损失。它不仅仅考虑一对程序的相对顺序，还考虑了这个顺序在最终排序列表中的位置影响（通过NDCG指标）。简单来说，它更加强调把真正高性能的程序排到列表的最前面。因为我们的目标就是找到Top-K的最快程序，至于第100名和第101名谁快谁慢，其实没那么重要。

2.3 持续知识蒸馏框架：让编译器学会“举一反三”

这是TCL实现“跨硬件”优化的核心。传统方法要么是为每个硬件单独训练一个模型（参数爆炸），要么是用多任务学习同时训练（需要所有硬件数据同时可得，不灵活）。TCL借鉴了持续学习的思想，设计了一个知识库（Knowledge Base, KB） + 活跃列（Active Column, AC）的架构。

你可以把知识库（KB）想象成编译器积累的“通用优化经验手册”。它从之前学过的所有硬件（比如先学了V100，又学了A100）中，提炼出那些硬件无关的通用规律，例如：“把循环拆分到这个粒度通常有利于缓存”，“这种数据布局转换在这里能减少访存冲突”。

活跃列（AC）则是针对当前正在学习的新硬件（比如一块新的华为昇腾芯片）的“专项学习模块”。它的任务是快速掌握这块新硬件的特性。

这个框架的学习过程是交替进行的两个阶段：

1. 持续学习（CL）阶段：学习新硬件。此时，知识库（KB）被冻结，防止遗忘旧知识。新硬件的训练数据只用于训练活跃列（AC）。但关键来了，AC不是从零开始学！它通过一个“渐进式网络”连接，能够利用KB的输出。公式（见论文）描述了AC第i层的输出，它同时依赖于AC自身前一层的输出和KB对应层的输出。这相当于AC一边学新知识，一边随时查阅那本“通用经验手册”获得提示，学得更快更好。
2. 知识蒸馏（KD）阶段：消化吸收新知识。当AC在新硬件上学有所成后，我们启动蒸馏阶段。此时，AC被冻结，扮演“老师”的角色；KB被解冻，扮演“学生”的角色。KB的目标是既学会AC从新硬件上学到的独特知识，又不能忘记自己以前从V100、A100那里学来的老知识。为了防止这种“灾难性遗忘”，TCL引入了弹性权重巩固（EWC）技术。EWC会给KB中那些对记住旧任务很重要的参数加上“保护锁”，在蒸馏更新时，这些参数的变化会受到惩罚，从而被保护起来。

通过“CL -> KD -> CL -> KD ...”的不断循环，KB就像一个不断成长的中央大脑，每接触一个新硬件，它的通用优化知识就增长一分。而当我们面对一个全新的、从未见过的硬件时，只需要为其创建一个新的AC，然后让它连接上这个已经饱经风霜的KB，就能极大地加速学习过程。

注意事项：在实现这个框架时，KB和AC的融合方式是关键。简单的特征拼接可能不够，TCL采用的是一种门控（gated）连接方式（公式中的 α_i 是可训练的门控向量），让AC网络动态决定从KB中汲取多少、以及汲取哪些信息。这个门控机制需要仔细调优，否则可能效果不佳。

3. 实现细节与实操要点

理解了原理，我们来看看如何将其落地。这里我会结合论文和工程经验，拆解几个关键的实现环节。

3.1 数据管道与RDU采样器的工程实现

首先，你需要一个初始的“种子数据集”。通常可以从公开数据集（如Tenset）中获取，或者用Ansor/TVM在自己的硬件上先跑一小轮自动调优，生成一批（比如1万个）(程序特征, 真实耗时)对。

RDU采样器的实现步骤：

1. 初始化：从海量未标注池中随机选取一小批（例如100个）程序，运行它们获取真实耗时，构成初始的已标注集D_l，并用它训练一个初始的成本模型f。
2. 计算算子分布：遍历整个未标注池D_u，统计所有程序中不同算子类型（Conv, MatMul, Pooling等）的出现频率，得到概率分布pro_map。这个分布将决定每类算子的采样预算。
3. 迭代采样循环： a.评分：对于D_u中的每一个程序x_i，用当前模型f预测其性能，并计算其与D_l中所有程序的最小性能差异作为多样性分数d_s。同时，计算将x_i加入D_l后，集合预测值的方差作为不确定性分数u_s。 b.综合排序：按公式总分数 = f(x_i) * d_s + u_s计算每个程序的最终得分，并降序排列。 c.预算控制：从高分往低分依次选择程序。检查该程序的算子类型，如果该类型已选中的数量还未达到其预算（总预算 * 该类型概率），则选中它，将其从D_u移入D_l，并支付一次“真实运行”的成本去获取其标签。如果该类型预算已满，则跳过此程序，检查下一个。 d.模型更新：当选中一批新程序（例如50个）并获取其真实标签后，用扩大的D_l重新训练或微调成本模型f。
4. 终止：当真实运行的预算（例如，只允许运行全数据集的10%）耗尽，或者模型性能在验证集上收敛时，停止采样。此时得到的D_l就是精选的高价值数据集，f也是基于此数据训练好的高效模型。

3.2 Mamba成本模型的搭建与训练

模型搭建（以PyTorch风格伪代码示意）：

import torch import torch.nn as nn from mamba_ssm import Mamba class MambaCostModel(nn.Module): def __init__(self, input_dim, mamba_d_model, mamba_d_state): super().__init__() # 编码器 self.encoder = nn.Sequential( nn.Linear(input_dim, 64), nn.SiLU(), nn.Linear(64, 128), nn.SiLU(), nn.Linear(128, mamba_d_model), nn.SiLU() ) # Mamba 块 (替代Transformer) self.mamba = Mamba( d_model=mamba_d_model, # 例如 128 d_state=mamba_d_state, # 状态维度，例如 16 d_conv=4, # 卷积核大小 expand=2 # 扩展因子 ) self.norm = nn.RMSNorm(mamba_d_model) # 解码器 self.decoder = nn.Sequential( nn.Linear(mamba_d_model, 64), nn.SiLU(), nn.Linear(64, 32), nn.SiLU(), nn.Linear(32, 1) # 输出一个预测值 ) def forward(self, x): # x: [batch_size, seq_len, feature_dim] x = self.encoder(x) x, _ = self.mamba(x) # Mamba处理序列 x = self.norm(x) # 可以取序列最后一个位置的输出，或做池化 x = x.mean(dim=1) # 全局平均池化 prediction = self.decoder(x) return prediction.squeeze(-1)

训练要点：

1. 输入准备：你的输入数据是三维张量[batch_size, seq_len, feature_dim]。seq_len是调度序列的长度（需填充到固定长度），feature_dim是特征维度（调度特征+硬件特征）。
2. 标签归一化：真实的程序耗时（latency）差异可能巨大，从微秒到毫秒。直接回归会导致模型被大值主导。建议使用对数缩放或分位数归一化。
3. LambdaRank实现：LambdaRank损失需要你实现一个lambda_weight的计算。你可以借助开源库（如lightgbm中的lambda rank目标函数思路），或者自己实现。核心是根据程序真实性能的排序，计算每一对程序的ΔNDCG和σ（排序敏感度参数）。
4. 评估指标：不要只看预测值和真实值的MSE或MAE。更重要的是看排序相关性，例如斯皮尔曼等级相关系数（Spearman's ρ），或者Top-K命中率（模型预测的前K个程序，有多少个在真实排序中也位于前K）。

3.3 持续知识蒸馏框架的整合

这是系统集成的难点。你需要维护两个模型实例：KB和AC。它们结构相同，但参数独立。

训练流程伪代码：

# 假设我们已经有一个在硬件A上训练好的KB knowledge_base = MambaCostModel(...) knowledge_base.load_state_dict(torch.load('kb_hardware_A.pth')) # 为新硬件B创建活跃列AC，其参数随机初始化 active_column = MambaCostModel(...) # 阶段1: 持续学习 (CL) - 在新硬件B的数据上训练AC def continual_learning_phase(ac, kb, data_b): # 冻结KB的参数 for param in kb.parameters(): param.requires_grad = False # AC的参数可训练 ac.train() # 定义融合层（公式7中的W, U, V, alpha等） # 在AC的前向传播中，每一层不仅接收自己上一层的输出，还接收KB对应层的输出，并通过门控融合 # 损失函数仅为在新数据上的排序损失 L_Rank optimizer = torch.optim.Adam(ac.parameters(), lr=1e-4) # ... 训练循环 ... return trained_ac # 阶段2: 知识蒸馏 (KD) - 将AC的知识蒸馏到KB中 def knowledge_distillation_phase(ac, kb, data_old=None): # 冻结AC的参数 for param in ac.parameters(): param.requires_grad = False # 解冻KB的参数，但要为EWC做准备 for param in kb.parameters(): param.requires_grad = True # EWC需要计算KB在旧任务上重要参数的Fisher信息矩阵（通常在CL阶段前计算并保存） # 定义蒸馏损失：L_KD = MSE(ac_features, kb_features) + λ * L_ewc # L_ewc 是EWC惩罚项，防止重要参数偏离太多 optimizer = torch.optim.Adam(kb.parameters(), lr=1e-5) # 使用更小的学习率 # ... 训练循环 ... return updated_kb # 交替执行 active_column = continual_learning_phase(active_column, knowledge_base, data_hardware_B) knowledge_base = knowledge_distillation_phase(active_column, knowledge_base, data_hardware_A) # 传入旧数据用于EWC计算

实操心得：EWC的实现需要仔细。你需要在开始学习新硬件B之前，在旧硬件A的数据上跑一遍KB，计算每个参数的重要性（Fisher信息）。这个矩阵会很大（参数数量级），通常只对角线上元素，或者只对某些层的参数应用EWC约束，以节省内存。λ是控制“不忘旧知识”力度的超参数，需要调优。

4. 效果评估与对比分析

论文中进行了详尽的实验，这里我提炼关键结论，并补充一些工程视角的解读。

实验设置：

• 硬件：Intel i7-12700F CPU, NVIDIA GeForce RTX 3080Ti GPU。
• 数据集：基于Tenset，并自行扩展收集的大规模张量程序数据集。
• 基线模型：Tenset-MLP（经典MLP方法）、TLP（基于Transformer的SOTA方法）。
• 评估指标：
  1. 调优时间加速比：达到相同或更优模型推理延迟，所需调优时间与基线相比的倍数。
  2. 推理延迟降低：最终生成的程序，其实际运行时间与基线方法生成程序相比的降低比例。
  3. 模型大小与训练开销：参数量、训练时间、内存占用。

核心结果：

深入分析：

1. RDU采样器的效率：实验显示，仅使用10%的数据，TCL训练出的成本模型，其预测排序质量能达到使用100%数据的95%以上。这意味着数据收集成本直接降了一个数量级。这对于需要支持众多边缘设备或新型AI芯片的编译器团队来说，是巨大的福音。
2. Mamba的成本效益：相比TLP的Transformer，Mamba模型在保持相当预测精度的同时，训练速度提升了3-5倍，推理速度提升更明显。这使得将成本模型集成到编译器的实时搜索循环中变得更加可行。
3. CKD的扩展性：论文模拟了从硬件A->B->C的顺序学习过程。TCL的KB在学完C后，在A和B上的性能几乎没有衰减（EWC的作用），同时在C上达到了接近单独为C训练一个模型的性能。这验证了其持续学习和抗遗忘的能力。

5. 潜在挑战与未来方向

尽管TCL展示了强大的潜力，但在实际工业部署中，我们仍需关注以下几个挑战：

1. 特征工程的通用性：当前硬件感知特征是针对CPU/GPU设计的。对于更特殊的硬件（如NPU、存算一体芯片），需要设计新的、能反映其计算架构特点的特征。这需要深入的硬件架构知识。
2. 超参数调优：RDU中三个指标的权重、Mamba模型的结构参数（状态维度、扩展因子）、CKD中EWC的惩罚系数λ等，都需要针对不同的硬件平台和算子集进行调优。这可能会带来一定的工程负担。
3. 对极端稀疏和动态形状的支持：当前研究大多基于静态形状、稠密算子。现实中的模型可能包含动态形状、稀疏计算。如何让成本模型适应这些情况，是一个开放问题。
4. 与编译器搜索算法的更深集成：TCL目前主要改进成本模型。未来的工作可以将RDU的主动学习思想与搜索算法本身结合，实现“学习”与“搜索”的闭环联动，动态决定下一步探索哪个区域，可能取得更好效果。

从我个人的工程经验来看，TCL框架最大的价值在于提供了一套系统性的方法论。它告诉我们，编译器优化不是一个单纯的搜索问题，而是一个数据、模型、知识联合优化的问题。对于工业界团队，可以分阶段采纳：

• 初级阶段：可以先实现RDU采样器，用于大幅降低在现有硬件上的数据收集成本。
• 中级阶段：引入Mamba成本模型，替换掉原有的XGBoost或笨重的Transformer模型，提升预测效率和精度。
• 高级阶段：当需要支持多种硬件家族时，再引入持续知识蒸馏框架，构建可积累、可迁移的优化知识中台。

这个框架的模块化设计，使得我们可以根据自身需求和资源，灵活地选用和组合这些技术，逐步构建起智能、高效的深度学习编译基础设施。