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.

翻译：资源调度与分配是诸多高影响力系统（从拥塞控制到云计算）的关键组成部分。为这些问题寻找更优的解决方案，通常能显著节约资源与时间、减少设备损耗，甚至可能降低碳排放。本文聚焦于调度问题的一个具体实例，即机器学习程序编译过程中出现的内存映射问题：将张量映射到不同内存层级以优化执行时间。我们提出一种基于强化学习解决内存映射问题的方法。强化学习作为一种解决方案范式，特别适用于具有可规划性的序列决策问题，以及带有高维数据输入的组合搜索空间。我们将该问题建模为单人博弈游戏，并称之为mallocGame，其中高回报的博弈轨迹对应目标硬件上高效的内存映射。我们还引入了一个强化学习智能体mallocMuZero，并证明其能够通过参与该博弈发现新的、更优的内存映射方案，从而在机器学习加速器上的真实工作负载中实现更快的执行时间。我们将mallocMuZero与加速线性代数（XLA）编译器在基准测试中针对真实机器学习工作负载所使用的默认求解器进行了性能对比。此外，我们证明mallocMuZero能够改进近期发表的AlphaTensor矩阵乘法模型的执行时间。