Canvas: End-to-End Kernel Architecture Search in Neural Networks
This work addresses the need for efficient and accurate neural network kernels for machine learning practitioners, representing a novel integration of existing techniques rather than a purely incremental advance.
The paper tackles the problem of optimizing neural networks for both performance and accuracy by introducing Kernel Architecture Search (KAS), which combines tensor compilation and Neural Architecture Search to generate fine-grained neural kernels; the result is Canvas, an end-to-end framework that achieves average 1.5x speedups compared to previous state-of-the-art methods with acceptable accuracy loss.
The demands for higher performance and accuracy in neural networks (NNs) never end. Existing tensor compilation and Neural Architecture Search (NAS) techniques orthogonally optimize the two goals but actually share many similarities in their concrete strategies. We exploit such opportunities by combining the two into one and make a case for Kernel Architecture Search (KAS). KAS reviews NAS from a system perspective and zooms into a more fine-grained level to generate neural kernels with both high performance and good accuracy. To demonstrate the potential of KAS, we build an end-to-end framework, Canvas, to find high-quality kernels as convolution replacements. Canvas samples from a rich set of fine-grained primitives to stochastically and iteratively construct new kernels and evaluate them according to user-specified constraints. Canvas supports freely adjustable tensor dimension sizes inside the kernel and uses two levels of solvers to satisfy structural legality and fully utilize model budgets. The evaluation shows that by replacing standard convolutions with generated new kernels in common NNs, Canvas achieves average 1.5x speedups compared to the previous state-of-the-art with acceptable accuracy loss and search efficiency. Canvas verifies the practicability of KAS by rediscovering many manually designed kernels in the past and producing new structures that may inspire future machine learning innovations. For source code and implementation, we open-sourced Canvas at https://github.com/tsinghua-ideal/Canvas.