A Deep Learning Based Cost Model for Automatic Code Optimization

Speaker: Riyadh Baghdadi

High-level aim: automatic code optimisation - a hard problem!
Predicting the performance of optimisations is hard & they have complex interactions.

Optimisation Selection

E.g. code + List of loop optimisations → optimisation selection
Becomes a search problem, wrt a cost model
2 problems then:
  1. Space exploration
  1. Candidate evaluation
    1. Cost estimation (how much saving?)
    2. Correctness check (is this still valid)
Aim here is 2.1.; For 1. we use MCTS or beam search

Goal

Build deep learning model for predicting speedup estimation of selection of optimisations.

Representations

Example features for a computation:
  1. Loop i
  1. Transformations applied on loop i
  1. Loop j
  1. ...
  1. Mem access 1
  1. ...
  1. Operations count
Uses these to create embeddings
Combine these hierarchically using LSTMs to create final embedding.
Feedforward NN applied to end, to predict single speedup value.

Training Dataset

Problem: not enough data for training
Solution: artificially generate random programs and optimisations and run them

Results

Pretty strong predicted speedup (16% MAE)
How does it compare to the optimal combination (found using search) → very close on most tasks! And better than Hailide AutoScheduler.