Mirage: A Multi-Level Superoptimizer for Tensor Programs
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2024
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866918047029133312 |
|---|---|
| author | Wu, Mengdi Cheng, Xinhao Liu, Shengyu Shi, Chunan Ji, Jianan Ao, Kit Velliengiri, Praveen Miao, Xupeng Padon, Oded Jia, Zhihao |
| author_facet | Wu, Mengdi Cheng, Xinhao Liu, Shengyu Shi, Chunan Ji, Jianan Ao, Kit Velliengiri, Praveen Miao, Xupeng Padon, Oded Jia, Zhihao |
| contents | We introduce Mirage, the first multi-level superoptimizer for tensor programs. A key idea in Mirage is $μ$Graphs, a uniform representation of tensor programs at the kernel, thread block, and thread levels of the GPU compute hierarchy. $μ$Graphs enable Mirage to discover novel optimizations that combine algebraic transformations, schedule transformations, and generation of new custom kernels. To navigate the large search space, Mirage introduces a pruning technique based on abstraction that significantly reduces the search space and provides a certain optimality guarantee. To ensure that the optimized $μ$Graph is equivalent to the input program, Mirage introduces a probabilistic equivalence verification procedure with strong theoretical guarantees. Our evaluation shows that Mirage outperforms existing approaches by up to 3.3$\times$ even for DNNs that are widely used and heavily optimized. Mirage is publicly available at https://github.com/mirage-project/mirage. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_05751 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Mirage: A Multi-Level Superoptimizer for Tensor Programs Wu, Mengdi Cheng, Xinhao Liu, Shengyu Shi, Chunan Ji, Jianan Ao, Kit Velliengiri, Praveen Miao, Xupeng Padon, Oded Jia, Zhihao Machine Learning Artificial Intelligence Programming Languages We introduce Mirage, the first multi-level superoptimizer for tensor programs. A key idea in Mirage is $μ$Graphs, a uniform representation of tensor programs at the kernel, thread block, and thread levels of the GPU compute hierarchy. $μ$Graphs enable Mirage to discover novel optimizations that combine algebraic transformations, schedule transformations, and generation of new custom kernels. To navigate the large search space, Mirage introduces a pruning technique based on abstraction that significantly reduces the search space and provides a certain optimality guarantee. To ensure that the optimized $μ$Graph is equivalent to the input program, Mirage introduces a probabilistic equivalence verification procedure with strong theoretical guarantees. Our evaluation shows that Mirage outperforms existing approaches by up to 3.3$\times$ even for DNNs that are widely used and heavily optimized. Mirage is publicly available at https://github.com/mirage-project/mirage. |
| title | Mirage: A Multi-Level Superoptimizer for Tensor Programs |
| topic | Machine Learning Artificial Intelligence Programming Languages |
| url | https://arxiv.org/abs/2405.05751 |