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副作用引起的图断裂：使6个模型的断裂计数降为0，并将另一模型的断裂数从5减少至2。在NVIDIA RTX 3090和A40 GPU上，GRAPHMEND最高可实现75%的延迟缩减与8%的端到端吞吐量提升。这些结果表明，高层代码变换是PyTorch动态JIT编译管线的有效补充，显著提升了可用性与性能。