GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2 (arxiv.org)

GraphMend is a high-level compiler that automatically removes FX graph breaks in PyTorch 2 programs by transforming source code before execution. PyTorch 2’s TorchDynamo and TorchInductor enable JIT-style compilation, but dynamic control flow and certain Python constructs (notably I/O) still fragment models into multiple FX graphs, triggering fallbacks to eager execution, costly CPU↔GPU synchronizations, and lost optimization opportunities. Built on the Jac compilation framework, GraphMend implements two targeted code transformations that rewrite control-flow and Python I/O patterns so the PyTorch compilation pipeline can capture larger, uninterrupted FX graphs without manual refactoring. On eight Hugging Face models GraphMend eliminated all fixable graph breaks for dynamic control flow and Python I/O (zero breaks in six models, and reduced breaks from five to two in another). On NVIDIA RTX 3090 and A40 GPUs the tool yielded up to 75% latency reduction and up to 8% higher end-to-end throughput. The work shows that high-level static transformations—applied ahead of TorchDynamo/TorchInductor—are a practical complement to dynamic JIT approaches, improving both usability and runtime performance for PyTorch 2 workloads by enabling more aggressive kernel fusion and fewer runtime fallbacks.