Gensor: A Graph-based Construction Tensor Compilation Method for Deep Learning

TL;DR

Gensor is a graph-based tensor compilation method that enhances deep learning performance by rapidly exploring a larger optimization space, leading to significant speedups on GPUs for various models.

Contribution

Gensor introduces a novel graph-based approach for tensor program optimization, expanding the search space and improving performance with faster compilation times.

Findings

Achieves 18-30% performance improvement over state-of-the-art methods.

Generates operator kernels in seconds, enabling rapid optimization.

Provides an average 20% acceleration for models like ResNet-50 and GPT-2.

Abstract

High-performance deep learning depends on efficient tensor programs. In recent years, automatic tensor program optimization, also known as tensor compilation, has emerged as the primary approach to generating efficient tensor programs. However, how to generate kernels with higher performance in a shorter time is still the key challenge. In this paper, we present Gensor, a graph-based construction tensor compilation method for deep learning, to further improve the performance of construction tensor compilation. Unlike existing tree-based methods, Gensor abstracts construction space into a graph structure. Gensor then explores the construction space with Markov analysis. Gensor takes tensor programs as states and models scheduling primitives as transition actions between these states. Therefore, the process of tensor program construction optimization is abstracted as a graph traversal process. This approach expands the optimization space, improving operator performance while ensuring rapid optimization. Extensive experiments with typical operators demonstrate that Gensor significantly outperforms the state-of-the-art methods on GPUs for both cloud servers and edge devices. As a result, Gensor can generate operator kernels in seconds, with performance increasing by 18\% on average, reaching a maximum of 30\%. It also achieves high speedup for end-to-end models like ResNet-50 and GPT-2, with an average acceleration of 20\%.