由大型语言模型（LLMs）驱动的生成式人工智能（AI）的兴起正在通过实现半自动代码生成来深刻改变开发者的编程方式。虽然这促进了软件的创建，但将代码片段（如计算内核）优化并映射到特定硬件目标的问题仍然是一个未解决的挑战。实际上，现代计算系统通常具有异构架构，包含多种硬件单元，如CPU、GPU（图形处理单元）和TPU（张量处理单元）。这些设备可以加速某些类型的操作，但传统的编译器启发式方法可能难以适应各种可能的映射选项，从而无法充分利用所有平台的潜力[1]。此外，只有在严格探索设备配置参数后，才能在给定计算单元上实现最佳性能，这进一步增加了这项任务的复杂性[2]、[3]。
这一挑战可以通过人工智能（AI）来解决，其中创建的模型旨在对源代码进行分类[4]，从而输出关于代码片段如何最佳映射到计算资源的决策[5]。值得注意的是，语言模型（LMs）和其他自然语言处理（NLP）技术经常因其能够识别与其指定任务相关的有用模式而被应用。有关NLP在软件工程领域应用的全面回顾，请参考Zhang等人的研究[6]。在源代码领域，这些模式可能包含了有关良好软件工程实践的知识，其获取有助于开发更智能、数据驱动的软件工程工具。
基于NLP的技术将源代码文件视为语言制品，这是基于自然语言和编程语言之间的相似性[5]。然而，源代码本身携带了与数据（例如变量和常量）和指令之间的因果和时间依赖关系相关的模式，这些模式在序列格式中难以识别。
事实上，最近的研究表明，广泛用于分析自然语言的大型语言模型（LLMs）在处理源代码优化和分析时可能效果不佳[7]、[8]。目前尚不清楚为源代码的LLMs设计的现有训练方法是否能够有效利用从这种数据类型中获得的知识[9]。相反，基于图的表示方法可以增强对高级代码概念的理解[10]。几项最近的研究加强了这种方法的合理性。在[4]中，作者展示了通过基于树的神经模型结合程序结构信息可以显著提高分类准确性。Wu等人和Jia等人的补充发现[11]、[12]强调了图神经架构在提取各种结构化数据集中的复杂关系模式方面的有效性。来自基于知识图的分类工作的类似证据进一步证实了学习基于图的表示方法的好处。因此，基于图的LM可能能够识别更有效的模式并利用所有可用信息。
近年来，深度学习社区从端到端模型转向了迁移学习技术[13]。在这种方法中，模型首先在一个大型数据集上进行预训练，然后针对特定任务进行微调。这种方法将问题解决方案的学习与数据中固有模式的提取区分开来，使得在预训练期间获得的一般知识可以在多个任务中重复使用[14]。
遵循这一趋势，我们旨在通过采用针对图神经网络（GNNs）定制的自监督预训练技术来实证验证基于图的学习方法的有效性，并测量预训练模型在几项下游任务中带来的改进：（i）将计算内核映射到计算单元（异构设备映射）；（ii）确定GPU内核的最有效线程块大小（线程块大小预测）；（iii）根据程序解决的问题对其进行分类（算法分类）。
为此，我们开发了“DeepCodeGraph”（DCG），这是一种用于构建基于图的源代码表示的通用语言模型的程序（见图1）。我们方法的关键组成部分是对掩码图自编码（MGAE）框架的改编。该技术涉及从输入图中随机选择一些元素并进行屏蔽，然后利用剩余信息让模型重建这些元素。尽管MGAE在其他场景中是一种流行的技术[15]，但据我们所知，先前的研究尚未探索将其作为源代码分析语言模型预训练方法的应用。
DeepCodeGraph的性能优于之前的基于图的技术和为代码优化设计的预训练程序，在上述下游任务上取得了业界领先的结果。特别是在异构设备映射任务上，我们的平均准确率超过了87%，比之前的基于图的模型提高了3.2%。此外，在线程块大小预测任务上，我们的准确率提高了4.9%，根据LS-CAT基准测试的性能指标达到了0.997。
本文的结构如下：第2节回顾了现有文献中与源代码分析和预训练技术相关的主要问题和解决方案。第3节描述了我们的贡献，包括一个基于图的LM、一个受MGAE领域最新方法启发的自监督学习框架，以及一个允许进行稳健预训练的大规模数据集。第4节检验了预训练和微调的结果，并评估了我们方法带来的好处。最后，第5节总结了这项工作并指出了未来的研究方向。
