OLLIE: Derivation-based Tensor Program Optimizer

TL;DR

OLLIE is a novel derivation-based tensor program optimizer that explores a broader expression space for DNNs, leading to significant performance improvements over existing methods.

Contribution

OLLIE introduces a derivation-based approach that enables optimization over general tensor algebra expressions, expanding the search space beyond fixed operator sets.

Findings

Up to 2.73× performance improvement on A100 GPU

Average 1.46× speedup across tested DNNs

Effective optimization on multiple GPU architectures

Abstract

Boosting the runtime performance of deep neural networks (DNNs) is critical due to their wide adoption in real-world tasks. Existing approaches to optimizing the tensor algebra expression of a DNN only consider expressions representable by a fixed set of predefined operators, missing possible optimization opportunities between general expressions. We propose OLLIE, the first derivation-based tensor program optimizer. OLLIE optimizes tensor programs by leveraging transformations between general tensor algebra expressions, enabling a significantly larger expression search space that includes those supported by prior work as special cases. OLLIE uses a hybrid derivation-based optimizer that effectively combines explorative and guided derivations to quickly discover highly optimized expressions. Evaluation on seven DNNs shows that OLLIE can outperform existing optimizers by up to 2.73$\times$ (1.46$\times$ on average) on an A100 GPU and up to 2.68$\times$ (1.51$\times$) on a V100 GPU, respectively.