This paper presents GRAPHMEND, a high-level compiler technique that eliminates FX graph breaks in PyTorch 2 programs. Although PyTorch 2 introduced TorchDynamo and TorchInductor to enable just-in-time graph compilation, unresolved dynamic control flow and unsupported Python constructs often fragment models into multiple FX graphs. These fragments force frequent fallbacks to eager mode, introduce costly CPU-to-GPU synchronizations, and reduce optimization opportunities. GRAPHMEND addresses this limitation by analyzing and transforming source code before execution. Built on the Jaseci compilation framework, GRAPHMEND introduces two code transformations that remove graph breaks due to dynamic control flow and Python side effects. This design allows PyTorch's compilation pipeline to capture larger, uninterrupted FX graphs without requiring manual refactoring by developers. Evaluation across eight Hugging Face models shows that GRAPHMEND removes graph breaks due to dynamic control flow and Python side effects, reducing the break count to 0 in 6 models and reducing it from 5 to 2 in another model. On NVIDIA RTX 3090 and A40 GPUs, GRAPHMEND achieves up to 75% latency reductions and up to 8% higher end-to-end throughput. These results demonstrate that high-level code transformation is an effective complement to PyTorch's dynamic JIT compilation pipeline, substantially improving both usability and performance.

翻译：本文提出 GRAPHMEND，一种用于消除 PyTorch 2 程序中 FX 图中断的高级编译技术。尽管 PyTorch 2 引入了 TorchDynamo 和 TorchInductor 以支持即时图编译，但未解决的动态控制流及不支持的 Python 结构仍常将模型分割为多个 FX 图。这些片段迫使系统频繁回退至即时执行模式，引入昂贵的 CPU 至 GPU 同步开销，并减少优化机会。GRAPHMEND 通过在执行前分析并转换源代码来解决此限制。该技术基于 Jaseci 编译框架构建，引入了两种代码转换方法，以消除由动态控制流和 Python 副作用导致的图中断。此设计使得 PyTorch 的编译流水线能够捕获更大且连续的 FX 图，无需开发者手动重构代码。在八个 Hugging Face 模型上的评估表明，GRAPHMEND 成功消除了由动态控制流和 Python 副作用引起的图中断：在六个模型中将中断数降为 0，在另一个模型中将中断数从 5 降至 2。在 NVIDIA RTX 3090 和 A40 GPU 上，GRAPHMEND 实现了高达 75% 的延迟降低和高达 8% 的端到端吞吐量提升。这些结果表明，高级代码转换技术可有效补充 PyTorch 的动态即时编译流水线，显著提升其可用性与性能。