Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning
Authors: Shauharda Khadka
Estelle Aflalo
Mattias Marder
Avrech Ben-David
Santiago Miret
Shie Mannor
Tamir Hazan
Hanlin Tang
Somdeb Majumdar
Venue: ICLR 2021 Poster (2021)
Keywords: Reinforcement Learning
Memory Mapping
Device Placement
Evolutionary Algorithms

Abstract: For deep neural network accelerators, memory movement is both energetically expensive and can bound computation. Therefore, optimal mapping of tensors to memory hierarchies is critical to performance. The growing complexity of neural networks calls for automated memory mapping instead of manual heuristic approaches; yet the search space of neural network computational graphs have previously been prohibitively large. We introduce Evolutionary Graph Reinforcement Learning (EGRL), a method designed for large search spaces, that combines graph neural networks, reinforcement learning, and evolutionary search. A set of fast, stateless policies guide the evolutionary search to improve its sample-efficiency. We train and validate our approach directly on the Intel NNP-I chip for inference. EGRL outperforms policy-gradient, evolutionary search and dynamic programming baselines on BERT, ResNet-101 and ResNet-50. We additionally achieve 28-78% speed-up compared to the native NNP-I compiler on all three workloads.  

Published as a conference paper at ICLR 2021
OPTIMIZING MEMORY PLACEMENT USING EVOLUTIONARY GRAPH REINFORCEMENT LEARNING
Shaubarda Khadka
Estelle Aflalo
Mattias Marder
Avrech Ben-David Technion
Intel Labs
Intel Israel
Intel Israel
Santiago Miret
Shie Mannor Technion
Tamir Hazan Technion
Hanlin Tang Intel Labs
Somdeb Majumdar Intel Labs
Intel Labs
ABSTRACT
For deep neural network accelerators, memory movement is both energetically expensive and can bound computation. Therefore, optimal mapping of tensors to memory hierarchies is critical to performance. The growing complexity of neural networks calls for automated memory mapping instead of manual heuristic approaches; yet the search space of neural network computational graphs have pre- viously been prohibitively large. We introduce Evolutionary Graph Reinforcement Learning (EGRL), a method designed for large search spaces, that combines graph neural networks, reinforcement learning, and evolutionary search. set of fast, stateless policies guide the evolutionary search to improve its sample-efficieng We train and validate our approach directly on the Intel NNP-I chip for inference. EGRL outperforms prrdcy,gradient, evolutionary search and dynamic programming baselines on BERT, ResNet- 101 and ResNet-50. We additionally achieve 28-78% speed-up compared to the native NNP-I compiler on all three workloads.
INTRODUCTION
The proliferation of deep learning (DL) has been fueled, part, by a rapid growth in the size and complexity of deep neural networks (DNN) (Dean et al. 2001 2 Ying et al. 2018). This has spurred the rapid development of hardware (Wang et al.] [0166] Jouppouuppi et al. 2017) and software (Abadi et al.
2016 Paszke et al.) 2018 Cyphers et al. 2018 dedicated to deep learning workloads that seek to optimize critical performance metrics like throughput and power efficiency (Mattson et al. 2020). Producing compiler optimizations that map the tensors of neural network's computational graph to the memory units on host hardware is critical challenge. Since different memory types trade off bandwidth and capacity differently, sub-optimal mapping could significantly increase latency.
For DL inference, the computational graph is static and placement can be pre-planned instead of relying on online cache management (Zhang al., 2020 Shi al.) 2019). However, this is especially challenging with DNNs due to the high dimensional search space. For example. ResNet-50 (He et al.) 2016) has 57 operational layers. Mapping each activation and weight tensor to, for example, three (DRAM, LLC, and SRAM) memory caches represents 3(2+57) x 1054 possible decisions. BERT Devlln et al. 2018) has 376 operational layers, and search space of ~ 10 Since optimizing this mapping intractable with traditional approaches, such dynamic programming (Bellman) 1954). current solutions primarily rely on manually-tunee heuristic rules encoded in a compiler.
Because of the large search space, prior reinforcement learning (RL) algorithms for automating mappings have relied on manually-designel grouping (Mirhoseini et al., 2017] 7 Addanki et al.) 2018)
or learned grouper whose hierarchical structure is domain dependent (Mirhoseini et al. | [0188). In
addition to the extremely large action space, the large number of nodes render the reward sparse and noisy, and thus further unsuitable for gradient-basee Deep RL algorithms. This sparsity stems from the fact that overall performance metric can only be measured after all nodes have been processed.
"quual Contribution Correspondence <sooddb.mmjjumde @ntel.com>>

---

Published as a conference paper at ICLR 2021
Workload
Graph
RL Agent
Memory Map
Node
DRAM
Conv
A: SRAM
Node
W: SRAM
ReLu
GNN
A: DRAM
DRAM
Node 3
Linear
A: SRAM
Graph Representation of Workload
Graph Representation of Memory Map
Figure 1: Workflow of Graph RI agent mapping weights (W) and activations (A) of each layer of a trained neural network workload various on-board memory components (e.g. DRAM, SRAM).
In this paper, we present Evolutionary Graph Reinforcement Learning (EGRL), a hybrid approach of evolutionary search with gradient-based learning, that is able to natively search in high-dimensional space that is wrders-of-magnitus
larger than previous approaches. EGRL is an extension of CERL (Khadka al. 2019), a population based method for sparse- reward tasks that combines fast policy gradient (PG) leanning with a stable evolutionary algorithm (EA). Since the action spaces explored this paper are several orders of magnitude larger than those explored CERL,, we introduce Boltzmann chromosomes a set of fast, stateless policies that accelerate evolution by providing partially optimized solutions as anchors. This mechanism is necessary to improve the sample- efficiency of the slow EA component for this large action space. Further, we employ graph neural network (GNN) (Wu et al. 2020 Scarselli et al., 2008) to represent our policy. This allows our agent to natively process computational graphs representing deep learning workloads, enabling generalization over workloads of varying size and connectivity.
We demonstrate solution the Intel Neural Network Processor for Inference (NNP-I), deep learning accelerator, to map modern neural networks on one of the three memory hierarchies on the chip. Each memory level in this chip has trade-offs in memory size and bandwidth, as detailed in Weehsler et al.| (2019). This additionally differentiates our work from prior works such as REGAL (Paliwal et al. (2200) that assume infinite bandwidth and memory that are not practical on real hardware. Additionally, we consider single-batch inference, an important industry benchmark Mattson et al.] (2020). While large batch sizes have greater computational efficiency (.o.dokuu et (l.20)) on NÑP-I), they are sub-optimal for a given inference example due to the latency associated with queuing up batch. Therefore, single-batch inference key to many time-critical applications (Park et al. 2018) where individual inference query needs to be processed in real-time.
Results on ResNet-50, ResNet-101 (He et al., 2016) and BERT, show that EGRL significantly outperforms the chipset's native compiler across all workloads, and exceeds the performance of dynamic programming, evolutionary search and policy-gradiess approaches.
Specifically, the contributions of this work are:
1. A generalized GNN-based policy that can natively accept a computational graph and produce a corresponding graph representation with the optimal memory maps. This eliminates the need for serialized, layer-dependent representations.
EGRL, a scalable population-based algorithm that can effectively train on sparse and noisy feedback from the host hardware in large search spaces.
3. An RL agent that trains directly the hardware, with feedback mechanism for constraint violation, and thus allowing direct deployment and testing on hardware.
2

---

Published as a conference paper at ICLR 2021
2 RELATED WORK
Optimizing Hardware using Machine Learning: this work, we study the problem of mapping tensors to memory components a device. Several recent works have studied the use of machine learning to optimize the execution of computation graphs on hardware -a similar combinatorics problem. Mirhoseini al. (2017) designed a policy gradient (PG) based Placer policy to map parts of neural models on hardware. However, it relied on a heuristic grouping strategy to significantly reduce the action space. Mirhoseini et al. (2018) improved this architecture by replacing the heuristic module with a Grouper policy. While this does represent an end-to-end PG approach, it is significantly more than EGRL's PG network. Specifically, their Placer is LSTM based Seq-to-Seq model with attention. The hierarchical structure is domain dependent specific to operation grouping. Mirhoseini et al. (2020) also applied Deep RL to learn subsystem placement to optimize power, performance and chip area but relied on similar heuristic grouping strategies to significantly reduce the action space seen by the agent.
complex and hyperparameter-hean
A closely related work is Placeto (Addanki et al. 2018)) where the nodes of a computation graph are sequentially placed on hardware. Similar to EGRL, they also operate on a GNN representation. EGRL primarily differs from Placeto in simultaneously mapping all nodes. We found empirically that a sequential mapping strategy was significantly more sample-inefficient and could not scale to larger workloads. Sequential mapping strategies have the additional disadvantage of not being able to exploit parallelism during policy training. Placeto also adopts the manual grouping strategy from Mirhoseini al. (018)) and adds additional Merge and Collocate heuristic.
In contrast, EGRL has a simpler architecture using generic PG and EA to scale to large search spaces. For example, previous work with manual grouping operate at most in 5280 = 10 dimensional action space (Mirhoseini et al. 2017), compared to ~ 10358 for our BERT problem. Compared to pure PG based approaches, EGRL has significantly fewer hyperparameters tune primarily due to the reduced dependency on PG learning by using a population based search to handle sparse rewards.
Another closely related work is REGAL (Paliwal et al.) 2020), which optimizes run-time and peak- memory via hardware placement. It also utilizes graph representation with genetic algorithm (GA) guided by RL. The RL agent predicts the parameters of GA a form of indirect information transfer while GA directly optimizes the final strategy. In contrast, our RL and EA components co-optimize the mapping strategies via direct information transfer (policy migration) and shared replay buffer. REGAL assumes infinite bandwidth, whereas we train and validate entirely on physical hardware introducing specific mechanisms to incentivize compiler-valid mappings. This ensures that solutions are performant under real-works operating conditions and closer to production-use.
As a relatively new research field, we are challenged by the unavailability of reproducible code for prior work. The domain specific heuristics as described above render it difficult to apply algorithms designed for, say, chip placement to our memory mapping problem. Therefore, we adopt a state- of-the-art PG method as baseline, since PG is common central component of the above prior work.
Classical Search Methods: Classical methods such as Simulated Annealing (SA) (Kirkpatrick et al. 1983) and genetic algorithms (GA) have also been studied for problems that have a similar combina- torial search complexity memory placement. SA evolves new solution via small perturbations on existing solutions and retaining solutions that yield improved performance. A temperature param- eter drives the exploration into new solution spaces. Our evolutionary algorithms (EA) (Floreano) et al. 2008 Lüders et al. 2017 Fogel 2006 Spears et al., 1993) improve on SA by systematically evolving new solutions within a population of solutions by performing mutations (similar to SA) and cross-over between pairs of solutions. Both methods are known to produce highly performant and stable solutions but are also significantly slower compared to Deep RI..
In this work, we use EA both as a component of EGRL and also as a baseline. The PG components of EGRL produce fast, neri-performant solutions which then become anchors in the EA module. This essentially "guides" the EA to performant solution by providing better anchors to search around. We demonstrate this via ablation studies that isolate the EA and PG components of EGRL. We also introduce Boltzmann chromosomes in the EA population perturb action proposals to accelerate exploration with a temperature term that balances exploration and exploitation. This component is motivated by SA.
set of stateless policies that directly
3

---

Published as a conference paper at ICLR 2021
Evolutionary RL: Our overall architecture builds top of CERL (Khadka and Tumer 2018 Khadka et al. 2019) which combines EA and PG. It diversifies exploration by allowing a population of EA policies to add data to a central replay buffer shared by population of PG learners. We directly build on CERL because it has been shown to be effective in optimizing sparse feedback signals. Our memory mapping solution inherently relies on optimizing a very sparse feedback signal (e.g., latency that is obtained at the end of an inference run on workload.
For the PG portion of our architecture, we adopt Soft-Actor Critic (SAC) (Haarnoj et al., 2018). the current state-of-the-a model-fire algorithm developed for continuous high-dimensional settings. SAC uses an actor-critic architecture with separate networks for the policy and the Q-value function. A stochastic Gaussian policy enables it to use a maximum entropy objective (Ziebart et al.) 2008) through which demonstrates state-of-the-ar results. We modify SAC to be compatible with our discrete action space.
METHOD
We formulate the hardware mapping problem as a Markov Decision Process (MDP) and apply RL to train a solution. Figure illustrates the high-level formulation of our workflow. A trained neural network (c.g., ResNet-50) is referred to as the workload. We convert the network to a graph representation which acts the input to GNN-based RL policy. The RL agent takes this graph input and proposes a memory location for the weight and activation tensors corresponding to each layer in the input workload, with the policy goal being to maximize performance metric. The RL agent is implemented as a Graph U-Net policy based on Gao and Ji (2019). Algorithm [ details the interaction of the agent with the environment.
In order to evaluate any policy, we run an inference call on hardware using the memory map proposals by replacing the compiler's 's native memory mapping module with our policy.
State: We convert the workload to a directed graph representation where each node represents an operational layer (conv, pooling etc) and the edges denote the connectivity between them. A detailed description of the node features can be found in Appendix A Since all the outgoing edges of a node denote the same output tensor, their associated information are encapsulated in their source node features, leaving the edges featureless. The graph represents all mappable tensors (nodes) simultaneously and thus allows us to have parallelized training pipeline. This compute advantage is primary reason for this design choice compared to serialized representations adopted in some prior works.
Each node contains information about two tensors, for the weights and the activations, that each must be individually mapped to one of three memory units (DRAM, LLC, SRAM) on the chip. The RL agent accepts the input graph and produces the node features of the output graph represent the two memory location proposals to map the weight and activation tensors corresponding to the node. The memory hierarchy of NNP-I is detailed in (Wechsler et al. 2019), with a high capacity (32GB) DRAM, a lower capacity (24MB) but faster (~> 10x) LLC,, and small SRAM (4MB) per chip that is very fast (~ 100x) compared to DRAM.
Actions:
output graph with identical topology, where
The agent's complete memory map M₁ is then sent to the compiler. If any of the mapping decisions cannot be executed on the hardware (i.e., invalid mapping), the compiler rectifies them and outputs a modified map, Mc. that is fully executable (Line 6 in Algorithm I].
Rewards: In standard RL setting, one can generally constrain the action space to avoid invalid actions. However, in our problem setting, constraining the action explicitly requires reproducing the compiler's logic for valid mappings, which would vary across hardware and compiler versions. In order to keep the RL algorithm independent of the compiler logic, we formulate separate reward domains for invalid and valid mappings. If the agent produces any invalid mapping that the compiler re-assigns, we do not execute an inference. Instead, we formulate negative reward as a function of the re-assigned bytes ratio, to quantify the extent of the invalidity (Line 12 in Algorithm []. This formulation allows us to avoid implementing the compiler logic explicitly instead relying on negative feedback that enables the agent to implicitly learn the rules for valid mappings. When the agent does produce fully valid mapping, execute inference and compute a positive reward. This reward is a function of the agent performance score normalized by that of the native compiler (Line 10 in Algorithm I]. While the normalization is not necessary when training on single workload,
4

---

Published as a conference paper at ICLR 2021
Evolutionary Algorithm
Potley
Polleles
Expertences
Seeding Priers
Champion Palicy
Shared Replay Buffer
Bottomoon Pollcic
Figure 2: EGRL Architecture: EA and PG operate concurrently via shared replay buffer. EA comprises sub- uub-populations of GNN and Boltzmann policies. PG policy periodically migrates to EA.
allows for flexibility concurrently train on multiple workloads that might have a wide range of scores. For our application, we maximize the reciprocal of latency as it is an intuitive measure of speed and easily implemented.
Training: Our training algorithm, EGRL, builds on the CERL framework (Khadka al., 2019) to tackle variable-sized, multi- discrete action settings. Figure 2 Iillssrates the high level architecture of EGRL. It is comprised a single PG learner, with a
Algorithm Agent's Interaction with the Environment
1: Initialize workload f. policy я, perf. metric
2: Initialize compiler and graph transform G
3: 9() +
Workload to graph
GNN architecture, and an EA population containing mixture of GNN and stateless Boltzmann policies. Each individual in the population provides a different mapping proposal for given input workload. These proposals are evaluated by performing an inference run on the input workload using the proposed mappings and measuring the resultant latency. During the evaluation 12:
4: for each iteration i do
M2 (9)
Agent's map Compile agent's 's map D Mapping error
6:
Me = C(M
7:
CM = MM
8:
if E = 0 then
9:
S2 I(Mc)
Run inference
10:
MM = Thaseline
2
Positive reward
11:
else
TM -EM
Negative reward
process, all data generated by all policies 13:
Ti++1 <- Ti
Update policy
stored the PG learner's is replay buffer.
The population then undergoes standard EA processes to produce generation of candidate policies.
Concurrently, the PG learner updates its actor and critic by sampling from the shared replay buffer, as typical for off-policy algorithms. The PG learner is periodically copied into the EA population form of information transfer and is allowed to participate in evolution. At any given time, the top-ranked policy in the EA population is chosen for deployment. We adopt SAC as our PG algorithm and modify for our discrete action space setting. A detailed description of our algorithm can be found in Appendix C with our modifications to SAC described in Appendix D
The Boltzmann chromosome is additional policy representation introduced into the population. Each Boltzmann chromosome is parameterized by set of prior probabilities (P) and temperature (T) for each node. To compute an action for each node, we sample from the Boltzmann softmax function (Asadi and Littman 2017) using that node's P and T, making it significantly faster to compute an action compared to neural network policy. higher temperature T results in higher exploration by selecting decisions farther from P. Crucially, T is learned (via evolution) for each node independently which allows for varying degrees of exploration -orrattioo-epllitatiion across different mapping decisions simultaneously. A longer discussion the Boltzmann chromosome is included in Appendix E
The EA population concurrently holds both GNN and Boltzmann policies. All policies share data and benefit from the joint exploration. The PG based GNN policy can directly leverage the states explored by the Boltzmann policy to compute gradients. Conversely, shown in Figure the Boltzmann
5

---

Published as a conference paper at ICLR 2021
policy's prior P is periodically seeded using the GNN policy's posterior probability distribution thus enabling it to directly bootstrap from the GNN population.
4 EXPERIMENTS
We evaluated the performance of EGRL and baseline implementations Intel NNP-I hardware. For a given DNN workload, our agents controlled how their intermediate tensors are mapped to memory units the chip. We then report the resulting latency measured directly in the hardware. We conduct both training and testing entirely on the physical hardware. A complete set of hyperparameter details can be found in Appendix B and code will be open-soureed.
Workloads Tested: We benchmarked our algorithms on three popular neural network workloads. ResNet-50, with 57 nodes, is widely used for benchmarks such as MI.Perf (Reddi et al., (019)). ResNet-101, with 108 nodes, allowed us to test our algorithms at greater scale. Lastly, BERT, with 376 nodes, is a state-of-the-arr natural language processing model. This allowed to test for scale and generalization of our approach. Since the action space for workload with N nodes is 3(2N), the corresponding sizes of the action spaces are 3114 x 1004 (ResNer50), 3216 as 10203 (ResNer101) and 3752
10358 (BERT) respectively.
Metrics Reported: We define speedup as the relative improvement in latency achieved by the agent's mapping versus that of the compiler. score greater than 1 indicates an improvement in latency while a score between O and indicates a degradation. A score of O indicates an invalid mapping. We conduct independent statistical runs and report the mean and standard deviation. Further, we report all speedups against iterations where an iteration refers to an inference process on the physical hardware. We report the iterations cumulatively across the population. Since policies are updated after each iteration, the number of iterations is equivalent to sample-complexig
Baselines: We use the Intel NNP-I's default compiler mapping as our primary baseline. The compiler consists of a collection of heuristic rules specific to the compute and memory capacity of the hardware. This is common to many other hardware subsystems. In all our runs, we simply replace the compiler's memory mapping module with our policy while keeping all other processes, including low-level cache management, the same.
We evaluate pure PG approach by testing the modified SAC-discrete algorithm. This is motivated by prior works like Mirhoseini et al. (2018) and Mirhoseini et al. (2017) which have shown effective results similar hardware mapping problems using PG. We also implement EA baseline given its importance as a classical search method. Since PG and EA are also core components of EGRL, these baselines also serve as ablation studies distil the relative contribution of each component to the final performance improvements.
Finally, we implement Greedy Dynamic Programming (DP) agent, inspired by classical DP methods for optimization (Andonov et al. 2000: Bertsimas and Thicle 2004). It makes layer-wise greedy decisions directly on the workload. It tries all possible maps for the first node choices each for 2 types of tensors), keeping all other mappings static, and chooses the action that maximizes the reward. It repeats this serially for each node in the workload and then conducts several passes. It essentially assumes conditional independence of mapping across N nodes to reduce the solution space from g - 9 s N. While this is a fairly naive assumption, running multiple passes through the graph produces reasonable solution.
RESULTS
We show the speedup achieved, relative to the compiler and measured directly on the NNP-I chip, for the various agents tested the ResNet-50, ResNet- 101 and BERT in Figure EA and EGRL significantly outperform the compiler across all three workloads. Greedy-Di approaches baseline performance while the PG agent in isolation fails to reach it at all.
On 50,, EGRL and EA significantly outperform the baseline compiler and all other agents reaching final speedup of 1.28 and 1.06. respectively. Greedy-DP underperforms the compiler at 0.72 while PG converges to 0.29. On ResNet-1011 101, EGRL again significantly outperforms the baseline compiler and all other agents reaching final speedup of 1.78 while EA converges 1.47.

---

Published as a conference paper at ICLR 2021
Rasnel-50
Rasnel-101
BERT
2.0
1.5
Speedup
1.0
1.5
1.0
1.0
0.5
0.5
0.0
0.0
0.0
0.00
0.75
1.50
2.25
3.00
3.75
0.00
0.75
1,50
2.25
3.00
0.00
0.75
1,50
2.25
3.00
Iteration (thousands)
Iteration (thousends)
lteration (thousands)
Figure 3: Speedup for different workloads, normalized to the heuristic compiler performance. EGRL consistently outperforms all baselines. Error bars indicate standard deviation over runs.
This performance gap demonstrates the role played by the collaborative learning using the shared replay buffer in EGRL. While the PG learner fails to find full mapping solutions by itself, the partial solutions finds carry vital information to the EA population. Greedy-DP outperforms the compiler, converging to 1.27, while PG alone converges to 0.23.
BERT, EGRL and EA significantly outperform the compiler and all other agents reaching a final speedup of 1.66 and 1.64, respectively. Greedy-DD converges to speedup of 0.67, greatly underperforming the compiler. This is unsurprising as BERT is comparatively much larger in size than the two ResNet models which breaks the Greedy-DP conditional independence assumption. PG fails to find good mappings and converges to 0.21.
Generalizability Properties:
While EGRL performs well when trained on each workload individu- ally, we investigated the policy embeddings derived from learning on one workload were performant on second workload. We tested this by taking policy trained on workload and evaluating it a second workload with no further training. This simply involves taking the GNN policy learned on the first workload and replacing the input and output graph representations to be compatible with the topology of the second workload keeping the hidden layers and thus the learned embeddings- untouched. We tested this various stages of training on given first workload.
Figure 4 shows how poli- cies trained on BERT and ResNet-50 transferred to the held-out workloads at different points in their training. While the trans- ferred policies clearly un- derperform those trained from scratch, it is no- table that they still outper- formed the baseline com- piler. While a full study of generalization and transfer learning out of the scope this work, this indicates that the GNN embeddings can be domain invariant. We also observe that the zero-shot performance drops after certain point in training, which indicates the point at which the policy overfits to training workload.
Trained BERT
Trainec on Resnet-50
1.6
1.2
Speedup
Speedup
1.2
0.8
Resnet-5 50
FERT
0.4
0.4
0.75
1.50
2.25
0.00
0.75
1.50
2.25
3.75
lteration (thousands)
Iteration (thousands)
Figure 4: Zero-shot Generalization of policy embeddings: policies
trained on one workload are tested on the others without fine-tuning
We also studied the differences between the mapping strategies of the compiler and EGRL at different stages of training. For ease of analysis, we convert the mappings to two-dimensional UMAP embedding (Mcinnes et al. 2) as shown in Fig For each workload, we collected mappings twice- first when the agent's 's mappings approximately reach the compiler's speedup performance (~~ 1) denoted as compilee-compptiiii-mmappen and second when the agent reaches its best recorded speedup denoted as best-mappings.
Visualizing Memory Mappings:
The
mappings and best-mappings are generally well-separable across all three workloads. Further, the compiler's 's mapping also fell within the cluster of compiler-comper mappings across all three workloads. This suggests that the agents learn to mimic the compiler's mappings at some point in their training. This is unsurprising as the reward we use to train the agents before they find valid mappings is based differences with the compiler. Interestingly, the

---

Published as a conference paper at ICLR 2021
Resnet-50
Resnet-101
BERT
compiar-compedivv--neee masppniinnne
best-mappings
15
complee-mapping
-5
10
-15
-20
-10
10
Figure 5: UMAP projection illustrating mappings that achieve compiler-competiinee performance
(speedup of ~ 1), the best mappings, and the compiler's mapping (highlighted with a red arrow).
Resnet-50
Resreet-11
BERT
0.00
32
0.32
SUAM
Weights
Top:
SRAM
Weights
ReoNett-11: Compllen EORL
Activations
DRAM
For each memory unit the y-axis, the corresponding row shows how EGRL changed the distribution of tensors originally mapped to it by the compiler. Bottom: Memory maps from the compiler vs the best ones found by EGRL for ResNet-50 and ResNet-101. Each bar denotes a tensor operation.
Figure 6: Memory map shifts Top:
intra-cluster spread for all three workloads. This indicates that the mappings associated with higher speedups more self-similar than those that are less performant, which also unsurprising since the number of inferior mappings is higher than that of the superior ones.
comperiiivv-mappiiin is markedly higher than best mast-mappnngs across
Figure 6 illussrates the differences in raw mappings between the compiler and EGRL. The transition matrices on top show how the distribution of tensors to the different memories shifted. Each row corresponds to memory unit. The corresponding columns indicate how EGRL fractionally re- distributed tensors originally mapped to that unit into all available memories. At the bottom, illustrate how each tensor in workload was mapped by the compiler and by EGRL, Each band represents either a weight an activation tensor.
While it is difficult to semantically interpret the mapping decisions reliably, we observe that EGRL generally found memory maps that avoided the slower, but higher-capacity DRAM. This difference particularly prominent for the weight tensors. EGRL also favored contiguity where tensors from neighboring layers generally got mapped to the same type of memory. Both are performant strategies optimize latency but not trivial to achieve using heuristics that need to trade-off speed and capacity for a large number tensors. One hypothesis is that EGRL's graph-basee global view of the workloads enables it to make globally optimal allocations compared to the sequential decision making of the compiler.
DISCUSSION AND FUTURE WORK
This paper introduced EGRL, a hybrid framework to learn effective memory mapping solutions for large deep learning workloads. We train policies end-to-end the NNP-I chipset ensure that
8

---

Published as a conference paper at ICLR 2021
the solutions are robust to the real-world constraints and uncertainties of the chip. We show that EGRL scales effectively across varying sizes and operational types of DL workloads, and exhibits some zero-shot transferability across workloads. Results show that EGRL outperforms several learning and classical search and optimization methods as well as the heuristic logic of the compiler. By combining evolutionary and gradient-based approaches, and including stateless policies accelerate evolution, we efficiently tackle the large search space of this problem. This scalability paves the way for learning-base baaein agents to tackle other hardware mapping problems. Specifically, future work will expand the action space of the EGRL agent to control multivariate settings like batch size, ring frequencies, power efficiency and data decomposition.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In 12th (USENIX) Symposium on Operating Systems Design and Implementation ({OSDI) 16), pages 085-283,12 2016.
Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan. Shreyan Gupta, Hongzi Mao, and Moham- mad Alizadeh. Placeto: Efficient progressive device placement optimization. In NIPS Machine Learning for Systems Workshop, 2018
Byung Hoon Ahn, Jinwon Lee, Jamie Menjay Lin, Hsin-Pa Cheng, Jilei Hou, and Hadi Esmacilzadeh. Ordering chaos: Memory- aware scheduling of irregularly wired neural networks for edge devices. Proceedings of the 3rd MLSys Conference, 2020.
Rumen Andonov, Vincent Poirriez, and Sanjay Rajopadhye. Unbounded knapsack problem: Dynamic programming revisited. European Journal of Operational Research, 123(2):394-40) 407, 2000.
Kavosh Asadi and Michael L Littman. An alternative softmax operator for reinforcement learning. In Proceedings of the 34th International Conference on Machine Leurning-Vont Volume 70, pages 243-252. JMLR. org, 2017.
Richard Bellman. The theory of dynamic programming. Bull. Amer: Math. Soc., 60(6):503-555, 11 1954.
9147
Dimitris Bertsimas and Aurélie Thiele. A robust optimization approach supply chain management. In International Conference on Integer Programming and Combinatorial Optimization, pages 86-100. Springer, 2004.
Guy Boudoukh, Eli Kfir, Ofir Zafrir, Uzi Sarel, Michael Behar, Wasserblat, Galina Ryvchin, Peter Adams, and Kiran Atmakuri.
Moshe
Intel c
Nervana NNP-I shows best-in-class throughput on BERT NLP model.
TM
www
intel. com wwtt///tttttt/ccc http:///ww/////..ataaaaaaaaaaaaa
nervana nnp- i shows wettbestinn-oude wtw.iasssttthouuuuu-uuuuu. bennbberttoo nlp model
html 2020.
Scott Cyphers, Arjun K Bansal, Anahita Bhiwandiwalla, Jayaram Bobba, Matthew Brookhart, Avijit Chakraborty, Will Constable, Christian Convey, Leona Cook, Omar Kanawi, et al. Intel ngraph: An intermediate representation, compiler, and executor for deep learning, arXiv preprint arXiv:1801.08058 2018.
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Mare' aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. In Advances in neural information processing systems, pages 1223-1233 2012.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800. 2018.
Dario Floreano, Peter Dürr, and Claudio Mattiussi. Neuroevolution: from architectures learning.
Evolutionary Intelligence, 1(1):47-62, 2008.
David B Fogel. Evolutionary computation: toward a new philosophy of machine intelligence,
volume 1. John Wiley & Sons, 2006.

---

Published as a conference paper at ICLR 2021
Scott Fujimoto, Herke van Hoof, and Dave Meger. Addressing function approximation error in actor-critic methods. arXiv preprint 10X1::180..09477,230 2018.
Hongyang Gao and Shuiwang Ji. Graph u-nets. arXiv preprint
2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- Off-policy maxi- mum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv: 18X1::8801.01290 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture, pages 1- 2017.
Shauharda Khadka and Kagan Tumer. ivolution-guides policy gradient in reinforcement learning. In Advances in Neural Information Processing Systems, pages 1198-120, 2018.
Shauharda Khadka, Somdeb Majumdar, Tarek Nassar, Zach Dwiel, Evren Tumer, Santiago Miret, Yinyin Liu, and Kagan Tumer. Collaborative evolutionary reinforcement learning. arXiv preprint arXiv:190..0096622, 2019.
Scott Kirkpatrick, C. Gelatt, and Mario Vechhi. Optimization by simulated annealing. In Science, 1983.
Benno Lüders, Mikkel Schläger, Aleksandra Korach, and Sebastian Risi. Continual and one-shot learning through neural networks with dynamic external memory. In European Conference on the Applications of Evolutionary Computation, pages 886-901. Springer, 2017.
Peter Mattson, Vijay Janapa Reddi, Christine Cheng, Cody Coleman, Greg Diamos, David Kanter, Paulius Micikevicius, David Patterson, Guenther Schmuelling. Hanlin Tang, et al. Mlperf: An industry standard benchmark suite for machine learning performance. IEEE Micro, 40(2):8-16. 2020.
Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger. Umap: Uniform manifold approximation and projection. Journal of Open Source Software, 3(29), 2018.
Azalia Mirhoseini, Hieu Pham, Quoe Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement optimization with 2017.
reinforcement learning. arXiv preprint 1rXiv:1706.00
Azalia Mirhoseini, Anna Goldie, Hieu Pham, Benoit Steiner, Quoe V. Le, and Jeff Dean. A hierarchi- cal model for device placement. In International Conference on Learning Representations, 2018. URL httph:tt///cccct httpps://ppenteview.nees. wtt.://wwwwweee/ennnnte..e forum? -k- 0w
Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Jiang, Ebrahim Songhori, Shen Wang, Young- Joon Lee, Eric Johnson, Omkar Pathak, Sungmin Bae, Azade Nazi, Jiwoo Pak, Andy Tong, Kavya Srinivasa, William Hang, Emre Tuncer, Anand Babu, Quoe V. Le, James Laudon, Richard Ho, Roger Carpenter, and Jeff Dean. Chip placement with deep reinforcement learning. arXiv preprint WrXiv:2004.10744 2020.
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. arXiv preprins 10tp.1900502444, 2020.
Jongsoo Park, Maxim Naumov, Protonu Basu, Summer Deng, Aravind Kalaiah, Daya Khudia, James Law, Parth Malani, Andrey Malevich, Satish Nadathur, et al. Deep learning inference in facebook data centers: Characterization, performance optimizations and hardware implications. arXiv preprint arXiv: 001////1.9998 09886, 2018.
10

---

Published as a conference paper at ICLR 2021
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In Advances in Neural Information Processing Systems, 2018.
Vijay Janapa Reddi, Christine Cheng, David Kanter, Peter Mattson, Guenther Schmuelling, Carole- Jean Wu, Brian Anderson, Maximilien Breughe, Mark Charlebois, William Chou, et al. MIperf inference benchmark. arXiv preprint
2019.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80) 2008
Zhan Shi, Xiangru Huang, Akanksha Jain, and Calvin Lin. Applying deep learning to the cache replacement problem. In Proceedings of the 52nd Annual JEEE/ACM International Symposium on Microarchiteetture, pages 413 425, 2019.
William M Spears, Kenneth A De Jong, Thomas Bäck, David B Fogel, and Hugo De Garis. An overview of evolutionary computation. In European Conference on Machine Learming, pages Springer, 1993.
Chao Wang, Lei Gong, Qi Yu, Li Yuan Xie, and Xuehai Zhou. Dlau: scalable deep learning accelerator unit fpga. IEEE Transactions Computer-Aided Design of Integrated Circuits and Systems, (6(3):513-5177 2016.
O. Wechsler, M. Behar, and B. Daga. Spring hill (nap-i 1000) intel's data center inference chip. In 2019 IEEE Hot Chips 3/ Symposium (HCS), pages 1-12. IEFE, 2019.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 2020.
Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec. Graph convolutional neural networks for web-scale recommender systems. In Proceedings the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 974-983, 2018.
Lei Zhang, Reza Karimi, Irfan Ahmad, and Ymir Vigfusson. Optimal data placement for heteroge- neous cache, memory, and storage systems. Proceedings of the ACM on Measurement and Analysis of Computing Systems, 2020.
Brian Ziebart, Andrew Maas, Andrew Bagnell, and Anind Dey. Maximum entropy inverse reinforcement learning. In AAAI, volume 8, pages 1433-1438. Chicago, IL, USA, 2008

---

Published as a conference paper at ICLR 2021
APPENDIX
A GRAPH EMBEDDING
Table T details the features we used for the node embedding. These features encapsulate encapsulate information about the input and output tensors of the given operation, well as summary information about future layers.
Node Features
Description Operation id
Optd weight_size ifms ifmy
Size in bytes of the weights if exist, otherwise Input feature map size on the x axis Input feature map size on the y axis Input feature map size on the z axis Output feature map size on the axis Output feature map size on the y axis Output feature map size on the 7. axis Total size the input feature map (ifmz * iffy ifms) Total size of the onput feature map (ofm, of ofm,) Total number operations after current node Total number of weights from current node to the last node Number of groups Convolution related parameter, set to otherwise Kernel size on x axis Convolution related parameter, set O otherwise Kernel size on y axis Convolution related parameter, set to 0 otherwise Stride size Convolution related parameter,
ifmz
ofms
ofmy
ofmz
ifm_size
ofm_size
n_ops p-ops_lef left
n a____ left
groups kernelz
kernely
stride
otherwise Padding size Convolution related parameter, set to 0 otherwise Dilation Convolution related parameter, set otherwise Input batch size
pad
dilation
batch.
Table 1: GNN Node Features
B HYPERPARAMETERS
Table details the hyperparameters used in the paper.
Hyperparameter
Range explored
Value used
GNN hidden layer size
[32, 64, 128]
128
GNN output layer size
[32, 64, 128]
128
GNN depth
4
4
Number of GNN attention heads
[1,4]
# Steps per Episode
51 10]
1
Initial mapping action
['DRAM']
'DRAM'
Reward for invalid mapping
[-10, -1]
-1 0.99
Discount Rate
[0.9, 0.97, 0.99]
EA population size
[10, 20]
20
PG Rollout size
[0, 10
I
Fraction of EA population that Boltzmann
10.1, 0.5]
0.2
Total steps in the environment
[4000, 10000]
4000
Replay buffer size Critic learning rate
[100000]
100000
[le-3, le-4]
le-3
Actor learning rate
[1e-3, 1e-4]
le-3
Alpha (Entropy Coefficient)
[0.05, 0.1, 0.2]
0.05
Tau (Double-Q Network synchronization rate)
[le-3]
le-3
Batch size for PG
24
24
Reward scaling multiplier
5
Gradients steps per environment step
1
1
Table 2: Hyperparameters
12

---

Published as a conference paper at ICLR 2021
C EGRL
Algorithm 2 EGRL Algorithm
Initialize a mixed population of k policies pops
2: Initialize an empry cyclic replay buffer R.
3: Define a random number generator r() E [0,1)
4: for generation = I, o0 do
5: for actor x E pop, do
6: fitness, Experiences = Rollout(z)
7: Add experiences to
8: Rank the population based on fitness scores
9: Select the first actors pop= elites
10:
Select (A- 1/-cc c) actors π from popx. to form Set S using tournament selection with replacement while [S] < (k e) do
12:
and πı E if na and иe are of the same encoding type then Use single-point crossover and append to $
Select nT
14:
15:
else
16:
Sample random state and get action from the GNN policy Use a to encode the prior of the Boltzmann chromosome
17:
18:
for Actor Set S do
if r() < prob then Mutate(f") by adding noise ~ A/
20:
21:
ups = # of environment steps taken this generation
22:
for =i= ups do
23:
Sample random minibatch of T transitions (si, a4, 8441) from R Update the critic via Bellman update
24:
25:
L, = 1 E,(,: (ys Qi(s;, )) where y/ ri + min =
26:
+ H((((15444)
j=1,2
27:
where a; = a, + e, clip(e ~ N(u, 0²) - c,
28:
Update Ln using the sampled policy gradient with noisy actions Soft update target networks:
29:
Lov + TLa= + (1 Ly= and
31:
Lge' + TLee + (1 Lge
Copy L into the population: for weakest E pops 0T Lg-
EGRL incorporates EA's population-basee search with powerful gradient- based methods from DRL to expedite learning. In this work, we instantiate the EA population to use both the GNN encodings well Boltzmann chromosome encoding direct its search. Concurrently, SAC Haarnoja et al. (2018) algorithm as our gradient- -bssed technique in training the GNN policies. Algorithm details the EGRL algorithm.
modified
A general flow of the EGRL algorithm proceeds as follow: mixed population of GNN-policies and Boltzmann-based policies is initialized with random weights. In addition the population, one additional actor network (referred to as PJgan henceforth) is initialized alongside critic network. The population is then evaluated the enviromment by allowing for the specified workload in the given hardware. A selection operator then selects a portion of the population for survival with probability commensurate on their relative performance. The population is then probabilistically perturbed through mutation and crossover operations create the next generation. A select portion top performers are preserved as clites and are shielded from the mutation step.
control the memory mapping
Unlike traditional evolutionary population, each individual (whether GNN Botlzmann- based) stores its experience defined by the tuple (current state, action, next state, reward) in a globally shared replay buffer. This is done for every interaction that takes place with the hardware maximize data efficiency. The critic samples random minibatch from this shared replay
Shared Replay Buffer:
13

---

Published as a conference paper at ICLR 2021
buffer and uses it update its parameters using gradient descent. The critic is then used to train the PGGNN using the sampled policy gradient.
The shared replay buffer is key mechanism that enables the sharing of information across the varying learning methods. In contrast to a standard search method which would extract the performance score and disregard the underlying data immediately, EGRL retains every interaction in the global buffer and engages the PNNNN and critic to learn from them repeatedly using powerful gradient- based methods. This enables maximal information extraction from each individual experiences as interfacing with the hardware is an expensive operation.
Mixed Exploration: A noisy version of the PGGNN using Gaussian noise generator is used to generate additional experiences for the replay buffer. In contrast to the population of GNN-actors which explore by noise their neural weights, the PGGNN actors explore through noise its action space. Boltzmann chromosomes tread this line in between where they explore in the parameters space more directly connected the action space. Overall, each exploration technique are complementary and collectively lead to effective exploration of the solution space.
Periodically, the PGANN network's ' weights are copied into the evolutionary population. This process enables the evolutionary framework to directly leverage the information learned through gradient descent. This process also serves to stabilize learning and make more robust to deception. If the policy learned by the PGENN is favorable, it will be selected to survive and extend its influence to the population over subsequent generations. However, in case is bad, it will be selected against and discarded. This mechanism ensures that the flow of information from the PGGNN the evolutionary population is constructive.
Migration:
D POLICY GRADIENT MODIFICATIONS TO SAC
Policy Gradient Algorithm: We build on SAC (Harmoja et al. 2018) to tackle our large multi- discrete actions space. Since our policy is discrete, we compute entropy directly as
H(T(sss) E.-D --E (-|s) log (-|s)]
We then average over all nodes to compute the overall entropy of the policy. Further, we use a noisy version of the one-hot encoded behavioral action to compute our Bellman update
L + (v: (,(s;,a))) where y4 = r. + min (5(41,,, a411) + H(x(((s44))
use the minimum of two heads from the Q-Network based on (Fujimoto et al. 2018). The noisy action ă is computed by adding Gaussian noise clipped between -c and
= + clip(e ~ N(H, o²),-c,c)
This noisy action smoothens the value estimate towards similar state-action value estimates by the policy. It serves to make the policy smooth and addresses overfitting the one-hot encoded behavioral output. The actor is trained using the sampled policy gradient.
E BOLTZMANN CHROMOSOME
Figure illustrates the operation of the Boltzmann chromosome for particular action choice in one node. Parameters for prior (p1, p2, p3) and tempera- ture fully encode the chromosome's 's policy. To com- pute an action, we first compute the probabilities by applying the Boltzmann softmax operation with the associated prior and temperature. Action is the sam- pled from this probability distribution. The choice
Prior Probabilities
Temperature
T
Pt
P2
P3
Boltzmann Sampling
14
P1
P2
P3
Action Probabilities

---

Published as a conference paper at ICLR 2021
of temperature t directly modulates the exploration- exploitation knob of decision making. A higher tem- perature leads higher entropy probability distribu- tion enabling higher exploration. value of temperature will lead to lower entropy in the probability distribution enabling exploitation of the prior information.
contrast, a lower
For the agent policy described in this paper, Boltz- mann chromosome solution comprises of priors and temperature parameters for each node and action choice in the computational graph. Learning either through seeding, mutation or crossover involves direct update of these parameters. Importantly, these parameters are learned independently within the context of each node allowing for varying degrees of pellorationexploitaaaie
position across nodes. For instance, the agent could be very confident about mapping specific node while concurrently be unsure for different node of the same workload at the same time. The enables the agent to systematically balance the exploration-exxlootaaion tradeoff at the resolution of individual node actions.
F MAPPING VISUALIZATIONS
Resnet-50
Resnet-101
B:RT
0.20
0.31
0.0L
0.03
1.00
0.00
0.00
3.2
0.45
32
-10
DRAM
SRAM
MHI
LLC
DRAM
Figure F.2: Memory map shifts Top: For each memory unit on the y-axis, the corresponding row shows how EGRL changed the distribution of tensors originally mapped to it by the compiler. Bottom: Memory maps from the compiler vs the best ones found by EGRL for ResNet-50, ResNet-101 and BERT. Each bar denotes tensor operation.
G NNP-I INFERENCE ACCELERATOR
NNP-I is hardware accelerator dedicated neural network inference and was put in production in late 2019. It was designed to handle inference workloads at data center scale achieving 4.8 TOPs/W. The memory organization in NNP-I is shown Figure G.3
NNP-I includes twelve inference compute engines (ICEs) each having a 4MB Deep SRAM. A 24MB shared memory cache (LLC) accessible by all ICEs, Additionally, a 32 GB DRAM is also accessible the ICFs through the LLC. In terms of bandwidth, the DRAM is the slowest, while the LLC. and Deep SRAM 10x and 100x faster than the DRAM respectively.
In NNP-I, the specific mapping of each network layer to DRAM, LLC or SRAM is pre-determined statically during compilation. In our work, we intercept this capability of the compiler and design a reinforcement learning policy that performs this mapping. In this work. our RL agent only makes a gross determination if tensor should be mapped to the DRAM, LLC or Deep SRAM memory components. If our agent specifies Deep SRAM as the recommended target for a tensor, the specific target out of the 12 possible Deep SRAM locations (one for each ICE) is left to the further logic in the compiler.
15

---

Published as a conference paper at ICLR 2021
Technology
Dynamic power management FIVE technology
Intel with AVI and VIEE
SNC IA Core
ICE
ICE
ICE
50/px 43 9e9 abx
ICE
ICE
ICE
Coherency Fabric
4032, 2d64 LPODRAX GTA,
Driouuti di 8OO
data data
SNC
ICE
ICE
ICE
Feature
SoC
Core
ICE
ICE
ICE
TOPs (NTT))
48-98 92
Engines (C2)
TDP
10-50w
TOPs/w
2.0 TOPs/w
IW baned syne ICE to ICE communication
Inference Engines (CE
10-12 ICEs
Total SRAM
75MB
DRAM BW
68 GB/s
Spring Hill
ICE
DL compute grid
~ 10100X
Compute
1536 KB
3072
~ 1000
Deep SRAM
4MB 12
100
B/W - 1X
DRAM
Figure G.3: Schematic overview of the NNP-I HW architecture. Characteristic for the NNP-I are its 12 Inference Computation Engines (ICE) each with its own 4MB Deep SRAM, the 24MB shared cache memory (LLC) accessible directly from all of the ICES. Outside the LLC the NNP-I has DRAM of 32GB accessible to the ICEs through the LLC.
G.1 EGRL'S INTERACTION WITH NNP-1
Based on the architecture described above, EGRL can be utilized to dictate memory mapping at various control hierarchies. In this work, we designed EGRL to make mapping proposals for each tensor to the DRAM. LLC or SRAM classes of memory. It does not specify which of the multiple SRAMs tensor should be mapped to. We the downstream compiler optimization, including cache management strategies, control this data flow. An alternative would have been to expand action space for EGRL to control for the number active ICE cores, mapping tensors to specific ICE addresses, branch scheduling and other control logic typically residing in the compiler. While increasing this complexity has the potential for better overall performance, also increases convergence time for the policy and, more pertinently, requires more domain knowledge of the underlying architecture. Thus, we chose the configuration that required minimal domain knowledge while discovering strategies that substantially outperformed the baseline.
16

---

Published as a conference paper at ICLR 2021
Importantly, the GNN in EGRL only encodes the imput workload explicitly. Since EGRL trains itself on the resultant latency, one could argue that it learns to model the consequence of the additional compiler optimizations implicitly but this is never modeled explicitly. This is designed to require minimal understanding of the underlying architecture and data flows control.
Additionally; EGRL does not modify any aspect of the input workload. In that respect it is com- plementary to other works that optimize the graph itself for example via quantization or model compression graph re-writing methods studied in (Ahn al. (0200) where operational sequences that are sub-optimal for memory bandwidth are re-wired for optimal throughput. In all of these scenarios, EGRL is expected achieve incremental performance improvements conditioned on the pepiication-speciiie metric (e.g., throughput) and hardware-specinn constraints.
17


Reviews:
[{"review_id": "g07Ws9sL6kr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThis paper proposes a new algorithm called EGRL to improve computation graph running time by optimizing placement of the graph's components on memory. Specifically, the authors demonstrate the algorithm on the Intel Neural Networks Processor for Inference (NNP-I), which allows them to map neural network components on one of three memory hierarchies, each with different tradeoffs. The authors demonstrate that this technique provides speedups on BERT, Resnet-50 and Resnet-101.\n\n\nPros:\n\n- Some past papers (for eg. [1]) in this domain evaluate their work in simulators instead of real hardware, and often, the simulators make assumptions that are not realisitc. The paper tests its technique on actual hardware, and this is definitely a plus.\n- The authors promise that they will open-source their code. This is important since many of the efforts in this domain remain fragemented and difficult to reproduce. This is primarily due to the lack of open source code or a standard benchmark.\n- The paper is well written.\n- The visualisations of the learned policy vs the baseline in Figure 6 are quite good.\n- EGRL directly builds upon CERL so it is not very novel, but it has not been applied before to this domain.\n\nCons:\n\n- The paper evaluates the technique on just 3 workloads. This is in contrast to [1] who evaluate on 372 different workloads and [2] who evaluate on 96 synthetic graphs. [3] and [4] also evaluate on a very small number of workloads, but I believe they probably got a freepass since they were the earliest works in their domain.\n- The baseline that the experiments are being compared against might be weak - in figure 3, it looks the policy from both EGRL and EA in iteration 0 itself beat the baseline for Resnet-101 and BERT! \n- How long does it take to perform one iteration? And how long does it take to train the policy? This would be useful to get an idea of how EGRL fairs against [3] and [4] which also trained on real hardware and took many hours to finish training the policy.\n- The demonstration of generalizability is insufficient - it is difficult to conclude that EGRL can generalize to other workloads. For eg. in Figure 4 (left), the policy performs worse than the baseline for Resnet-50. Moreover, two of the workloads are from the Resnet family.\n- If it takes a long time to train each policy and if the model also shows poor zero-shot generalizability, it makes me question if this approach is practical for a compiler setting where a user would typically want the compilation to be completed quickly.\n\nOverall:\n\nI felt that the paper has some interesting ideas but needs more experiments.\n\n\nQuestions and Clarifications:\n\n- I believe that the related work section should add a clarification - [1], [2], [3] and [4] primarily deal with device placement, i.e., placing components of computation graph on different CPUs/GPUs to optimize run time via better parallelization. While this work is concerned with mapping components to different memory hierarchies on the same device. \n\n- While EGRL's action space is larger than [5], the action space in [1] is much larger - for a graph with 2000 nodes to be placed on 2 devices, there are 2^2000 possible choices ~ 10^603\n\n- In Figure 6 (Bottom), is there any reason why you didn't show the result for BERT?\n\n\nReferences:\n\n[1] Aditya Paliwal,  Felix Gimeno,  Vinod Nair,  Yujia Li,  Miles Lubin,  Pushmeet Kohli,  and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. arXiv preprint arXiv:1905.02494, 2020\n\n[2] Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Moham-mad Alizadeh. Placeto: Efficient progressive device placement optimization. In NIPS MachineLearning for Systems Workshop, 2018\n\n[3] Azalia Mirhoseini, Anna Goldie, Hieu Pham, Benoit Steiner, Quoc V. Le, and Jeff Dean. A hierarchical model for device placement. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=Hkc-TeZ0W.\n\n[4] Azalia Mirhoseini, Hieu Pham, Quoc V. Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, NaveenKumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device placement optimization with reinforcement learning. arXiv preprint arXiv:1706.04972, 2017.\n\n[5] Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Jiang, Ebrahim Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Sungmin Bae, Azade Nazi, Jiwoo Pak, Andy Tong, KavyaSrinivasa, William Hang, Emre Tuncer, Anand Babu, Quoc V. Le, James Laudon, Richard Ho,Roger Carpenter, and Jeff Dean. Chip placement with deep reinforcement learning. arXiv preprint arXiv:2004.10746, 2020.", "extra_scores": {}, "extra_text": {}}, {"review_id": "dCCxpALWJ4", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Optimizing the execution of deep neural networks has tremendous impact on the cost and performance in many industries due to the proliferation of \"Deep Learning\". There has recently been an interesting line of work of using learning to optimize policies related to placement and scheduling of the neural network computation graph outperforming tediously hand-crafted heuristics. The proposed paper would be a nice extension along this line.\nThe impact of memory placement for DNN has been clearly motivated and is easy to appreciate.\nThe paper overall is clear and easy to follow. The methodology is sound and justified. Experiments and results make a compelling case in supporting the claims.\n\nI think one aspect that is insufficiently motivated is the need to use a hybrid approach between RL and evolutionary algorithms. Results show improvement in performance but it is not clear to me why. Perhaps this is addressed in the CERL paper which I am not familiar with.\n\nIn the results section, it would be nice to see the sample complexity of each of the methods. I see that the number of iterations are shown in Figure 3 but it is not clear to me if that also corresponds to the number of samples consumed by each of the approaches before they converge.", "extra_scores": {}, "extra_text": {}}, {"review_id": "XlRlf2Pb_M", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper describes a machine learning method for mapping computational nodes in a neural network onto different levels of the memory hierarchy (DRAM, LLC, and SRAM) to minimize latency of inference. The proposed approach builds on CERL, combining policy gradient, reinforcement learning, and graph neural network, achieving 28-78% speed-up over the native NNP-I compiler on vision and language benchmarks (ResNet-50, ResNet-101, and BERT).\n\nOverall, the paper was well-written, targets an impactful problem, and the reported improvements (28-78% over native compiler) are impressive.\n\nIn the related work section, I did have a concern, as the authors state “For example, previous work with manual grouping (sic) operate at most in 5^280 \\~= 10^196 dimensional action space (Mirhoseini et al., 2020), compared to 10~ 20^358 for our BERT problem”. However, Mirhoseini et al., 2020 (“Chip placement with deep reinforcement learning”) places “a few thousand clusters” (>=2000 nodes) onto a grid with “an average of 30 rows and columns” (~900 cells), so wouldn’t the action space be at least 900^2000? Also, didn’t that work use a heuristic grouper (hMETIS), but maybe that’s close enough to “manual”? \n\nThe authors only look at three benchmarks, but they were well-chosen (two representative vision models and one large language model). It’s also good that they compare against PG and EA alone as a form of ablation, given that their method is effectively a combination of these two. It would have been better if they also had compared with prior state-of-the-art (e.g. HDP, REGAL, Placeto, or (the unmentioned) GDP / GO), but it is somewhat understandable given that their code does not seem to be open-sourced.  \n\nI liked that the authors report mean and standard deviation for the five runs, and measured “true” reward by running on hardware. I also thought they did a good job motivating their method (aside from the questionable statements about action spaces in prior work), and of analyzing and visualizing its performance. \n\nNits:\nIn the Method section, “the compiler rectifies them and outputs a modified map, M_c, that is fully executable (Line 6).” It would probably be good to add “in Algorithm 1” so as not to confused the reader.\n\n“It comprises of a single PG learner” -> “It is comprised of…”\n\n“Both methods are known to produce highly performant and stable solutions but are also significantly slow compared to Deep RL” (“significantly slower than”?)\n\n“While the transferred policies are clearly underperform those from scratch” -> “underperforming”", "extra_scores": {}, "extra_text": {}}, {"review_id": "CvfwPhx8Jky", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes Evolutionary Graph Reinforcement Learning to solve the memory placement problem. Main ideas are using GNN as the network architecture for reinforcement learning agents that look for more informed priors for evolutionary algorithms. Overall novelty of the paper comes from the neat combination of RL, EA, and GNN, and applying it to memory placement (ML for Systems).\n\nThe paper indeed tackles an important problem that can affect the overall performance and efficiency of the hardware. I believe the reorganization of various off-the-shelf ML techniques to solve real problems in the systems domain marks a large contribution, hence the positive overall rating.\n\nOne of the main drawbacks of the paper is that the paper only tests on a single type/configuration of hardware. While this is fine to some extent, this makes it hard to get confirmation about the generality of the overall method considering the large variance of the speedup.\n\nAnother related question comes from how this work relates to the optimizations of the dataflows [1,2]. As it is difficult to evaluate the overall memory communication without considering the order of operations, etc. the work in turn neglects the big question and focuses on only the partial view of the problem. It would provide a nice reference point if some of these points are discussed in the paper.\n\nLast question comes from the baselines. While the previous works on tensor optimizations [3,4] are very closely related and many of the ideas provide a good comparison point, these have not been discussed nor cited. For example, I guess AutoTVM's way of approximating the search space using TreeGRU or XGBoost can help. Also, Chameleon's way of sampling the examples using adaptive sampling may provide an interesting reference point in terms of reduction of number of samples.\n\nOverall, I have enjoyed reading the paper and I find the ideas in the paper interesting. I am currently weakly pro for the paper, and look forward to the authors' response :)\n\nQuestions\n1. Could you provide and overview of the NNP-I's architecture in the appendix? Also, possibly ablation studies over different configurations of the hardware.\n2. What are the relative communication speed of SRAM, LLC, and DRAM? \n3. It is well discussed in the computer architecture community that the memory communication is very much determined by the dataflows of the architecture. How are the results affected by these dataflows?\n4. How does the work compare to the methods described in [3,4]?\n\n[1] \"Eyeriss: A spatial architecture for energy-efficient dataflow for convolutional neural networks\", ISCA 2016\n\n[2] \"Interstellar: Using Halide's Scheduling Language to Analyze DNN Accelerators\", ASPLOS 2020\n\n[3] \"Learning to optimize tensor programs\", NeurIPS 2018\n\n[4] \"Chameleon: Adaptive code optimization for expedited deep neural network compilation\", ICLR 2020", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Most of the reviewers agree that this paper presents interesting ideas for an important problem. The paper could be further improved by having a thorough discussion of related works (e.g. Placeto) and construct proxy baselines that reflect these approaches.


The meta-reviewer decided to accept the paper given the positive aspects, and encourages the author to further improve the paper per review comments.


Thank you for submitting the paper to ICLR. 

Decision: Accept (Poster)