title: Optimizing Memory Mapping Using Deep Reinforcement Learning
summary: 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.
paper_authors._id: 6461cce85ea1bc9a9da954a0
paper_authors.name: Pengming Wang
paper_authors._id: 6461cce85ea1bc9a9da954a1
paper_authors.name: Mikita Sazanovich
paper_authors._id: 6461cce85ea1bc9a9da954a2
paper_authors.name: Berkin Ilbeyi
paper_authors._id: 6461cce85ea1bc9a9da954a3
paper_authors.name: Phitchaya Mangpo Phothilimthana
paper_authors._id: 6461cce85ea1bc9a9da954a4
paper_authors.name: Manish Purohit
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-05-26T09:41:42.213Z
paper_authors._id: 6122fde3d09e016ab5187a5f
paper_authors.avatarUrl: /avatars/366e14ff8768e8deaee0ff7b21b9f76a.svg
paper_authors.fullname: manish purohit
paper_authors.type: user
paper_authors.user: manishpurohit
paper_authors._id: 6461cce85ea1bc9a9da954a5
paper_authors.name: Han Yang Tay
paper_authors._id: 6461cce85ea1bc9a9da954a6
paper_authors.name: Ngân Vũ
paper_authors._id: 6461cce85ea1bc9a9da954a7
paper_authors.name: Miaosen Wang
paper_authors._id: 6461cce85ea1bc9a9da954a8
paper_authors.name: Cosmin Paduraru
paper_authors._id: 6461cce85ea1bc9a9da954a9
paper_authors.name: Edouard Leurent
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-05-26T09:42:32.843Z
paper_authors._id: 633e093315033639f991760f
paper_authors.avatarUrl: /avatars/0339ec428a8d1fb961cb0b095e0839bd.svg
paper_authors.fullname: Edouard Leurent
paper_authors.type: user
paper_authors.user: eleurent
paper_authors._id: 6461cce85ea1bc9a9da954aa
paper_authors.name: Anton Zhernov
paper_authors._id: 6461cce85ea1bc9a9da954ab
paper_authors.name: Julian Schrittwieser
paper_authors._id: 6461cce85ea1bc9a9da954ac
paper_authors.name: Thomas Hubert
paper_authors._id: 6461cce85ea1bc9a9da954ad
paper_authors.name: Robert Tung
paper_authors._id: 6461cce85ea1bc9a9da954ae
paper_authors.name: Paula Kurylowicz
paper_authors._id: 6461cce85ea1bc9a9da954af
paper_authors.name: Kieran Milan
paper_authors._id: 6461cce85ea1bc9a9da954b0
paper_authors.name: Oriol Vinyals
paper_authors._id: 6461cce85ea1bc9a9da954b1
paper_authors.name: Daniel J. Mankowitz
paper_title: Optimizing Memory Mapping Using Deep Reinforcement Learning
paper_summary: 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.
paper_ai_keywords: Reinforcement Learning RL memory mapping problem mallocGame mallocMuZero high-reward trajectories single-actor game combinatorial search spaces sequential decision making problems execution time high-dimensional data inputs Accelerated Linear Algebra (XLA) compiler AlphaTensor matrix multiplication model