Vectorizing Pytorch for RISC-V RVV

Abstract

In this internship we explore avenues for the vectorized execution of Pytorch models on RISC-V CPUs with Vector support. We identify 3 areas where Pytorch would benefit from vectorization: 1. the ATen computation backend, 2. the BLAS library, 3. the oneDNN compute library. Our contributions are as follows: we implement the vectorized class of ATen using RVV intrinsics, and we integrate vectorized version of BLAS and oneDNN into the Pytorch build process. This required us to setup an advanced, custom cross-compilation toolchain, including automated assembly modifications. Finally we evaluation the performance gained in elementary functions, fundamental building blocks of Deep Learning models (Linear Layers, Attention Layer and Convolutional Layers) and full AI models on our target hardware system, which is the EPAC (European Processor Accelerators) design, which is part of the European Processor Initiative.