MIREncoder: Multi-modal IR-based Pretrained Embeddings for Performance Optimizations

TL;DR

MIREncoder is a multi-modal IR-based auto-encoder that pre-trains embeddings capturing code semantics and structure, enabling transfer learning for performance optimization with reduced overhead in high-performance computing.

Contribution

The paper introduces MIREncoder, a novel multi-modal IR-based auto-encoder that enables effective pre-training of embeddings for code analysis and optimization tasks, reducing overhead and improving performance.

Findings

Outperforms state-of-the-art methods in code optimization tasks.

Reduces overhead compared to large language model approaches.

Effectively captures code semantics and structure for downstream use.

Abstract

One of the primary areas of interest in High Performance Computing is the improvement of performance of parallel workloads. Nowadays, compilable source code-based optimization tasks that employ deep learning often exploit LLVM Intermediate Representations (IRs) for extracting features from source code. Most such works target specific tasks, or are designed with a pre-defined set of heuristics. So far, pre-trained models are rare in this domain, but the possibilities have been widely discussed. Especially approaches mimicking large-language models (LLMs) have been proposed. But these have prohibitively large training costs. In this paper, we propose MIREncoder, a M}ulti-modal IR-based Auto-Encoder that can be pre-trained to generate a learned embedding space to be used for downstream tasks by machine learning-based approaches. A multi-modal approach enables us to better extract features from compilable programs. It allows us to better model code syntax, semantics and structure. For code-based performance optimizations, these features are very important while making optimization decisions. A pre-trained model/embedding implicitly enables the usage of transfer learning, and helps move away from task-specific trained models. Additionally, a pre-trained model used for downstream performance optimization should itself have reduced overhead, and be easily usable. These considerations have led us to propose a modeling approach that i) understands code semantics and structure, ii) enables use of transfer learning, and iii) is small and simple enough to be easily re-purposed or reused even with low resource availability. Our evaluations will show that our proposed approach can outperform the state of the art while reducing overhead.