Neural Topological Ordering for Computation Graphs
Anonymous Author(s)
AffiliationAddress email
Abstract
Recent works on machine learning for combinatorial optimization have shown thatlearning based approaches can outperform heuristic methods in terms of speed andperformance.
In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimizationproblem which arises in compilers.
We propose an end-to-end machine learningbased approach for topological ordering using an encoder-decoder framework.
Our encoder is a novel attention based graph neural network architecture called
Topoformer which uses different topological transforms of a DAG for messagepassing.
The node embeddings produced by the encoder are converted into nodepriorities which are used by the decoder to generate a probability distribution overtopological orders.
We train our model on a dataset of synthetically generatedgraphs 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 ofreal-world computation graphs, showing performance improvements.
Introduction
Many problems in computer science amount to finding the best sequence of objects consistent withsome precedence constraints.
An intuitive example comes from routing problems, where we wouldlike to find the shortest route between cities but we have requirements (i e. for example to pick upand subsequently deliver a package) on the order in which the cities should be visited [1].
Anothercase is found in compiler pipelines wherein the "cities" become operations to be executed and theconstraints come from the data dependencies between these operations, such as when the result ofan operation is an operand in a subsequent one.
In this case, the metric to be optimized can be therun time of the compiled program, or the memory required to execute the program [2].
Commonacross this class of problems is their formulation in term of finding the optimal topological orderof the Directed Acyclic Graph (DAG) that encodes the precedence constraints, which induces a
Combinatorial Optimization [3] (CO) problem which is in general computationally hard [4].
Already from the two examples above, one can immediately grasp the relevance of such problemsfor industrial Operations Research, which has prompted various actors to invest in the developmentof efficient CO solvers; these solvers usually encapsulate heuristic methods whose design typicallyrequires extensive use of domain-specific and problem-specific knowledge, across decades of development.
In recent years, considerable interest has emerged in the possibility of replacing suchhandcrafted heuristics with ones learned by deep neural nets [5] (machine learning for combinatorialoptimization, MLCO).
As a matter of fact, both of our two examples of DAG-based CO problemshave indirectly been object of study in the Machine Learning literature.
References [6, 7, 8, 9] takeinto consideration Routing Problems, especially the Traveling Salesperson Problem (TSP) which, onaccount of its richness, complexity and long history of mathematical study [10], has attained the statusof a standard benchmark for MLCO [8].
Conversely, less attention has been devoted to operations sequencing likely due to the proprietary and sensitive nature of compiler workflows, which hampersthe definition of public benchmarks.
References [11, 12] both consider the task of optimizing therun time of a neural network’s forward pass by optimizing the ordering and device assignment ofits required operations.
However, in this last case the sequencing stage is only one part of a largercompiler pipeline, and as a result of this both the performance metrics and the datasets employedcannot be made available for reproduction by third parties.
This makes it both hard to assess theresults therein, and to draw general conclusions and guidelines for the advancement of MLCO, whichstill suffers from a lack of commonly accepted and standard datasets and benchmarks.
In this work, we address the problem of finding optimal topological orders in a DAG using deeplearning, focusing on the compiler task of optimizing the peak local memory usage during execution.
We make the following contributions:
• We present a neural framework to optimize sequences on direct acyclic graphs.
Mindful of
the need for scalability, we consider a non-auto-regressive (NAR) scheme for parametrizingthe probability distribution of topological orders.
This allows our method to attain anextremely favorable performance vs. run time tradeoff: it always outperforms fast baselines,and is only matched or outperformed by those requiring a much longer (in one case 4000x
more) run time.
• We address the problem of how to perform meaningful message-passing on DAGs, a graphtype which has received comparatively less attention in the literature on Graph NeuralNetworks.
We introduce Topoformer, a flexible, attention-based architecture wherein
messages can be passed between each and every pair of nodes, with a different set oflearnable parameters depending on the topological relation between these.
• To test our method, we introduce an algorithm for the generation of synthetic , layered,
Neural Net-like computation graphs, allowing any researcher to generate a dataset of asmany as desired graphs of any desired size .
These graphs are a more faithful model ofreal NN workflows, and allow us to prove our method on a much larger and varied dataset,
than previous
efforts [11].
To our knowledge, this is the first public algorithm of this kind.
Nevertheless, we also test our method on proprietary graphs to illustrate its relevance torealistic compiler workflows.
Related work
Machine Learning for Combinatorial optimization Combinatorial optimization as a use case fordeep learning poses interesting technical challenges.
First, the combinatorial nature of the problemconflicts with the differentiable structure of modern deep neural networks; and second, the modelsneed to be run at large scale to solve real world instances, exacerbating the challenges in trainingdeep learning models.
Given the discrete nature of CO problems, a natural approach is to pose them as reinforcementlearning (RL) problems [13].
The aim is then to learn a policy that selects the best actions tomaximize a reward directly related to the optimization objective.
Algorithms then differ in the waythe policy is parameterized: either in an end-to-end manner where the actions directly correspond tosolutions of the optimization problem [12, 6, 8, 14], or in a hybrid manner, where the policy augmentsparts of a traditional solver, e.g. by replacing heuristics used in setting parameters of an algorithm, seee.g.
Our approach follows an end-to-end design philosophy, which, not having to relyon an external algorithm, affords better control of post-compile run time and facilitates applicationon edge devices [2].
Furthermore, RL has the advantage of being useful as a black box optimizer,when no handcrafted heuristics can be designed.
Sequence optimization via ML
Within MLCO, much effort has been devoted to the task ofpredicting optimal sequences [15, 6, 16, 17, 18].
The end-to-end nature of our method places it closeto the one proposed in [6], although to the best of our knowledge, our work is the first to tackle thechallenge of enforcing precedence constraints in the network predictions.
As we shall see in moredetail below, this generalization is non-trivial: already counting the number of topological ordersbelongs to the hardest class of computational problems [4].
This has to be contrasted with the factthat number of sequences without topological constraints is simply n ! for n objects.
Besides as pointed out in [ no MLCO method has so far been able to convincingly tackle TSPs of sizes above a few hundred nodes, when it comes to zero-shot generalization to unseen problem instances,i.e. when no fine tuning on the test set is done.
It is also therein pointed out how an auto-regressiveparametrization of the sequence (which was the method used in ref.
[6]) appears to be necessary toachieve acceptable performance even at those small sizes.
Conversely, in the present work we showcompelling zero-shot performance on DAGs of sizes up to thousands of nodes, while nonethelessgenerating our sequences in a fully non-auto-regressive (NAR) way and maintaining a strong runtime advantage over classical sequencing algorithms.
Our results can then also be interpreted ascautioning against the idea of using the TSP as the sole, paradigmatic test-bed for MLCO research,as [8] remarks.
ML for compiler optimization The DAG sequencing task we consider is an omnipresent stagein compiler workflows, which usually also include such tasks as device assignment and operationsfusion [12].
In such a setting, jointly optimizing these tasks to reduce the run time of a certainworkflow (such as the forward pass of a Neural Net) is a common objective, which in refs [11, 12, 19]is tackled with ML methods.
In this work we focus on the task of minimizing the peak local memoryusage during execution, which does not require a performance model or simulator as well as beingrelevant to applications on edge devices [2].
In [11], the ML solution leans on an existing geneticalgorithm, whilst our solution is end-to-end, much like that proposed in [12].
Another characteristicof the solution proposed in [12] is the idea of interpolating between AR and NAR via an iterativerefinement scheme, in which sequences are generated in one pass but subsequently refined during anuser-defined number of subsequent passes; conversely, we generate all our sequences in a single pass.
While in [11] the run time optimization is studied on both real-world and synthetic random graphs –the latter being relatively small (up to about 200 nodes), the peak memory optimization is studied onlyon a proprietary dataset augmented via perturbation of the node attributes.
In [12] the authors trainand test their method on a relatively small set of six proprietary workflows which are not disclosed tothe reader, and out of those six, only the size of the largest instance is mentioned.
Deep Graph Neural Networks Given that our problem is specified as a DAG, it is a logical choiceto parametrize our sequence-generation policy with a Graph Neural Network architecture [20].
Thebasic idea of every GNN architecture is to update graph and edge representations by passing messagesbetween the graph nodes along the graph edges [21].
However, this can be too restrictive whenit comes to sequence generation on DAGs.
For example, nodes that come after each other in thesequence might not be linked by an edge in the graph, and therefore are unable to directly influenceeach other’s representation.
Notice how this difficulty is another consequence of the presence ofprecedence constraints in our problem, which conversely was not an issue in e.g. [6] where the graphis fully connected and no constraints are present.
Relatively few efforts (see e.g. [22, 23, 24]) havebeen devoted to devise a way to perform meaningful message passing on DAGs.
As a matter offact, the quest for expressive GNN architectures is at the center of intense theoretical investigation[25, 26].
Background
3.1 Topological orders and DAGs
We here introduce the mathematical background, starting with a few definitions.
A partial order is anirreflexive transitive relation < between certain pairs of a set V .
We call a pair (
x, y ) ∈
V × V that isrelated by < comparable, and incomparable otherwise.
A Directed Acyclic Graph (DAG)
G
=
( V, E )is a directed graph with no directed loops.
We can map a DAG G =
( V, E ) to a partially ordered set( V, < ) where x < y
if there is a directed path from node x to node y .
Multiple DAGs map to the samepartial order.
For example, the DAGs with vertex set { x, y, z } and edge sets E = { x → y, y → z }and E ′ =
{ x → y, y → z, x → z } , where s → t denotes a directed edge from s to t , correspondto the same partial order x < y <
z .
We define the transitive closure (TC) of a DAG as the graphwith most edges that has the same underlying partial order, so that there exists a directed edge ( x, y )whenever x < y .
Conversely, the transitive reduction (TR) is the graph with least edges that resultsin the same partial order.
We denote the order induced by a DAG by < G .
A topological order or sorting of a DAG G is a bijection σ : V → { 1 , . .
, | V |} such that σ ( x ) < σ ( y )whenever x < G y .
The set T G of topological orders of G is a subset of the permutation group of thevertices and coincides with total orders on V that respect < G , called linear extensions of the partial order.
While there are several well-known algorithms to compute a topological order of a DAG, e.g.breadth first search and depth first search, counting the number of topological orders is one of thehardest computational problems, being # P complete [4].
In this work we develop a general machinelearning method to find a topological order that minimizes a given cost function on a DAG, which wedefine in the next section.
3.2 Peak Memory Minimization
Deciding the best way to schedule operations in a computational graph representing a neural networkis a central problem in compilers [11, 12, 2].
We can associate a DAG to a computational graphin such a way that nodes represent operations ("ops"), and incoming/outgoing edges representoperands/results of these operations.
Every time one executes an op, the inputs 1 to that op need tobe in memory, and memory for the outputs needs to be allocated.
Therefore, each node of the DAGcarries a label m : V → N specifying the memory required to store the output of that op.
A typicalfirst step in scheduling a DAG is to identify topological orders to execute operations.
Compilersfor edge devices, which have limited memory, aim at choosing the optimal topological order thatminimizes the peak memory footprint [2].
We focus therefore on the peak local memory usageminimization task, which can be formulated as the following combinatorial optimization problem ona labeled DAG G =
( V, E, m ) :
[Equation]
with the definitions
[Equation]
[Equation]
[Equation]
S l and ∀ ( i
[Equation]
j
[Equation]
∈ E
[Equation]
j ∈ σ 1: t (cid:41) [Equation] [Equation] i.e. the memory usage at time t is given by the memory usage I t − 1 of the outputs which have not yetbeen consumed, at time t − 1 , by downstream operations, plus the memory requirement of the outputof operation σ ( t ) .
I t is in turn obtained by subtracting from M t the memory costs of nodes whoseoutgoing edges only connect to already scheduled nodes, i.e. nodes whose output was only requiredby already scheduled operations.
Naturally, I 0 = 0 , S 0 = ϕ .
Method
We use an encoder-decoder architecture whose schematic is shown in figure 1.
Our encoder is
Topoformer , a novel GNN architecture, which derives an embedding for each node of the graph.
Theembeddings are used by the decoder which generates a distribution in the sequence space and finallythe distribution can be converted to a sequence via different inference methods like sampling, greedyinference or beam search.
Next, we describe each of the component in detail.
4.1 Topoformer: topologically masked attention
A Graph Neural Network (GNN) is a natural choice to encode our scheduling problem via embeddingof the DAG nodes.
All canonical GNN architectures operate by updating these embeddings via theaggregation of "messages" sent from the other nodes, usually in the form of some function of theirown embedding [20].
Architectures mainly differ in how the set of sender nodes is constructed andthe aggregation function is chosen.
In a Graph Convolutional Network [27], the senders are the firstneighbors of a node and the aggregation function is a weighted average, whilst in a vanilla Graph
Attention Network [28], the senders are all the other nodes, but their contributions are aggregatedvia averaging with learned weights so as to account for their degree of relevance.
When trying toapply such mechanisms on DAGs, a common point of contention is whether, and how in practice, thepartial ordering encoded by it should reflect in the direction of travel of the messages [29, 24, 22].
While disregarding the DAG structure entirely (as one would do in a vanilla GAT), does not appearwise, it might be too restrictive when it comes to our task.
For example, nodes that are next to each other in the sequence might well be incomparable, and thus lack a path for messages between them
[Figure]
The combinatorial nature of the task also poses requirements; it is known [25, 5] that reasoningabout CO problems on a graph requires the capacity to reason about the global structure of it,whilst architectures such as those proposed in [29, 24, 22] limit the set of sender nodes to a localneighborhood of the receiver node.
In summary, our architecture must strike a compromise betweenaccounting for global structure and local partial ordering information.
Our Topoformer architecture meets these requirements.
A vector x i of input features (see the appendixfor details about its definition and dimensionality) is first turned into an initial node embedding h (0) ivia a node-wise linear transformation, h (0) i
= W x i + b .
Subsequently, a succession of L attentionlayers, each of them consisting of a Multi-Head Attention (MHA) [28] sub-layer followed by onemore node-wise MLP, updates these embeddings, similar to a vanilla Transformer [30]; however, weconfer a topological inductive bias to these updates by having a separate group of attention headsmasked by each of the following graphs induced by the original DAG:
• Its transitive reduction (TR).
• The directed graph obtained by removing the TR edges from the DAG:
G \ E TR ( G ) .
• The directed graph obtained by removing the edges of the DAG from its TC: TC ( G ) \ E .•
The backwards versions (i.e. with flipped edges) of each of the three above.
• The undirected graph obtained by joining all incomparable pairs.
By adding together these graphs, one would obtain the fully connected graph relative to the nodeset V , whereupon all nodes would attend to all nodes.
Then effectively, the propagation rules of
Topoformer are same as those of a vanilla transformer encoder,
[Equation]
save for the presence of the mask M j , which ensures that head j only attends to its assigned graphamong the seven listed above.
Following [31], we also apply layer normalization [32] to the MHAand MLP inputs.
The number of heads assigned to each graph can be chosen independently (settingit to zero means to not message-pass along the edges of the respective graph), or parameters canbe tied among different MHAs.
One should also remark how the MLP sub-layer allows the flowof information between different attention heads.
All nodes are then able to influence each other’srepresentation, while anyway injecting a strong inductive bias based on the DAG structure.
Werefer the reader to the appendix for information about the Topoformer configurations we use in ourexperiments.
4.2 Decoder
Once the embeddings of the nodes are generated, the decoder’s task is to derive a stochastic policyp ( σ | G ) over the valid topological orders of the graph.
The most straightforward way is to takeadvantage of the chain rule of conditional probability to decompose the policy as a product
[Equation]
We could then sample a complete sequence by autoregressively choosing a new node at each step asdone e.g. in [6].
This scheme is the most principled and expressive; however, when a NN is used as afunction approximator for p θ , it also requires that | V | calls to this NN be performed, which limits itsfeasibility to relatively small graphs due to the amount of computation required.
This makes it acceptable to sacrifice expressivity for run time, by employing a Non-Auto-Regressive(NAR) scheme which decouples the number of NN calls from the graph size.
Similar to the approachof [12], we assign scheduling priorities y i ∈ R to the nodes, rather than scheduling probabilities.
Thepriority for node i is derived by passing its final embedding through an MLP:
[Equation]
These priorities are assigned with a single NN inference.
The sequence itself is subsequentlyconstructed by adding a new node at each step.
Given the partial sequence σ 1: i − 1 , the next node canonly be selected from a subset S ( σ 1: i − 1 , G ) of schedulable nodes, due to both the graph topologyand choices made earlier in the sequence.
Then, the distribution of the next node to be added at step iis given as follows:
[Equation]
Decoding methods : We use the following three methods to obtain the next node in the partialsequence from the distribution p ( σ t | σ 1: t − 1 , h , G ) :
Greedy : At each step t , select the node with the highest probability i
e.
[Equation]
[Equation]
Sampling :
At each step t , sample from the next node distribution i.e.
[Equation]
[Equation] 1: t − 1 , h , G )
Beam search with state-collapsing : We can also expand the
partial sequences by using abeam search method where the score function is total probability of the partial sequence.
We improve our beam search routine by making the following observation: suppose thereare two partial sequences in consideration, σ 1: t and ˜ σ 1: t , such that both have scheduledthe same set of nodes so far (but different order), and C ( σ 1: t ) <
C (˜ σ 1: t ) .
Then, we canignore the partial sequence ˜ σ 1: t and only keep σ 1: t in the beam search.
This is because bothpartial sequences must schedule the same set of remaining nodes, and hence the set of futurememory costs are identical for both σ 1: t and ˜ σ 1: t , but the current peak memory cost is higherfor ˜ σ 1: t .
Thus, σ 1: t dominates ˜ σ 1: t in terms of achievable minimal peak memory usage.
4.3 Training
Our encoder-decoder architecture induces a distribution p θ ( σ | G ) on the set of topological orders fora given DAG G .
The expected cost incurred is given by J ( θ | G ) =
E p θ ( σ | G )
[ C ( σ ( θ ))] .
We minimizethe cost J ( θ ) =
E G
[ J ( θ | G )] via gradient descent using the REINFORCE gradient estimator [33, 13]as follows
[Equation]
where b ( G ) is a baseline meant to reduce the variance of the estimator.
We follow [6] in setting itequal to the cost of a greedy rollout of a baseline policy on the graph G
[Equation]
Experiments
We conduct experiments on a synthetic dataset of graphs which we refer to as "layered graphs", aswell as a set of real-world computation graphs.
We compare our approach with the following classictopological ordering baselines:
• Depth/Breadth first sequencing - Find the topological order by traversing the graph indepth/breadth first manner according to the layout of the graph generated using pygraphviz.• Depth-first dynamic programming (DP) - Depth-first DP is a global depth-first methodfor searching the optimal sequence, with automatic backtracking when equivalent partialsequences are found; it retains the full sequence with minimum cost so far, and returns it ifsearch does not complete before the prescribed timeout.• Approximate DP - In approximate DP, a beam of partial sequences are considered at eachstep. For each beam in the subsequent step, the top- K partial sequences with the lowest costsare retained.
This DP is also able to find the optimal sequence given enough memory andcompute resources (with the unlimited beam size K ), but we only consider an approximateversion with K = 10 5 in this work.
Note that approximate DP uses the state-collapsing andparallelism to improve its efficiency.
• Random order - We generate 100 random topological orders, and pick the one with smallestcost.
Please see the appendix for more detailed description of the baselines.
Therein, we also report someablation studies on various neural baselines including ablation studies on both components (encoderand decoder) by considering other neural architectures, as well as a comparison with an end to endbaseline adapted from ref.
Neural topo order greedy, sample and BS denote the performance ofour model in greedy, sampling and beam search inference mode respectively.
We use a sample sizeand beam size of 16 sequences, of which the best one is subsequently picked, for all our experiments.
Next, we describe in detail the results of the two experiments.
5.1 Layered graphs
In order to generate a large corpus of training data we come up with a way to synthetically generategraphs of a given size which have similar structure to the computation graphs of feed-forward neuralnetworks.
We call our synthetic graph family layered graphs , as these graphs comprise of well-definedlayers of nodes.
The nodes in a layer have connections to the nodes in the subsequent layer and canalso have skip connections with nodes in layers farther down.
The number of layers, number of nodeper layer, number of edges between subsequent layers, number of skip connections and memoryutilization of the nodes are all generated randomly, and can be controlled by setting appropriateparameters.
We refer the reader to the appendix for more details on layered graphs, including theirgeneration algorithm and some visual examples.
We train our model on 500 -node layered graphs for 325 epochs, where in each epoch we generate atraining set of 1000 new graphs.
We test the performance of our model on a set of 300 unseen graphsof the same size, generated with the same method.
We also evaluate the cross-size generalizationperformance of our trained model by testing it on graphs of size 1000 and 2000 .
We refer the readerto the appendix for a comprehensive description of the training algorithm and model configuration.
Figure 2 shows the performance vs. run time plot on layered graphs of size | V | = 500 , 1000 , and2000 .
We report the performance in terms of the % gap of peak memory utilization from the peakmemory obtained via approximate DP, which we consistently observed to be the best-performingbaseline.
Note that the run time is plotted on a log-scale.
We can observe that for 500 -node graphs,our model beats all the baselines except approximate DP in terms of both the memory usage and runtime.
Our model is slightly worse than approximate DP from the memory usage perspective, but itruns 100x faster.
We also observe that our model generalizes well to larger sized graphs.
For thecase of 2000 -node graphs our model performs better than approximate DP in terms of peak memoryusage, while being 4000 x times faster.
This shows that while approximate DP performs more poorlyas graph size increases, our model is able to generalize to larger graphs by learning meaningfulembeddings of the topological structure thanks to Topoformer, and to be extremely fast thanks to our
NAR decoding scheme.
[Table]
[Table]
[Table]
Depth
[Table]
[Table]
[Table]
[Table]
[Table]
Breadth-first seq.
[Table]
[Table]
[Table]
[Table] Neural Topo Order ✓ Greedy 4. 2.44 ✓ Sample 3. 2.87 ✓
[Table]
search 3. 5.2 Real-world graphs While our synthetic layered graphs are convenient for experimentation, we see value in also presentingresults obtained from neural computation graphs used for commercial development of our artificialintelligence hardware and software products.
Here we sample 115 representative graphs that havediverse architectures (classifiers, language processors, denoisers, etc.) and size (from a few dozen to1k nodes).
We split this dataset into a training set and test set via a random 80 − 20 split.
We trainour model for 500 epochs and report the performance on the unseen test set at the end of training intable 2.
In order to ensure fair comparison of run times, we stratify the test set into 3 categories basedon the graph size.
Figure 3 shows the performance vs run time plot on the test set of real graphs.
We observe thatfor real graphs the performance gap between the best baseline (approximate DP) and our model isremarkable.
We can obtain sequences which are 50% better than approximate DP on average whilealso being almost 1000x faster on average.
This proves the capability of our model to generalize andperform well on real-world computation workflows.
[Table]
Conclusions
In this work we propose an end-to-end machine learning method for the task of optimizing topologicalorders in a directed acyclic graph.
Two key elements in our design are: (1) an attention-based GNNarchitecture named Topoformer that employs message passing that is both global and topologicallyaware in directed acyclic graphs, (2) a non-autoregressive parametrization of the distribution ontopological orders that enables fast inference.
We demonstrated, for both synthetic and real-worldgraphs, the effectiveness of the method in tackling the problem of minimizing peak local memoryusage for a compute graph – a canonical task in compiler pipelines.
Said pipelines also includeother tasks [12], chief amongst them the one of assigning operations to devices for execution.
Atthe present stage, our method and dataset cannot be leveraged for solving these, or for end-to-endoptimization of a whole pipeline.
Extending our method to this more challenging setting is thereforea natural direction for future research.