[论文解读] ProGraML: Graph-based Deep Learning for Program Optimization and Analysis

Chris Cummins, Zacharias V. Fisches+3 more|arXiv (Cornell University)|Mar 23, 2020

Software Engineering Research参考文献 55被引用 24

一句话总结

ProGraML 引入了一种基于图的程序表示方法，该方法源自编译器中间表示（IR），将控制、数据和调用依赖关系表示为有向带属性的多重图。在该结构上应用消息传递神经网络（MPNN），其在核心编译器分析任务上实现了最先进的性能——F1得分为94.0，并在异构设备映射和程序分类任务上优于先前方法。

ABSTRACT

The increasing complexity of computing systems places a tremendous burden on optimizing compilers, requiring ever more accurate and aggressive optimizations. Machine learning offers significant benefits for constructing optimization heuristics but there remains a gap between what state-of-the-art methods achieve and the performance of an optimal heuristic. Closing this gap requires improvements in two key areas: a representation that accurately captures the semantics of programs, and a model architecture with sufficient expressiveness to reason about this representation. We introduce ProGraML - Program Graphs for Machine Learning - a novel graph-based program representation using a low level, language agnostic, and portable format; and machine learning models capable of performing complex downstream tasks over these graphs. The ProGraML representation is a directed attributed multigraph that captures control, data, and call relations, and summarizes instruction and operand types and ordering. Message Passing Neural Networks propagate information through this structured representation, enabling whole-program or per-vertex classification tasks. ProGraML provides a general-purpose program representation that equips learnable models to perform the types of program analysis that are fundamental to optimization. To this end, we evaluate the performance of our approach first on a suite of traditional compiler analysis tasks: control flow reachability, dominator trees, data dependencies, variable liveness, and common subexpression detection. On a benchmark dataset of 250k LLVM-IR files covering six source programming languages, ProGraML achieves an average 94.0 F1 score, significantly outperforming the state-of-the-art approaches. We then apply our approach to two high-level tasks - heterogeneous device mapping and program classification - setting new state-of-the-art performance in both.

研究动机与目标

• 解决机器学习用于程序分析时，序列化和非结构化程序表示方法的局限性。
• 通过提升表示的保真度和模型的表达能力，弥合基于机器学习的启发式方法与最优编译器优化之间的性能差距。
• 开发一种与编译器无关、与语言无关的程序表示方法，以捕捉丰富的语义结构，供深度学习模型使用。
• 利用图神经网络实现对复杂程序分析（如活跃变量分析、支配树和公共子表达式检测）的端到端学习。
• 在高层次任务（如异构设备映射和程序分类）上展示该方法的泛化能力，实现新的最先进性能。

提出的方法

• 从LLVM和XLA IR中构建ProGraML，将其表示为有向带属性的多重图，编码指令之间的控制、数据和调用依赖关系。
• 通过丰富的特征（包括操作码、数据类型、修饰符和顺序）表示每条指令和操作数，以保留语义信息。
• 应用消息传递神经网络（MPNN）在图中传播潜在表示，支持对每条指令和每个变量的分类。
• 在全程序图上端到端训练模型，以执行下游任务，如控制流可达性分析、数据依赖检测和活跃变量分析。
• 采用统一的图表示方式，同时支持低级IR和高级程序理解任务，实现跨领域的迁移学习。
• 利用ProGraML与传统编译器IR之间的结构相似性，确保与现有编译流水线和工具的兼容性。

实验结果

研究问题

• RQ1 基于编译器IR构建的图表示方法，是否能在学习程序语义方面优于基于序列或手工特征的方法？
• RQ2 消息传递神经网络在无需手工设计启发式规则的情况下，能在多大程度上近似经典编译器分析（如支配树和活跃变量分析）？
• RQ3 与现有的基于图的表示方法（如CDFG和IR2Vec）相比，ProGraML在建模程序行为方面的表达能力如何？
• RQ4 ProGraML能否泛化到高层次优化任务（如异构设备映射和程序分类），并实现最先进性能？
• RQ5 与先前方法相比，ProGraML在真实编译器工作负载中的可扩展性和推理效率如何？

主要发现

• 在涵盖六种编程语言的25万个LLVM-IR文件基准测试中，ProGraML在传统编译器分析任务上实现了94.0的F1得分，显著优于当前最先进方法。
• 该模型在变量活跃性分析和公共子表达式检测任务上表现更优，而此前方法因未能充分表达操作数语义而表现不佳。
• ProGraML在异构设备映射任务上创下新的最先进性能，表明基于图的学习方法能有效推理数据局部性和执行目标。
• 与同类基于CDFG的方法相比，该方法的推理吞吐量提高了76.6倍，表明其在计算复杂度相近的情况下仍具备出色的实用可扩展性。
• 图表示能够准确建模大型代码库中长距离依赖关系和语义关联，克服了序列化模型的局限性。
• 该框架在多种编程语言和优化任务中均表现出良好泛化能力，表明其应用范围远超特定编译器优化阶段。