Original: R EINFORCED G ENETIC A LGORITHM L EARNING FOR O PTIMIZING C OMPUTATION G RAPHS
Revised: R EINFORCED G ENETIC A LGORITHM L EARNING FOR O PTIMIZING C OMPUTATION G RAPHS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler.
Revised: We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler.

Original: Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer ofﬂine and then generalizes to previously unseen graphs without further training.
Revised: Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer ofﬂine and then generalizes to previously unseen graphs without further training.

Original: This allows our approach to produce high-quality execution decisions on real-world
Revised: This allows our approach to produce high-quality execution decisions on real-world

Original: TensorFlow graphs in seconds instead of hours.
Revised: TensorFlow graphs in seconds instead of hours.

Original: We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage.
Revised: We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage.

Original: In comparison to an extensive set of baselines, our approach achieves signiﬁcant improvements over classical and other learning-based methods on these two tasks.
Revised: In comparison to an extensive set of baselines, our approach achieves signiﬁcant improvements over classical and other learning-based methods on these two tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep Learning frameworks such as MXNet (Chen et al., 2015), PyTorch (Paszke et al., 2017), and TensorFlow (TensorFlow Authors, 2016a) represent neural network models as computation graphs.
Revised: Deep Learning frameworks such as MXNet (Chen et al., 2015), PyTorch (Paszke et al., 2017), and TensorFlow (TensorFlow Authors, 2016a) represent neural network models as computation graphs.

Original: Efﬁciently executing such graphs requires optimizing discrete decisions about how to map the computations in a graph onto hardware so as to minimize a relevant cost metric (e.g., running time, peak memory).
Revised: Efﬁciently executing such graphs requires optimizing discrete decisions about how to map the computations in a graph onto hardware so as to minimize a relevant cost metric (e.g., running time, peak memory).

Original: Given that execution efﬁciency is critical for the success of neural networks, there is growing interest in the use of optimizing static compilers for neural network computation graphs, such as Glow (Rotem et al., 2018), MLIR (MLIR Authors, 2018), TVM (Chen et al., 2018a), and XLA (XLA team, 2017).
Revised: Given that execution efﬁciency is critical for the success of neural networks, there is growing interest in the use of optimizing static compilers for neural network computation graphs, such as Glow (Rotem et al., 2018), MLIR (MLIR Authors, 2018), TVM (Chen et al., 2018a), and XLA (XLA team, 2017).

Original: Here we consider the model parallelism setting where a computation graph can be executed using multiple devices in parallel.
Revised: Here we consider the model parallelism setting where a computation graph can be executed using multiple devices in parallel.

Original: Nodes of the graph are computational tasks, and directed edges denote dependencies between them.
Revised: Nodes of the graph are computational tasks, and directed edges denote dependencies between them.

Original: We consider here jointly optimizing over placement , i.e., which nodes are executed on which devices, and schedule , i.e., the node execution order on each device.
Revised: We consider here jointly optimizing over placement , i.e., which nodes are executed on which devices, and schedule , i.e., the node execution order on each device.

Original: These decisions are typically made in either one or two passes in the compiler.
Revised: These decisions are typically made in either one or two passes in the compiler.

Original: We consider two different objectives: 1) minimize running time, subject to not exceeding device memory limits, and 2) minimize peak memory usage.
Revised: We consider two different objectives: 1) minimize running time, subject to not exceeding device memory limits, and 2) minimize peak memory usage.

Original: In the optimization literature, such problems are studied under the class of task scheduling , which is known to be NP-hard in typical settings (Sinnen, 2007; Kwok & Ahmad, 1999).
Revised: In the optimization literature, such problems are studied under the class of task scheduling , which is known to be NP-hard in typical settings (Sinnen, 2007; Kwok & Ahmad, 1999).

Original: As scheduling and placement are just a few of the many complex decisions made in a compiler, it is essential in a production setting that a solution 1) produce solutions of acceptable quality fast, even on large graphs (e.g., thousands of nodes) and decision spaces, and 2) handle diverse graphs from various types of applications, neural network architectures, and users.
Revised: As scheduling and placement are just a few of the many complex decisions made in a compiler, it is essential in a production setting that a solution 1) produce solutions of acceptable quality fast, even on large graphs (e.g., thousands of nodes) and decision spaces, and 2) handle diverse graphs from various types of applications, neural network architectures, and users.

Original: In this work we consider learning an optimizer that satisﬁes these requirements.
Revised: In this work we consider learning an optimizer that satisﬁes these requirements.

Original: Crucially, we aim to learn an optimizer that generalizes to a broad set of previously unseen computation graphs, without the need for training on such graphs, thus allowing it to be fast at test time.
Revised: Crucially, we aim to learn an optimizer that generalizes to a broad set of previously unseen computation graphs, without the need for training on such graphs, thus allowing it to be fast at test time.

Original: Previous works on learning to optimize model parallelism decisions (Mirhoseini et al., 2017; 2018;
Revised: Previous works on learning to optimize model parallelism decisions (Mirhoseini et al., 2017; 2018;

Original: Addanki et al., 2019) have not considered generalization to a broad set of graphs nor joint optimization of placement and scheduling.
Revised: Addanki et al., 2019) have not considered generalization to a broad set of graphs nor joint optimization of placement and scheduling.

Original: In Mirhoseini et al.
Revised: In Mirhoseini et al.

Original: (2017; 2018), learning is done from scratch for each computation graph and for placement decisions only, requiring hours (e.g., 12 to 27 hours per graph).
Revised: (2017; 2018), learning is done from scratch for each computation graph and for placement decisions only, requiring hours (e.g., 12 to 27 hours per graph).

Original: This is too slow to be broadly useful in a general-purpose production compiler.
Revised: This is too slow to be broadly useful in a general-purpose production compiler.

Original: We propose an approach that takes only seconds to optimize similar graphs.
Revised: We propose an approach that takes only seconds to optimize similar graphs.

Original: In concurrent work to ours, Addanki et al.
Revised: In concurrent work to ours, Addanki et al.

Original: (2019) shows generalization to unseen graphs, but they are generated artiﬁcially by architecture search for a single learning task and dataset.
Revised: (2019) shows generalization to unseen graphs, but they are generated artiﬁcially by architecture search for a single learning task and dataset.

Original: In contrast, we collect real user-deﬁned graphs spanning a broad set of tasks, architectures, and datasets.
Revised: In contrast, we collect real user-deﬁned graphs spanning a broad set of tasks, architectures, and datasets.

Original: In addition, both Mirhoseini et al.
Revised: In addition, both Mirhoseini et al.

Original: (2017; 2018) and
Revised: (2017; 2018) and

Original: [Figure]
Revised: [Figure]

Original: BRKGA requires proposal distributions for each node in the graph to generate candidate solutions in its search loop.
Revised: BRKGA requires proposal distributions for each node in the graph to generate candidate solutions in its search loop.

Original: The default choice is agnostic to the input graph: uniform distribution over [0 , 1] at all nodes.
Revised: The default choice is agnostic to the input graph: uniform distribution over [0 , 1] at all nodes.

Original: We use a graph neural network policy to predict node-speciﬁc non-uniform proposal distribution choices (parameterized as beta distributions over [0 , 1] ).
Revised: We use a graph neural network policy to predict node-speciﬁc non-uniform proposal distribution choices (parameterized as beta distributions over [0 , 1] ).

Original: BRKGA is then run with those choices and outputs the best solution found by its iteration limit.
Revised: BRKGA is then run with those choices and outputs the best solution found by its iteration limit.

Original: By controlling the non-uniformity of the distributions, the policy directs how BRKGA’s search effort is allocated such that a better solution can be found with the same search budget.
Revised: By controlling the non-uniformity of the distributions, the policy directs how BRKGA’s search effort is allocated such that a better solution can be found with the same search budget.

Original: Addanki et al.
Revised: Addanki et al.

Original: (2019) consider only placement decisions and rely on TensorFlow’s dynamic scheduler; they do not address the static compiler setting where it is natural to jointly optimize scheduling and placement.
Revised: (2019) consider only placement decisions and rely on TensorFlow’s dynamic scheduler; they do not address the static compiler setting where it is natural to jointly optimize scheduling and placement.

Original: The key idea of our approach (Figure 1) is to learn a neural network that, conditioned on the input graph to be optimized, directs an existing optimization algorithm’s search such that it ﬁnds a better solution in the same search budget.
Revised: The key idea of our approach (Figure 1) is to learn a neural network that, conditioned on the input graph to be optimized, directs an existing optimization algorithm’s search such that it ﬁnds a better solution in the same search budget.

Original: We choose the Biased Random-Key Genetic Algorithm (BRKGA (Gonçalves & Resende, 2011)) as the optimization algorithm after an extensive evaluation of several choices showed that it gives by far the best speed-vs-quality trade-off for our application.
Revised: We choose the Biased Random-Key Genetic Algorithm (BRKGA (Gonçalves & Resende, 2011)) as the optimization algorithm after an extensive evaluation of several choices showed that it gives by far the best speed-vs-quality trade-off for our application.

Original: BRKGA produces good solutions in just a few seconds even for real-world TensorFlow graphs with thousands of nodes, and we use learning to improve the solution quality signiﬁcantly at similar speed.
Revised: BRKGA produces good solutions in just a few seconds even for real-world TensorFlow graphs with thousands of nodes, and we use learning to improve the solution quality signiﬁcantly at similar speed.

Original: We train a graph neural network (Battaglia et al., 2018) to take a computation graph as input and output node-speciﬁc proposal distributions to use in the mutant generation step of BRKGA’s inner loop.
Revised: We train a graph neural network (Battaglia et al., 2018) to take a computation graph as input and output node-speciﬁc proposal distributions to use in the mutant generation step of BRKGA’s inner loop.

Original: BRKGA is then run to completion with those input-dependent distribution choices, instead of inputagnostic default choices, to compute execution decisions.
Revised: BRKGA is then run to completion with those input-dependent distribution choices, instead of inputagnostic default choices, to compute execution decisions.

Original: The distributions are predicted at each node, resulting in a high-dimensional prediction problem.
Revised: The distributions are predicted at each node, resulting in a high-dimensional prediction problem.

Original: There is no explicit supervision available, so we use the objective value as a reward signal in a contextual bandit approach with REINFORCE (Williams, 1992).
Revised: There is no explicit supervision available, so we use the objective value as a reward signal in a contextual bandit approach with REINFORCE (Williams, 1992).

Original: Our approach, “Reinforced Genetic Algorithm Learning” (REGAL), uses the network’s ability to generalize to new graphs to signiﬁcantly improve the solution quality of the genetic algorithm for the same objective evaluation budget.
Revised: Our approach, “Reinforced Genetic Algorithm Learning” (REGAL), uses the network’s ability to generalize to new graphs to signiﬁcantly improve the solution quality of the genetic algorithm for the same objective evaluation budget.

Original: We follow the static compiler approach of constructing a coarse static cost model to evaluate execution decisions and optimizing them with respect to it, as done in (Addanki et al., 2018; Jia et al., 2018).
Revised: We follow the static compiler approach of constructing a coarse static cost model to evaluate execution decisions and optimizing them with respect to it, as done in (Addanki et al., 2018; Jia et al., 2018).

Original: This is in contrast to evaluating the cost by executing the computation graph on hardware (Mirhoseini et al., 2017; 2018).
Revised: This is in contrast to evaluating the cost by executing the computation graph on hardware (Mirhoseini et al., 2017; 2018).

Original: A computationally cheap cost model enables fast optimization.
Revised: A computationally cheap cost model enables fast optimization.

Original: It is also better suited for distributed training of RL policies since a cost model is cheap to replicate in parallel actors, while hardware environments are not.
Revised: It is also better suited for distributed training of RL policies since a cost model is cheap to replicate in parallel actors, while hardware environments are not.

Original: Our cost model corresponds to classical NP-hard scheduling problems, so optimizing it is difﬁcult.
Revised: Our cost model corresponds to classical NP-hard scheduling problems, so optimizing it is difﬁcult.

Original: In this paper we focus fully on learning to optimize this cost model, leaving integration with a compiler for future work.
Revised: In this paper we focus fully on learning to optimize this cost model, leaving integration with a compiler for future work.

Original: We structure the neural network’s task as predicting proposal distributions to use in the search over execution decisions, rather than the decisions themselves directly.
Revised: We structure the neural network’s task as predicting proposal distributions to use in the search over execution decisions, rather than the decisions themselves directly.

Original: Empirically we have found the direct prediction approach to be too slow at inference time for our application and generalizes poorly.
Revised: Empirically we have found the direct prediction approach to be too slow at inference time for our application and generalizes poorly.

Original: Our approach potentially allows the network to learn a more abstract policy not directly tied to detailed decisions that are speciﬁc to particular graphs, which may generalize better to new graphs.
Revised: Our approach potentially allows the network to learn a more abstract policy not directly tied to detailed decisions that are speciﬁc to particular graphs, which may generalize better to new graphs.

Original: It can also make the learning task easier as the search may succeed even with sub-optimal proposal distribution predictions, thus smoothening the reward function and allowing the network to incrementally learn better proposals.
Revised: It can also make the learning task easier as the search may succeed even with sub-optimal proposal distribution predictions, thus smoothening the reward function and allowing the network to incrementally learn better proposals.

Original: The node-speciﬁc proposal distribution choices provide a rich set of knobs for the network to ﬂexibly direct the search.
Revised: The node-speciﬁc proposal distribution choices provide a rich set of knobs for the network to ﬂexibly direct the search.

Original: Combining learning with a search algorithm has been shown to be successful (e.g., (Silver et al., 2017; 2018)), and our work can be seen as an instance of the same high-level idea.
Revised: Combining learning with a search algorithm has been shown to be successful (e.g., (Silver et al., 2017; 2018)), and our work can be seen as an instance of the same high-level idea.

Original: This paper makes several contributions:
Revised: This paper makes several contributions:

Original: • We are the ﬁrst to demonstrate learning a policy for jointly optimizing placement and scheduling that generalizes to a broad set of real-world TensorFlow graphs.
Revised: • We are the ﬁrst to demonstrate learning a policy for jointly optimizing placement and scheduling that generalizes to a broad set of real-world TensorFlow graphs.

Original: REGAL signiﬁcantly outperforms all baseline algorithms on two separate tasks of minimizing runtime and peak memory usage (section 5.3) on datasets constructed from 372 unique real-world TensorFlow graphs, the largest dataset of its kind in the literature and at least an order of magnitude larger than the ones in previous works ( Mirhoseini et al., 2017; 2018; Chen et al., 2018b; Addanki et al., 2018; 2019).
Revised: REGAL signiﬁcantly outperforms all baseline algorithms on two separate tasks of minimizing runtime and peak memory usage (section 5.3) on datasets constructed from 372 unique real-world TensorFlow graphs, the largest dataset of its kind in the literature and at least an order of magnitude larger than the ones in previous works (Mirhoseini et al., 2017; 2018; Chen et al., 2018b; Addanki et al., 2018; 2019).

Original: • We use a graph neural network to predict mutant sampling distributions of a genetic algorithm, speciﬁcally BRKGA, for the input graph to be optimized.
Revised: • We use a graph neural network to predict mutant sampling distributions of a genetic algorithm, speciﬁcally BRKGA, for the input graph to be optimized.

Original: This directs BRKGA’s search in an input-dependent way, improving solution quality for the same search budget.
Revised: This directs BRKGA’s search in an input-dependent way, improving solution quality for the same search budget.

Original: • We compare extensively to classical optimization algorithms, such as enumerative search, local search, genetic search and other heuristics, and analyze room-for-improvement in the objective value available to be captured via learning.
Revised: • We compare extensively to classical optimization algorithms, such as enumerative search, local search, genetic search, and other heuristics, and analyze room-for-improvement in the objective value available to be captured via learning.

Original: Both are missing in previous works.
Revised: Both are missing in previous works.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Learning to optimize computation graphs: AutoTVM (Chen et al., 2018b) applies learning to the very different problem of optimizing low-level implementations of operators in a tensor program, while we focus on optimizing higher-level decisions such as placement and scheduling of ops.
Revised: Learning to optimize computation graphs: AutoTVM (Chen et al., 2018b) applies learning to the very different problem of optimizing low-level implementations of operators in a tensor program, while we focus on optimizing higher-level decisions such as placement and scheduling of ops.

Original: Mao et al.
Revised: Mao et al.

Original: (2019) use graph neural nets and RL to learn a scheduling policy for data processing jobs on clusters.
Revised: (2019) use graph neural nets and RL to learn a scheduling policy for data processing jobs on clusters.

Original: These works are conceptually similar to ours in their use of learning, applied to a different domain.
Revised: These works are conceptually similar to ours in their use of learning, applied to a different domain.

Original: Learning for combinatorial optimization: Our work is an instance of applying learning for combinatorial optimization (Bengio et al., 2018).
Revised: Learning for combinatorial optimization: Our work is an instance of applying learning for combinatorial optimization (Bengio et al., 2018).

Original: Previous works on learning graph combinatorial optimization algorithms (e.g., Li et al.
Revised: Previous works on learning graph combinatorial optimization algorithms (e.g., Li et al.

Original: ; Khalil et al.
Revised: ; Khalil et al.

Original: (2017)) have focused on problems such as Minimum
Revised: (2017)) have focused on problems such as Minimum

Original: Vertex Cover, Maximum Clique, Maximum Independent Set, etc.
Revised: Vertex Cover, Maximum Clique, Maximum Independent Set, etc.

Original: The task scheduling problem we consider is signiﬁcantly different in that the objective value is a more complex function on node-level decisions.
Revised: The task scheduling problem we consider is signiﬁcantly different in that the objective value is a more complex function on node-level decisions.

Original: Also, we focus on large-scale, real-world TensorFlow graphs, while e.g., Khalil et al.
Revised: Also, we focus on large-scale, real-world TensorFlow graphs, while e.g., Khalil et al.

Original: (2017) uses small-scale, synthetic graph distributions.
Revised: (2017) uses small-scale, synthetic graph distributions.

Original: Learning a proposal distribution for stochastic search: Bunel et al.
Revised: Learning a proposal distribution for stochastic search: Bunel et al.

Original: (2017) learns a policy for predicting instance-dependent proposal distributions to be used in the stochastic optimizer STOKE (Schkufza et al., 2013) for superoptimizing programs.
Revised: (2017) learns a policy for predicting instance-dependent proposal distributions to be used in the stochastic optimizer STOKE (Schkufza et al., 2013) for superoptimizing programs.

Original: However, it uses handcrafted instance features and shows results on relatively simple, small programs.
Revised: However, it uses handcrafted instance features and shows results on relatively simple, small programs.

Original: In contrast, we automatically learn the instance representations and show results on real-world graphs.
Revised: In contrast, we automatically learn the instance representations and show results on real-world graphs.

Original: An earlier work by Paige &
Revised: An earlier work by Paige &

Original: Wood (2016) similarly learns a neural network to predict input-dependent proposal distributions for sequential Monte Carlo search for inference in a graphical model.
Revised: Wood (2016) similarly learns a neural network to predict input-dependent proposal distributions for sequential Monte Carlo search for inference in a graphical model.

Original: Optimization without learning: Parallel task scheduling (Sinnen, 2007; Kwok & Ahmad, 1999) is a classical problem for scheduling ops in a computational graph to minimize runtime.
Revised: Optimization without learning: Parallel task scheduling (Sinnen, 2007; Kwok & Ahmad, 1999) is a classical problem for scheduling ops in a computational graph to minimize runtime.

Original: Learning is not traditionally a part of the approaches proposed in this literature.
Revised: Learning is not traditionally a part of the approaches proposed in this literature.

Revised: Mayer et al.

Revised: (2017) studies greedy task scheduling approaches for TensorFlow.

Original: Jia et al.
Revised: Jia et al.

Original: (2018) develops a simulation-based optimizer for deep learning computation graphs that uses a larger decision space by combining data, model, and attribute parallelism.
Revised: (2018) develops a simulation-based optimizer for deep learning computation graphs that uses a larger decision space by combining data, model, and attribute parallelism.

Original: Our approach can potentially be extended to such larger decisions spaces to achieve even bigger improvements in execution cost.
Revised: Our approach can potentially be extended to such larger decisions spaces to achieve even bigger improvements in execution cost.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Figure 1 shows an overview of our approach.
Revised: Figure 1 shows an overview of our approach.

Original: Given an input graph to optimize, instead of applying
Revised: Given an input graph to optimize, instead of applying

Original: BRKGA directly with the default uniform distribution at all nodes, a graph neural network predicts beta distribution choices at each node.
Revised: BRKGA directly with the default uniform distribution at all nodes, a graph neural network predicts beta distribution choices at each node.

Original: BRKGA is run with these choices to optimize placement and scheduling decisions with respect to the objective deﬁned by the performance model.
Revised: BRKGA is run with these choices to optimize placement and scheduling decisions with respect to the objective deﬁned by the performance model.

Original: We ﬁrst explain the performance model and BRKGA in this section, and the learning component in the next.
Revised: We ﬁrst explain the performance model and BRKGA in this section, and the learning component in the next.

Original: 3. P ERFORMANCE MODEL
Revised: 3. P ERFORMANCE MODEL

Original: A computation graph has a set of ops to run.
Revised: A computation graph has a set of ops to run.

Original: Each op produces zero or more tensors and requires zero or more tensors as input.
Revised: Each op produces zero or more tensors and requires zero or more tensors as input.

Original: The runtime of each op is known and ﬁxed (e.g., given by a simulator as in Jia et al.
Revised: The runtime of each op is known and ﬁxed (e.g., given by a simulator as in Jia et al.

Original: The memory use of each tensor is known (an assumption that holds in static compilers like XLA).
Revised: The memory use of each tensor is known (an assumption that holds in static compilers like XLA).

Original: We assume a collection of d homogeneous devices that have separate local memory and can run at most one op at a time.
Revised: We assume a collection of d homogeneous devices that have separate local memory and can run at most one op at a time.

Original: An op can run only when its input tensors are present in the local memory.
Revised: An op can run only when its input tensors are present in the local memory.

Original: Tensors can be transferred across devices by synchronous (blocking) transfers.
Revised: Tensors can be transferred across devices by synchronous (blocking) transfers.

Original: Tensors are freed from local memory after all local consumers have run.
Revised: Tensors are freed from local memory after all local consumers have run.

Original: In this setting, we consider the problem of ﬁnding an assignment of ops to devices and an overall schedule such that each op is run once with the objectives of (1) minimizing the peak local memory use across devices (e.g., to ﬁnd a feasible way to run a large computation graph), or (2) minimizing the runtime subject to a constraint on the peak memory used on any device.
Revised: In this setting, we consider the problem of ﬁnding an assignment of ops to devices and an overall schedule such that each op is run once with the objectives of (1) minimizing the peak local memory use across devices (e.g., to ﬁnd a feasible way to run a large computation graph), or (2) minimizing the runtime subject to a constraint on the peak memory used on any device.

Original: The performance model does not consider rematerialization of tensors, fragmentation when computing memory use, and asynchronous transfers between devices.
Revised: The performance model does not consider rematerialization of tensors, fragmentation when computing memory use, and asynchronous transfers between devices.

Original: Despite these simpliﬁcations, the model yields slight variants of problems that are known to be NP-hard (Eyraud-Dubois et al., 2015) and therefore remains a challenging setting in which to study how to learn an optimizer.
Revised: Despite these simpliﬁcations, the model yields slight variants of problems that are known to be NP-hard (Eyraud-Dubois et al., 2015) and therefore remains a challenging setting in which to study how to learn an optimizer.

Original: See section A. for more details of the model.
Revised: See section A. for more details of the model.

Original: 3. B IASED RANDOM - KEY GENETIC ALGORITHM
Revised: 3. B IASED RANDOM - KEY GENETIC ALGORITHM

Original: Biased random-key genetic algorithm (BRKGA) is a meta-heuristic framework that has been successful in a wide array of applications for solving hard combinatorial optimization problems (Gonçalves & Resende, 2011).
Revised: Biased random-key genetic algorithm (BRKGA) is a meta-heuristic framework that has been successful in a wide array of applications for solving hard combinatorial optimization problems (Gonçalves & Resende, 2011).

Original: In BRKGA, chromosomes in a population are encoded as n -dimensional vectors with entries in [0 , 1] for some ﬁxed n .
Revised: In BRKGA, chromosomes in a population are encoded as n -dimensional vectors with entries in [0 , 1] for some ﬁxed n .

Original: This random-key encoding decouples the application from the genetic algorithm, speciﬁcally the crossover and mutant generation procedures (Bean, 1994).
Revised: This random-key encoding decouples the application from the genetic algorithm, speciﬁcally the crossover and mutant generation procedures (Bean, 1994).

Original: The BRKGA variant we use is speciﬁed by (1) a ﬁtness evaluation function f : [0 , 1] n → R , (2) scalar integer parameters π , π e , and π c representing the population size, number of elites, and number of children, resp.
Revised: The BRKGA variant we use is speciﬁed by (1) a ﬁtness evaluation function f : [0 , 1] n → R , (2) scalar integer parameters π , π e , and π c representing the population size, number of elites, and number of children, resp.

Original: , (3) an elite bias ρ ∈ [0 . 5 , 1 . 0) , and (4) a mutant generation distribution D over [0 , 1] n .
Revised: , (3) an elite bias ρ ∈ [0 . 5 , 1 . 0) , and (4) a mutant generation distribution D over [0 , 1] n .

Original: The procedure aims to ﬁnd a chromosome that maximizes f .
Revised: The procedure aims to ﬁnd a chromosome that maximizes f .

Original: The initial population (a collection of π chromosomes) is created by sampling from D .
Revised: The initial population (a collection of π chromosomes) is created by sampling from D .

Original: (Known good solutions may also be used to initialize a population.)
Revised: (Known good solutions may also be used to initialize a population.)

Original: One evolution step is completed as follows.
Revised: One evolution step is completed as follows.

Original: Sort the chromosomes in order of decreasing ﬁtness using f .
Revised: Sort the chromosomes in order of decreasing ﬁtness using f .

Original: Denote the ﬁrst π e chromosomes as elites and the remaining chromosomes as nonelites .
Revised: Denote the ﬁrst π e chromosomes as elites and the remaining chromosomes as nonelites .

Original: Construct the next generation from three different sources of chromosomes: (a) Copy the elite chromosomes unmodiﬁed from the last generation.
Revised: Construct the next generation from three different sources of chromosomes: (a) Copy the elite chromosomes unmodiﬁed from the last generation.

Original: (b) For each of the π c new children, select two parent chromosomes uniformly at random, one from the nonelites and one from the elites.
Revised: (b) For each of the π c new children, select two parent chromosomes uniformly at random, one from the nonelites and one from the elites.

Original: Apply the crossover procedure (described below) to generate a new chromosome given the two parents.
Revised: Apply the crossover procedure (described below) to generate a new chromosome given the two parents.

Original: (c) Generate the remaining π − π e − π c by sampling from D .
Revised: (c) Generate the remaining π − π e − π c by sampling from D .

Original: We continue the evolution procedure for a ﬁxed number of evaluations, i.e., calls to f .
Revised: We continue the evolution procedure for a ﬁxed number of evaluations, i.e., calls to f .

Original: Given an elite and nonelite chromosome a , b ∈ [0 , 1] n (resp.), the crossover procedure produces a child chromosome c by independently combining entries from the parents.
Revised: Given an elite and nonelite chromosome a , b ∈ [0 , 1] n (resp.), the crossover procedure produces a child chromosome c by independently combining entries from the parents.

Original: Speciﬁcally, for each index i ∈ 1 , . . .
Revised: Speciﬁcally, for each index i ∈ 1 , . . .

Original: , n independently, let c i
Revised: , n independently, let c i

Original: = a i with probability ρ
Revised: = a i with probability ρ

Original: and c i =
Revised: and c i =

Original: b i with probability 1 − ρ .
Revised: b i with probability 1 − ρ .

Original: Our use of BRKGA is standard except for the mutant-sampling distribution D , which is usually ﬁxed to the uniform distribution.
Revised: Our use of BRKGA is standard except for the mutant-sampling distribution D , which is usually ﬁxed to the uniform distribution.

Original: We generalize BRKGA for instance-speciﬁc learning by sampling from n independent beta distributions, whose parameters can vary by index.
Revised: We generalize BRKGA for instance-speciﬁc learning by sampling from n independent beta distributions, whose parameters can vary by index.

Original: Beta ﬂexibly admits non-uniform distribution choices and also subsumes the uniform choice.
Revised: Beta ﬂexibly admits non-uniform distribution choices and also subsumes the uniform choice.

Original: Given a computation graph, let d be the number of devices, o the number of ops, and t the number of tensors.
Revised: Given a computation graph, let d be the number of devices, o the number of ops, and t the number of tensors.

Original: We deﬁne the chromosome encoding a scheduling solution to have three distinct parts: (1) o × d entries specifying op-to-device afﬁnities; (2) o entries specifying scheduling priorities for each op; (3) t × d entries specifying tensor-to-device priorities for transfers that may be needed.
Revised: We deﬁne the chromosome encoding a scheduling solution to have three distinct parts: (1) o × d entries specifying op-to-device afﬁnities; (2) o entries specifying scheduling priorities for each op; (3) t × d entries specifying tensor-to-device priorities for transfers that may be needed.

Original: Given a chromosome, op placements are picked by maximum afﬁnity.
Revised: Given a chromosome, op placements are picked by maximum afﬁnity.

Original: Transfer ops are created as implied by the placements.
Revised: Transfer ops are created as implied by the placements.

Original: We then obtain a schedule by performing a topological sort over the ops given their tensor dependencies, breaking ties by using the corresponding node priorities.
Revised: We then obtain a schedule by performing a topological sort over the ops given their tensor dependencies, breaking ties by using the corresponding node priorities.

Original: Once a schedule is constructed, the performance model is used to evaluate its peak memory and/or runtime.
Revised: Once a schedule is constructed, the performance model is used to evaluate its peak memory and/or runtime.

Original: When enforcing a memory constraint, the ﬁtness of a schedule is encoded such that all memory-feasible schedules have better ﬁtness than infeasible schedules.
Revised: When enforcing a memory constraint, the ﬁtness of a schedule is encoded such that all memory-feasible schedules have better ﬁtness than infeasible schedules.

Original: An example is provided in section A.5.
Revised: An example is provided in section A.5.

Original:  REGAL 4. C ONTEXTUAL B ANDIT F ORMULATION
Revised:  REGAL 4. C ONTEXTUAL B ANDIT F ORMULATION

Original: We train a contextual bandit policy that predicts beta distribution choices for each of the nodes of a computation graph to be used by BRKGA to optimize it.
Revised: We train a contextual bandit policy that predicts beta distribution choices for each of the nodes of a computation graph to be used by BRKGA to optimize it.

Original: For each round in the bandit setting, ﬁrst the context is observed by drawing a computation graph G as an i.i.d.
Revised: For each round in the bandit setting, ﬁrst the context is observed by drawing a computation graph G as an i.i.d.

Original: sample from a distribution
Revised: sample from a distribution

Original: D (e.g., a distribution over TensorFlow graphs).
Revised: D (e.g., a distribution over TensorFlow graphs).

Original: G has a set of nodes V and a set of edges E , with features associated with the nodes and edges.
Revised: G has a set of nodes V and a set of edges E , with features associated with the nodes and edges.

Original: A policy p ( a | G ) is applied to make a set of decisions at each node.
Revised: A policy p ( a | G ) is applied to make a set of decisions at each node.

Original: These decisions, denoted a v for each v ∈ V , across all nodes form one action a = { a v ∈ V } .
Revised: These decisions, denoted a v for each v ∈ V , across all nodes form one action a = { a v ∈ V } .

Original: One decision in a corresponds to playing one arm of the bandit, and specifying the entire a corresponds to playing several arms together in a single round.
Revised: One decision in a corresponds to playing one arm of the bandit, and specifying the entire a corresponds to playing several arms together in a single round.

Original: This can be viewed as a combinatorial multi-armed bandit problem (Chen et al., 2013).
Revised: This can be viewed as a combinatorial multi-armed bandit problem (Chen et al., 2013).

Original: The action a speciﬁes all the node-speciﬁc beta distributions BRKGA needs to optimize placement and scheduling decisions for G .
Revised: The action a speciﬁes all the node-speciﬁc beta distributions BRKGA needs to optimize placement and scheduling decisions for G .

Original: To enable a policy over discrete choices, we quantize the mean and variance parameters of the beta distribution.
Revised: To enable a policy over discrete choices, we quantize the mean and variance parameters of the beta distribution.

Original: The environment then runs BRKGA with those distribution choices with a ﬁxed iteration limit.
Revised: The environment then runs BRKGA with those distribution choices with a ﬁxed iteration limit.

Original: The ﬁnal objective value is used to compute the reward.
Revised: The ﬁnal objective value is used to compute the reward.

Original: To make the reward values comparable across different graphs, we divide the objective value o a ( G ) achieved on a graph G with action a by the objective value o s ( G ) achieved by standard
Revised: To make the reward values comparable across different graphs, we divide the objective value o a ( G ) achieved on a graph G with action a by the objective value o s ( G ) achieved by standard

Original: BRKGA using uniform distributions.
Revised: BRKGA using uniform distributions.

Original: Since we want to minimize the objective (e.g., runtime or peak memory), we deﬁne the reward as r ( a , G ) =
Revised: Since we want to minimize the objective (e.g., runtime or peak memory), we deﬁne the reward as r ( a , G ) =

Original: − o a ( G ) o s ( G ) .
Revised: − o a ( G ) o s ( G ) .

Original: So a reward > − 1 corresponds to an action that achieves a better objective value than standard BRKGA on a graph.
Revised: So a reward > − 1 corresponds to an action that achieves a better objective value than standard BRKGA on a graph.

Original: We maximize the expected reward L = E G [ (cid:80) a p ( a | G ) r ( a , G )] , where E G is an expectation over graphs in our training set.
Revised: We maximize the expected reward L = E G [ (cid:80) a p ( a | G ) r ( a , G )] , where E G is an expectation over graphs in our training set.

Original: Learning is done by REINFORCE (Williams, 1992).
Revised: Learning is done by REINFORCE (Williams, 1992).

Original: We added a scalar baseline b ( G ) for reducing the variance of the gradient estimates.
Revised: We added a scalar baseline b ( G ) for reducing the variance of the gradient estimates.

Original: 4. G RAPH NEURAL NETWORK POLICY
Revised: 4. G RAPH NEURAL NETWORK POLICY

Original: From computation graphs, we derive multigraphs with attributed nodes and directed edges.
Revised: From computation graphs, we derive multigraphs with attributed nodes and directed edges.

Original: Denote a multigraph G =
Revised: Denote a multigraph G =

Original: ( V, E ) .
Revised: ( V, E ) .

Original: In our setup, the nodes V correspond 1:1 to the ops.
Revised: In our setup, the nodes V correspond 1:1 to the ops.

Original: An edge e ∈ E exists from u to v for each tensor that op v requires that is produced by op u .
Revised: An edge e ∈ E exists from u to v for each tensor that op v requires that is produced by op u .

Original: As a tensor can be required by multiple ops, the correspondence from edges to tensors may be many to one.
Revised: As a tensor can be required by multiple ops, the correspondence from edges to tensors may be many to one.

Original: Each node v ∈ V and edge e ∈ E has an attribute vector x v and x e .
Revised: Each node v ∈ V and edge e ∈ E has an attribute vector x v and x e .

Original: The attributes contain respective features, e.g., sizes of the tensors.
Revised: The attributes contain respective features, e.g., sizes of the tensors.

Original: We learn a model that predicts good mutant sampling distributions for BRKGA given this multigraph.
Revised: We learn a model that predicts good mutant sampling distributions for BRKGA given this multigraph.

Original: Each node has d + 1 independent beta distributions, corresponding to device afﬁnities and scheduling priorities, whose parameters are represented as a vector a v .
Revised: Each node has d + 1 independent beta distributions, corresponding to device afﬁnities and scheduling priorities, whose parameters are represented as a vector a v .

Original: These are the model’s actions in RL terms, and our model speciﬁes a distribution over actions a = { a v } v ∈ V for each graph, p ( a | G ) .
Revised: These are the model’s actions in RL terms, and our model speciﬁes a distribution over actions a = { a v } v ∈ V for each graph, p ( a | G ) .

Original: Note the action space is different from graph to graph.
Revised: Note the action space is different from graph to graph.

Original: We use Graph Neural Networks (GNNs) (Scarselli et al., 2009;
Revised: We use Graph Neural Networks (GNNs) (Scarselli et al., 2009;

Original: Li et al., 2015; Gilmer et al., 2017;
Revised: Li et al., 2015; Gilmer et al., 2017;

Original: Battaglia et al., 2018) to learn representations for computation graphs.
Revised: Battaglia et al., 2018) to learn representations for computation graphs.

Original: Given a (multi)graph G , a
Revised: Given a (multi)graph G , a

Original: GNN computes representation vectors h v for each node through an iterative message passing process as follows:
Revised: GNN computes representation vectors h v for each node through an iterative message passing process as follows:

Original: [Equation]
Revised: h

Original: [Equation]
Revised: [Equation]

Original: where e s is the source node of edge e and e t is the target node.
Revised: where e s is the source node of edge e and e t is the target node.

Original: In our formulation, MLP n and MLP e are multilayer perceptrons (MLPs) that encode node and edge attributes, MLP msg and MLP (cid:48)
Revised: In our formulation, MLP n and MLP e are multilayer perceptrons (MLPs) that encode node and edge attributes, MLP msg and MLP (cid:48)

Original: msg compute messages along the edges in the edge direction ( m ( t ) e ) and the opposite direction ( m ( t ) (cid:48) e ),
Revised: msg compute messages along the edges in the edge direction ( m ( t ) e ) and the opposite direction ( m ( t ) (cid:48) e ),

Original: MLP node updates node representations and [ . ] represents ﬂat vector concatenation.
Revised: MLP node updates node representations and [ . ] represents ﬂat vector concatenation.

Original: After T rounds of message passing, the representation for each node h v = h ( T ) v will contain information from the
Revised: After T rounds of message passing, the representation for each node h v = h ( T ) v will contain information from the

Original: T -hop neighborhood around v in the graph.
Revised: T -hop neighborhood around v in the graph.

Original: Given the h v ’s, we produce a v ’s through conditionally independent predictions, where the prediction for one node v does not depend on the predictions of other nodes given the computed representations:
Revised: Given the h v ’s, we produce a v ’s through conditionally independent predictions, where the prediction for one node v does not depend on the predictions of other nodes given the computed representations:

Original: [Equation]
Revised: [Equation]

Original: MLP a is shared across all nodes for predicting the parameters of the output distributions.
Revised: MLP a is shared across all nodes for predicting the parameters of the output distributions.

Original: In our experiments, we quantize the continuous beta distribution parameters and use a discrete action space.
Revised: In our experiments, we quantize the continuous beta distribution parameters and use a discrete action space.

Original: The outputs are therefore categorical, and we use the MLP to compute the logits of the corresponding softmax distributions.
Revised: The outputs are therefore categorical, and we use the MLP to compute the logits of the corresponding softmax distributions.

Original: More details are included in section A.6.
Revised: More details are included in section A.6.

Original: The baseline is computed using a separate GNN, where after we obtained the node representations h v , we aggregate across nodes and compute b ( G ) as b ( G ) =
Revised: The baseline is computed using a separate GNN, where after we obtained the node representations h v , we aggregate across nodes and compute b ( G ) as b ( G ) =

Original: MLP b (cid:16) 1 | V | (cid:80) v MLP g
Revised: MLP b (cid:16) 1 | V | (cid:80) v MLP g

Original: ( h v ) (cid:17) .
Revised: ( h v ) (cid:17) .

Original:  E XPERIMENTAL RESULTS 5. T ASKS AND DATASETS
Revised:  E XPERIMENTAL RESULTS 5. T ASKS AND DATASETS

Original: We consider two tasks, one is minimizing peak memory and the other is minimizing running time, both on two homogeneous devices with 16 GiB of memory each and synchronous tensor transfers with zero cost (zero latency and inﬁnite bandwidth).
Revised: We consider two tasks, one is minimizing peak memory and the other is minimizing running time, both on two homogeneous devices with 16 GiB of memory each and synchronous tensor transfers with zero cost (zero latency and inﬁnite bandwidth).

Original: We train a separate neural network for each task-dataset pair for the case of two devices.
Revised: We train a separate neural network for each task-dataset pair for the case of two devices.

Original: We have collected a dataset of 372 topologically-unique real-world TensorFlow graphs by mining machine learning jobs on a company’s internal compute cluster (see A.1.2).
Revised: We have collected a dataset of 372 topologically-unique real-world TensorFlow graphs by mining machine learning jobs on a company’s internal compute cluster (see A.1.2).

Original: These jobs are from a wide range of production and research use cases.
Revised: These jobs are from a wide range of production and research use cases.

Original: The dataset is split into {train, valid, test} sets containing {60%, 20%, 20%} graphs, respectively.
Revised: The dataset is split into {train, valid, test} sets containing {60%, 20%, 20%} graphs, respectively.

Original: These sets are disjoint with respect to graph topology, so at test time the policy needs to generalize to new topologies.
Revised: These sets are disjoint with respect to graph topology, so at test time the policy needs to generalize to new topologies.

Original: We augment the dataset by applying multiplicative noise to tensor sizes and op running times to create several variants per graph.
Revised: We augment the dataset by applying multiplicative noise to tensor sizes and op running times to create several variants per graph.

Original: Even though the variants of the same graph share the same topology, they represent different optimization problem instances.
Revised: Even though the variants of the same graph share the same topology, they represent different optimization problem instances.

Original: We create separate datasets for minimizing runtime and peak memory.
Revised: We create separate datasets for minimizing runtime and peak memory.

Original: The TF runtime dataset has 16329 training, 5470 validation, andtest graphs.
Revised: The TF runtime dataset has 16329 training, 5470 validation, andtest graphs.

Original: The TF peak memory dataset has 22400 training, 7400 validation, and 7400 test graphs.
Revised: The TF peak memory dataset has 22400 training, 7400 validation, and 7400 test graphs.

Original: For reproducibility, we have released 1 a synthetic dataset of computation graphs with 10000 training, 1000 validation, and 1000 test cases.
Revised: For reproducibility, we have released 1 a synthetic dataset of computation graphs with 10000 training, 1000 validation, and 1000 test cases.

Original: The graph topologies are generated from several classical random graph models, and the op running times and tensor sizes are sampled from Gaussian distributions (see A.1.4).
Revised: The graph topologies are generated from several classical random graph models, and the op running times and tensor sizes are sampled from Gaussian distributions (see A.1.4).

Original: On this dataset we minimize running time without a memory constraint (e.g., on two homogeneous devices with inﬁnite memory).
Revised: On this dataset we minimize running time without a memory constraint (e.g., on two homogeneous devices with inﬁnite memory).

Original: 5. B ASELINES
Revised: 5. B ASELINES

Original: Graph Partitioning + Depth First Search (GP+DFS): Combines a graph partitioning (GP) baseline for device placement to minimize communication across devies and a Depth-First Search heuristic similar to the one implemented in XLA TensorFlow Authors (2016b) to compute per-device schedules given placements.
Revised: Graph Partitioning + Depth First Search (GP+DFS): Combines a graph partitioning (GP) baseline for device placement to minimize communication across devies and a Depth-First Search heuristic similar to the one implemented in XLA (TensorFlow Authors, 2016b) to compute per-device schedules given placements.

Original: This is representative of the XLA compiler’s solution for model parallelism.
Revised: This is representative of the XLA compiler’s solution for model parallelism.

Original: Local Search : The method starts with a random placement and schedule and greedily improve it by moving an op across devices or changing an op’s order in the current schedule.
Revised: Local Search : The method starts with a random placement and schedule and greedily improves it by moving an op across devices or changing an op’s order in the current schedule.

Original: Graph-As-Sequence Model (GAS) :
Revised: Graph-As-Sequence Model (GAS) :

Original: Like Mirhoseini et al.
Revised: Like Mirhoseini et al.

Original: (2017; 2018), we convert the graph into a sequence using a topological sort and apply a recurrent neural network to predict node-level  distributions to be used by BRKGA.
Revised: (2017; 2018), we convert the graph into a sequence using a topological sort and apply a recurrent neural network to predict node-level [Table] distributions to be used by BRKGA.

Original: This comparison measures the usefulness of graph structure for learning.
Revised: This comparison measures the usefulness of graph structure for learning.

Original: [Table] BRKGA XK : Run BRKGA for X thousand iterations with uniform sampling distributions using default hyperparameters consistent with Gonçalves & Resende(2011).
Revised: BRKGA XK : Run BRKGA for X thousand iterations with uniform sampling distributions using default hyperparameters consistent with Gonçalves & Resende

Original: This comparison measures the performance of the default version of BRKGA.
Revised: This comparison measures the performance of the default version of BRKGA.

Original: [Table]

Original:  BRKGA : Apply grid search to BRKGA’s hyperparameters on the training set and pick the best.
Revised: Tuned BRKGA : Apply grid search to BRKGA’s hyperparameters on the training set and pick the best.

Original: This represents how well BRKGA performs by customizing it to the distribution of computation graphs, but without instance-dependent customization.
Revised: This represents how well BRKGA performs by customizing it to the distribution of computation graphs, but without instance-dependent customization.

Original: Instance-dependent Random Search (IDRS) :
Revised: Instance-dependent Random Search (IDRS) :

Original: Same as REGAL, but BRKGA is replaced with random search.
Revised: Same as REGAL, but BRKGA is replaced with random search.

Original: This is done by running BRKGA for only one generation using the proposal distributions computed by the neural network.
Revised: This is done by running BRKGA for only one generation using the proposal distributions computed by the neural network.

Original: Additionally, we use a Constraint Programming (CP) approach with the CP-SAT solver of Google ORtools (Google, 2019) to establish a provably global optimum for each computation graph optimization problem instance by running for up to 24 hours.
Revised: Additionally, we use a Constraint Programming (CP) approach with the CP-SAT solver of Google ORtools (Google, 2019) to establish a provably global optimum for each computation graph optimization problem instance by running for up to 24 hours.

Original: As an enumerative algorithm, it is generally not competitive when run only for seconds.
Revised: As an enumerative algorithm, it is generally not competitive when run only for seconds.

Original: For a fair comparison, we ﬁx the number of performance model evaluations allowed per graph to be the same across algorithms.
Revised: For a fair comparison, we ﬁx the number of performance model evaluations allowed per graph to be the same across algorithms.

Original: (Except GP+DFS, which does not allow ﬁxing it.)
Revised: (Except GP+DFS, which does not allow ﬁxing it.)

Original: Given typical
Revised: Given typical

Original: TensorFlow graph sizes and compiler running time constraints, we estimate that a budget of 5, evaluations is feasible in practice, so we use that in the experiments.
Revised: TensorFlow graph sizes and compiler running time constraints, we estimate that a budget of 5, evaluations is feasible in practice, so we use that in the experiments.

Original: Learning to directly predict a solution: We have explored two more approaches for training a graph neural network to predict placement and scheduling solutions directly, without BRKGA.
Revised: Learning to directly predict a solution: We have explored two more approaches for training a graph neural network to predict placement and scheduling solutions directly, without BRKGA.

Original: We used supervised learning to train a network to predict BRKGA’s solutions.
Revised: We used supervised learning to train a network to predict BRKGA’s solutions.

Original: The best accuracy was achieved by predicting the solution autoregressively, one variable at a time conditioned on previously predicted variables.
Revised: The best accuracy was achieved by predicting the solution autoregressively, one variable at a time conditioned on previously predicted variables.

Original: We also used RL to learn a policy with IMPALA (Espeholt et al., 2018) to optimize the objective value by incrementally predicting the solution one variable at a time, and once complete, iteratively improving it with a learned local search policy.
Revised: We also used RL to learn a policy with IMPALA (Espeholt et al., 2018) to optimize the objective value by incrementally predicting the solution one variable at a time, and once complete, iteratively improving it with a learned local search policy.

Original: The inference cost for both approaches is quadratic in the number of nodes (the graph net is applied a linear number of times, each with linear cost), while REGAL’s inference cost is linear, making them orders of magnitude slower than REGAL at test time.
Revised: The inference cost for both approaches is quadratic in the number of nodes (the graph net is applied a linear number of times, each with linear cost), while REGAL’s inference cost is linear, making them orders of magnitude slower than REGAL at test time.

Original: An evaluation on a test set of small graphs showed that neither approach improves on BRKGA5K. Improving the scalability and the generalization of these approaches is left as future work, and do not present their results here.
Revised: An evaluation on a test set of small graphs showed that neither approach improves on BRKGA5K. Improving the scalability and the generalization of these approaches is left as future work, and do not present their results here.

Original: 5. C OMPARISON TO BASELINE ALGORITHMS
Revised: 5. C OMPARISON TO BASELINE ALGORITHMS

Original: We use two metrics to compare algorithms.
Revised: We use two metrics to compare algorithms.

Original: 1) Average percent improvement over BRKGA 5K: For a given graph, compute the percent improvement in the objective achieved by an algorithm relative to BRKGA with evaluation limit set to 5,000.
Revised: 1) Average percent improvement over BRKGA 5K: For a given graph, compute the percent improvement in the objective achieved by an algorithm relative to BRKGA with evaluation limit set to 5,000.

Original: BRKGA 5K is a natural reference for measuring the effect of learning approaches that predict proposal distributions for it.
Revised: BRKGA 5K is a natural reference for measuring the effect of learning approaches that predict proposal distributions for it.

Original: Average percent gap from best known solution: Compute the best known objective value among all the algorithms.(This will [Figure] P
Revised: Average percent gap from best known solution: Compute the best known objective value among all the algorithms.

Original: [Figure]

Original: r

Original: [Figure] n t a g e o f t e s t s e t [Figure] be found by CP-SAT if it ﬁnishes within the time limit.)
Revised: (This will be found by CP-SAT if it ﬁnishes within the time limit.)

Original: Compute the percent difference between an algorithm’s solution and the best known objective value.
Revised: Compute the percent difference between an algorithm’s solution and the best known objective value.

Original: We report averages over test set graphs.
Revised: We report averages over test set graphs.

Original: Training set results are similar and reported in section A.3.
Revised: Training set results are similar and reported in section A.3.

Original: Table 1 compares REGAL to other algorithms on the two TensorFlow test sets and the synthetic dataset.
Revised: Table 1 compares REGAL to other algorithms on the two TensorFlow test sets and the synthetic dataset.

Original: REGAL outperforms all the baselines on all three tasks.
Revised: REGAL outperforms all the baselines on all three tasks.

Original: It gives 1 .
Revised: It gives 1 .

Original: 9 × and 4 .
Revised: 9 × and 4 .

Original: 4 × bigger improvements than the next best algorithm on runtime and peak memory minimization tasks, respectively.
Revised: 4 × bigger improvements than the next best algorithm on runtime and peak memory minimization tasks, respectively.

Original: The percent improvement over GP + DFS is 44.4% and 10.1% for runtime and peak memory, respectively.
Revised: The percent improvement over GP + DFS is 44.4% and 10.1% for runtime and peak memory, respectively.

Original: REGAL reduces the average percent gap from the best known solution by about 1 . 8 × with respect to
Revised: REGAL reduces the average percent gap from the best known solution by about 1 . 8 × with respect to

Original: BRKGA 5K on both TensorFlow test sets, and by about 6 × and 3 .
Revised: BRKGA 5K on both TensorFlow test sets, and by about 6 × and 3 .

Original: 3 × with respect to GP + DFS on the TensorFlow Runtime and Peak Memory test sets, respectively.
Revised: 3 × with respect to GP + DFS on the TensorFlow Runtime and Peak Memory test sets, respectively.

Original: (For an XLA user GP + DFS is the current, albeit weak, state-of-the-art algorithm.)
Revised: (For an XLA user, GP + DFS is the current, albeit weak, state-of-the-art algorithm.)

Original: The synthetic test set shows similar results.
Revised: The synthetic test set shows similar results.

Original: The learned policy successfully generalizes to previously unseen graphs, to the extent that a large fraction of the estimated room for improvement over BRKGA 5K is captured using the same evaluation limit.
Revised: The learned policy successfully generalizes to previously unseen graphs, to the extent that a large fraction of the estimated room for improvement over BRKGA 5K is captured using the same evaluation limit.

Original: To further test the limits of generalization for the policies learned using REGAL, we evaluate them on XLA graphs from a production compiler team’s internal performance benchmark.
Revised: To further test the limits of generalization for the policies learned using REGAL, we evaluate them on XLA graphs from a production compiler team’s internal performance benchmark.

Original: XLA uses a different set of ops from TensorFlow, and the benchmark graphs on average have about an order of magnitude more nodes and edges than the TensorFlow graphs in our training set, so this is a difﬁcult generalization challenge.
Revised: XLA uses a different set of ops from TensorFlow, and the benchmark graphs on average have about an order of magnitude more nodes and edges than the TensorFlow graphs in our training set, so this is a difﬁcult generalization challenge.

Original: REGAL achieves 0 .
Revised: REGAL achieves 0 .

Original: 58% average runtime improvement over BRKGA 5K on 94 graphs, and
Revised: 58% average runtime improvement over BRKGA 5K on 94 graphs, and

Original: 74% average peak memory improvement on 32 graphs.
Revised: 74% average peak memory improvement on 32 graphs.

Original: It is promising that any improvements are possible at all despite training only on TensorFlow graphs, and points to the possibility of bigger improvements by training directly on XLA graphs.
Revised: It is promising that any improvements are possible at all despite training only on TensorFlow graphs, and points to the possibility of bigger improvements by training directly on XLA graphs.

Original: Optimizer running times: BRKGA 5K takes on average 0.89 seconds on the TensorFlow Peak
Revised: Optimizer running times: BRKGA 5K takes on average 0.89 seconds on the TensorFlow Peak

Original: Memory test set to optimize a computation graph, while REGAL takes 1.04 seconds.
Revised: Memory test set to optimize a computation graph, while REGAL takes 1.04 seconds.

Original: (The times are similar on the Runtime test set.)
Revised: (The times are similar on the Runtime test set.)

Original: Instead of taking hours to compute a solution per graph (e.g.,
Revised: Instead of taking hours to compute a solution per graph (e.g.,

Original: Mirhoseini et al.
Revised: Mirhoseini et al.

Original: (2017; 2018)), REGAL produces solutions in orders of magnitude less time, while still being better than all the baselines.
Revised: (2017; 2018)), REGAL produces solutions in orders of magnitude less time, while still being better than all the baselines.

Original: 5. C OMPARING REGAL VS . BRKGA
Revised: 5. C OMPARING REGAL VS . BRKGA

Original: Figure 2 shows histograms of percent improvements in runtime (left) and peak memory (right) achieved by REGAL over BRKGA 5K on the test sets.
Revised: Figure 2 shows histograms of percent improvements in runtime (left) and peak memory (right) achieved by REGAL over BRKGA 5K on the test sets.

Original: Green bars correspond to graphs on which
Revised: Green bars correspond to graphs on which

Original: REGAL improved over BRKGA 5K, while red bars correspond to graphs on which REGAL was worse.
Revised: REGAL improved over BRKGA 5K, while red bars correspond to graphs on which REGAL was worse.

Original: (Ties have been omitted for clarity.)
Revised: (Ties have been omitted for clarity.)

Original: REGAL matches or beats BRKGA 5K on 87.4% of the runtime test set, and 88.9% of the peak memory test set.
Revised: REGAL matches or beats BRKGA 5K on 87.4% of the runtime test set, and 88.9% of the peak memory test set.

Original: The highest improvement is 26.0% for run time and 54.3% for peak memory, while the worst regression is 24.0% for run time and 17.9% for peak memory.
Revised: The highest improvement is 26.0% for run time and 54.3% for peak memory, while the worst regression is 24.0% for run time and 17.9% for peak memory.

Original: To assess whether the improvements provided by REGAL’s policy generalize to evaluation limits other than the one for which it was trained, we varied the evaluation limit used by both BRKGA [Figure] t [Figure]
Revised: To assess whether the improvements provided by REGAL’s policy generalize to evaluation limits other than the one for which it was trained, we varied the evaluation limit used by both BRKGA and REGAL at test time.

Original: g e i m p [Figure] m e n t ov [Figure]
Revised: The results are shown in ﬁgure 3.

Original: r [Figure]
Revised: REGAL’s performance improves with more evaluations, conﬁrming that the policy generalizes to higher evaluation limits.

Original: KGA [Figure] [Figure] BRKGA that perform well regardless of the evaluation limit, and REGAL learns to predict those choices.
Revised: In other words, there exist node-level choices for the distributions used in BRKGA that perform well regardless of the evaluation limit, and REGAL learns to predict those choices.

Original: This is particularly useful in cases where the actual evaluation limit to use will be known only at test time, so that the same policy can be applied without re-training.
Revised: This is particularly useful in cases where the actual evaluation limit to use will be known only at test time, so that the same policy can be applied without re-training.

Original: Interestingly, even with 50,000 evaluations, BRKGA is not able to match REGAL’s performance with just 5,000 evaluations!
Revised: Interestingly, even with 50,000 evaluations, BRKGA is not able to match REGAL’s performance with just 5,000 evaluations!

Original: 5. G RAPH - DEPENDENT POLICY
Revised: 5. G RAPH - DEPENDENT POLICY

Original: The RL agents’ actions are instance dependent.
Revised: The RL agents’ actions are instance dependent.

Original: The agent that performs the best on the TF Runtime dataset has a choice of 16 different node placement actions for each node in a graph.
Revised: The agent that performs the best on the TF Runtime dataset has a choice of 16 different node placement actions for each node in a graph.

Original: For each graph in the TF Runtime test set, we compute the entropy of the distribution of the node placement actions taken by the agent and plot a histogram of these entropies in Figure 4(a).
Revised: For each graph in the TF Runtime test set, we compute the entropy of the distribution of the node placement actions taken by the agent and plot a histogram of these entropies in Figure 4(a).

Original: The mean of this distribution  is 1.71 nats which implies that the actions are neither uniform random, nor constant, and vary from graph to graph.
Revised: The mean of this distribution [Figure] [Figure] is 1.71 nats which implies that the actions are neither uniform random, nor constant, and vary from graph to graph

Original: Furthermore, the agent’s performance overall gets better with more graph message passing iterations
Revised: Furthermore, the agent’s performance overall gets better with more graph message passing iterations

Original: T .
Revised: T .

Original: Figure 4(b) shows the peak validation reward reached within a hyperparameter sweep for each T for the TF runtime optimization task.
Revised: Figure 4(b) shows the peak validation reward reached within a hyperparameter sweep for each T for the TF runtime optimization task.

Original: Models that utilize the GNN with message passing ( T > 0 ) reach higher performance than T = 

Revised: Models that utilize the GNN with message passing ( T > 0 ) reach higher performance than T =

Original: (i.e., ignoring the graph structure).
Revised: (i.e., ignoring the graph structure).

Original:  C ONCLUSIONS AND FUTURE WORK
Revised:  C ONCLUSIONS AND FUTURE WORK

Original: By training a graph neural network policy to predict graph-conditional node-level distributions for
Revised: By training a graph neural network policy to predict graph-conditional node-level distributions for

Original: BRKGA, REGAL successfully generalizes to new graphs, signiﬁcantly outperforms all baselines in solution quality, and computes solutions in about one second on average per TensorFlow test set graph.
Revised: BRKGA, REGAL successfully generalizes to new graphs, signiﬁcantly outperforms all baselines in solution quality, and computes solutions in about one second on average per TensorFlow test set graph.

Original: REGAL’s speed and generalization make it a strong choice for use in a production compiler that needs to handle a diverse set of graphs under a limited time budget.
Revised: REGAL’s speed and generalization make it a strong choice for use in a production compiler that needs to handle a diverse set of graphs under a limited time budget.

Original: We foresee several extensions.
Revised: We foresee several extensions.

Original: Integrating REGAL into a neural network compiler would allow us to evaluate the end-to-end gains due to better placement and scheduling decisions.
Revised: Integrating REGAL into a neural network compiler would allow us to evaluate the end-to-end gains due to better placement and scheduling decisions.

Original: To further improve
Revised: To further improve

Original: REGAL’s own performance, one could use a Mixture of Experts architecture.
Revised: REGAL’s own performance, one could use a Mixture of Experts architecture.

Original: Given the diversity of graphs, a mixture model can train specialized sub-models on different types of graphs (e g., convolutional networks, recurrent networks, etc.).
Revised: Given the diversity of graphs, a mixture model can train specialized sub-models on different types of graphs (e.g., convolutional networks, recurrent networks, etc.).

Original: Another is to replace BRKGA with alternatives, e.g., combining learned neural policies with local search
Revised: Another is to replace BRKGA with alternatives, e.g., combining learned neural policies with local search. [Figure] [Figure] [Figure]

Revised: [Figure]

Original: A.1. T ENSOR F LOW DATASET
Revised: [Figure] [Figure] [Figure] V a li d a ti on r e w a r d

Revised: [Figure]

Revised: . D ATASETS A.1. D ATASET STATISTICS

Revised: Figure 5 and ﬁgure 6 give statistics for the number of nodes and edges in the datasets.

Revised: The broad range of graph sizes indicates the diversity of the datasets.

Revised: A.1. T ENSOR F LOW DATASET

Original: We collected a dataset by mining TensorFlow jobs running in a shared production cluster and extracting computation graphs in the MetaGraphDef 2 format.
Revised: We collected a dataset by mining TensorFlow jobs running in a shared production cluster and extracting computation graphs in the MetaGraphDef 2 format.

Original: As lots of computation graphs were repeated due to device/machine/job replicas, we deduplicate the dataset by graph topology (speciﬁcally node in-degree sequence).
Revised: As lots of computation graphs were repeated due to device/machine/job replicas, we deduplicate the dataset by graph topology (speciﬁcally node in-degree sequence).

Original: We have not applied any other kind of ﬁltering to restrict the dataset in any way (e.g., by architecture, input modality, learning task, dataset, etc.).
Revised: We have not applied any other kind of ﬁltering to restrict the dataset in any way (e.g., by architecture, input modality, learning task, dataset, etc.).

Original: Since the organization from which the graphs were collected has a large and diverse set of production and research use cases across input modalities, learning types, and datasets, we strongly believe our dataset is representative of a broad, real-world distribution of TensorFlow graphs.
Revised: Since the organization from which the graphs were collected has a large and diverse set of production and research use cases across input modalities, learning types, and datasets, we strongly believe our dataset is representative of a broad, real-world distribution of TensorFlow graphs.

Original: Computational costs for these computation graphs are simulated with an in-house simulator (based on
Revised: Computational costs for these computation graphs are simulated with an in-house simulator (based on

Original: Grappler 3 ) that outputs memory and running time proﬁled information in the CostGraphDef 4 format.
Revised: Grappler 3 ) that outputs memory and running time proﬁled information in the CostGraphDef 4 format.

Original: The simulator TensorFlow Op coverage didn’t include custom kernels or complicated control ﬂow Yu et al.
Revised: The simulator TensorFlow Op coverage didn’t include custom kernels or complicated control ﬂow Yu et al.

Original: (2018) like cycles (e.g. tf.while_loop ).
Revised: (2018) like cycles (e.g. tf.while_loop ).

Original: The train-validation-test set split is made by selecting a set of 5 graphs from a list of graphs sorted by number of nodes, and splitting them as 3-1-1 across the three sets, respectively.
Revised: The train-validation-test set split is made by selecting a set of 5 graphs from a list of graphs sorted by number of nodes, and splitting them as 3-1-1 across the three sets, respectively.

Original: This ensures that the distribution of the number of nodes is similar for the three sets.
Revised: This ensures that the distribution of the number of nodes is similar for the three sets.

Original: For each graph in the train/validation/test sets, we make 99 copies of it and multiply each tensor size and each Op running time cost with a uniform sampled number in the interval (0 . 5 , 1 . 5) (one sample 2 https://github . com/tensorflow/tensorflow/blob/master/tensorflow/core/ protobuf/meta_graph . proto 3 https://github .
Revised: For each graph in the train/validation/test sets, we make 99 copies of it and multiply each tensor size and each Op running time cost with a uniform sampled number in the interval (0 . 5 , 1 . 5) (one sample per tensor size per copy plus one sample per TF Op per copy).

Original: com/tensorflow/tensorflow/tree/master/tensorflow/core/ grappler 4 https://github .

Original: com/tensorflow/tensorflow/blob/master/tensorflow/core/ framework/cost_graph .

Original: proto per tensor size per copy plus one sample per TF Op per copy).

Original: The modiﬁed copies are added back to the respective set so that the graph topologies in train/validation/test do not overlap.
Revised: The modiﬁed copies are added back to the respective set so that the graph topologies in train/validation/test do not overlap.

Original: Graphs with no relevant cost information or no room for improvement (e.g. a graph with a single node, a chain in the minimizing running time task) are ﬁltered.
Revised: Graphs with no relevant cost information or no room for improvement (e.g. a graph with a single node, a chain in the minimizing running time task) are ﬁltered.

Original: This results in two different datasets for the two objective functions, one for runtime and another for peak memory.
Revised: This results in two different datasets for the two objective functions, one for runtime and another for peak memory.

Original: The encoding is described in A.
Revised: The encoding is described in A.

Original: A.1. XLA DATASET
Revised: A.1. XLA DATASET

Original: We also collected a dataset by dumping CostGraphDefs during the XLA compilation of several benchmark graphs and extracted the greatest-size control-ﬂow-free subgraph.
Revised: We also collected a dataset by dumping CostGraphDefs during the XLA compilation of several benchmark graphs and extracted the greatest-size control-ﬂow-free subgraph.

Original: After deduplication 94 graphs remained.
Revised: After deduplication 94 graphs remained.

Original: The encoding is described in A.
Revised: The encoding is described in A.

Original: A.1. S YNTHETIC DATASET
Revised: A.1. S YNTHETIC DATASET

Original: We sample synthetic graphs from a set of classic random graph models, including the Erdos-Renyi model Erdos & Rényi (1960), the Barabasi-Albert model Barabási & Albert (1999), the WattsStrogatz model Watts & Strogatz (1998) and the stochastic block model Holland et al.
Revised: We sample synthetic graphs from a set of classic random graph models, including the Erdos-Renyi model Erdos & Rényi (1960), the Barabasi-Albert model Barabási & Albert (1999), the WattsStrogatz model Watts & Strogatz (1998) and the stochastic block model Holland et al.

Original: The parameters we used for each of the random graph models are listed inTable 2. [Table] , with probability p = [Table]
Revised: The parameters we used for each of the random graph models are listed in

Original: 3 to

Original: [Table]

Original: an

Original: [Table]

Original: with a random

Original: [Table]
Revised: [Table]

Original: Stochastic

Original: [Table] Model [Table] The number of nodes in a graph is sampled uniformly from range [50 , 200] .
Revised: The number of nodes in a graph is sampled uniformly from range [50 , 200] .

Original: Note that all the above random graph models generate undirected graphs.
Revised: Note that all the above random graph models generate undirected graphs.

Original: To convert the graphs into directed acyclic graphs, for each graph we sample a random ordering of nodes π , and then set the direction of all edges ( i, j ) such that π ( i ) <
Revised: To convert the graphs into directed acyclic graphs, for each graph we sample a random ordering of nodes π , and then set the direction of all edges ( i, j ) such that π ( i ) <

Original: π ( j ) , i.e. node i comes before node j in the ordering π .
Revised: π ( j ) , i.e. node i comes before node j in the ordering π .

Original: After setting all the edge directions, we locate all the head nodes that don’t have any predecessors and all the tail nodes that don’t have any successors, and then create a source node and a sink node, and create edges from the source node to all the head nodes, and from all the tail nodes to the sink node.
Revised: After setting all the edge directions, we locate all the head nodes that don’t have any predecessors and all the tail nodes that don’t have any successors, and then create a source node and a sink node, and create edges from the source node to all the head nodes, and from all the tail nodes to the sink node.

Original: Real TensorFlow graphs all contain one source and one sink node.
Revised: Real TensorFlow graphs all contain one source and one sink node.

Original: Examples of synthetic graphs generated from therandom graph models are shown in Figure 6.
Revised: Examples of synthetic graphs generated from therandom graph models are shown in Figure 7.

Original: {+Erdos-Renyi      8        42     34      40         
+}
Revised: {+Erdos-Renyi      8        42     34      40         
+}

Original: {+Barabasi-Albert 0      494 5 6  8         38   402034 22  25 2933 35 
+}
Revised: {+Barabasi-Albert 0      494 5 6  8         38   402034 22  25 2933 35 
+}

Original: {+Watts-Strogatz             12 16   32       
 [Figure]+} Each edge ( i, j ) in the graph represents either a control dependency (op j can run only when op i is ﬁnished) or a data dependency (op i can run only when some output tensor(s) produced by op j is available).
Revised: {+Watts-Strogatz             12 16   32       
 [Figure]+} Each edge ( i, j ) in the graph represents either a control dependency (op j can run only when op i is ﬁnished) or a data dependency (op i can run only when some output tensor(s) produced by op j is available).

Original: We assign a probability of 0.1, 0.8 and 0.1 for each op to produce 0, 1 or 2 output tensors.
Revised: We assign a probability of 0.1, 0.8 and 0.1 for each op to produce 0, 1 or 2 output tensors.

Original: When op i produces 0 output tensors, any other op j that depends on it can only be through a control dependency.
Revised: When op i produces 0 output tensors, any other op j that depends on it can only be through a control dependency.

Original: Otherwise, op j will have a probability of 0.2 to make the dependency a control dependency, and probability of 0.8 to make this a data dependency, in which case an output tensor is picked according to an uniform distribution.
Revised: Otherwise, op j will have a probability of 0.2 to make the dependency a control dependency, and probability of 0.8 to make this a data dependency, in which case an output tensor is picked according to an uniform distribution.

Original: We ﬁll in the memory cost for each tensor by sampling from a Gaussian distribution with meanand standard deviation 10.
Revised: We ﬁll in the memory cost for each tensor by sampling from a Gaussian distribution with meanand standard deviation 10.

Original: The time cost for each op is computed as the sum of all the input and output tensor memory costs plus a random noise that is a fraction r of the total memory cost, with r sampled from a Gaussian with 0 mean and standard deviation 0.1.
Revised: The time cost for each op is computed as the sum of all the input and output tensor memory costs plus a random noise that is a fraction r of the total memory cost, with r sampled from a Gaussian with 0 mean and standard deviation 0.1.

Original: The source and sink nodes do not have any memory or time costs.
Revised: The source and sink nodes do not have any memory or time costs.

Original: To make sure the generated synthetic graphs are interesting, we apply an additional ﬁltering step by running BRKGA 1k and BRKGA 10k, and keeping only the graphs whose runtime improved by at [-Erdos-Renyi      8        42     34      40         
-]
Revised: To make sure the generated synthetic graphs are interesting, we apply an additional ﬁltering step by running BRKGA 1k and BRKGA 10k, and keeping only the graphs whose runtime improved by at [-Erdos-Renyi      8        42     34      40         
-]least 18%.

Original: [-Barabasi-Albert 0      494 5 6  8         38   402034 22  25 2933 35 
-]
Revised: [-Barabasi-Albert 0      494 5 6  8         38   402034 22  25 2933 35 
-]

Original: [-Watts-Strogatz             12 16   32       
-]
Revised: [-Watts-Strogatz             12 16   32       
-]

Original: [Figure]

Original: least 18%.

Original: This gives us a dataset of graphs that on average improve runtime by 20% from running
Revised: This gives us a dataset of graphs that on average improve runtime by 20% from running

Original: BRKGA 1k to 10k.
Revised: BRKGA 1k to 10k.

Original: The synthetic data in CostGraphDef format is available at https://drive .
Revised: The synthetic data in CostGraphDef format is available at https://drive .

Original: google .
Revised: google .

Original: com/ drive/folders/1lxRl1ocsWu-POwbdEY06Mrzq6Ot99j7N .
Revised: com/ drive/folders/1lxRl1ocsWu-POwbdEY06Mrzq6Ot99j7N .

Original: The encoding is described in A.
Revised: The encoding is described in A.2.

Original: A. D ATA ENCODING
Revised: A. D ATA ENCODING

Original: We consider “control dependencies” as being tensors of size zero.
Revised: We consider “control dependencies” as being tensors of size zero.

Original: Each triple (op producer, tensor, op consumer) in the CostGraphDef is encoded as a separate edge.
Revised: Each triple (op producer, tensor, op consumer) in the CostGraphDef is encoded as a separate edge.

Original: Each of these edges e is associate with three features denoted by x e : • The size of the tensor that is associated with the edge. • A one hot feature indicating if the edge is a control edge. • The normalized index of the hyperedge to which the edge belongs.
Revised: Each of these edges e is associated with three features denoted by x e : the size of the tensor that is associated with the edge, a one-hot feature indicating if the edge is a control edge, and normalized index of the hyperedge to which the edge belongs.

Original: This means that the graph neural network’s input is a directed graph with multiple edges.
Revised: This means that the graph neural network’s input is a directed graph with multiple edges.

Original: (There are alternative encodings like a bipartite graph where both TF ops and TF tensors are nodes, and edges exist only between ops and tensors when the op consumes or produces the Tensor.)
Revised: (There are alternative encodings like a bipartite graph where both TF ops and TF tensors are nodes, and edges exist only between ops and tensors when the op consumes or produces the tensor.)

Original: Each node v in the computation graph is associate with node features x v .
Revised: Each node v in the computation graph is associated with node features x v .

Original: There are 11 node features in total which can be grouped into three categories: memory-based, runtime-based (not used in the peak memory task) and BRKGA-based (task dependent).
Revised: There are 11 node features in total, which can be grouped into three categories: memory-based, runtime-based (not used in the peak memory task) and BRKGA-based (task dependent).

Original: As memory node features, we use the sum of input tensor sizes, the sum of output tensor sizes, the extra internal memory of the TensorFlow op, and a binary number indicating whether the op is the one which uses the greatest memory in the graph.
Revised: As memory-based node features, we use the sum of input tensor sizes, the sum of output tensor sizes, the extra internal memory of the TensorFlow op, and a one-hot indicator of whether the op is the one which uses the greatest memory in the graph.

Original: As runtime node features, we use the sum of direct predecessor nodes’ running times, the sum of direct successor nodes’ running times, the running time cost of the TensorFlow op, and a binary number indicating whether the op is the one with the greatest runtime cost in the graph.
Revised: As runtime-based node features, we use the sum of direct predecessor nodes’ running times, the sum of direct successor nodes’ running times, the running time cost of the op itself, and one-hot indicator of whether the op is the one with the greatest runtime cost in the graph.

Original: As BRKGA node features, we have a node aggregation (the expectation of the placement per device and the schedule order for each node) of the chromosomes found by BRKGA (minimizing peak memory for the peak memory dataset and minimizing runtime for the runtime dataset) running for 400 evaluations with uniform random distributions.
Revised: As BRKGA-based node features, we have a node aggregation (the expectation of the placement per device and the schedule order for each node) of the chromosomes found by BRKGA (minimizing peak memory for the peak memory dataset and minimizing runtime for the runtime dataset) running for 400 evaluations with uniform random distributions.

Original: To make comparisons fair, REGAL with K ﬁtness evaluations means 400 evaluations to compute features, and K − 400 ﬁtness evaluations for
Revised: To make comparisons fair, REGAL with K ﬁtness evaluations means 400 evaluations to compute features, and K − 400 ﬁtness evaluations for

Original: BRKGA using the instance-speciﬁc distributions.
Revised: BRKGA using the instance-speciﬁc distributions.

Original: For each graph, all node or edge features relating to memory size are normalized by the greatest memory size number in that graph and all features relating to runtime are are normalized by the greatest op runtime cost.
Revised: For each graph, all node or edge features relating to memory size are normalized by the greatest memory size number in that graph and all features relating to runtime are are normalized by the greatest op runtime cost.

Original: To break symmetry and reduce a single degree of freedom without loss in performance, we ﬁx the placement of the node with highest memory for the memory task and runtime for the runtime task to the ﬁrst device. 
Revised: To break symmetry and reduce a single degree of freedom without loss of performance, we ﬁx the placement of the node with highest memory for the memory task and runtime for the runtime task to the ﬁrst device. [Figure] e r [Figure]

Revised: i n i

Original: A. T RAINING SET RESULTS
Revised: [Figure] s [Figure] t

Original:  Figure 7 shows the reward curves on the training set for the runtime minimization (left) and peak memory minimization (right).
Revised: [Figure] [Figure] T RAINING SET RESULTS Figure 8 shows the reward curves on the training set for the runtime minimization (left) and peak memory minimization (right).

Original: Each point in the curves is the average reward achieved on a minibatch of training graphs at the corresponding training step.
Revised: Each point in the curves is the average reward achieved on a minibatch of training graphs at the corresponding training step.

Original: The ﬁnal average percent improvement over
Revised: The ﬁnal average percent improvement over

Original: BRKGA5K on the training set is similar to that of the test set for both tasks: 7 . 25% on the training set vs. 7 . 09% on the test set for runtime minimization, and 4 .
Revised: BRKGA5K on the training set is similar to that of the test set for both tasks: 7 . 25% on the training set vs. 7 . 09% on the test set for runtime minimization, and 4 .

Original: 36% on the training set vs. 3 .
Revised: 36% on the training set vs. 3 .

Original: 56% on the test set for peak memory minimization.
Revised: 56% on the test set for peak memory minimization.

Original: The small gap between train and test results shows that the policy is able to generalize successfully to unseen graphs at test time.[Figure] [Figure]
Revised: The small gap between train and test results shows that the policy is able to generalize successfully to unseen graphs at test time.

Original: A. P ERFORMANCE MODEL
Revised: A. P ERFORMANCE MODEL

Original: The scheduling problem is speciﬁed by a set of devices D and a computation graph.
Revised: The scheduling problem is speciﬁed by a set of devices D and a computation graph.

Original: The computation graph has the list of ops j ∈ N and tensors τ ∈ T , the tensors produced by each op
Revised: The computation graph has the list of ops j ∈ N and tensors τ ∈ T , the tensors produced by each op

Original: I ( j ) ⊆
Revised: I ( j ) ⊆

Original: T , the tensors consumed by each
Revised: T , the tensors consumed by each

Original: op C ( j ) ⊆
Revised: op C ( j ) ⊆

Original: T , the memory used by each tensor m τ [Figure] and the execution time of each op r j .
Revised: T , the memory used by each tensor m τ , and the execution time of each op r j .

Original: A tensor is produced by exactly one op but can be consumed by many.
Revised: A tensor is produced by exactly one op but can be consumed by many.

Original: Solutions to the scheduling problem are constructed as follows.
Revised: Solutions to the scheduling problem are constructed as follows.

Original: A placement is an assignment pl : N → D from ops to devices.
Revised: A placement is an assignment pl : N → D from ops to devices.

Original: Given a placement we deﬁne ˜ N pl as the set of ops N extended with synchronous inter-device transfer operations.
Revised: Given a placement we deﬁne ˜ N pl as the set of ops N extended with synchronous inter-device transfer operations.

Original: A transfer operation consumes a tensor on the device where it was created and produces it on a device where it is needed.
Revised: A transfer operation consumes a tensor on the device where it was created and produces it on a device where it is needed.

Original: Given a placement pl , a schedule is a total ordering s : ˜ N → { 1 , 2 , . . .
Revised: Given a placement pl , a schedule is a total ordering s : ˜ N → { 1 , 2 , . . .

Original: , | ˜ N |} on ops in ˜ N
Revised: , | ˜ N |} on ops in ˜ N

Original: pl .
Revised: pl .

Original: We say that op j runs at simulation time step s ( j ) .
Revised: We say that op j runs at simulation time step s ( j ) .

Original: We model the schedule execution as follows.
Revised: We model the schedule execution as follows.

Original: At each simulation time step t ∈ { 1 , 2 , . . .
Revised: At each simulation time step t ∈ { 1 , 2 , . . .

Original: , | ˜ N |} , each device d has a list l d,t of tensors currently in memory.
Revised: , | ˜ N |} , each device d has a list l d,t of tensors currently in memory.

Original: A tensor is added to the list when produced by an op that runs on the device or by a transfer op that receives the tensor on the device.
Revised: A tensor is added to the list when produced by an op that runs on the device or by a transfer op that receives the tensor on the device.

Original: A tensor is removed immediately after all of its consumers on the device have run.
Revised: A tensor is removed immediately after all of its consumers on the device have run.

Original: A schedule is valid if for each op j , all the input tensors are available on the corresponding device at simulation time step s ( j ) .
Revised: A schedule is valid if for each op j , all the input tensors are available on the corresponding device at simulation time step s ( j ) .

Original: See Section A.9 for an example schedule.
Revised: See Section A.9 for an example schedule.

Original: The memory used on a device at simulation time step t is the sum of the memory used by each tensor that is in memory, i.e., (cid:80) τ ∈ l d,t m τ .
Revised: The memory used on a device at simulation time step t is the sum of the memory used by each tensor that is in memory, i.e., (cid:80) τ ∈ l d,t m τ .

Original: The peak memory of a schedule is the maximum value of the memory used at any time and on any device.
Revised: The peak memory of a schedule is the maximum value of the memory used at any time and on any device.

Original: The runtime of a schedule is computed by stepping through the simulation time steps in order and accounting for the execution time of each op on each device.
Revised: The runtime of a schedule is computed by stepping through the simulation time steps in order and accounting for the execution time of each op on each device.

Original: Synchronous transfers block until both the sender and the receiver have completed their preceding tasks, and their execution time depends on a known bandwidth between devices.
Revised: Synchronous transfers block until both the sender and the receiver have completed their preceding tasks, and their execution time depends on a known bandwidth between devices.

Original: A. BRKGA CHROMOSOME ENCODING
Revised: A. BRKGA CHROMOSOME ENCODING

Original: Let the input graph G contain o ops and t tensors which must be placed over d devices.
Revised: Let the input graph G contain o ops and t tensors which must be placed over d devices.

Original: Then, the
Revised: Then, the

Original: BRKGA chromosome for this graph is a vector c ∈ [0 , 1] o × d
Revised: BRKGA chromosome for this graph is a vector c ∈ [0 , 1] o × d

Original: + o + t ×
Revised: + o + t ×

Original: d composed of the following parts
Revised: d composed of the following parts

Original: The ﬁrst o × d entries in c represent the node-device afﬁnities, one value for
Revised: Node-Device Affinities Node Scheduling Priorities Tensor Transfer Priorities

Original: each
Revised: Node Node Node Node

Revised: Node

Original: node
Revised: Node

Revised: Tensor A

Original: , device).
Revised: Tensor B Tensor C Dev Dev

Revised: Dev

Revised: Dev

Revised: Dev

Original: Each node is
Revised: Dev 2 Dev 1 Dev 2 Dev 1 Dev 2 Dev 1 Dev0. A B 3 C The ﬁrst o × d entries

Original:  assigned to the device for which it has the highest value in the chromosome.
Revised: c the node- afﬁnities, one value for each (node, device). Each node is assigned to the device for which it has the highest value in the chromosome.

Original: The next o entries represent the node scheduling priorities.
Revised: The next o entries represent the node scheduling priorities.

Original: A valid schedule of the computation graph is obtained by performing a topological sort over the nodes of the graph, breaking ties using the node scheduling priorities.
Revised: A valid schedule of the computation graph is obtained by performing a topological sort over the nodes of the graph, breaking ties using the node scheduling priorities.

Original: Nodes with higher priority are scheduled ﬁrst. 3.
Revised: Nodes with higher priority are scheduled ﬁrst. 3.

Original: The ﬁnal t × d entries represent the tensor transfer priorities, one entry for each (tensor, device) pair.
Revised: The ﬁnal t × d entries represent the tensor transfer priorities, one entry for each (tensor, device) pair.

Original: These priorities determine the order in which tensors are transferred across devices.
Revised: These priorities determine the order in which tensors are transferred across devices.

Original: An example of a chromosome encoding is shown in Figure 8 for a graph with o = 3 nodes, t =tensors and d = 2 devices.
Revised: An example of a chromosome encoding is shown in Figure 9 for a graph with o = 3 nodes, t =tensors and d = 2 devices.

Original: As per the example, nodes 1 and 3 are placed on device 1 while node 2 is placed on device 2.
Revised: As per the example, nodes 1 and 3 are placed on device 1 while node 2 is placed on device 2.

Original: The scheduling order over the nodes is 1, 2, 3.
Revised: The scheduling order over the nodes is 1, 2, 3.

Original: Since nodes 1 and 2 are placed on different devices, tensors A and C must be transferred from device 1, where they are produced, to their consumer, node 2, which is on device 2.
Revised: Since nodes 1 and 2 are placed on different devices, tensors A and C must be transferred from device 1, where they are produced, to their consumer, node 2, which is on device 2.

Original: As per the tensor transfer priorities, tensor C is transferred before tensor A since tensor C has a higher priority to get transferred to device 2.
Revised: As per the tensor transfer priorities, tensor C is transferred before tensor A since tensor C has a higher priority to get transferred to device 2.

Original: A. A CTION DEQUANTIZATION
Revised: A. A CTION DEQUANTIZATION

Original: Each real number in a BRKGA chromosome is sampled from its own Beta distribution, which is parameterized by two real numbers α and β .
Revised: Each real number in a BRKGA chromosome is sampled from its own Beta distribution, which is parameterized by two real numbers α and β .

Original: To be more precise, if we denote the chromosome by c ∈ R L , then c i ∼ D ( α i , β i ) ∀ 1 ≤ i ≤ L where D ( α i , β i ) is a Beta distribution with parameters
Revised: To be more precise, if we denote the chromosome by c ∈ R L , then c i ∼ D ( α i , β i ) ∀ 1 ≤ i ≤ L where D ( α i , β i ) is a Beta distribution with parameters

Original: α i and β i .
Revised: α i and β i .

Original: To be able to run BRKGA, REGAL must propose the values for α i and β i for each i .
Revised: To be able to run BRKGA, REGAL must propose the values for α i and β i for each i .

Original: As described in A.5, the BRKGA chromosome consists of three parts.
Revised: As described in A.5, the BRKGA chromosome consists of three parts.

Original: In REGAL, we optimize the
Revised: In REGAL, we optimize the

Original: RL agents over choices of α i and β i for the ﬁrst two parts of the chromosome, i.e., the parts of the chromosome corresponding to the node placement and scheduling decisions.
Revised: RL agents over choices of α i and β i for the ﬁrst two parts of the chromosome, i.e., the parts of the chromosome corresponding to the node placement and scheduling decisions.

Original: The Beta distribution parameters for the tensor transfer priorities are ﬁxed to α
Revised: The Beta distribution parameters for the tensor transfer priorities are ﬁxed to α

Original: i =
Revised: i =

Original: β i = 1 which correspond to the uniform random distribution.
Revised: β i = 1 which correspond to the uniform random distribution.

Original: Thus, for a graph with o ops and with d devices, the RL agent must propose ( d + 1) × 2 values for each of the o ops in the graph.
Revised: Thus, for a graph with o ops and with d devices, the RL agent must propose ( d + 1) × 2 values for each of the o ops in the graph.

Original: To make the learning task easier, rather than directly predicting values of α and β , we quantize the output space of the RL agent’s actions such that each action uniquely maps to a Beta distribution.
Revised: To make the learning task easier, rather than directly predicting values of α and β , we quantize the output space of the RL agent’s actions such that each action uniquely maps to a Beta distribution.

Original: This mapping is done as follows:
Revised: This mapping is done as follows:

Original: • For each node in the graph 1 ≤ i ≤ o , and for each entry 1 ≤ j ≤ ( d + 1) correponding to node i in the BRKGA chromosome, the agent performs the set of actions 
Revised: • For each node in the graph 1 ≤ i ≤ o , and for each entry 1 ≤ j ≤ ( d + 1) correponding to node i in the BRKGA chromosome, the agent performs the set of actions [Equation] [Equation] ,d

Original: [Equation]
Revised: [Equation]

Revised: 1 , v i,d +1 } .

Original: • m ij and v ij ∈ { 0 , 1 , . .
Revised: • m ij and v ij ∈ { 0 , 1 , . .

Original: , k − 1 } where k is some ﬁxed constant greater than 1, and these represent the quantized mean µ ij and variance σ 2 ij of the Beta distribution which are related to each other as follows:
Revised: , k − 1 } where k is some ﬁxed constant greater than 1, and these represent the quantized mean µ ij and variance σ 2 ij of the Beta distribution which are related to each other as follows:

Original: [Equation]
Revised: [Equation]

Original: • µ ij and σ 2 ij can be mapped to α ij and β ij for a Beta distribution as follows:
Revised: • µ ij and σ 2 ij can be mapped to α ij and β ij for a Beta distribution as follows:

Original: [Equation]
Revised: [Equation]

Original: • The values m ij and v ij are sampled from a Categorical distribution whose logits are determined by MLP a ( h i )
Revised: • The values m ij and v ij are sampled from a Categorical distribution whose logits are determined by MLP a ( h i )

Original: [ k ∗ ( j − 1) :
Revised: [ k ∗ ( j − 1) :

Original: k ∗ j ] .
Revised: k ∗ j ] .

Original: We use a similar quantization strategy for the BRKGA crossover probabilities.
Revised: We use a similar quantization strategy for the BRKGA crossover probabilities.

Original: For every crossover probability, we sample an integer c ∈ { 0 , 1 , . . .
Revised: For every crossover probability, we sample an integer c ∈ { 0 , 1 , . . .

Original: , k − 1 } from a Categorical distribution for some ﬁxed integer constant k , and the dequantized crossover probability is given by 0 .
Revised: , k − 1 } from a Categorical distribution for some ﬁxed integer constant k , and the dequantized crossover probability is given by 0 .

Original: 5 ∗ (cid:0) 1 + c +1 k (cid:1)
Revised: 5 ∗ (cid:0) 1 + c +1 k (cid:1)

Original: A. E XTRA MODEL DETAILS
Revised: A. E XTRA MODEL DETAILS

Original: MLPs Multi-layer perceptrons, or multi-layer fully connected neural networks are models that map input vectors to output vectors through layers of linear transformations and nonlinear activation functions, like the following:
Revised: MLPs Multi-layer perceptrons, or multi-layer fully connected neural networks are models that map input vectors to output vectors through layers of linear transformations and nonlinear activation functions, like the following:

Original: [Equation]
Revised: [Equation]

Original: where x is an input vector, ( W i , b i ) are the parameters for the i th layer, and h is the output vector.
Revised: where x is an input vector, ( W i , b i ) are the parameters for the i th layer, and h is the output vector.

Original: σ is a nonlinear scalar function applied element-wise to the input vectors.
Revised: σ is a nonlinear scalar function applied element-wise to the input vectors.

Original: Typical choices include the logistic sigmoid function σ ( x ) =
Revised: Typical choices include the logistic sigmoid function σ ( x ) =

Original: [Equation]
Revised: [Equation]

Original: − x , tanh function σ ( x ) =

Original: [Equation]

Original: x and the ReLU function σ ( x ) = max { 0 , x } .
Revised: and the ReLU function σ ( x ) = max { 0 , x } .

Original: RNNs and LSTMs Recurrent neural networks (RNNs) are good sequence models.
Revised: RNNs and LSTMs Recurrent neural networks (RNNs) are good sequence models.

Original: Typical RNNs contains a recurrent memory c t that is updated recursively by taking some input at each step t as the following:
Revised: Typical RNNs contains a recurrent memory c t that is updated recursively by taking some input at each step t as the following:

Original: [Equation]
Revised: [Equation]

Original: where x t is the input at step t .
Revised: where x t is the input at step t .

Original: The simplest RNN cell has the following form
Revised: The simplest RNN cell has the following form

Original: [Equation]
Revised: [Equation]

Original: where W , b are the parameters and σ is a nonlinearity.
Revised: where W , b are the parameters and σ is a nonlinearity.

Original: Long-short term memory (LSTM) models are a type of RNNs that uses explicit gating to control the access to the memory.
Revised: Long-short term memory (LSTM) models are a type of RNNs that uses explicit gating to control the access to the memory.

Original: LSTMs distinguish the memory c t and the output of the LSTM h t as two sets of vectors, and compute the update at step t as
Revised: LSTMs distinguish the memory c t and the output of the LSTM h t as two sets of vectors, and compute the update at step t as

Original: [Equation]
Revised: [Equation]

Original: Here i , f , o are the input, forget and output gates and (cid:12) is element-wise multiplication.
Revised: Here i , f , o are the input, forget and output gates and (cid:12) is element-wise multiplication.

Original: The carefully designed memory access control through gating makes LSTMs better at modeling long-term dependencies.
Revised: The carefully designed memory access control through gating makes LSTMs better at modeling long-term dependencies.

Original: A. A UTOREGRESSIVE PREDICTION MODELS
Revised: A. A UTOREGRESSIVE PREDICTION MODELS

Original: We can use an autoregressive model to capture some structure in the outputs.
Revised: We can use an autoregressive model to capture some structure in the outputs.

Original: Given the node representations h v for each of the nodes from the GNN, we can utilize an ordering of the nodes, e.g. from a topological sort, and treat the node representations as a sequence, and then use an LSTM Hochreiter & Schmidhuber (1997) to predict the outputs y v sequentially.
Revised: Given the node representations h v for each of the nodes from the GNN, we can utilize an ordering of the nodes, e.g. from a topological sort, and treat the node representations as a sequence, and then use an LSTM (Hochreiter & Schmidhuber, 1997) to predict the outputs y v sequentially.

Original: We tried this approach but found that using an LSTM on top of the h v ’s to predict y v ’s did not perform as well as the conditionally independent model.
Revised: We tried this approach but found that using an LSTM on top of the h v ’s to predict y v ’s did not perform as well as the conditionally independent model.

Original: The reasons for this might be: (1) the autoregressive approach relies on a sequential ordering of the nodes, and this ordering might not be reliable nor consistent across graphs; (2) the number of nodes in the computation graphs can be large, and learning recurrent models on long sequences is known to be challenging Pascanu et al.
Revised: The reasons for this might be: (1) the autoregressive approach relies on a sequential ordering of the nodes, and this ordering might not be reliable nor consistent across graphs; (2) the number of nodes in the computation graphs can be large, and learning recurrent models on long sequences is known to be challenging Pascanu et al.

Original: (3) the noisy training signal in our REINFORCE-based training setup makes this model even more difﬁcult to train.
Revised: (3) the noisy training signal in our REINFORCE-based training setup makes this model even more difﬁcult to train.

Original: A. P LACEMENT AND SCHEDULING EXAMPLE
Revised: A. P LACEMENT AND SCHEDULING EXAMPLE

Original: Figure 9 illustrates a computation graph, a valid schedule, and how we account for which tensors are in memory at a given time under the model presented in sec. 3.1.
Revised: Figure 10 illustrates a computation graph, a valid schedule, and how we account for which tensors are in memory at a given time under the model presented in sec. 3.1.

Original: A
Revised: A

Original: B
Revised: B

Original: C
Revised: C

Original: D
Revised: D

Original: E
Revised: E

Original: Op
Revised: Op

Original: Inmem.
Revised: Inmem.

Original: (dev.
Revised: (dev.

Original: Inmem.
Revised: Inmem.

Original: (dev.
Revised: (dev.

Original: Run
Revised: Run

Original: A, B ∅
Revised: A, B ∅

Original: Transfer B
Revised: Transfer B

Original: A, B
Revised: A, B

Original: B
Revised: B

Original: Run
Revised: Run

Original: A, C
Revised: A, C

Original: B
Revised: B

Original: Run
Revised: Run

Original: C
Revised: C

Original: B, D
Revised: B, D

Original: Run
Revised: Run

Original: C, E
Revised: C, E

Original: D
Revised: D

Original: Transfer D
Revised: Transfer D

Original: E, D
Revised: E, D

Original: D
Revised: D

Original: Run
Revised: Run

Original: D, E ∅
Revised: D, E ∅

Original: A. B ASELINES
Revised: A. B ASELINES

Original: Here we provide supplementary information about our baselines.
Revised: Here we provide supplementary information about our baselines.

Original: • CP SAT:
Revised: • CP SAT:

Original: The multi-device peak memory minimization problem is formulated for the CP solver using a model of the operation execution order, tensor lifetimes, and cumulative constraints to evaluate peak memory usage.
Revised: The multi-device peak memory minimization problem is formulated for the CP solver using a model of the operation execution order, tensor lifetimes, and cumulative constraints to evaluate peak memory usage.

Original: The solver is guaranteed to ﬁnd the globally optimal solution given sufﬁcient time.
Revised: The solver is guaranteed to ﬁnd the globally optimal solution given sufﬁcient time.

Original: • Graph partition + DFS: The graph partitioning objective is to minimize data transferred across devices.
Revised: • Graph partition + DFS: The graph partitioning objective is to minimize data transferred across devices.

Original: We use the modiﬁed implementation of the Kernighan-Lin algorithm Kernighan & Lin (1970) used by XLA for device placement in some settings.
Revised: We use the modiﬁed implementation of the Kernighan-Lin algorithm (Kernighan & Lin, 1970) used by XLA for device placement in some settings.

Original: This implementation is generally slower than heuristics implemented in popular libraries like METIS Karypis & Kumar (1998) although it tends to ﬁnd better quality solutions.
Revised: This implementation is generally slower than heuristics implemented in popular libraries like METIS (Karypis & Kumar, 1998) although it tends to ﬁnd better quality solutions.

Original: • Local Search:
Revised: • Local Search:

Original: The initial schedule is a topological sort order of the ops, and the initial placement selects devices uniformly randomly.
Revised: The initial schedule is a topological sort order of the ops, and the initial placement selects devices uniformly randomly.

Original: A local move either changes the device assignment of the op, or changes the op order in the current schedule.
Revised: A local move either changes the device assignment of the op, or changes the op order in the current schedule.

Original: The hyperparameters (e.g., number of random restarts) are set to values that perform the best on a sample of 10,000 graphs in the training set as found by grid search.
Revised: The hyperparameters (e.g., number of random restarts) are set to values that perform the best on a sample of 10,000 graphs in the training set as found by grid search.

Original: • Tuned BRKGA:
Revised: • Tuned BRKGA:

Original: The following hyperparameters of BRKGA using grid search: the Beta distribution parameters (two scalars), and the number of chromosomes, elites, mutants, and populations.
Revised: The following hyperparameters of BRKGA using grid search: the beta distribution parameters (two scalars), and the number of chromosomes, elites, mutants, and populations.

Original: The grid search tries 648 hyperparameter settings and picks the best one as evaluated on 10,000 training set graphs.
Revised: The grid search tries 648 hyperparameter settings and picks the best one as evaluated on 10,000 training set graphs.

Original: • REGAL: The performance of REGAL is stochastic both because the actions are stochastic and because BRKGA itself depends on random samples for mutation and crossover.
Revised: • REGAL: The performance of REGAL is stochastic both because the actions are stochastic and because BRKGA itself depends on random samples for mutation and crossover.

Original: We estimated the standard deviation of the percent improvement statistics with respect to these sources of randomness as below 0.1%, which is small compared to the differences we observe.
Revised: We estimated the standard deviation of the percent improvement statistics with respect to these sources of randomness as below 0.1%, which is small compared to the differences we observe.

Original: Hence we have omitted the error bars from ﬁgures 3.
Revised: Hence we have omitted the error bars from ﬁgures 3.

Original: A. R UNNING TIME COMPARISON FOR ALGORITHMS
Revised: A. R UNNING TIME COMPARISON FOR ALGORITHMS

Original: Table 3 shows the average running times of the various algorithms on the TensorFlow test set and the XLA dataset, as measured on an Intel Xeon E5-1650 3.60GHz machine.
Revised: Table 3 shows the average running times of the various algorithms on the TensorFlow test set and the XLA dataset, as measured on an Intel Xeon E5-1650 3.60GHz machine.

Original: The times are averaged over the unaugmented graphs in the test set.
Revised: The times are averaged over the unaugmented graphs in the test set.

Original: REGAL provides both fast running time as well as high solution quality.
Revised: REGAL provides both fast running time as well as high solution quality.

Original: For a slightly higher running time than BRKGA 5K, REGAL improves the solution quality signiﬁcantly.
Revised: For a slightly higher running time than BRKGA 5K, REGAL improves the solution quality signiﬁcantly.

Original: Almost all of the added running time is due to extra cost of sampling beta distributions by REGAL compared to uniform distributions by BRKGA.
Revised: Almost all of the added running time is due to extra cost of sampling beta distributions by REGAL compared to uniform distributions by BRKGA.

Original: This can be seen from  the nearly identical running times of REGAL and Tuned BRKGA, which also uses beta distributions, but without the neural network policy.
Revised: This can be seen from [Table] [Table] the nearly identical running times of REGAL and Tuned BRKGA, which also uses beta distributions, but without the neural network policy.

Original: The local search heuristic runs slowly because it was not implemented efﬁciently, e.g., with incremental evaluations of the objective; we show its timing results for completeness only.[Table] [Table]
Revised: The local search heuristic runs slowly because it was not implemented efﬁciently, e.g., with incremental evaluations of the objective; we show its timing results for completeness only.

Original: A. A BLATION ANALYSIS OF AGENT ACTION TYPES
Revised: A. A BLATION ANALYSIS OF AGENT ACTION TYPES

Original: REGAL can train a policy to generate any subset of the following actions for BRKGA: 1) Actions for node placement priorities, and 2) Actions for node scheduling priorities.
Revised: REGAL can train a policy to generate any subset of the following actions for BRKGA: 1) Actions for node placement priorities, and 2) Actions for node scheduling priorities.

Original: We train REGAL with various subsets of these actions and compare their performance against each other in table 4.
Revised: We train REGAL with various subsets of these actions and compare their performance against each other in table 4.

Original: We observe that on the validation set, REGAL performs best when it has to learn actions for both placement and scheduling compared to just scheduling or placement alone.
Revised: We observe that on the validation set, REGAL performs best when it has to learn actions for both placement and scheduling compared to just scheduling or placement alone.

Original: A. A NALYSIS OF THE STRUCTURE OF PLACEMENT AND SCHEDULING DECISIONS ,
Revised: A. A NALYSIS OF THE STRUCTURE OF PLACEMENT AND SCHEDULING DECISIONS ,

Original: Our best model trained on the TF peak memory dataset is capable of generating 16 different kinds of actions for node placement decisions and 4 different kind of actions for node scheduling decisions.
Revised: Our best model trained on the TF peak memory dataset is capable of generating 16 different kinds of actions for node placement decisions and 4 different kind of actions for node scheduling decisions.

Original: Each of these actions determine the shape of the Beta distributions from which we sample the node-device afﬁnities and the node scheduling priorities.
Revised: Each of these actions determine the shape of the Beta distributions from which we sample the node-device afﬁnities and the node scheduling priorities.

Original: In this section we attempt to gain insights into the structure of these actions.
Revised: In this section we attempt to gain insights into the structure of these actions.

Original: We divide our node placement decisions into three categories:
Revised: We divide our node placement decisions into three categories:

Original: • Actions that give a node a higher probability to be placed on device• Actions that give a node a higher probability to be placed on device• Actions that give equal preference to the two devices.
Revised: • Actions that give a node a higher probability to be placed on device• Actions that give a node a higher probability to be placed on device• Actions that give equal preference to the two devices.

Original: Similarly, we divide our node scheduling decisions in two categories:
Revised: Similarly, we divide our node scheduling decisions in two categories:

Original: • Actions that give nodes a “high” scheduling priority.
Revised: • Actions that give nodes a “high” scheduling priority.

Original: • Actions that give node a “low” scheduling priority.
Revised: • Actions that give node a “low” scheduling priority.

Original: Finally, we aggregate the average relative memory consumption of all nodes that were assigned the same set of actions, where memory consumption of a node is deﬁned as the sum of the memory uses of all its input and output tensors.
Revised: Finally, we aggregate the average relative memory consumption of all nodes that were assigned the same set of actions, where memory consumption of a node is deﬁned as the sum of the memory uses of all its input and output tensors.

Original: The relative memory usage is the memory usage normalized by the largest memory usage of a node in the graph.
Revised: The relative memory usage is the memory usage normalized by the largest memory usage of a node in the graph.

Original: We plot this data in Figure 10.

Original: On the right, each cell represents the average relative memory consumption of the nodes that were assigned a particular placement and scheduling decision (darker cells indicate nodes with higher average relative memory consumption).

Original: On the left, each cell represents the frequency of the nodes that were assigned a particular placement and scheduling decision (darker cells indicate higher frequency).

Original: With these we can make the following observations:

Original: L o w
Revised: L o w

Original: H i gh
Revised: H i gh

Original: Node scheduling priority
Revised: Node scheduling priority

Original: Same
Revised: Same

Original: Device Device
Revised: Device Device

Original: D e v
Revised: D e v

Original: i ce
Revised: i ce

Original: a ff i n it y
Revised: a ff i n it y

Original: L o w
Revised: L o w

Original: H i gh
Revised: H i gh

Original: Node scheduling priority
Revised: Node scheduling priority

Original: Same
Revised: Same

Original: Device Device
Revised: Device Device

Original: D e v
Revised: D e v

Original: i ce
Revised: i ce

Original: a ff i n it y 0.
Revised: a ff i n it y 0.

Original:  • On average, nodes with higher normalized memory consumption are assigned lower scheduling priorities.
Revised: in • On average, nodes with higher normalized memory consumption are assigned lower scheduling priorities.

Original: • Most of the nodes with the highest relative memory consumption have no afﬁnity for either of the two devices.
Revised: • Most of the nodes with the highest relative memory consumption have no afﬁnity for either of the two devices.

Original: • For nodes with the lowest relative memory consumption, most of them have an afﬁnity to be placed on device 2, while a smaller but still signiﬁcant number of them prefer device 1.
Revised: • For nodes with the lowest relative memory consumption, most of them have an afﬁnity to be placed on device 2, while a smaller but still signiﬁcant number of them prefer device 1.

Original: This implies that the node placement strategy is more complicated than a trivial "place lighter nodes on device 2” strategy and REGAL’s actions are non-trivially dependent on the input.
Revised: This implies that the node placement strategy is more complicated than a trivial "place lighter nodes on device 2” strategy and REGAL’s actions are non-trivially dependent on the input.

Original: A. H YPERPARAMETERS OF THE BEST AGENT FOR PEAK MEMORY TF
Revised: A. H YPERPARAMETERS OF THE BEST AGENT FOR PEAK MEMORY TF

Original: The graph neural network had a state size of 32 for each node and edge, 16 propagations, all networks MLP n MLP e MLP node
Revised: The graph neural network had a state size of 32 for each node and edge, 16 propagations, all networks MLP n MLP e MLP node

Original: MLP msg MLP (cid:48)
Revised: MLP msg MLP (cid:48)

Original: msg being two layers of size 32, the aggregation used was mean pooling.
Revised: msg being two layers of size 32, the aggregation used was mean pooling.

Original: For faster training, the reward of the training set was made with 1000 ﬁtness evaluations for REGAL and BRKGA ( 4600 for REGAL and 5000 for BRKGA for the validation and test sets).
Revised: For faster training, the reward of the training set was made with 1000 ﬁtness evaluations for REGAL and BRKGA ( 4600 for REGAL and 5000 for BRKGA for the validation and test sets).

Original: Training lasted 100000 gradient steps with each step having a mini-batch of size 4 and with gradient clipping by L 2 norm with value 10 .
Revised: Training lasted 100000 gradient steps with each step having a mini-batch of size 4 and with gradient clipping by L 2 norm with value 10 .

Original: The baseline mean squared error term’s contribution to the overall loss was weighted by 0 . 0001 .
Revised: The baseline mean squared error term’s contribution to the overall loss was weighted by 0 . 0001 .

Original: The optimizer was Adam with beta1 0 .
Revised: The optimizer was Adam with beta1 0 .

Original: 9 beta2 0 .epsilon
Revised: 9 beta2 0 .epsilon

Original: 1 e − 8 and learning rate 0 . 0001 .
Revised: 1 e − 8 and learning rate 0 . 0001 .

Original: The number of devices (for the memory model) was 2.
Revised: The number of devices (for the memory model) was 2.

Original: A.H YPERPARAMETERS OF THE BEST AGENT FOR RUNTIME TF
Revised: A.

Revised: H YPERPARAMETERS OF THE BEST AGENT FOR RUNTIME TF

Original: The graph neural network had a state size of 32 for each node and edge, 16 residual graph propagations, all networks MLP n MLP e MLP node
Revised: The graph neural network had a state size of 32 for each node and edge, 16 residual graph propagations, all networks MLP n MLP e MLP node

Original: MLP msg MLP (cid:48)
Revised: MLP msg MLP (cid:48)

Original: msg being two layers of size 32, the aggregation used was sum.
Revised: msg being two layers of size 32, the aggregation used was sum.

Original: Training lasted 100000 gradient steps with each step having a mini-batch of size 4 and with gradient clipping by L 2 norm with value 10 .
Revised: Training lasted 100000 gradient steps with each step having a mini-batch of size 4 and with gradient clipping by L 2 norm with value 10 .

Original: The baseline mean squared error term’s contribution to the overall loss was weighted by 0 . 0001 .
Revised: The baseline mean squared error term’s contribution to the overall loss was weighted by 0 . 0001 .

Original: The optimizer was Adam with beta1 0 .beta2 0 .
Revised: The optimizer was Adam with beta1 0 .beta2 0 .

Original: 999 epsilon 1 e − 8 and learning rate 0 . 0001 .
Revised: 999 epsilon 1 e − 8 and learning rate 0 . 0001 .

Original: With k = 16 for scheduling and k = 2 for placement (k being the quantization level deﬁned in A.6). 
Revised: With k = 16 for scheduling and k = 2 for placement (k being the quantization level deﬁned in A.6). l ll l ll l lll lll

Revised: lllllllllllllllllllllllll

Revised: l l lll ll lll ll llll ll ll l l ll l lllll l llll l llllllllllll l ll lll l ll ll l l l lll l l l lll

Revised: lll l l l l l l l l l ll l l l l lll lll l l llll

Original: A. H YPERPARAMETERS OF THE BEST AGENT FOR RUNTIME S YNTHETIC
Revised: [Figure] R e w a r d ( h i ghe r i s be tt [Figure]

Revised: Test graphs, by unique topology

Revised: A. H YPERPARAMETERS

Original: Same as A.15 but with 2 graph propagations and gru node updates and aggregation using mean.
Revised: Same as A.15 but with 2 graph propagations and gru node updates and aggregation using mean.

Revised: A. P ERFORMANCE BY GRAPH ARCHITECTURE

Revised: Figure 12 shows the distribution of rewards on the TF Runtme test set broken down by graph topology.

Revised: As described in Section 5.1, for each unique topology we augment the dataset by perturbing the tensor sizes and op running times.

Revised: This generates a distribution of 100 rewards per topology.

Revised: The variance for a ﬁxed topology is typically relatively small.