Swift：用于加速深度学习推理的张量运算符的高并行程序生成

《ACM Transactions on Architecture and Code Optimization》：Swift: High Parallelism Program Generation of Tensor Operators for Accelerating Deep Learning Inference

优化小批量推理中GPU资源利用的Swift编译器提出了一种搜索空间扩展方法，通过结合tiling结构和reduction并行化，显著提升了深度学习模型执行效率，实验在三种GPU架构上平均加速比达1.19倍，优于现有编译器和手工优化库2.40倍。

　　

摘要

优化深度学习推理过程，尤其是降低小批量数据下张量计算的执行延迟，对于深度神经网络（DNN）模型的成功和广泛应用至关重要。然而，当前的深度学习编译器和手工调优的库在执行小批量任务时往往无法实现高效的硬件利用。主要原因在于运算过程中的序列化特性（例如，在LLM解码过程中扁平化后的GEMM运算沿隐藏维度进行的运算），这种特性难以并行化，因此无法充分利用可用的硬件资源。

在本文中，我们提出了一种基于搜索的新方法Swift，该方法通过最大化硬件利用率来高效生成适用于GPU的高性能程序。其核心思想是将运算并行化与现有的 tile 结构相结合，从而显著扩大高性能程序的搜索空间。具体而言，我们首先枚举所有可能的循环并行映射方式，生成一个包含高性能程序的庞大搜索空间；然后通过子空间移位探索技术来识别有潜力的区域，从而有效剔除大部分潜力较低的搜索结果。

我们在三种不同的GPU架构上进行了实验，使用了能够代表典型应用场景的多种基准测试工具。实验结果表明，Swift相比现有的基于编译器的方法平均加速比达到了1.19倍；与厂商提供的人工调优库相比，Swift的平均加速比达到了2.40倍。