Relax: Composable Abstractions for End-to-End Dynamic Machine Learning
This addresses the problem of efficiently deploying dynamic machine learning models across diverse hardware environments for practitioners and researchers.
The paper tackles the challenge of optimizing dynamic shape computations in modern machine learning workloads, particularly large language models, by introducing Relax, a compiler abstraction that enables dynamic shape-aware cross-level optimizations. Experimental results show Relax delivers performance competitive with state-of-the-art systems across various GPUs and enables deployment to diverse environments like mobile phones and web browsers.
Dynamic shape computations have become critical in modern machine learning workloads, especially in emerging large language models. The success of these models has driven the demand for their universal deployment across 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 a cross-level abstraction that encapsulates computational graphs, loop-level tensor programs, and external library calls in a single representation. Relax also introduces first-class symbolic shape annotations to track dynamic shape computations globally across the program, enabling dynamic shape-aware cross-level optimizations. We build an end-to-end compilation framework using the proposed approach to optimize dynamic shape models. Experimental results on LLMs show that Relax delivers performance competitive with state-of-the-art systems across various GPUs and enables deployment of emerging models to a broader set of emerging environments, including mobile phones, embedded devices, and web browsers.