Optimal Kernel Orchestration for Tensor Programs with Korch
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2024
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866909223533674496 |
|---|---|
| author | Hu, Muyan Venkatram, Ashwin Biswas, Shreyashri Marimuthu, Balamurugan Hou, Bohan Oliaro, Gabriele Wang, Haojie Zheng, Liyan Miao, Xupeng Zhai, Jidong |
| author_facet | Hu, Muyan Venkatram, Ashwin Biswas, Shreyashri Marimuthu, Balamurugan Hou, Bohan Oliaro, Gabriele Wang, Haojie Zheng, Liyan Miao, Xupeng Zhai, Jidong |
| contents | Kernel orchestration is the task of mapping the computation defined in different operators of a deep neural network (DNN) to the execution of GPU kernels on modern hardware platforms. Prior approaches optimize kernel orchestration by greedily applying operator fusion, which fuses the computation of multiple operators into a single kernel, and miss a variety of optimization opportunities in kernel orchestration.
This paper presents Korch, a tensor program optimizer that discovers optimal kernel orchestration strategies for tensor programs. Instead of directly fusing operators, Korch first applies operator fission to decompose tensor operators into a small set of basic tensor algebra primitives. This decomposition enables a diversity of fine-grained, inter-operator optimizations. Next, Korch optimizes kernel orchestration by formalizing it as a constrained optimization problem, leveraging an off-the-shelf binary linear programming solver to discover an optimal orchestration strategy, and generating an executable that can be directly deployed on modern GPU platforms. Evaluation on a variety of DNNs shows that Korch outperforms existing tensor program optimizers by up to 1.7x on V100 GPUs and up to 1.6x on A100 GPUs. Korch is publicly available at https://github.com/humuyan/Korch. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2406_09465 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Optimal Kernel Orchestration for Tensor Programs with Korch Hu, Muyan Venkatram, Ashwin Biswas, Shreyashri Marimuthu, Balamurugan Hou, Bohan Oliaro, Gabriele Wang, Haojie Zheng, Liyan Miao, Xupeng Zhai, Jidong Data Structures and Algorithms Machine Learning Kernel orchestration is the task of mapping the computation defined in different operators of a deep neural network (DNN) to the execution of GPU kernels on modern hardware platforms. Prior approaches optimize kernel orchestration by greedily applying operator fusion, which fuses the computation of multiple operators into a single kernel, and miss a variety of optimization opportunities in kernel orchestration. This paper presents Korch, a tensor program optimizer that discovers optimal kernel orchestration strategies for tensor programs. Instead of directly fusing operators, Korch first applies operator fission to decompose tensor operators into a small set of basic tensor algebra primitives. This decomposition enables a diversity of fine-grained, inter-operator optimizations. Next, Korch optimizes kernel orchestration by formalizing it as a constrained optimization problem, leveraging an off-the-shelf binary linear programming solver to discover an optimal orchestration strategy, and generating an executable that can be directly deployed on modern GPU platforms. Evaluation on a variety of DNNs shows that Korch outperforms existing tensor program optimizers by up to 1.7x on V100 GPUs and up to 1.6x on A100 GPUs. Korch is publicly available at https://github.com/humuyan/Korch. |
| title | Optimal Kernel Orchestration for Tensor Programs with Korch |
| topic | Data Structures and Algorithms Machine Learning |
| url | https://arxiv.org/abs/2406.09465 |