Towards a high-performance AI compiler with upstream MLIR
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_ | 1866917648018702336 |
|---|---|
| author | Golin, Renato Chelini, Lorenzo Siemieniuk, Adam Madhu, Kavitha Hasabnis, Niranjan Pabst, Hans Georganas, Evangelos Heinecke, Alexander |
| author_facet | Golin, Renato Chelini, Lorenzo Siemieniuk, Adam Madhu, Kavitha Hasabnis, Niranjan Pabst, Hans Georganas, Evangelos Heinecke, Alexander |
| contents | This work proposes a compilation flow using open-source compiler passes to build a framework to achieve ninja performance from a generic linear algebra high-level abstraction. We demonstrate this flow with a proof-of-concept MLIR project that uses input IR in Linalg-on-Tensor from TensorFlow and PyTorch, performs cache-level optimizations and lowering to micro-kernels for efficient vectorization, achieving over 90% of the performance of ninja-written equivalent programs. The contributions of this work include: (1) Packing primitives on the tensor dialect and passes for cache-aware distribution of tensors (single and multi-core) and type-aware instructions (VNNI, BFDOT, BFMMLA), including propagation of shapes across the entire function; (2) A linear algebra pipeline, including tile, fuse and bufferization strategies to get model-level IR into hardware friendly tile calls; (3) A mechanism for micro-kernel lowering to an open source library that supports various CPUs. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2404_15204 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Towards a high-performance AI compiler with upstream MLIR Golin, Renato Chelini, Lorenzo Siemieniuk, Adam Madhu, Kavitha Hasabnis, Niranjan Pabst, Hans Georganas, Evangelos Heinecke, Alexander Programming Languages Artificial Intelligence Hardware Architecture Distributed, Parallel, and Cluster Computing Machine Learning This work proposes a compilation flow using open-source compiler passes to build a framework to achieve ninja performance from a generic linear algebra high-level abstraction. We demonstrate this flow with a proof-of-concept MLIR project that uses input IR in Linalg-on-Tensor from TensorFlow and PyTorch, performs cache-level optimizations and lowering to micro-kernels for efficient vectorization, achieving over 90% of the performance of ninja-written equivalent programs. The contributions of this work include: (1) Packing primitives on the tensor dialect and passes for cache-aware distribution of tensors (single and multi-core) and type-aware instructions (VNNI, BFDOT, BFMMLA), including propagation of shapes across the entire function; (2) A linear algebra pipeline, including tile, fuse and bufferization strategies to get model-level IR into hardware friendly tile calls; (3) A mechanism for micro-kernel lowering to an open source library that supports various CPUs. |
| title | Towards a high-performance AI compiler with upstream MLIR |
| topic | Programming Languages Artificial Intelligence Hardware Architecture Distributed, Parallel, and Cluster Computing Machine Learning |
| url | https://arxiv.org/abs/2404.15204 |