[Paper Review] Neural Topological Ordering for Computation Graphs

Mukul Gagrani, Corrado Rainone+8 more|arXiv (Cornell University)|Jul 13, 2022

Advanced Graph Neural Networks4 citations

TL;DR

This paper proposes a novel end-to-end deep learning framework, Topoformer, for learning optimal topological orders in directed acyclic graphs (DAGs) to minimize peak memory usage in compiler workloads. Using a non-autoregressive, attention-based graph neural network with topologically-aware message passing, the method achieves state-of-the-art or competitive performance with significantly faster inference than baselines, including those requiring up to 4000x more runtime.

ABSTRACT

Recent works on machine learning for combinatorial optimization have shown that learning based approaches can outperform heuristic methods in terms of speed and performance. In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimization problem which arises in compilers. We propose an end-to-end machine learning based approach for topological ordering using an encoder-decoder framework. Our encoder is a novel attention based graph neural network architecture called \emph{Topoformer} which uses different topological transforms of a DAG for message passing. The node embeddings produced by the encoder are converted into node priorities which are used by the decoder to generate a probability distribution over topological orders. We train our model on a dataset of synthetically generated graphs called layered graphs. We show that our model outperforms, or is on-par, with several topological ordering baselines while being significantly faster on synthetic graphs with up to 2k nodes. We also train and test our model on a set of real-world computation graphs, showing performance improvements.

Motivation & Objective

• To address the combinatorial optimization problem of finding optimal topological orders in DAGs for memory minimization in compiler pipelines.
• To develop a scalable, end-to-end deep learning framework that replaces heuristic methods with learned policies for topological ordering.
• To introduce a novel graph neural network architecture, Topoformer, that enables topologically-aware, global message passing on DAGs.
• To create a publicly available synthetic dataset of layered, neural network-like computation graphs for reproducible benchmarking in MLCO.
• To demonstrate performance gains on both synthetic and real-world computation graphs while maintaining fast inference speed.

Proposed method

• The method employs a non-autoregressive (NAR) decoder to efficiently parametrize the probability distribution over topological orders, enabling fast inference.
• The encoder is Topoformer, a novel attention-based graph neural network that performs message passing using different learnable parameters based on the topological relationship (e.g., ancestor, descendant, unrelated) between node pairs.
• Topoformer processes the DAG through seven distinct topological transforms, including forward and backward edge propagation, to capture global structural dependencies.
• Node embeddings from Topoformer are converted into node priorities to guide the decoder’s generation of topological sequences.
• The model is trained on a synthetic dataset of layered computation graphs, which are designed to mimic real neural network workloads and allow scalable data generation.
• Training uses supervised learning with approximate dynamic programming (DP) solutions as targets to supervise the policy network.

Experimental results

Research questions

• RQ1 Can a learnable, end-to-end deep learning framework outperform traditional heuristic methods in minimizing peak memory usage on DAGs?
• RQ2 How effective is topologically-aware, global message passing in a GNN architecture for DAGs compared to local or edge-only message passing?
• RQ3 Can a non-autoregressive decoding scheme achieve competitive performance with drastically reduced inference time compared to autoregressive baselines?
• RQ4 To what extent can a synthetic, publicly available dataset of layered DAGs enable reproducible benchmarking in machine learning for combinatorial optimization?
• RQ5 Does the proposed method generalize to real-world compiler workloads beyond synthetic graphs?

Key findings

• On 500-node synthetic graphs, the full Topoformer model achieves a mean % gap of 3.08 ± 0.51 from approximate DP using beam search, outperforming GAT (forward+backward) and other baselines.
• On 1000-node graphs, Topoformer with beam search achieves a % gap of 0.05 ± 0.36, significantly outperforming GAT (forward+backward) and other models.
• The full Topoformer model reduces the % gap from approximate DP by 1.5–2.5 percentage points compared to its forward+backward variant, demonstrating the benefit of global message passing.
• Topoformer achieves state-of-the-art performance with a mean inference time of 1.62s for beam search on 500-node graphs, while outperforming baselines that require up to 4000x more time.
• On real-world computation graphs, the model shows consistent performance improvements over strong baselines, validating its relevance to industrial compiler workflows.
• The proposed synthetic graph generation algorithm enables scalable, reproducible benchmarking and is publicly available for the research community.