Tiga compiler speeds up graph learning on memory-limited devices
Tiga: Compiling Graph Message Passing at Scale
Distributed, Parallel, and Cluster Computing
Summary
Graph message passing is a way to handle complex data relationships, but running these programs can be slow and use too much memory, especially on smaller devices like laptops. This paper introduces Tiga, a tool that separates the program’s logic from how it actually processes data, allowing it to run more efficiently on CPUs and GPUs. Tiga can handle huge graphs by streaming data from disk and can work across multiple machines, all while integrating smoothly with Python and PyTorch. The authors show that Tiga reduces runtime and memory use compared to existing tools and can process huge graphs on limited hardware.
What this means in practice
- •For machine learning engineers: Run large-scale graph learning models efficiently on devices with limited memory by compiling message passing into optimized data traversal.
- •For simulation software developers: Accelerate simulation workflows involving graph-based calculations using Tiga’s fused traversal and memory management techniques on heterogeneous hardware.
Authors
Mingyuan Chi
Abstract
Graph message passing offers a common way to express learning algorithms, physical simulations, and numerical solvers. Efficient execution depends on interaction structure and data movement, which can be obscured when a program is expressed as a sequence of tensor operations. On memory-constrained systems such as laptops, materializing connectivity and intermediate messages can also exhaust device memory. We present Tiga, a just-in-time compiler that separates the definition of a message-passing program from how its interactions are traversed, computed, and stored. Tiga preserves graph relations and reducer algebra in a multi-level intermediate representation, enabling traversal specialization, fusion of relation generation with aggregation, and reverse-mode automatic differentiation. Backend-specific lowering targets CPUs and GPUs. Its runtime extends execution beyond device-memory capacity by streaming graph partitions from disk through host memory with page-sized device staging buffers; partition ownership and halo exchange extend the same programming model to distributed execution. A Python interface interoperates with ordinary PyTorch tensors and autograd for forward and backward computation. Numerical checks validate outputs and gradients for differentiable workloads. A differentiated geometric workload characterizes the memory--time tradeoff of fused forward and backward execution. Evaluation against matched Torch and PyTorch Geometric baselines demonstrates reduced runtime and device-memory use for generated-relation workloads, while offloaded forward execution processes billion-edge graphs on a single memory-limited GPU. Measurements on heterogeneous devices further characterize the communication and load-balance costs of distributed execution.