Syno: Structured Synthesis for Neural Operators

TL;DR

Syno introduces an automated framework for neural operator synthesis, enabling the discovery of novel neural operators that improve accuracy and speed, surpassing traditional methods constrained by manual operator design.

Contribution

We propose Syno, a novel end-to-end neural operator synthesis framework utilizing fine-grained primitives and guided search to efficiently generate superior neural operators.

Findings

Achieves 1.37x to 2.06x speedups on various hardware.

Maintains less than 1% accuracy loss on NAS-optimized models.

Effectively discovers novel operators outperforming manually designed ones.

Abstract

The desires for better prediction accuracy and higher execution performance in neural networks never end. Neural architecture search (NAS) and tensor compilers are two popular techniques to optimize these two goals, but they are both limited to composing or optimizing existing manually designed operators rather than coming up with completely new designs. In this work, we explore the less studied direction of neural operator synthesis, which aims to automatically and efficiently discover novel neural operators with better accuracy and/or speed. We develop an end-to-end framework Syno, to realize practical neural operator synthesis. Syno makes use of a novel set of fine-grained primitives defined on tensor dimensions, which ensure various desired properties to ease model training, and also enable expression canonicalization techniques to avoid redundant candidates during search. Syno further adopts a novel guided synthesis flow to obtain valid operators matched with the specified input/output dimension sizes, and leverages efficient stochastic tree search algorithms to quickly explore the design space. We demonstrate that Syno discovers better operators with average speedups of $1.37\times$ to $2.06\times$ on various hardware and compiler choices, while keeping less than 1% accuracy loss even on NAS-optimized models.