Enabling Polyhedral optimizations in TensorFlow through Polly

annanay25 cs14btech11001@iith.ac.in

Annanay Agarwal, IITHMichael Kruse, Polly LabsBrian Retford, Vertex.aiTobias Grosser, ETHZ, Polly LabsRamakrishna Upadrasta, IITH

1

Fast Machine Learning for everyone.

2

A wild appears!

❏ TensorFlow - Open Source Deep learning

Framework by Google.

❏ Built-in cross platform support for

writing Machine learning code.

❏ Numerical computation using data flow

graphs.

3

XLA - X(Acc)elerated Linear Algebra

❏ Recently open sourced (Jan’17)!

❏ JIT (Just In Time) -

❏ Runtime Compilation!

❏ Know the size of dataset you are dealing with!

❏ JIT compiles subgraphs of the TensorFlow computation!

❏ Uses LLVM as a backend!

4

The Force Awakens.

5

Architecture Diagram.

6

Polyhedral Compilation

❏ for ( i = 0; i < N; i++ ) ❏ for ( j = 0; j < M; j++ )❏ Stmt ( i, j )

❏ for ( i = 0; i < N; i++ ) ❏ for ( j = i ; j < M + i ; j++ )❏ Stmt ( i, j - i )

7

Pollyyy!

8

Polly in TensorFlow.

Polly Passes9

Behold!

10

Convolutional Neural Networks

11

Results

SCoP DetectionPolly’s SCoP detection was modified to

detect the convolution kernel from the

LLVM IR generated for tf.conv2d()

operation in XLA.

12

Results

PerformanceReduced cache misses - advanced data

locality optimizations like tiling.

Also performs operator fusion.

13

Results

PerformanceBetter overall runtime for convolution

kernel despite having a greater compilation

time.

14

Future Work

Thank you!

❏ SCoP detection and pattern optimization does not work for other deep

learning kernels like Recurrent Neural Networks (RNNs).

❏ Expand support to more Deep Learning kernels.

❏ Polly is capable of generating GPGPU code.

❏ Polly as a backend.

annanay25 cs14btech11001@iith.ac.in 15