Towards Just-in-Time Instruction Generation for Accelerated Sparse Matrix-Matrix Multiplication on GPUs

Sparse Matrix-Matrix Multiplication (SpMM) is a fundamental operation in neural networks and high performance graph algorithms. Most popular solutions to S pMM adhere to ahead-of-time (AOT) compilation, where the entire program is compiled before execution. With the introduction of GPU threading in SpMM there have been significant leaps in performance. However, AOT compilation for threaded SpMM has four major limitations: unnecessary memory access, greater branch overhead, unnecessary retained memory, and slow memory transfers between the host and the device. These limitations are due to the nature of AOT compilations where key information is unaccessible during the compilation time but later becomes available during the programs runtime. In this paper, we propose a SIMD GPU threaded just-in-time (JIT) compilation solution to SpMM, improving SpMM resource usage and run time on GPUs. The most significant advantage of our framework is that JIT compilation provides our solution with runtime information. This information is used to dynamically create precise thread structures based on input matrices. Operations are evenly distributed to each thread in their smallest denominations to achieve maximum efficiency and a balanced workload distribution. Runtime information is also used to manage a self-maintained memory pool which increases memory efficiency, bandwidth, and SpMM speed. Improved forms of memory allocation and storage are used, decreasing the amount of transfers needed between the device and host, as well as increasing GPU memory bandwidth and decreasing memory access times. We conduct a performance evaluation of our framework and compare it to a AOT baseline provided by NVIDIA, a threaded matrix multiplier using the CUDA framework, adjusted for SpMM using an array in compressed sparse row (CSR) format. Our JIT framework consistently delivers a significant performance enhancement over the native CUDA solution, achieving on average, a 1.3x improvement in execution time and a 3x improvement in GPU memory footprint.