Classifying Speedup of Loop Optimizations

Project Overview

For my Machine Learning final project, I worked within an existing research codebase created by Marco Alvarez and under the guidance of Christian Esteves and investigated how machine-learning models could classify the performance impact of compiler optimizations applied to loops in source code.

The broader problem is useful in high-performance computing: a compiler transformation may improve one loop while slowing down another, and determining that outcome normally requires actually compiling and benchmarking the transformed program. The research instead used source-code embeddings generated from the reference and transformed code and trained a classifier to predict whether the resulting program experienced a slowdown or speedup.

My work focused on understanding why the existing model struggled to generalize. That required studying the dataset, extending the experiment infrastructure, comparing stratification strategies, narrowing the data at the benchmark and application levels, and running hyperparameter sweeps.

Problem Setup

The dataset was built from benchmark applications whose loops had different combinations of compiler transformations applied to them. Every transformed loop had an associated measured speedup value. We bucketized those values into four classes from high slowdown to high speedup.

I worked with source-code embeddings generated using CodeRank. The classifier itself was a multilayer perceptron using linear layers, batch normalization, ReLU activations, dropout, and a four-class output.

The first major issue was that the dataset was not balanced. It contained 68,427 samples across 1,591 loop groups, but the classes were heavily skewed toward slowdowns and the number of samples per loop group varied heavily.

That meant a small number of applications and loop groups could dominate training even if the classification labels were stratified.

Building the Experiment Workflow

I created an Experiment Runner around the existing training and evaluation functions. It handled model initialization and configuration, experiment execution, Weights & Biases logging, and later hyperparameter sweeps and leave-one-out cross-validation.

I also extended the data-loading workflow so experiments could filter the dataset by benchmark and application. This made it possible to more easily investigate how the model behaved on different subsets of the data and to isolate where generalization was failing.

It was very useful for me to work on a problem of this kind because we had discussed the various methods used in classical machine learning, and now I got to apply them to a real problem that doesn’t behave nicely.

Stabilizing the Problem

I narrowed the problem to the NAS Parallel Benchmarks (NPB), which dominated the available data, and implemented leave-one-application-out cross-validation across the NPB applications. For each fold, the model trained on the remaining applications and validated against the one left out.

Training folds were generally large and similar in size, while some validation applications contained only a few dozen examples. Performance was therefore extremely mixed: larger applications produced useful validation results, while tiny folds gave unstable or poor accuracy.

There was a second issue. Even with enough data, an application performing one kind of computation may not be representative of another. Training on applications with one memory-access or arithmetic pattern does not guarantee that their embeddings contain enough information to generalize to a substantially different workload.

After the NPB experiments, I narrowed the dataset again to SP, the largest individual NPB application. This intentionally made the task simpler: before trying to generalize across heterogeneous applications, I wanted to verify how well the model behaved on a comparatively stable subset.

I compared random stratification against majority stratification. Across these SP runs, majority stratification produced higher validation accuracy than the random split, reaching roughly 93% compared with about 89% for random stratification. I treated this as a result specific to the stabilized SP experiment rather than evidence that majority stratification would generalize across the full dataset.

Hyperparameter Sweeps

With the SP subset behaving more consistently, I used Weights & Biases sweeps to explore the effect of:

• number of hidden layers;
• embedding latent dimension;
• dropout;
• learning rate;
• batch vs. layer normalization;
• weight decay.

One of the most useful findings was that the model was relatively robust to moderate hyperparameter changes. I could improve or degrade performance somewhat, but poor parameter choices were not enough to explain the large failures seen in the earlier cross-application experiments.

Normalization was one of the clearest patterns. Runs without normalization generally performed worse, while both batch and layer normalization produced stronger results. The later experiments also suggested that two or three hidden layers and 128- or 256-dimensional latent embeddings were better directions than making the network unnecessarily large.

My strongest runs reached validation accuracy around 93%, with relatively small gaps between training and validation on the better configurations.

←→

Weights & Biases sweeps let me trace model configuration through to training and validation accuracy rather than comparing isolated runs manually.

Showing image 1 of 2: Parallel-coordinate plot of model hyperparameters and validation accuracy

Results & Next Steps

The most important conclusions were:

• dataset composition was a larger source of instability than small hyperparameter changes;
• application and loop-group sizes needed to be considered explicitly when constructing train/validation splits;
• majority stratification was promising on the stabilized SP experiments;
• normalization consistently improved the MLP experiments;
• progressively narrowing the dataset was useful for separating model behavior from cross-application generalization problems.

The next research direction I identified was to build the dataset back up gradually rather than jump immediately from one application to every benchmark. In particular, I proposed a three-fold experiment using BT, SP, and LU, since they were among the largest NPB applications and shared more similar computational characteristics.

I also wanted to investigate the input representation itself. The active compiler transformations were represented with a long sparse binary vector, which may not have been ideal when only a small number of transformations were active. Improving that representation—and potentially the source-code embeddings themselves—could provide the model with more meaningful features before adding additional benchmarks.

Reflection

This project was useful because the research process did not have an obvious path to a better result.

I entered an existing codebase, had to learn how its data, embeddings, model, and experiment infrastructure fit together, and then determine what was actually preventing the classifier from improving.

It gave me a much more realistic experience of machine-learning research: the data is not necessarily clean, a higher-capacity model is not automatically the answer, and useful progress can mean understanding why an experiment fails well enough to design the next one.

Thank you to Professor Alvarez for allow me to work on this as my final project and for Christian Esteves for mentoring me through the process.

Connection to Published Work

This project was a small exploratory final project within a larger research effort on machine learning for program analysis and compiler optimization. My experiments focused specifically on the dataset and evaluation side of the problem: how class imbalance, application imbalance, stratification, and cross-application validation affected the apparent performance of the classifier. The related research continued and resulted in the paper “LLM-Based Embeddings for Program Analysis and Optimization,” published at IJCNN 2025 by Calvin Higgins and Marco Alvarez.

One important conclusion is that conventional stratified splits could place identical or closely related programs across training and testing sets, inflating apparent performance, while grouped evaluation exposed a much harder generalization problem.