Patent Description:
Compiler optimization techniques aim to make executing a computer program more efficient, which can mean running in a shorter amount of time or requiring fewer computing resources, to name just two examples. Compilers often rely on hand-tuned heuristics that indicate under what conditions certain optimizations, e.g., loop unrolling, should be performed. In this sense, optimization means applying transformations to how a program operates, but does not imply that the result will be optimal or provide the best possible solution.

Machine learning (ML) frameworks have their own compiler optimization considerations. Many ML frameworks use compilers that operate on computer programs that define a graph structure that will be executed on a distributed system having many computers. In mapping a computational graph to machine code that executes on a collection of devices, ML compilers often need to solve many optimization problems, including graph rewriting, assignment of operations to devices, operation fusion, layout and tiling of tensors, and scheduling.

ML compilers typically apply heuristics to solve these problems individually, an approach that suffers from two key limitations. First, these heuristics often lead to sub-optimal configurations, especially for previously unseen model architectures. Second, by solving these problems in isolation, the compiler misses out on opportunities for joint optimizations across tasks.

Because of the complexity of such machine learning models, and the computing resources required to train such models, which compiler optimizations are employed and how can have a dramatic impact on the run time of both the training and inference processes.

"<NPL>) discloses an efficient end-to-end method of automated device placement based on a scalable sequential attention mechanism over a graph neural network that is transferable to new graphs. To further reduce the computation cost, the policy network is pretrained on a set of dataflow graphs and a superposition network used to fine-tune it on each individual graph.

This specification describes techniques for using a compiler optimization network to learn compiler graph optimizations that are generalizable to unseen computation graphs. The compiler optimization network is a machine learning model having an architecture as described in more detail below.

A compiler optimization network as described in this specification can learn policies for a variety of compiler optimizations, including device assignment, scheduling, and operation fusion. These techniques are particularly useful for optimizing computer programs that themselves implement machine learning models because many machine learning models have or can be converted into a computational graph structure.

In order to make the resulting network generalizable, the network can be trained jointly over a set of computation graphs, which can be computation graphs from different machine learning domains. For example, recurrent networks typically have very different operation types and connection patterns compared to convolutional networks. The techniques described in this specification can be used to jointly train a compiler optimization network using input graphs from both domains. The resulting the resulting network is therefore generalizable to new computation graphs.

One example implementation of a compiler optimization model includes an end-to-end deep reinforcement learning (RL) model, where the learned policy is generalizable to new graphs and transferable across multiple tasks. Specifically, the example RL model includes an inductive graph-embedding network that encodes operation features and dependencies in a trainable graph representation, followed by a policy network of segmented recurrent attention layers. The policy network transforms the graph representations into an optimization decision with soft attention. These two networks can be jointly trained in an end-to-end fashion using a supervised reward. To generalize to arbitrary and held-out graphs, the RL model can be trained jointly over a set of computation graphs (instead of one at a time) and then fine-tuned on new graphs. By transferring the learned graph embeddings and optimization policies, the RL model can converge faster using fewer resources.

On empirical data, the resulting compiler optimization networks can outperform state-of-the-art compiler optimizations in terms of the resulting run time of the input programs. In addition, they generalize well to unseen graphs.

The training process tends to converge faster as well because the compiler optimization network can generate decisions on the entire graph in a single-shot fashion rather than operating on each individual node or group of nodes in the graph. The training process has been demonstrated in real-world experiments to converge 15x faster, or more, than other state-of-the-art techniques.

The techniques described in this specification can be implemented using a system with computer programs on one or more computers in one or more locations that performs compiler optimizations for a machine learning task configured to process a network input.

The machine learning task can be any appropriate machine learning task that (i) operates on a network input that is an input sequence, (ii) generates a network output that is an output sequence, or (iii) both. In some implementations, the machine learning task is defined by a computation graph.

Some examples of machine learning tasks that the system can be configured to perform compiler optimizations for include the following examples.

As one example, if the input to the neural network is a sequence of text in one language, the output generated by the neural network may be a score for each of a set of pieces of text in another language, with each score representing an estimated likelihood that the piece of text in the other language is a proper translation of the input text into the other language.

As another example, the task may be an audio processing task. For example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network may be a score for each of a set of pieces of text, each score representing an estimated likelihood that the piece of text is the correct transcript for the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can indicate whether a particular word or phrase ("hotword") was spoken in the utterance. As another example, if the input to the neural network is a sequence representing a spoken utterance, the output generated by the neural network can identify the natural language in which the utterance was spoken.

As another example, the task can be a natural language processing or understanding task, e.g., an entailment task, a paraphrase task, a textual similarity task, a sentiment task, a sentence completion task, a grammaticality task, and so on, that operates on a sequence of text in some natural language.

As another example, the task can be a text to speech task, where the input is text in a natural language or features of text in a natural language and the network output is a spectrogram, a waveform, or other data defining audio of the text being spoken in the natural language.

As another example, the task can be a health prediction task, where the input is a sequence derived from electronic health record data for a patient and the output is a prediction that is relevant to the future health of the patient, e.g., a predicted treatment that should be prescribed to the patient, the likelihood that an adverse health event will occur to the patient, or a predicted diagnosis for the patient. Such electronic health record data may comprise physiological data relating to a patient, such as blood glucose, blood pressure, body temperature, heartrate or the like. Examples of adverse health events comprise hypo- and/or hyper-glycemic events, heart attacks, strokes or the like.

As another example, the task can be a text generation task, where the input is a sequence of text, and the output is another sequence of text, e.g., a completion of the input sequence of text, a response to a question posed in the input sequence, or a sequence of text that is about a topic specified by the first sequence of text. As another example, the input to the text generation task can be an input other than text, e.g., an image, and the output sequence can be text that describes the input.

As another example, the task can be an image generation task, where the input is a conditioning input and the output is a sequence of intensity value inputs for the pixels of an image.

As another example, the task can be an agent control task, where the input is a sequence of observations or other data characterizing states of an environment and the output defines an action to be performed by the agent in response to the most recent data in the sequence. The agent can be, e.g., a real-world or simulated robot, a control system for an industrial facility, or a control system that controls a different kind of agent. Examples of observations or other data characterizing states of an environment comprise sensor data related to the agent, such as visual data, LIDAR readings, sonar readings, or the like. The output defines an action to be performed by the agent may comprise control signals for controlling the agent to perform the action, or data that may be converted to such signals.

As another example, the task can be a genomics task, where the input is a sequence representing a fragment of a DNA sequence or other molecule sequence and the output is either an embedding of the fragment for use in a downstream task, e.g., by making use of an unsupervised learning technique on a data set of DNA sequence fragments, or an output for the downstream task. Examples of downstream tasks include promoter site prediction, methylation analysis, predicting functional effects of non-coding variants, and so on.

As mentioned above, the compiler optimization network can be trained on a graph for one of these machine learning task examples or jointly on multiple graphs for multiple types of these machine learning task examples.

The techniques described in this specification enable the creation of a compiler optimization network that outperforms state-of-the-art methods in terms of the resulting runtime of the input programs. In particular, the techniques described below outperform a network having only a graph embedding network in placement performance. In addition, the techniques described below are faster than processing large sequence lengths compared to networks having LSTM policy networks instead of attention layers. Lastly, the techniques described below handle large graphs much better than graph attention networks.

When input programs have better optimizations, a system can learn better and more sophisticated models in less time. In addition, the techniques described below generalize well to unseen graphs, and the training process converges faster than known methods. Furthermore, the optimizations generated by the network can be output in a one-shot fashion. The compiler optimization network can also be augmented to jointly solve multiple optimization problems without incurring significantly more parameters or training time. In addition, unlike other existing approaches, the techniques described below are scalable to very large computation graphs that include <NUM>, <NUM>, or <NUM> or more computation nodes.

This specification describes techniques for training a compiler optimization network. These techniques can be used to train the network to learn how to utilize a number of different candidate computational graph optimizations. The following example graph optimization problems will be discussed in detail.

A first example optimization is device placement. Given a computational graph, the goal of device placement is to learn a policy that assigns a device D for all nodes in a given graph G, to maximize a reward rG,D, which can be defined based on run time of the resulting optimized graph. Such devices may include, for example, CPU cores, GPUs, TPUs, FGPAs or the like.

A second example optimization is operation scheduling. An operation ("op") in a dataflow graph is ready to run when its incoming tensors are present in the device memory. A frequently used scheduling strategy is to maintain a ready queue of operations for each device and schedule operations in first-in-first-out order. However, such schedules can be sub-optimal especially when the operations are distributed across a set of devices. Sub-optimal schedules typically exhibit underutilized devices since ops for these devices are blocked on producer ops in ready queues of other devices. We propose a priority-based scheduler where the scheduler maintains a priority queue of ready operations for each device and schedules the highest priority operation in the queue first. Similar to device placement, operation scheduling can be formulated as the problem of learning a policy that assigns a scheduling priority P for all ops in the graph to maximize a reward rG,P , which can also be defined based on run time of the resulting optimized graph.

A third example optimization is operation fusion ("op fusion"). Op fusion is the process of merging multiple ops into a single op.

<FIG> illustrates an example of an op fusion optimization. In <FIG>, op K1 <NUM> is producing an output which is consumed by op K2 <NUM>. If these two ops are fused into a fused kernel <NUM>, the intermediate data produced by K1 <NUM> is immediately used for K2 <NUM> when K1 <NUM> finishes on the device, without the need to perform read and write transactions with the global memory, thereby reducing the communication overhead.

<FIG> illustrates an example of how the decision of which ops to fuse can change application performance. In this case, the graph includes an element-wise multiplication (Mul), a reduction (Reduce), and a sigmoid function (Sigmoid) connected to each other as shown.

Should an optimization process choose Reduce and Sigmoid for fusion (left), the performance will not improve much, since the amount of intermediate values transferred to or from the memory will not change significantly.

On the other hand, if an optimization process selects Mul and Reduce (right) to be fused, the intermediate tensor after the multiplication will potentially stay in faster memory for the Reduce op. Therefore, the amount of transferred data from or to global memory has decreased dramatically. Simple strategies like considering ops in topological order can make inefficient fusion decisions and lead to suboptimal performance.

To address this, op fusion can be reformulated as a priority-based ordering problem: learning a policy that assigns a fusion priority F for all ops in the graph to maximize a reward rG,F defined based on run time. Each node can be associated with a fusion priority score F, where the size of F can be a hyperparameter of the network. The scores can then be used to determine the order in which nodes are fused.

Another advantage of the techniques described in this specification is that instead of being limited to learning the features of a single graph, the objective for the compiler optimization network is to simultaneously reduce the expected run time of the optimization over multiple dataflow graphs: <MAT> where T ∈ {D, P, F} denotes the task, θ denotes the parameters of the RL policy, and <IMG> denotes an empirical distribution of dataflow graphs.

In addition, these techniques can be extended to handle multiple tasks jointly. The multi-task objective can be expressed as: <MAT> Parameters across the three tasks can be partially shared: θT = (φ,ψT), where T ∈ {D, P, F}, φ denotes the shared parameters and φT denotes task-specific parameters. The shared policies can be parameterized in the form of multiple recurrent attention layers.

<FIG> is a diagram that illustrates an example architecture for a compiler optimization network <NUM>. The architecture is an example of a network that can be used to learn how to solve a number of different computational graph optimization problems in a way that generalizes to unseen graphs.

The compiler optimization network <NUM> includes a graph embedding network <NUM> followed by a policy network <NUM>, which can be jointly trained in an end-to-end fashion. The graph embedding network <NUM> can learn a graph representation hG of a computational graph, and the decision policy <NUM> can learn a optimization strategy p(y I G) over a given graph representation.

The goal for the compiler optimization network can to compute an action distribution for a node in the graph based on the actions of all previous nodes in an auto-regressive manner, which can be expressed as: <MAT> where yi is the action for the ith node.

If N is very large, e.g., more than <NUM>, <NUM>, or <NUM> nodes, which is a common situation for modem computational graphs, it may be infeasible to compute all the yi's sequentially.

Therefore, an iterative and non-autoregression process can be used as an approximation, which can be expressed as: <MAT>.

The N sampling procedures can now be carried out in parallel within each iteration t. Decisions over the N nodes are allowed to mutually influence each other because the process above can be repeated for T times, where T << N. The distribution of y(t) can be informed by y(t-<NUM>), the actions made over all the nodes in the previous iteration.

In terms of network implementation, the temporal dependencies between actions in multiple RL steps can be carried out by a state embedding <NUM> such that the state of the previous step is used as an input to the network in the next step. The actions are then produced by the final layer <NUM> in <FIG>. In addition, the network can include one or more Transformer layers with residual connections as task heads.

The architecture of the compiler optimization network <NUM> is designed to be invariant over the underlying graph topology, enabling the same learned policy to be applied to a wide range of input graphs. In some implementations, the compiler optimization network <NUM> can optimize the objective function above described in Section <NUM> using a Proximal Policy Optimization for improved sample efficiency.

The graph embedding network <NUM> is designed to capture topological information encoded in the computation graph and can be implemented, for example, using a graph neural network (GNN). One example of a suitable graph neural network is GraphSAGE, which is an inductive network that leverages node attribute information to generalize to previously unseen data. In addition, the system can use the feature aggregation scheme proposed in GraphSAGE and to a general, end-to-end graph optimization method for a wide range of computation graphs.

Nodes and edges in the computation graph can be represented as the concatenation of their meta features, e.g., operation type and output shape. The nodes and edges are then further encoded by the graph embedding network into a trainable representation. The graph embedding process can take multiple iterations, and the computation procedure for the /-th iteration can proceed as follows. First, each node v∈V aggregates the feature representations of its neighbors <MAT>, into a single vector <MAT>.

This aggregation outcome is a function of all previously generated representations, including the initial representations defined based on the input node features. The aggregation function can use the following form with max pooling: <MAT> where (W(l),b(l)) define an affine transform and σ stands for the sigmoid activation function. The node's current representation, <MAT>, is then concatenated with the aggregated neighborhood vector, <MAT>, and the concatenated vector is processed through a fully connected layer f(l+<NUM>).

Unlike other graph neural networks including GraphSAGE, parameters in the graph embedding network <NUM> can be trained jointly with a decision network via stochastic gradient descent with PPO, in a supervised fashion. In addition, when enhanced with a global attention network discussed below, the optimal layers are smaller.

In graph optimization problems, the optimal actions for a node are often influenced by actions for other nodes in the graph. While a graph neural network is a suitable feature aggregation network, GNNs often lack the ability to track global node dependencies in a scalable fashion. In contrast, an attention network can learn this dependency and the relative importance of dependencies across an entire graph.

In addition to tracking dependencies in a scalable fashion, handling large computational graphs in the compilation stacks is another practical concern when designing the compiler optimization network, which should be able to handle computational graphs from realistic workloads having over <NUM>,<NUM> nodes.

LSTM-based models proposed for language tasks usually target a shorter sequence length, incurring vanishing (and exploding) gradients or substantially longer training time. Although hierarchical grouping has been used to address longer sequences in a LSTM-based network, the proposed grouper network comes with limited flexibility and generality. And the non-differentiable grouping procedure prevents training the networks end-to-end.

Instead, a Transformer-based attentive network can be used to generate the optimization decision in an end-to-end fashion as shown in <FIG>. Because the graph embedding already includes topological information for each node, the positional embedding can be removed in the original transformer to prevent the model from overfilling node identifications.

In order to capture long-term dependencies efficiently among a large set of nodes, segment-level recurrence introduced in Transformer-XL can be used, where hidden states computed for the previous set of nodes are cached with gradient flows disabled and reused as an extended context during the training of the next segment. Besides achieving extra long context, the segment-level recurrent attention is much faster than a LSTM-based method.

One of the key goals of this work is to ensure the generalizability of the network over a wide variety of graphs from potentially different application domains (e.g. computer vision, language, and speech). Not only do these graphs vary in the number of operations from a few thousand to a million, but they also have drastically different network architectures, in terms of computational operations, data shape, and network topology. As an example, recurrent networks have completely different operation types and connections compared to convolutional networks that are widely used in computer vision. A naive strategy of training the shared policy network with batches of heterogeneous graphs is unlikely to perform as well as networks exclusively trained for a particular type of graph.

To overcome this limitation, we propose a feature modulation mechanism similar to parameter superposition. The idea is to re-weight network parameters by generating a feature modulation layer based on the input graph features as shown in <FIG>, to mitigate the potentially undesirable interference among different input graphs. The feature modulation layer is dot multiplied with all intermediate feature maps in the decision network: <MAT> where a(l) stands for an attention layer in the network, m stands for the feature modulation layer, and hG is the feature representation of the input graph generated by the graph-embedding network. The feature modulation layer can be implemented with minimum overhead by adding an additional transformer layer to the decision network. The output of the modulation layer can then be defined as the last-layer features of the additional transformer.

<FIG> illustrates how a compiler optimization network can be augmented to jointly solve multiple optimization problems. The generality of the compiler optimization network presented above allows jointly solving multiple dependent graph optimization problems without introducing significant parameters or training time. A recurrent attention policy network can be used that not only applies a segment-level recurrent attention to the graph representation spatially, but also generates recurrent actions for multiple optimization problems through residual connections and parameter sharing across multiple recurrent attention layers.

In <FIG>, the network <NUM> includes the policy network <NUM> described above that is augmented by three additional recurrent attention layers <NUM>, <NUM>, and <NUM>. In this example, the first recurrent attention layer <NUM> is for the fusion optimization problem, the second recurrent attention layer <NUM> is for the placement optimization problem, and the third recurrent attention layer <NUM> is for the scheduling optimization problem.

Compared to the single-problem policy net, <FIG> shows that the multi-problem network introduces just one recurrent attention layer for each additional optimization problem to be solved, and the parameters are shared across the different optimization problems. The recurrent attention layers with residual connections of actions enables tracking inter-problem dependencies.

<FIG> illustrates a multi-problem policy network including a spatial recurrent attention that attends to the input sequence (the recurrent attention policy network) and temporal recurrent attention that attends to the problem sequence (the multi-problem recurrent attention layers). For the t-th problem, the recurrent attention layer can compute the its hidden state Ht and action At as follows: <MAT> <MAT> where LN stands for layer normalization, and MultiHeadAttn stands for a multi-head attention layer whose parameters are shared across layers. FC is a fully-connected layer that can include a rectified-linear activation between two affine transformations, applied position-wise, that projects the hidden state to the action space of the problem.

<FIG> is a flowchart of an example process for using a compiler optimization network to generate an optimization plan for an input program. The process will be described as being performed by a system of one or more computers in one or more locations and programmed appropriately in accordance with this specification.

The system receives an input program that defines a graph of operations (<NUM>). As described above, a computation graph has nodes and edges with node in the graph representing an operation, which can be implemented in software by a particular operation module. Each edge between nodes in the graph representing one operation receiving the output generated by another operation.

The system can then provide the input program to a compiler optimization network having a graph-embedding network and a policy network. As described above, the compiler optimization network can be trained for a single computation graph or for multiple different computation graph representing different respective computing tasks, which for example, can be programs for training machine learning models, possibly from different machine-learning domains. In addition, the compiler optimization network described in this specification can be trained end-to-end over both the graph-embedding network and the policy network.

The system generates a graph embedding representation of the input program (<NUM>). As described above, the system can implement a graph-embedding network that is configured to encode operation features and operation dependencies of the operation modules of the input program into a graph embedding representation. For example, the system can use the graph-embedding network to iteratively aggregate respective feature representations for one or more neighboring nodes of each node in the graph.

The system generates an optimization action for each of one or more nodes encoded in the graph embedding representation (<NUM>). The system can use a policy network that is configured to generate an optimization action for each of one or more nodes encoded in the graph embedding representation. In some implementations, the policy network includes a feature modulation layer that is configured to re-weight network parameters. In some implementations, the system can use segment-level recurrent attention to solve multiple optimization problems concurrently.

The system obtains an output optimization plan having one or more optimization actions for the input program (<NUM>). The optimization actions that are generated depends on the optimization problem being solved. For example, if the optimization problem is scheduling, the output of the compiler optimization network can be a value representing a respective scheduling priority for each of one or more nodes in the graph. If the optimization problem is device placement, the output of the compiler optimization network can be an assignment of each operation in the graph to a particular device, with each device potentially being assigned multiple nodes in the graph. And if the optimization problem is operation fusion, each node can be associated with a respective fusion priority score, which represents an order in which nodes should be fused.

The system can then scheduling the input program for execution on the one or more devices using the output optimization plan generated by the compiler optimization network. For example, the system can assign nodes to devices according to the output of the network, schedule operations according to the output of the network, or fuse nodes in an order according to the output of the network.

Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory program carrier for execution by, or to control the operation of, data processing apparatus. The computer storage medium is not, however, a propagated signal.

A computer program (which may also be referred to or described as a program, software, a software application, a module, a software module, a script, or code) can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment.

Claim 1:
A computer-implemented method comprising:
receiving (<NUM>) an input program, wherein the input program defines a graph of operation modules, wherein each node in the graph is a respective operation module, and each edge between nodes in the graph represents one operation module receiving the output generated by another operation module;
providing a representation of the input program to a compiler optimization network comprising:
a graph-embedding network that is configured to encode (<NUM>) operation features and operation dependencies of the operation modules of the input program into a graph embedding representation, and
a policy network that is configured to generate (<NUM>) an optimization action for each of one or more nodes encoded in the graph embedding representation, comprising generating recurrent actions for multiple optimization problems using residual connections and parameter sharing across multiple recurrent attention layers.; and
obtaining (<NUM>), from the compiler optimization network, an output optimization plan comprising one or more optimization actions for the input program.