Resource scheduling and allocation is a critical component of many high
impact systems ranging from congestion control to cloud computing. Finding more
optimal solutions to these problems often has significant impact on resource
and time savings, reducing device wear-and-tear, and even potentially improving
carbon emissions. In this paper, we focus on a specific instance of a
scheduling problem, namely the memory mapping problem that occurs during
compilation of machine learning programs: That is, mapping tensors to different
memory layers to optimize execution time.
We introduce an approach for solving the memory mapping problem using
Reinforcement Learning. RL is a solution paradigm well-suited for sequential
decision making problems that are amenable to planning, and combinatorial
search spaces with high-dimensional data inputs. We formulate the problem as a
single-player game, which we call the mallocGame, such that high-reward
trajectories of the game correspond to efficient memory mappings on the target
hardware. We also introduce a Reinforcement Learning agent, mallocMuZero, and
show that it is capable of playing this game to discover new and improved
memory mapping solutions that lead to faster execution times on real ML
workloads on ML accelerators. We compare the performance of mallocMuZero to the
default solver used by the Accelerated Linear Algebra (XLA) compiler on a
benchmark of realistic ML workloads. In addition, we show that mallocMuZero is
capable of improving the execution time of the recently published AlphaTensor
matrix multiplication model.

通过采用强化学习技术，本文提出了一个解决编译机器学习程序中内存映射问题的新方法，称之为 mallocMuZero 算法，并将其应用在 ML 加速器上以缩短执行时间，表现优于目前的默认求解器。

使用深度强化学习优化内存映射

Optimizing Memory Mapping Using Deep Reinforcement Learning

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.

本文提出了一种基于进化图强化学习（EGRL）的方法，用于深度神经网络加速器的自动内存映射，该方法结合了图神经网络、强化学习和进化搜索等技术，并在 Intel NNP-I 芯片上进行了训练和验证，实验结果表明，EGRL 在 BERT、ResNet-101 和 ResNet-50 三种工作负载中表现良好，并相对于原生 NNP-I 编译器获得 28-78% 的加速。