1 of 36

DL-based Compiler Optimizations

-Dibyendu Das

Sr. Principal Engineer, Intel India

2 of 36

Agenda

  • Brief Background
  • DeepTune- End-to-end Deep Learning of Optimization Heuristics
  • Ithemal : Basic Block Throughput Prediction
  • Graph-Coloring based Register Allocation via LSTMs
  • Compiler Auto-Vectorization using Imitation Learning
  • Conclusion and Way Forward

LLVM Social Bangalore Sep 2021

2

9/18/2021

3 of 36

Brief Background

LLVM Social Bangalore Sep 2021

3

9/18/2021

Program

IR

IR

IR

Optimized program in target language

Optimized Machine Code

Analysis

Transformation

Code Generation

Compiler Optimization Pass

4 of 36

Brief Background (Cont’d)

​

​

  • Lots of compiler optimization depend on carefully constructed heuristics
    • Which optimization order to apply
    • How to inline functions, which functions should be inlined
    • Scheduling and register allocation
    • How much to unroll loops
  • Heuristics are developed and maintained by compiler engineers
    • Not always scalable
    • Difficult to adapt to different hardware
    • Fragile in many cases
    • The space of possible optimizations is also vast, making it very hard for compiler writers to design heuristics that take all these considerations into account
      • As a result, many compiler optimizations are out of date or poorly tuned

LLVM Social Bangalore Sep 2021

4

9/18/2021

5 of 36

Brief Background (Cont’d)

​

​

  • Given a program, compiler writers would like to know what compiler heuristic or optimization to apply in order to make the code better
    • Execute faster
    • Smaller code footprint
    • Reduced power
  • Machine learning can be used to build a model used within the compiler, that makes such decisions for any given program
    • Early work concentrated on feature engineering ( ex: Loop Unrolling )
    • More recent work applies DL, RL techniques

LLVM Social Bangalore Sep 2021

5

9/18/2021

* Figure taken from: Machine Learning in Compiler Optimisation, Zheng Wang and Michael O’Boyle

6 of 36

DeepTune- End-to-end Deep Learning of Optimization Heuristics

LLVM Social Bangalore Sep 2021

6

9/18/2021

7 of 36

DeepTune

​

  • Used for optimal mapping of heterogeneous parallelism
    • Should we run an OpenCL kernel on a CPU or a GPU ?
  • End2End DL of optimization heuristics
    • One of the earliest works to test this strategy

LLVM Social Bangalore Sep 2021

7

9/18/2021

OpenCL Code

Source Rewriter

Sequence Encoder

Embedding

LSTM

Auxiliary Inputs

Batchnorm

Dense

+

Predicted Device Mapping– {CPU,GPU}

Language Model

8 of 36

DeepTune Components

​

  • Language Model
  • Source Rewriter
  • Sequence Encoder

LLVM Social Bangalore Sep 2021

8

9/18/2021

​

​

__kernel void memset_kernel(__global char *mem_d,

short val, int num_bytes) {

const int thread_id = get_global_id(0);

mem_d[thread_id] = val;

}

​

​

​

__kernel void A(__global char *a, short b, int c) {

const int d = get_global_id(0);

a[d] = b;

}

​

9 of 36

DeepTune Components

​

  • Embedding
    • Converts sparse to dense representations
  • LSTM
  • Heuristic Model

LLVM Social Bangalore Sep 2021

9

9/18/2021

Auxiliary Inputs for

GPUs - work-group size and data-size

def init(self, seed: int):

from keras.layers import Input, Embedding, LSTM, Dense

from keras.layers.merge import Concatenate

from keras.layers.normalization import BatchNormalization

from keras.models import Model

np.random.seed(seed)

​

# Language model. Takes as inputs source code sequences.

code_in = Input(shape=(1024,), dtype="int32", name ="code_in")

x = Embedding(input_dim=atomizer.vocab_size + 1, input_length=1024, output_dim=64,

name="embedding")(code_in)

x = LSTM(64, implementation=1, return_sequences=True, name="lstm_1")(x)

x = LSTM(64, implementation=1, name="lstm_2")(x)

langmodel_out = Dense(2, activation="sigmoid")(x)

​

# Auxiliary inputs. wgsize and dsize.

auxiliary_inputs = Input(shape=(2,))

x = Concatenate()([auxiliary_inputs, x])

x = BatchNormalization()(x)

x = Dense(32, activation="relu")(x)

out = Dense(2, activation="sigmoid")(x)

self.model = Model(inputs=[auxiliary_inputs, code_in], outputs=[out, langmodel_out])

self.model.compile( optimizer="adam", metrics=['accuracy'], loss=["categorical_crossentropy",

"categorical_crossentropy"], loss_weights=[1., .2])

return self

https://github.com/ChrisCummins/paper-end2end-dl/blob/master/code/Case%20Study%20A.ipynb

10 of 36

Visualizing DeepTune

LLVM Social Bangalore Sep 2021

10

9/18/2021

* Figure taken from: End-to-end Deep Learning of Optimization Heuristics. Chris Cummins, Pavlos Petoumenos, Zheng Wang, Hugh Leather

11 of 36

Performance

LLVM Social Bangalore Sep 2021

11

9/18/2021

  • The optimal static mapping achieves 58% accuracy
  • Grewe et al. (hand-crafted features + ML) achieves 73% accuracy
  • DeepTune predictive models achieve accuracy of 84%

* Figure taken from: End-to-end Deep Learning of Optimization Heuristics. Chris Cummins, Pavlos Petoumenos, Zheng Wang, Hugh Leather

12 of 36

Ithemal : Basic Block Throughput Prediction

LLVM Social Bangalore Sep 2021

12

9/18/2021

13 of 36

Estimating Basic Block Throughput

​

  • The problem:
    • Given a basic-block of x86 instructions estimate the throughput in terms of clock cycles
  • Analytical models (llvm-mca, IACA) are usually used in such scenarios

LLVM Social Bangalore Sep 2021

13

9/18/2021

High-Level Code

Optimizing Compiler

lea r14, [rbx-0x40]

…

lea rdx, [rbp+0x48]

cmp rdi, rax

lea r14, [rbx-0x40]

…

lea rdx, [rbp+0x48]

cmp rdi, rax

lea r14, [rbx-0x40]

…

lea rdx, [rbp+0x48]

cmp rdi, rax

…

40 cycles

44 cycles

36 cycles

14 of 36

Accurate Modeling of Processor Core is Complex

​

  • Modeling the micro-architectural details of a complex core is a very hard problem
      • Very easy to omit details
      • Specs are not always accurate
      • Some details are proprietary

LLVM Social Bangalore Sep 2021

14

9/18/2021

lea r14, [rbx-0x40]

…

lea rdx, [rbp+0x48]

cmp rdi, rax

K cycles

15 of 36

Ithemal: A Data-driven approach

​

  • Run/train a DL model using many samples of x86 BBs and their corresponding cycle count
    • Measured on real hardware
    • A new hardware just means re-training the new model OR some form of transfer learning
  • High accuracy and ease of portability

​

LLVM Social Bangalore Sep 2021

15

9/18/2021

* Figure taken from: Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks, C. Mendis et al.

16 of 36

Ithemal DL model

​

  • Hierarchical LSTMs
    • 2-layers
    • Layer-1 for the sequence of operands of each instruction
    • Layer-2 for the sequence of instructions

LLVM Social Bangalore Sep 2021

16

9/18/2021

Layer-1 Instruction-level LSTM

Layer-2 BB-level LSTM

* Figure taken from: Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks, C. Mendis et al.

17 of 36

Ithemal Comparative Performance

​

LLVM Social Bangalore Sep 2021

17

9/18/2021

* Figure taken from: Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks, C. Mendis et al.

18 of 36

DL-based Graph Coloring Register Allocation

LLVM Social Bangalore Sep 2021

18

9/18/2021

19 of 36

Graph Coloring

​

  • Graph-coloring is an important problem in CS with numerous applications
  • A graph coloring is an assignment of labels, called colors, to the vertices of a graph such that no two adjacent vertices share the same color
  • The chromatic number X(G) of a graph G is the minimal number of colors for which such an assignment is possible
  • Solved using heuristics

LLVM Social Bangalore Sep 2021

19

9/18/2021

A graph coloring for a graph with 6 vertices. It is impossible to color the graph with 2 colors, so the graph has chromatic number 3.

20 of 36

Register Allocation as a Graph-Coloring Problem

​

  • Register Allocation is an important problem in the area of compiler code generation
  • Usually the number of registers available < number of variables used
  • Create what is known as the *interference graph* which models registers which need to be *live* at the same time

LLVM Social Bangalore Sep 2021

20

9/18/2021

g

k

j

h

Interference Graph of Live Ranges

Live Range

21 of 36

Modeling Graph Coloring using LSTMS

​

  • Viewed as a sequence-2-sequence translation via LSTMs
  • An input sequence where each item of the sequence corresponds to a node of the graph
  • The output sequence is of the same length as the input sequence (number of nodes of the graph)
  • Trained using random graphs

​

​

LLVM Social Bangalore Sep 2021

21

9/18/2021

1 1 0 0 …

0 1 0 1 …

1 0 0 1…

Adj Matrix

1024 hidden units per LSTM

22 of 36

Inference and Color-Correction phase

​

  • Difficult to encode constraints in LSTM that two adjacent nodes cannot have some color
  • *Invalid* edges may appear during inference
  • Rectify these edges using a post-inference color-correct pass

LLVM Social Bangalore Sep 2021

22

9/18/2021

  • Forest-Fire graph
      • has a chromatic number of 5
  • LSTM-based model colors with 4 colors resulting in 2 invalid edges
  • These edges are <v2,v3> and <v1,v5>
  • v5 can reuse the color c3 as none of its neighboring nodes use c3. v2 requires a new color c5 as both v2 and v3’s neighbors use all the colors
  • So finally we also get 5 colors ☺

​

​

​

23 of 36

DL-model vs LLVM’s Greedy Register Allocator (GRA)

​

  • Collected the interference graphs for the functions of certain SPEC CPU® 2017 benchmarks
    • Use these graphs to predict colors using the DL-model
  • Collect the actual register count of each function after codegen from LLVM
  • Comparison shows DL-model performing better than GRA

LLVM Social Bangalore Sep 2021

23

9/18/2021

LSTM-based model

Interference graph

Color/register assignment

https://github.com/dasdibye/DL4RegAlloc

24 of 36

Compiler Auto-vectorization using Imitation Learning

LLVM Social Bangalore Sep 2021

24

9/18/2021

25 of 36

Auto-Vectorization

​

  • The problem:
    • Solve the SLP (Superword-Level Parallelism) using ML
    • SLP is a superset of Loop Vectorization whereby stratight-line code can also be vectorized
  • Naïve ML strategies may lead to correctness issues as only isomorphic statements can be vectorized

LLVM Social Bangalore Sep 2021

25

9/18/2021

​

a[1] = b[1] + c[1]

a[2] = b[2] + c[2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

​

a[1:2] = b[1:2] + c[1:2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

Scalar Code

Vector Code

26 of 36

Auto-Vectorization

​

  • Naïve ML strategies may lead to correctness issues as only isomorphic statements can be vectorized

LLVM Social Bangalore Sep 2021

26

9/18/2021

​

a[1] = b[1] + c[1]

a[2] = b[2] + c[2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

​

a[1] = b[1] + c[1]

a[2] = b[2] + c[2]

a[3] = b[3] + c[3]

a[4:5] = b[4:5] * c[4:5]

​

Scalar Code

Vector Code

Neural Machine Translation

27 of 36

MDP Formulation

​

  • The pairwise instruction packing problem formulated as a MDP by iteratively forming one vector pack at a time following a particular instruction traversal policy – Bottom-Up or Top-Down

LLVM Social Bangalore Sep 2021

27

9/18/2021

​

a[1] = b[1] + c[1]

a[2] = b[2] + c[2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

​

a[1:2] = b[1:2] + c[1:2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

Choose a “valid” action

{a[1],a[2]},{a[2],a[3]},{a[3],a[4]}

State

New State

​

a[1:2] = b[1:2] + c[1:2]

a[3] = b[3] + c[3]

a[4] = b[4] + c[4]

a[5] = b[5] * c[5]

​

​

a[1:2] = b[1:2] + c[1:2]

a[3:4] = b[3:4] + c[3:4]

a[5] = b[5] * c[5]

​

Choose a “valid” action

{a[3],a[4]}

State

New State

Iterate

Get Reward

{a[1],a[2]}

28 of 36

MDP Formulation – State details

​

  • Nodes : 5 types of nodes to encode the graph features of a state:
    • Instruction Node: correspond to each instruction with at least one valid packing opportunity or instructions which are already packed
    • Pack Node: common node representing overhead packing instructions
    • Unpack Node: common node representing overhead unpacking instructions
    • Constant Node: common node representing any constant value used by instructions
    • Focus Node: connected to the instruction that is considered for packing
  • Edges: Following are the 4 types of edges connecting the above nodes:
    • Dependency Edge: encodes if an instruction must be executed after another
    • Possible Pack Edge: encodes whether two instructions can be packed together
    • Packed Edge: encodes instructions that are already packed together
    • Focus Edge: the focus edge connects the focus node to the instruction node that is considered for packing

LLVM Social Bangalore Sep 2021

28

9/18/2021

* Figure taken from: Compiler 2.0: How to Modernize Compiler Technology, Saman Amarasinghe et al.

29 of 36

Imitation Learning

​

  • goSLP solves the SLP problem *exactly* using ILP solvers
  • Use such a solution to imitate/mimic the action space
  • Use actual runtimes/estimates as cost/reward function
  • Gated Graph Neural Network (GGNN) used as part of the policy network modeling to make packing decisions for each state

​

LLVM Social Bangalore Sep 2021

29

9/18/2021

Choose a “valid” action

{a[1],a[2]},{a[2],a[3]},{a[3],a[4]}

State

New State

{a[1],a[2]}

GGNN

30 of 36

Related Work

  • End-to-end Deep Learning of Optimization Heuristics. Chris Cummins, Pavlos Petoumenos, Zheng Wang, Hugh Leather. PACT 2017.
  • Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks. Charith Mendis, Alex Renda, Saman Amarasinghe, Michael Carbin. ICML 2019.
  • Deep Learning-based Approximate Graph-Coloring Algorithm for Register Allocation. Dibyendu Das, Shahid Asghar Ahmad, Venkataramanan Kumar. LLVM-HPC 2020 Workshop, co-located with Supercomputing 2020.
  • Compiler Auto-Vectorization with Imitation Learning. Charith Mendis, Cambridge Yang, Yewen Pu, Saman Amarasinghe, Michael Carbin. NeurIPS 2019.
  • Compiler 2.0: How to Modernize Compiler Technology, Saman Amarasinghe, Charith Mendis, Sam Larsen, Ajay Brahmakshatriya, Yishen Chen, Alex Renda, Cambridge Yang, Yewen Pu, Michael Carbin.

​

LLVM Social Bangalore Sep 2021

30

9/18/2021

31 of 36

Some other work we did not cover

  • Static Neural Compiler Optimization via Deep Reinforcement Learning – R. Mammdil et al. LLVM-HPC Workshop (2020) co-located with Supercomputing 2020
    • Solving the phase-ordering problem for the LLVM compiler pass-pipeline
  • NeuroVectorizer: End-to-End Vectorization with Deep Reinforcement Learning – A Haj-Ali et al. CGO 2020.
    • Predict the most beneficial VFs(Vectorization Factor) for Loop Vectorization
  • MLGO: a Machine Learning Guided Compiler Optimizations Framework – M. Trofin et al. https://arxiv.org/pdf/2101.04808.pdf
    • Replacing the heuristics-based inlining-for-size optimization in LLVM with machine learned models
  • CompilerGym - Facebook AI. https://github.com/facebookresearch/CompilerGym
    • CompilerGym is a toolkit for exposing compiler optimization problems for reinforcement learning. It allows machine learning researchers to experiment with program optimization techniques without requiring any experience in compilers, and provides a framework for compiler developers to expose new optimization problems for AI

LLVM Social Bangalore Sep 2021

31

9/18/2021

32 of 36

Wrapping it Up

LLVM Social Bangalore Sep 2021

32

9/18/2021

33 of 36

Conclusion and Future Directions

  • We are just at the cusp of utilizing DL techniques for compiler analysis/transform
  • A big challenge is to ensure correctness of the transformed code
  • Some of the proposed work consists of a mix of NLP-like techniques and RL
  • No consensus yet on standardized program representations though techniques along the lines of word2vec seem to be gaining traction
  • No consensus yet on representing control and data-flow though Graph Neural Networks (GNNs) are being used in some cases

​

LLVM Social Bangalore Sep 2021

33

9/18/2021

34 of 36

Conclusion and Future Directions

  • Some authors like Cummins & Leather suggest OaaS
    • Optimization as a service
    • All optimizations implemented as some kind of an RL-based model
  • Personal opinion:
    • RL too costly to be introduced for all optimization passes
    • Compile time may blow up and we may need sophisticated HW resources
    • Optimizations of the future may be a mix of NLP-like and RL models
  • Open questions:
    • Can we generate *semantically correct* transformed code using a DL model ?
    • Would the optimizations of the future be a mix of DL models + Fast correction algorithms ?
      • Get most of it right using a DL model and then apply a quick and simple correction pass
    • Can these models be really portable – across compilers and across hardware ?

LLVM Social Bangalore Sep 2021

34

9/18/2021

* Figure taken from: Machine Learning in Compilers: Past, Present and Future, Chris Cummins and Hugh Leather

35 of 36

The Future ?

LLVM Social Bangalore Sep 2021

35

9/18/2021

Program

IR

IR

IR

Optimized program in target language

Optimized Machine Code

Learnt Analysis Model

Learnt Transformation Model

Code Generation

Learnt Compiler Optimization Pass

36 of 36

Thanks!

LLVM Social Bangalore Sep 2021

36

9/18/2021