We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. 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 offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.
Summary
This paper proposes an ML-based method to optimize TensorFlow Graph execution. Specifically, it combines graph neural networks (GNNs) and BRKGA (a genetic algorithm) to search over the joint space of TF node-device placement and scheduling. The core claims on the advantages of this method are that (1) it co-searches placement and scheduling space, (2) the trained model can generalize to different graphs and inference cost is very  small. The experimental results show that REGEL can outperform a few baseline methods on this problem.
Writing
- The paper is well-written and I enjoyed reading the paper.
- Some more descriptions about the BRKGA algorithm could be added in.
Method and Results
Some confusion if the authors could answer:
- I am very confused by one of the claims that “ the first work on learning a policy for jointly optimizing placement and scheduling”. I don’t see much evidence in the result section about showing the co-searching the joint space yield advantages? I am fairly familiar with the line of work on only optimizing device placement, but it would be good to see some ablation studies showing search over the joint space is advantageous. 
- The model is trained with standard REINFORCE -- how many training time and resources are needed to train a REGEL model for a task? How’re the training dynamics looking like (variance, convergence, etc?)? 
- In terms of the generalization ability of REGEL, the paper has clearly shown that REGEL is able to generalize to differently shaped graphs, with acceptable cost, but I am wondering for the same dataflow graph, how REGEL generalizes to different input data configurations (size, modality, etc.)? E.g. if the batch size of the input data is changed, the execution time of each kernel and their memory usage (in general, the system treatment) would change; Can a trained REGEL model on a data config A generalize to B? How would this affect the performance of REGEL?
- It seems the method and assumptions about graphs or training data are pretty coupled with TensorFlow and graph-mode execution, how could the method be generalized to other ML frameworks (e.g. frameworks with eager execution)
- Could the authors clarify why the two methods mentioned in “Learning to directly predict a solution” has quadratic complexity w.r.t. # of nodes and whereas REGEL is linear? 
- Confusion on Figure 4(b): Could some more critical statistics about the graphs in the training/test dataset be reported? e.g. what’s the average depth of the training graphs? When there are 32 MP layers a node’s feature will be passed across its 32-hop neighborhood, which seems surprising as it is common to observe GNN starts degenerating with increased depth (because all node features become similar during message passing)