Vectorization and Distribution with Deep RL
Predict LLVM optimization factors with reinforcement learning methods
LLVM’s baseline models for loop optimization arebheuristic-based in nature and often lead to suboptimal performance. We use Code2Vec to generate code embeddings that characterize loops. With embeddings plus static features generated from LLVM pass, we train a reinforcement learning (RL) model that predicts the loop optimization factors to minimize execution time. We train the model on the loop dataset extracted from LLVM test suite and tune hyperparameters. The inference results show the effectiveness of adding distribution as an optimization factor, but the static features are too trivial to be helpful.
Figure 1: Pipeline of the method
Figure 2: Reward of training
Figure 3: Reward of training when numerical features are added