Fusing the Polyhedral and Tensor Compilers to Accelerate Scientific Computing Kernels

Abstract

Polyhedral compilers and tensor compilers have achieved great success on accelerating scientific computing kernels and deep learning networks, respectively. Although much work has been done to integrate techniques of the polyhedral model to tensor compilers for accelerating deep learning, leveraging the powerful auto-tuning ability of modern tensor compilers to accelerate more general scientific computing kernels is challenging and is still at its dawn. In this work, we introduce a method to accelerate a family of basic scientific computing kernels by fusing the polyhedral compiler Pluto and the tensor compiler Tensor Virtual Machine (TVM) to generate efficient implementations targeting the heterogeneous CPU/GPU platform. The fusion is done in four steps: building a polyhedral model for the loop description of a given scientific kernel; designing schedules to transform the polyhedral model to new ones to enable rectangular tiling and expose explicit parallelism; selecting a new polyhedral model and converting it to the tensor compute representation; auto-tuning the tensor compute to generate efficient implementations on both CPUs and GPUs. Shifting and padding optimizations are also considered to avoid conditionals. Experiments on 30 typical scientific computing kernels show that our method achieves $3.31\times$ speedup on average over a typical polyhedral compiler PPCG on GPU.