Auto Tuning Quantized Graph Neural Networks on GPU Tensor Core via TVM

Research on quantized graph neural networks (QGNN) have driven the development of related inference acceleration techniques. Both lower memory footprint and lower inference latency have led more and more researchers to focus on these new techniques. Although there is work on how to boost the performance of QGNN on modern GPUs by significant manual effort, how to automate building high-performance kernels on modern GPUs is still our concern. We propose AQGN, an extension plug-in on the TVM compilation stack for automatically exploring the sparsity of quantized graph neural networks on new tensor core hardware. Combined with post training quantization method, AQGN does not require model retraining. We construct a extended configuration search space to automate model tuning. We craft a module in TVM that automatically search the quantization configuration space given the hardware’s feedback. Experimental data demonstrates that our implementation outperforms the current DGL framework in terms of efficiency.