自动优化动态张量运算的执行

深度学习模型核心依赖于涉及张量（矩阵的高维类比）的代数运算，这些运算可能重复数万次。高效学习需要优化频繁重复的张量运算，但涉及不同形状（32x32、64x64、128x128等）张量的运算必须单独优化。

自动调度器是学习优化形状的程序，这些形状在当前张量运算库中的实现可能不够优化。然而，现有自动调度器难以处理形状变化的工作负载。例如，许多自然语言处理应用接受任意长度的输入，这意味着任意形状的张量。

在今年的机器学习与系统会议（MLSys）上，团队展示了一种名为DietCode的新型自动调度器，其处理动态形状工作负载的效率远超前代方案。现有自动编码器必须单独优化每个可能形状，而DietCode构建了形状通用搜索空间，使其能够同时优化所有可能形状。

在自然语言处理（NLP）任务测试中，输入大小范围从1到128个token。当使用反映真实世界分布的随机输入大小采样时，优化过程速度比最佳现有自动调度器快近六倍。考虑所有可能形状时，加速比提高到94倍以上。

尽管速度大幅提升，DietCode还将生成代码的性能提高了70%（相对于先前自动调度器）和19%（相对于现有张量运算库中手动优化的代码）。这有望加速客户的动态形状机器学习工作负载。

动态工作负载

处理任意长度文本字符串的NLP模型是动态设计模型的示例，允许可变大小的输入。但其他应用也需要动态工作负载。例如，神经架构搜索通过不同形状组件构建尝试不同深度学习架构，这需要不同形状张量的运算。某些模型（如BERT语言模型）在网络不同层应用相同运算，这些层具有不同数量的节点。

微内核

自动调度器通常依赖计算内核（程序模板），其使用大大加速不同候选优化的评估速度。然而，奇异形状工作负载可能无法精确匹配内核。例如，如果张量在某个维度有513个元素，但内核容量仅为512，则必须将两个内核拼接以容纳张量。

但拼接内核在相关维度上的组合容量为1,024，而输入张量仅为513。因此必须对输入张量进行填充以填满内核。这种填充会显著减慢优化过程，因为它导致不必要的计算，随后必须从结果中修剪。

DietCode使用根据可用硬件（而非输入形状）调整大小的微内核，这有助于针对该硬件进行优化。对于给定硬件配置，DietCode还可以生成一系列不同的微内核形状和大小，这些可以组合使用。

微内核足够小，通常可以平铺在输入上以更精确地适应其形状。这可能仍然需要在边缘进行一些填充，但比大型内核需要的填充少得多。

然而，微内核的真正优势是使DietCode能够同时优化多个形状的运算符。标准自动调度器接受工作负载形状，根据需要填充以适应其平铺内核，然后使用提取程序特征（如循环结构和内存访问模式）的成本模型估计不同实现的效率。然后对下一个形状重复该过程。

相比之下，DietCode将运算符分解到微内核上。成本模型有两个组件：一个评估分配给每个微内核的部分运算特征，另一个评估将这些部分运算拼接在一起形成完整运算符的成本。

传统自动编码器（左）分别优化不同形状工作负载的张量运算实现。DietCode（右）同时优化多个工作负载的实现，节省时间并提高性能。

在这里实现了最大的效率提升，因为每个部分运算都是多个工作负载形状运算符的组成部分。与评估运算的计算成本（涉及实际硬件测量的机器学习过程）相比，将部分运算拼接在一起的成本较低。

借助优化的微内核，训练高效的决策树模型将工作负载形状映射到微内核。该决策树被纳入执行张量运算的二进制文件中，以将任意形状的输入路由到适当的微内核进行处理。

有关实验结果和更多细节，请参阅相关论文。

致谢：Cody Yu, Yizhi Liu, Gennady Pekhimenko