Data-efficient Performance Modeling via Pre-training

TL;DR

This work addresses the data hunger of deep learning–based performance models used in autoschedulers by introducing a self-supervised autoencoder pre-training pipeline that learns robust code representations from unlabeled random programs. The pre-trained encoder embeds computation vectors before speedup prediction, significantly reducing the labeled data required to train accurate models on Tiramisu’s autoscheduler. Empirical results show that, at smaller data budgets, the pre-trained approach achieves substantial gains in speedups and maintains competitive MAPE (e.g., $MAPE$ improvements at 5× data reductions), while imposing only modest autoscheduler overhead. Overall, the method lowers the barrier to deploying data-driven performance models in compilers, enabling faster development and broader adoption across domains that rely on code optimization. The work contributes a scalable, source-level pre-training paradigm and provides extensive ablations and design guidance for practitioners.

Abstract

Performance models are essential for automatic code optimization, enabling compilers to predict the effects of code transformations on performance and guide search for optimal transformations. Building state-of-the-art performance models with deep learning, however, requires vast labeled datasets of random programs -- an expensive and time-consuming process, stretching over months. This paper introduces a self-supervised pre-training scheme with autoencoders to reduce the need for labeled data. By pre-training on a large dataset of random programs, the autoencoder learns representations of code and transformations, which are then used to embed programs for the performance model. Implemented in the Tiramisu autoscheduler, our approach improves model accuracy with less data. For example, to achieve a MAPE of 20.72%, the original model requires 18 million data points, whereas our method achieves a similar MAPE of 22.44% with only 3.6 million data points, reducing data requirements by 5x.

Figures (6)

• Figure 1: An example nested for loop structure.
• Figure 2: DNN based performance model by Merouani_2024
• Figure 3: Architecture of the Encoder Part of the Autoencoder
• Figure 4: MAPE achieved after training the performance model with (OURS) and without (ORIGINAL) pre-trained autoencoder on different datasize
• Figure 5: MAPE achieved after training two performance models, both utilizing our encoder architecture, but one with pre-trained weights (pretrained) and the other with randomly initialized weights (complex-embed).