High Accuracy Matrix Computations on Neural Engines: A Study of QR Factorization and its Applications

Fueled by the surge of ever expanding successful applications of deep neural networks and the great computational power demanded, modern computer processors and accelerators are beginning to offer half precision floating point arithmetic support, and special units (neural engines) such as NVIDIA TensorCore on GPU and Google Tensor Processing Unit (TPU) to accelerate the training and prediction of deep neural networks. It remains unclear how neural engines can be profitably used in application other than neural networks. In this paper we present an endeavor of accelerating and stabilizing a fundamental matrix factorization on neural enginesthe QR factorization-which may open doors to much wider relevance to scientific, engineering, and data sciences. We show that traditional Householder QR algorithms and implementations do not have the necessary data locality, parallelism, accuracy and robustness on neural engines which are characterized by extreme speed and low precision/range.

We demonstrate that neural engines can be effectively used to accelerate matrix computations (QR 3.0x-14.6x speedup compared to cuSOLVER, reaching up to 36.6TFLOPS); however different algorithms (recursive Gram-Schmidt) are needed to expose more locality and parallelism, even at the cost of increased computations. Moreover, scaling, iterative refinement, and other safeguarding procedures are also needed to regain accuracy and avoid overflowing. Our experience seems to suggest that presently with neural engines the matrix factorizations (QR, LU, Cholesky) are best to be co-designed with its applications (linear solver, least square, orthogonalization, SVD etc) to achieve high performance and adequate accuracy and reliability, rather than used as a black-box.