Подход к построению компиляторов нейронных сетей с использованием инфраструктуры MLIR

Развитие матричных расширений процессорных архитектур, а также внедрение этих расширений в специализированные AI-процессоры, позволяет существенно повысить эффективность выполнения искусственных нейронных сетей. В работе выполнен обзор базовых функциональных возможностей некоторых популярных матричных расширений процессорных архитектур, в частности расширений ARM SME, RISC-V IME, RISC-V AME, а также процессорной архитектуры DaVinci. В результате проведенного анализа была предложена модель абстрактного матричного процессора, отражающая особенности современных процессорных архитектур, которые поддерживают матричное расширение. Для введенной модели матричного процессора разработано гетерогенное матричное промежуточное представление, которое может быть использовано для построения компиляторов нейронных сетей. Предложенное промежуточное представление было реализовано в инфраструктуре MLIR в виде диалекта heteroMx. В работе также описан подход к построению AI-компилятора с использованием разработанного диалекта heteroMx. Разработанное промежуточное представление может быть адаптировано или конкретизировано для других матричных процессорных архитектур.

1. Sousa R. et al. Tensor slicing and optimization for multicore NPUs //Journal of Parallel and Distributed Computing. – 2023. – Т. 175. – С. 66-79.

2. Jouppi N. P. et al. In-datacenter performance analysis of a tensor processing unit //Proceedings of the 44th annual international symposium on computer architecture. – 2017. – С. 1-12.

3. Lattner C. et al. MLIR: Scaling compiler infrastructure for domain specific computation //2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO). – IEEE, 2021. – С. 2-14.

4. Remke S., Breuer A. Hello SME! Generating Fast Matrix Multiplication Kernels Using the Scalable Matrix Extension //arXiv preprint arXiv:2409.18779. – 2024.

5. Stephens N. et al. The ARM scalable vector extension //IEEE micro. – 2017. – Т. 37. – №. 2. – С. 26-39.

6. The RISC-V IME Set Specification. https://github.com/space-mit/riscv-ime-extension-spec/releases/download/v0429/spacemit-ime-asciidoc.pdf.

7. H. Liao, J. Tu, J. Xia and X. Zhou, "DaVinci: A Scalable Architecture for Neural Network Computing," 2019 IEEE Hot Chips 31 Symposium (HCS), Cupertino, CA, USA, 2019, pp. 1-44, doi: 10.1109/HOTCHIPS.2019.8875654.

8. TOSA specification. https://www.mlplatform.org/tosa/tosa_spec.html, Accessed July 2023.

9. Goto K., Geijn R. A. Anatomy of high-performance matrix multiplication //ACM Transactions on Mathematical Software (TOMS). – 2008. – Т. 34. – №. 3. – С. 1-25.