title: Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
summary: Dynamic shape computations have become critical in modern machine learning
  workloads, especially in emerging large language models. The success of these
  models has driven demand for deploying them to a diverse set of backend
  environments. In this paper, we present Relax, a compiler abstraction for
  optimizing end-to-end dynamic machine learning workloads. Relax introduces
  first-class symbolic shape annotations to track dynamic shape computations
  globally across the program. It also introduces a cross-level abstraction that
  encapsulates computational graphs, loop-level tensor programs, and library
  calls in a single representation to enable cross-level optimizations. We build
  an end-to-end compilation framework using the proposed approach to optimize
  dynamic shape models. Experimental results on large language models show that
  Relax delivers performance competitive with state-of-the-art hand-optimized
  systems across platforms and enables deployment of emerging dynamic models to a
  broader set of environments, including mobile phones, embedded devices, and web
  browsers.
paper_authors._id: 6549afd6c8b11b62711aad01
paper_authors.name: Ruihang Lai
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:26:24.458Z
paper_authors._id: 6339c512ecbd8b193575aee4
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1678644922528-6339c512ecbd8b193575aee4.jpeg
paper_authors.fullname: Ruihang Lai
paper_authors.type: user
paper_authors.user: ruihanglai
paper_authors._id: 6549afd6c8b11b62711aad02
paper_authors.name: Junru Shao
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:26:39.294Z
paper_authors._id: 646db758075bbcc48ddcd28f
paper_authors.avatarUrl: /avatars/0aee3fac8609979b5837d696d18190a6.svg
paper_authors.fullname: Junru Shao
paper_authors.type: user
paper_authors.user: junrushao
paper_authors._id: 6549afd6c8b11b62711aad03
paper_authors.name: Siyuan Feng
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:27:39.006Z
paper_authors._id: 62061f8f03825909dcbeba27
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/62061f8f03825909dcbeba27/PgxdcYKIuMBIzf-FFHK5r.jpeg
paper_authors.fullname: Siyuan Feng
paper_authors.type: user
paper_authors.user: Siyuan
paper_authors._id: 6549afd6c8b11b62711aad04
paper_authors.name: Steven S. Lyubomirsky
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:27:50.189Z
paper_authors._id: 65031e902cec4c6d94ff7284
paper_authors.avatarUrl: /avatars/d1e77fe4aba07f347995ca52ef9e6dfa.svg
paper_authors.fullname: Steven Lyubomirsky (octoml)
paper_authors.type: user
paper_authors.user: slyubomirsky-octoml
paper_authors._id: 6549afd6c8b11b62711aad05
paper_authors.name: Bohan Hou
paper_authors._id: 6549afd6c8b11b62711aad06
paper_authors.name: Wuwei Lin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:28:22.889Z
paper_authors._id: 632d1184ea6e62428abb8b8b
paper_authors.avatarUrl: /avatars/8f29e0e221a7ec50e801f0670d84c2b9.svg
paper_authors.fullname: Wuwei Lin
paper_authors.type: user
paper_authors.user: vinx13
paper_authors._id: 6549afd6c8b11b62711aad07
paper_authors.name: Zihao Ye
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:28:30.326Z
paper_authors._id: 632b9f58ffe3618eff2ffb93
paper_authors.avatarUrl: /avatars/1ffaf2da94fb9f993aa0af4bc3808197.svg
paper_authors.fullname: Zihao Ye
paper_authors.type: user
paper_authors.user: zhye
paper_authors._id: 6549afd6c8b11b62711aad08
paper_authors.name: Hongyi Jin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:28:36.802Z
paper_authors._id: 644043227663594a126420a9
paper_authors.avatarUrl: /avatars/c0d0e0e5bac33de90c40d8d330dc0bc3.svg
paper_authors.fullname: Hongyi Jin
paper_authors.type: user
paper_authors.user: hongyij
paper_authors._id: 6549afd6c8b11b62711aad09
paper_authors.name: Yuchen Jin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:28:44.228Z
paper_authors._id: 6500bca7e850c952d4516c5a
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6500bca7e850c952d4516c5a/dT3Obwruf2ruZBNorc_aJ.jpeg
paper_authors.fullname: Yuchen Jin
paper_authors.type: user
paper_authors.user: yuchenj
paper_authors._id: 6549afd6c8b11b62711aad0a
paper_authors.name: Jiawei Liu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:31:47.648Z
paper_authors._id: 644b584a9279988e0cbeb664
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/644b584a9279988e0cbeb664/fz4pn7HeXXfpvz23v8E8t.png
paper_authors.fullname: Jiawei Liu
paper_authors.type: user
paper_authors.user: ganler
paper_authors._id: 6549afd6c8b11b62711aad0b
paper_authors.name: Lesheng Jin
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:32:05.983Z
paper_authors._id: 64308a5eab607482bb0466eb
paper_authors.avatarUrl: /avatars/21f8f9599ad6ad24ed81ec3e17f4cab0.svg
paper_authors.fullname: Lesheng Jin
paper_authors.type: user
paper_authors.user: leshenj15
paper_authors._id: 6549afd6c8b11b62711aad0c
paper_authors.name: Yaxing Cai
paper_authors._id: 6549afd6c8b11b62711aad0d
paper_authors.name: Ziheng Jiang
paper_authors._id: 6549afd6c8b11b62711aad0e
paper_authors.name: Yong Wu
paper_authors._id: 6549afd6c8b11b62711aad0f
paper_authors.name: Sunghyun Park
paper_authors._id: 6549afd6c8b11b62711aad10
paper_authors.name: Prakalp Srivastava
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:33:43.286Z
paper_authors._id: 65454441dbce6bd2beb57b50
paper_authors.avatarUrl: /avatars/ec70cdd46413f49780bb3c90ad518508.svg
paper_authors.fullname: Prakalp Srivastava
paper_authors.type: user
paper_authors.user: psrivas2
paper_authors._id: 6549afd6c8b11b62711aad11
paper_authors.name: Jared G. Roesch
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:33:53.270Z
paper_authors._id: 636bf40aaae2da3c76b7b5d2
paper_authors.avatarUrl: /avatars/41b3f81fe12349bcfa70eff20eaeb187.svg
paper_authors.fullname: Jared Roesch
paper_authors.type: user
paper_authors.user: jroesch
paper_authors._id: 6549afd6c8b11b62711aad12
paper_authors.name: Todd C. Mowry
paper_authors._id: 6549afd6c8b11b62711aad13
paper_authors.name: Tianqi Chen
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2023-11-07T11:35:12.707Z
paper_authors._id: 6407c3cb2e309e654528860f
paper_authors.avatarUrl: /avatars/b2177ae76ff8001564d0606752538884.svg
paper_authors.fullname: Tianqi Chen
paper_authors.type: user
paper_authors.user: tqchen
paper_title: Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
paper_summary: Dynamic shape computations have become critical in modern machine learning
  workloads, especially in emerging large language models. The success of these
  models has driven demand for deploying them to a diverse set of backend
  environments. In this paper, we present Relax, a compiler abstraction for
  optimizing end-to-end dynamic machine learning workloads. Relax introduces
  first-class symbolic shape annotations to track dynamic shape computations
  globally across the program. It also introduces a cross-level abstraction that
  encapsulates computational graphs, loop-level tensor programs, and library
  calls in a single representation to enable cross-level optimizations. We build
  an end-to-end compilation framework using the proposed approach to optimize
  dynamic shape models. Experimental results on large language models show that
  Relax delivers performance competitive with state-of-the-art hand-optimized
  systems across platforms and enables deployment of emerging dynamic models to a
  broader set of environments, including mobile phones, embedded devices, and web
  browsers.
paper_ai_keywords: symbolic shape annotations computational graphs loop-level tensor programs library calls cross-level optimizations