PaperlayerView exact revision on arXiv

Source-linked AI summary

NeuroVectorizer: End-to-End Vectorization with Deep Reinforcement Learning

Ameer Haj-Ali, Nesreen K. Ahmed, Ted Willke, Sophia Shao, Krste Asanovic, Ion Stoica

arXiv:1909.13639v4cs.DCcs.PFcs.PL

TL;DR

LLVM's heuristic cost models do not capture computation-graph structure, while manual vectorization requires substantial effort. The paper proposes an end-to-end code-embedding and deep RL framework for choosing vectorization factors, reporting 1.29×–4.73× average performance improvement over the LLVM baseline and results only 3% worse than brute-force search.

  • Problem

    LLVM's heuristic cost models do not consider computation graphs, while manual vectorization is labor-intensive and error-prone.

    Source: §1 · p. 1

  • Method

    The framework embeds loop code and uses deep RL to predict vectorization factors, with RL also supporting co-optimization of multiple objectives.

    Source: §1 · p. 2, §2.3 · p. 4

  • Results

    1.29×–4.73× average performance improvement over the baseline cost model implemented in LLVM was reported, with results only 3% worse than brute-force search.

    Source: §7 · p. 11

  • Takeaways & Limitations

    The authors report that RL can learn with fewer samples than supervised methods and co-optimize compilation time, code size, and execution time.

    Source: §2.3 · p. 4

    Interpretability remains a challenge for deep RL and neural networks generally.

    Source: §4.2 · p. 10

Abstract

from arXiv · show

One of the key challenges arising when compilers vectorize loops for today's SIMD-compatible architectures is to decide if vectorization or interleaving is beneficial. Then, the compiler has to determine how many instructions to pack together and how many loop iterations to interleave. Compilers are designed today to use fixed-cost models that are based on heuristics to make vectorization decisions on loops. However, these models are unable to capture the data dependency, the computation graph, or the organization of instructions. Alternatively, software engineers often hand-write the vectorization factors of every loop. This, however, places a huge burden on them, since it requires prior experience and significantly increases the development time. In this work, we explore a novel approach for handling loop vectorization and propose an end-to-end solution using deep reinforcement learning (RL). We conjecture that deep RL can capture different instructions, dependencies, and data structures to enable learning a sophisticated model that can better predict the actual performance cost and determine the optimal vectorization factors. We develop an end-to-end framework, from code to vectorization, that integrates deep RL in the LLVM compiler. Our proposed framework takes benchmark codes as input and extracts the loop codes. These loop codes are then fed to a loop embedding generator that learns an embedding for these loops. Finally, the learned embeddings are used as input to a Deep RL agent, which determines the vectorization factors for all the loops. We further extend our framework to support multiple supervised learning methods. We evaluate our approaches against the currently used LLVM vectorizer and loop polyhedral optimization techniques. Our experiments show 1.29X-4.73X performance speedup compared to baseline and only 3% worse than the brute-force search on a wide range of benchmarks.

1 Introduction

Loop vectorization requires choosing vectorization and interleaving factors, but existing cost models miss computation-graph structure and manual tuning is labor-intensive. The paper proposes an end-to-end learning framework and reports speedups over the LLVM baseline close to brute-force search.

1 Introduction

  • Loop vectorization converts scalar operations to SIMD operations, with VF determining how many instructions are packed and IF determining memory-access stride.IF supports vectorization of non-consecutive addresses.Source: §1 · p. 1
  • LLVM's heuristic cost models estimate instruction costs without considering the computation graph, while manual vectorization is labor-intensive, error-prone, and difficult to maintain or port.Prior machine-learning approaches use hand-engineered features that may not capture code functionality; supervised methods require labels and can struggle with large, multi-objective search spaces.Source: §1 · pp. 1–2
  • The proposed framework embeds loop code and maps the embedding to vectorization factors, aiming to co-optimize objectives while preserving code correctness.The authors describe the embedding as representing code in a way similar to how a human vectorization expert examines loop structure and operations.Source: §1 · p. 2
  • 1.29×–4.73× average performance speedup was reported, with results only 3% worse than the brute-force solution.The evaluation also compares the approaches with the LLVM cost model and Polly.Source: §1 · p. 2
  • The contributions include a dataset of more than 10,000 synthetic loops and an extensible framework supporting deep RL and multiple supervised learning methods.The paper also reports evaluations across learning hyperparameters and benchmark suites against the LLVM cost model and Polly.Source: §1 · p. 2

2 Motivation and Background

The section motivates learning loop vectorization factors by showing that optimal VF and IF choices depend on loop and hardware characteristics, and can outperform LLVM’s baseline. It then introduces deep RL as an end-to-end approach for mapping loop embeddings to vectorization factors.

2.1 Vectorization Characterization

  • Choosing VF and IF is cumbersome because the right values depend on loop instructions, stride, hardware, computation graph, and functionality.Source: §2.1 · p. 2
  • VF = 64, IF = 8 achieves up to 20% better performance than LLVM’s baseline on the dot product kernel, while brute-force search improves all tested LLVM vectorizer cases by up to 1.5×.The 1.5× gap occurs on more complicated tests.Source: §2.1 · p. 2, Figure 1 · p. 3, Figure 2 · p. 3

2.2 State-of-the-Art Auto-Vectorzation

  • Manual factor selection is error-prone and time-consuming; Polly uses tiling and fusion for locality, while prior learned approaches rely on hand-engineered features or labels.Supervised learning also faces challenges with label acquisition and large multi-objective search spaces.Source: §2.2 · pp. 2–3
  • The proposed framework learns loop embeddings and maps them to optimal VF and IF with RL, which can co-optimize objectives without brute-force search.The authors describe RL as more sample-efficient than supervised methods.Source: §2.2 · p. 3, §2.3 · p. 4

2.3 Deep Reinforcement Learning for Auto-Vectorization

  • The framework extracts loops, embeds their text, predicts vectorization factors, injects compiler pragmas, and uses execution-time improvements as rewards.Source: Figure 3 · p. 4
  • Contextual bandits reduce RL to selecting one action in the current state, and PPO updates policies while limiting deviation from the previous policy.Source: §2.3 · p. 4

3 The Proposed Framework for Automatic Vectorization

The framework extracts loops, embeds their code, and uses a learning agent to select vectorization factors. Its RL environment evaluates pragmas against LLVM’s baseline, while the dataset and compilation-time handling address practical training constraints.

3 The Proposed Framework for Automatic Vectorization

  • The framework extracts loops and their contexts, embeds the code, predicts vectorization factors, injects pragmas, and uses compiled execution-time improvements as RL rewards.Once trained, the model can be used for inference without further retraining; pragmas act as compiler hints that LLVM may ignore when infeasible.Source: §3 · pp. 4–5

3.1 Code Embedding

  • Code2vec represents each loop snippet as a fixed-length vector of 340 features intended to capture code characteristics such as semantic similarities, combinations, and analogies.The embedding generator aims to place similar loop code close together so the agent can make similar vectorization decisions.Source: §3.1 · p. 5

3.2 The RL Environment Definition

  • The agent’s reward compares execution time under its injected pragmas with execution time under LLVM’s baseline cost model, normalized by the baseline time.The action selects VF and IF; a joint neural network predicting both performed better than two independent agents, consistent with their direct correlation.Source: §3.2 · p. 5

3.3 Dataset Description

  • The authors generated more than 10,000 synthetic loop examples from LLVM’s vectorization test suite, varying parameters, stride, iteration count, functionality, instructions, and nesting.They varied parameter names to reduce embedding noise from name-based bias; available open-source benchmarks for training were limited.Source: §3.3 · p. 6
  • Different loops have different optimal VF and IF values; (VF = 4, IF = 2) is the most common pair in the dataset.The distribution comes from brute-force search with MAX_VF = 16 and MAX_IF = 8, though the loops do not represent all existing loops.Source: §3.3 · p. 6, Figure 5 · p. 6

3.4 Handling Long Compilation Time

  • To limit long compilation during training, programs exceeding ten times baseline compilation time receive a −9 penalty reward.The authors report that, for the slow-compiling programs they examined, waiting for compilation did not yield better performance than other vectorization configurations.Source: §3.4 · p. 6

4 Evaluation

The evaluation compares learning-based vectorization methods with LLVM’s baseline, Polly, and brute-force search, then tests training behavior and transfer to new benchmarks. It also discusses the practical challenge of deploying learning-based compiler optimization.

4 Evaluation

  • At a 5e-5 learning rate, reward mean reached its maximum fastest; 5e-3 did not reach a higher maximum, and smaller-than-500 batch sizes produced unstable policies.FCNN architecture differences were minor, and the policy converged with fewer samples as batch size decreased.Source: Figure 6 · p. 7, §4 · p. 8
  • The discrete action space performed best among the tested definitions, where the agent selected indices for vectorization and interleaving factors or rounded continuous encodings.The study compared one discrete and two continuous action-space definitions.Source: §4 · p. 8, Figure 7 · p. 9
  • 2.67× average improvement over the baseline and performance only 3% worse than brute-force search were achieved by RL across the evaluated benchmarks.NNS and decision trees also outperformed the baseline by 2.65× and 2.47×, respectively.Source: §4 · p. 8, Figure 8 · p. 9
  • Polly outperformed the baseline by 17% but performed 56% worse than RL; its loop transformations beyond vectorization suggest potential benefit from combining Polly and RL.For benchmark #10, Polly also outperformed brute-force search.Source: §4 · p. 8
  • With 5,000 compilations, deep RL reached high performance, only 5% below its peak and 1.26× better than supervised FCNN; FCNN needed 70,000 compilations for that level.The authors report FCNN as 14× less sample efficient; RL training is performed once, with inference requiring a single step.Source: §4 · p. 8, Figure 9 · p. 10

4.1 Transfer Learning

  • Deep RL improved average PolyBench performance 3.42× over the baseline and 1.33× over Polly, outperforming Polly on three of six benchmarks despite Polly’s PolyBench focus.Using best rather than average performance across multiple inferences raised the reported average improvement to 4.77×.Source: §4.1 · p. 9, Figure 10 · p. 10
  • On MiBench, deep RL outperformed Polly and the baseline on every benchmark, with average performance improvement of 1.1× over the baseline.The authors note that loops are a minor part of MiBench programs and execution time was measured for whole programs.Source: §4.1 · pp. 9–10, Figure 11 · p. 10

4.2 Discussion: Deployability

  • Compiler adoption of learning-based optimization faces demands for deterministic, simple, explainable methods; the authors identify interpretability as the main remaining challenge for deep RL.They also say sufficient training data can make deep RL deterministic and performant across large application sets, while inference is simple to use and deploy.Source: §4.2 · p. 10

5 Future Directions

The paper outlines extensions to the embedding, action, and optimization stages, alongside broader training and architecture coverage needed for standardization.

5 Future Directions

  • Polyhedral analysis could provide loop-focused code embeddings at lower computational cost.The authors also propose combining deep RL with Polly, potentially training the agent to predict whether Polly should be used.Source: §5 · p. 10
  • The vectorizer could operate at the intermediate-representation level, while different architectures require architecture features or separately fitted models.The authors say training across a wide range of applications and target architectures remains necessary for standardization in LLVM.Source: §5 · pp. 10–11
  • For newer compiler support of vectorization at different nested-loop levels, the agent could make multiple sequential decisions within one episode instead of one decision per loop nest.The current approach assumes a single decision per loop nest.Source: §5 · p. 11
  • The framework could tune other pragmas and compiler-wide choices by defining an action space and reward, including different phase orderings for different code sections.Examples include function inlining, loop unrolling, superword-level parallelism, and scatter/gather.Source: §5 · p. 11
  • A vanilla neural network could rank vectorization-factor and interleaving-factor choices by predicted normalized execution time, effectively learning a cost model trainable end-to-end.The authors describe this as a potential replacement for the baseline cost model.Source: §5 · p. 11

6 Related Work

Related work uses fixed cost models, heuristics, and machine learning for vectorization and compiler optimization. This work distinguishes itself through deep RL and end-to-end feature learning from code text.

6 Related Work

  • This work is described as the first to use deep RL to explore the vectorization space and compiler optimization in general.Prior deep supervised and RL approaches addressed phase ordering.Source: §6 · p. 11
  • Unlike approaches using engineered program features, this work learns useful features end-to-end from text code to vectorization factors.Prior work uses handcrafted features, while end-to-end supervised learning can require time-consuming labels and struggles with large multi-objective search spaces.Source: §6 · p. 11
  • LLVM and proposed cost models use predefined linear cost functions based on instruction distribution to estimate vectorized-loop execution time.Source: §6 · p. 11
  • Other approaches use heuristics or improve SLP, with SLP methods relying on baseline or fixed cost models rather than addressing loop vectorization directly.The work in does not address loop vectorization and relies on the baseline cost model to decide when code should remain unvectorized.Source: §6 · p. 11

7 Conclusion

The framework detects loops, learns their structures, and uses deep RL to inject vectorization pragmas. It improves performance by 1.29×–4.73× over LLVM’s baseline cost model and averages only 3% worse than brute-force search.

7 Conclusion

  • 1.29×–4.73× average performance improvement versus LLVM’s baseline cost model, with results averaging only 3% worse than the brute-force solution.Source: §7 · p. 11
  • The end-to-end framework automatically detects loops, learns their structures, and applies deep RL to inject vectorization pragmas into the compiler.Source: §7 · p. 11
  • The authors foresee potential for end-to-end code tuning and optimization using machine learning techniques such as deep RL.Source: §7 · p. 11
Loading 1909.13639v4…