Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
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.
6549afd6c8b11b62711aad01
Ruihang Lai
admin_assigned
2023-11-07T11:26:24.458Z
6339c512ecbd8b193575aee4
https://cdn-avatars.huggingface.co/v1/production/uploads/1678644922528-6339c512ecbd8b193575aee4.jpeg
Ruihang Lai
user
ruihanglai
6549afd6c8b11b62711aad02
Junru Shao
admin_assigned
2023-11-07T11:26:39.294Z
646db758075bbcc48ddcd28f
/avatars/0aee3fac8609979b5837d696d18190a6.svg
Junru Shao
user
junrushao
6549afd6c8b11b62711aad03
Siyuan Feng
admin_assigned
2023-11-07T11:27:39.006Z
62061f8f03825909dcbeba27
https://cdn-avatars.huggingface.co/v1/production/uploads/62061f8f03825909dcbeba27/byOi6gxlycUtsUPIUPscQ.jpeg
Siyuan Feng
user
Siyuan
6549afd6c8b11b62711aad04
Steven S. Lyubomirsky
admin_assigned
2023-11-07T11:27:50.189Z
65031e902cec4c6d94ff7284
/avatars/d1e77fe4aba07f347995ca52ef9e6dfa.svg
Steven Lyubomirsky (octoml)
user
slyubomirsky-octoml
6549afd6c8b11b62711aad05
Bohan Hou
6549afd6c8b11b62711aad06
Wuwei Lin
admin_assigned
2023-11-07T11:28:22.889Z
632d1184ea6e62428abb8b8b
/avatars/8f29e0e221a7ec50e801f0670d84c2b9.svg
Wuwei Lin
user
vinx13
6549afd6c8b11b62711aad07
Zihao Ye
admin_assigned
2023-11-07T11:28:30.326Z
632b9f58ffe3618eff2ffb93
/avatars/1ffaf2da94fb9f993aa0af4bc3808197.svg
Zihao Ye
user
zhye
6549afd6c8b11b62711aad08
Hongyi Jin
admin_assigned
2023-11-07T11:28:36.802Z
644043227663594a126420a9
/avatars/c0d0e0e5bac33de90c40d8d330dc0bc3.svg
Hongyi Jin
user
hongyij
6549afd6c8b11b62711aad09
Yuchen Jin
admin_assigned
2023-11-07T11:28:44.228Z
6500bca7e850c952d4516c5a
https://cdn-avatars.huggingface.co/v1/production/uploads/6500bca7e850c952d4516c5a/dT3Obwruf2ruZBNorc_aJ.jpeg
Yuchen Jin
user
yuchenj
6549afd6c8b11b62711aad0a
Jiawei Liu
admin_assigned
2023-11-07T11:31:47.648Z
644b584a9279988e0cbeb664
https://cdn-avatars.huggingface.co/v1/production/uploads/644b584a9279988e0cbeb664/fhWCI_Q26tTruhdFkjejw.jpeg
Jiawei Liu
user
ganler
6549afd6c8b11b62711aad0b
Lesheng Jin
admin_assigned
2023-11-07T11:32:05.983Z
64308a5eab607482bb0466eb
/avatars/21f8f9599ad6ad24ed81ec3e17f4cab0.svg
Lesheng Jin
user
leshenj15
6549afd6c8b11b62711aad0c
Yaxing Cai
6549afd6c8b11b62711aad0d
Ziheng Jiang
6549afd6c8b11b62711aad0e
Yong Wu
6549afd6c8b11b62711aad0f
Sunghyun Park
6549afd6c8b11b62711aad10
Prakalp Srivastava
admin_assigned
2023-11-07T11:33:43.286Z
65454441dbce6bd2beb57b50
/avatars/ec70cdd46413f49780bb3c90ad518508.svg
Prakalp Srivastava
user
psrivas2
6549afd6c8b11b62711aad11
Jared G. Roesch
admin_assigned
2023-11-07T11:33:53.270Z
636bf40aaae2da3c76b7b5d2
/avatars/41b3f81fe12349bcfa70eff20eaeb187.svg
Jared Roesch
user
jroesch
6549afd6c8b11b62711aad12
Todd C. Mowry
6549afd6c8b11b62711aad13
Tianqi Chen
admin_assigned
2023-11-07T11:35:12.707Z
6407c3cb2e309e654528860f
/avatars/b2177ae76ff8001564d0606752538884.svg
Tianqi Chen
user
tqchen
Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
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.
Relax is a compiler abstraction that optimizes dynamic shape computations in machine learning workloads, enabling competitive performance and broader deployment of large language models.
dynamic shape computations compiler abstraction symbolic shape annotations end-to-end compilation computational graphs loop-level tensor programs library calls cross-level optimizations large language models mobile phones embedded devices web browsers