Automatic Differentiation of Parallelised Convolutional
Neural Networks - Lessons from Adjoint PDE Solvers
Jan Hückelheim, Imperial College London
Paul Hovland, Argonne National Laboratory

December 9, 2017

Jan Hückelheim

Many-core adjoints

1About me• M.Sc. from RWTH Aachen, Germany, 2012
• PhD from Queen Mary University of London, 2017
• Research Associate at Imperial College London, present• Inria, work on Tapenade static analysis• Argonne National Laboratory, parallel AD• AD and verification in Computational Fluid Dynamics, Seismic Imaging

Jan Hückelheim

Many-core adjoints

2An example from PDE solvers: Seismic imaging• Seismic imaging: Explore the subsurface geological structure

ic
m

m
ic

ic

shot

m

ic
m

m

m

ic

ic

• In real life: Shots are being fired, and the reflections recorded

surfacesubsurface
structure

Jan Hückelheim

Many-core adjoints

3An example from PDE solvers: Seismic imaging• In simulation, the same experiment is conducted• Since we don’t know the subsurface yet, we assume something

surface?

?
?

Jan Hückelheim

Many-core adjoints?

?
?

unknownsubsurface
structure4

An example from PDE solvers: Seismic imaging• Back-propagate the mismatch between simulation and measurement• Minimise mismatch by updating assumed subsurface structuresurface?

?
?

Jan Hückelheim

Many-core adjoints?

?
?

unknownsubsurface
structure5

Back-propagation in CNNs
• Convolutional layers, subsampling layers, unknown weights everywhere• Models are ”trained” to minimise misclassifications

forward pass

output andtraining data
mismatchbackwards passJan Hückelheim

Many-core adjoints

6More similarities• Stencil computations in PDE solvers look like convolutions
Updatedwave
field

Features

Wave
field

Image

window

stencil

Note that there are also differences:• CNNs have few layers, compared to many iterations in PDE solvers• Loop bodies more complex in PDE solvers• Boundary treatment is different

Let’s see how much AD knowledge we can transfer.
Jan Hückelheim

Many-core adjoints

7

Algorithmic differentiation (AD)• Given a program (”primal”) that implements some function

J = F (α),
AD generates a program that implements the derivative

Tangent mode
• Computes the Jacobian-vector productJ̇ = (∇F (x)) · α̇.

Adjoint mode
• Computes the transpose Jacobian-vector productᾱ = (∇F (x))T · J̄.
Jan Hückelheim

Many-core adjoints

8Forward vs. reverse• Tangent mode is simple to understand and implement, but: Need to

re-run for every input.
• Adjoint mode is cheaper for many inputs and few outputs (run once,

get all directional derivatives).
Original program
alpha

Reversedifferentiationintermediate
valuesJ
Forward differentiation

Jan Hückelheim

Many-core adjoints

9AD approaches
There are at least two ways of implementing AD:

Source-to-source transformation• Creates code that computes partial derivative of each operation, andassembles them with chain-rule.
• Fast, efficient, but hard to get right. Mainly Fortran/C

Operator overloading• Trace the computation at runtime, compute adjoints based on trace.

Slow, huge memory footprint, easy to implement. Works for mosthigh-level languages.
Source transformation can lead to more efficient derivative codes,
Operator overloading is often easier to use, better language support.
Jan Hückelheim

Many-core adjoints

10Source transformation example• Each instruction is augmented by its derivative instruction• Variables are augmented by derivative variables• Data flow reversal: r receives from a and b, rb sends to ab and bb.
float f_d(float a, float ad, float b, float bd, float *f) {*f = a*b;
return ad*b + a*bd;}

forward mode
float f(float a, float b) {return a*b;}

reverse mode

void f_b(float a, float *ab, float b, float *bb, float fb) {float f;*ab = *ab + b*fb;*bb = *bb + a*fb;}
Jan Hückelheim

Many-core adjoints

11Why do we need AD for parallel code?
• We can’t wait for faster processors.

Image from https://en.wikipedia.org/wiki/File:Clock CPU Scaling.jpg
See also: Andrew Danowitz et.al., Recording Microprocessor History, Communications of theACM, Vol. 55 No. 4, Pages 55-63 10.1145/2133806.2133822
Jan Hückelheim

Many-core adjoints

12Parallelism has many dimensions• More compute nodes (each node with its own memory and processor)• More cores (each processor can do several unrelated things at once)• Vectors (each core can apply the same operation to multiple values)

Each of these lends itself to different programming models:• Message-passing (e.g. MPI)
• Shared-memory parallelism (Pthreads, OpenMP, OpenACC)• SIMD/SIMT vectorisation (intel intrinsics, OpenMP, CUDA, OpenCL)

There are also performance portability frameworks.

What can AD do?
• Best case: AD always generates efficient parallel codes (unrealistic)• Second-best case: AD generates efficient parallel codes if the input

was well parallelised (realistic?)
Jan Hückelheim

Many-core adjoints

13

AD for MPI• If the original code sends, the adjoint code must receive• If the original code receives, the adjoint code must send
• Remaining problems with non-blocking communication and othersubtleties• Adjoint MPI: libraries are available, and used in practiceeasy adjoints for blocking calls

forward

P1

RECV(c)

RECV(b)b=d;P2
SEND(d)a=a+tRECV(t)P1
b=0
SEND(b)c=0
SEND(c)adjoint

SEND(a)c=a;a=a+c; c=0;

P2
d=d+t
RECV(t)d=d+b; b=0;

Graphic: J. Utke, Adjoints of MPI programs, ECCO2 meeting slides, Argonne National
Laboratory, 2008
Jan Hückelheim

Many-core adjoints

14

Adjoint MPI: Some references• P. Hovland, Automatic differentiation of parallel programs, PhD thesis,

1997• J. Utke et al, Toward adjoinable MPI, IPDPS, 2009
• AdjointMPI, AMPI, with more references:https://www.stce.rwth-aachen.de/research/software/ampi• AdjoinableMPI, also with more references:https://trac.mcs.anl.gov/projects/AdjoinableMPI

What can AD do?
• AD can generally handle this well enough for practical use.

Jan Hückelheim

Many-core adjoints

15

The brutal way to adjoint MPI• In practice, AD tool support is often not necessary• Hand-differentiate the MPI layer, and apply AD only to some kernel

forward

P1

RECV(c)

RECV(b)b=d;P2
SEND(d)manuala=a+t
RECV(t)AD P1
manualb=0
SEND(b)c=0
SEND(c)adjoint

SEND(a)c=a;a=a+c; c=0;

P2
d=d+t
RECV(t)d=d+b; b=0;

• Just make sure that P1 and P2 don’t contain communication calls(”grep -ri MPI” is your friend)

Jan Hückelheim

Many-core adjoints

16

AD for multi-core/many-core/SIMD• Most processors today have multiple cores• Examples:• Intel Core i5, between 2 and 6 cores• Intel Xeon Platinum, up to 28 cores• Intel XeonPhi, up to 68 cores• Raspberry Pi: 4 core ARM Cortex-A53• iPhone X: 6 cores (4+2 different cores)• If we aren’t using the cores, we are wasting resources.
• If the original code is using all cores, the generated adjoint codeshould also use them!

Jan Hückelheim

Many-core adjoints

17

Shared-memory parallelism• Multiple threads run in parallel (e.g. on multi-core CPU)
• Memory visible to all threads, no explicit communication• Parallel read-access is fine, parallel write access is a problem
S

S
Thread 1

Thread 2

Thread 1

Thread 2

P

P

P

P

• Avoid parallel write access(if necessary, use atomic updates, critical sections or barriers)

Jan Hückelheim

Many-core adjoints

18

Reverse AD and OpenMP - the challenge• Situation: primal code is parallelised with OpenMP.
• Source-transformation used to generate adjoint code.
• AD support for OpenMP, Pthreads, CUDA, OpenCL etc is poor.
• Can we use the brutal method that worked with MPI?

parallel forP
end

parallel forP
end?
pthread_create(P1)pthread_create(P2)Jan Hückelheim

Many-core adjoints

pthread_create(P1)pthread_create(P2)19Example: a convolution• Let’s apply a filter to layer k, resulting in layer k + 1

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

20

Example: a convolution• We could do this in parallel, with two threads

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

21

Example: a convolution• Each thread writes to its own output index, no problem

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

22

Example: a convolution• What about the back-propagation?

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

23

Example: a convolution• Each thread reads from its own index...

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

24

Example: a convolution• And scatters the result to overlapping memory regions. Conflict!

Layer kweights

Layer k+1

Jan Hückelheim

Many-core adjoints

25

Why did this happen?
• Overlapping write access to ū happens if there was overlapping readaccess from u in primal.
• We can only easily parallelise adjoint if primal had exclusive read access• Reference for this: M Förster, Algorithmic Differentiation of

Pragma-Defined Parallel Regions: Differentiating Computer Programs
Containing OpenMP, PhD thesis, 2014

Jan Hückelheim

Many-core adjoints

26

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 1
real, dimension(10) :: b,c!$omp parallel dodo i=1,10b(i) = sin(c(i))
end doJan Hückelheim

Many-core adjoints

27

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 1
real, dimension(10) :: b,c!$omp parallel dodo i=1,10b(i) = sin(c(i))
end do• Answer: YescLoop 1b

Jan Hückelheim

Many-core adjoints

28

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 2:
real :: areal, dimension(10) :: b,c!$omp parallel dodo i=1,10b(i) = a+c(i)end doJan Hückelheim

Many-core adjoints

29

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 2:
real :: areal, dimension(10) :: b,c!$omp parallel dodo i=1,10b(i) = a+c(i)end do• Answer: NoacLoop 2b
Jan Hückelheim

Many-core adjoints

30

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 3:
size = read_from_command_line(1)!$omp parallel dodo i=1+size,10-sizeb(i) = 0
do j=i-size,i+size
b(i) = b(i) + c(j)end doend doJan Hückelheim

Many-core adjoints

31

Exclusive read access examples• Do these loops have exclusive read access?

! Example loop 3:
size = read_from_command_line(1)!$omp parallel dodo i=1+size,10-sizeb(i) = 0
do j=i-size,i+size
b(i) = b(i) + c(j)end doend do• Answer: Depends on size, unknown at compile timec?

Loop 3b
Jan Hückelheim

Many-core adjoints

32

The problem with exclusive read
• Any use of a global memory can become a problem• Exclusive read is undecidable in general.
• Can’t just use grep to find it.
• Are there heuristics? Maybe. (One example shown later, but mostly

unexplored question)• Can we rely on users giving pragmas?
• Can we generate several versions (efficient version, safe fallback) anddecide at runtime

Jan Hückelheim

Many-core adjoints

33

What if there’s no exclusive read?
• Or: what if we are not sure?
• Use atomic updates (potentially slow)• Atomic updates are acceptable if the computation is otherwise

expensive enough to hide the overhead of few atomic updates• Use OpenMP reduction (Taf does this)
a_adj

a_adj
thread 1

a_adj
thread 2

Loop 2b_adj

Jan Hückelheim

Many-core adjoints

34

Reduction memory footprint• Depending on OpenMP implementation, reduction may require

temporary private copy on every thread• What if the array is large, and we have dozens/hundreds of threads?
c_adj

c_adj
thread 2c_adj
thread 1?

Loop 3b_adj

Jan Hückelheim

Many-core adjoints

35Summary so far:• Primal parallelism does not imply adjoint parallelism (*)• ”Exclusive read” is a sufficient condition for (*) to hold.
• Exclusive read is impossible to detect in general.
• Can we detect it in practice?
• What if it doesn’t hold?

Jan Hückelheim

Many-core adjoints

36Detection of exclusive read
• Static control flow, indices affine functions of loop counter? Maybe.
• Indirections, non-affine indexing, pointer arithmetic, dependence onuser input? Maybe not.
• Special case where it works for complicated indexing withruntime-dependent indirections: set of read indices identical to set ofwrite indices. In this case, exclusive read property follows from correct
parallelisation of the primal (see our paper, Reverse-mode algorithmic
differentiation of an OpenMP-parallel compressible flow solver, 2017)

Jan Hückelheim

Many-core adjoints

37

What if exclusive read doesn’t hold?
• Traditional adjoint is not parallel.
• Can we do something non-traditional?
• TF-MAD: transposed forward-mode algorithmic differentiation,

combines forward and reverse mode to compute adjoints using theoriginal communication pattern.
• Idea: Split code up into segments where each segment writes to only

one index. The redistribute these segments so that everything thatwrites to the same index is collected in the same iteration.

Jan Hückelheim

Many-core adjoints

38

Forward stencil• Each point in the k layer influences 9 points in the k + 1 layer• Right weight goes to left neighbour, left weight goes to right neighbour

Jan Hückelheim

Many-core adjoints

39

Backward stencil• Each point in k layer receives from 9 points in k + 1 layer• Again, right neighbour goes through left weight, and vice versa

Jan Hückelheim

Many-core adjoints

40

TF-MAD• Implement this more efficiently:• Flip filter horizontally and vertically, switch from push to pull

Jan Hückelheim

Many-core adjoints

41

Parallelising the adjoint, step 1: Look at the primal• A stencil code that pulls data from neighbours to update some value.
• Outer loop: parallel loop over all nodes i.
• Inner loop: sequential loop, reading from all neighbours of i andupdating i (write/increment denoted by ↑).
• On the right: Small example mesh, we’ll come back to this.

2

1
4

Primal code
Jan Hückelheim

Many-core adjoints

42

3Parallelising the adjoint, step 2: Sequential adjoint• Primal outer loop is parallel, adjoint outer loop is not.
• Reason: Every inner iteration writes to ūj (some neighbour), andmaybe some other thread is writing to this at the same time.

Primal code

Jan Hückelheim

Sequential Adjoint codeMany-core adjoints

43

Parallelising the adjoint, step 3: Segmented adjoint• Loop body is split into two segments, each writes to only one index.
• Relis on multi-activity differentiation, see our paper, Algorithmic

Differentiation of Code with Multiple Context-Specific Activities, ACM
TOMS, 2017

Sequential Adjoint code

Jan Hückelheim

Many-core adjoints

Segmented Adjoint code44Parallelising the adjoint, step 3: Redistributed adjoint• “Transpose the off-diagonal term”• Why does this work? See next slides.

Segmented Adjoint code

Jan Hückelheim

Many-core adjoints

Redistributed Parallel Adjoint

45

Illustration: Standard adjoint• Every outer iteration writes almost everywhere ¯   2,1∂f11,2~¯r2(∂f1 + ∂f14,1 )~¯r1~u1
  1,23,24,2
~u¯2  
∂f22,1~¯r1 (∂f2 + ∂f2 + ∂f2 )~¯r2 
 =
+




3,2
~u¯3  0
 
∂f3 ~¯r2~u¯4∂f44,1~¯r1∂f44,2~¯r2

 

0
∂f11,4~¯r4 ∂f22,3~¯r3∂f22,4~¯r4 +  2,3
+
4,3
(∂f3 + ∂f3 )~¯r3  ∂f33,4~¯r4
4,3¯
1,4
2,4
3,4 ¯
∂f4 ~r3(∂f4 + ∂f4 + ∂f4 )~r4

Jan Hückelheim

Many-core adjoints

46

Illustration: Reorganised adjoint• Every outer iteration writes only to one indexi  h 2,1
~u¯1 = (∂f1 + ∂f14,1 )~¯r1 + ∂f11,2~¯r2 + ∂f11,4~¯r4i  h 2,1
~u¯2 = ∂f2 ~¯r1 + (∂f21,2 + ∂f23,2 + ∂f24,2 )~¯r2 + ∂f22,3~¯r3 + ∂f22,4~¯r4i  h 3,2~u¯3 = ∂f3 ~¯r2 + (∂f32,3 + ∂f34,3 )~¯r3 + ∂f33,4~¯r4i  h 4,1
~u¯4 = ∂f4 ~¯r1 + ∂f44,2~¯r2 + ∂f44,3~¯r3 + (∂f41,4 + ∂f42,4 + ∂f43,4 )~¯r4

Jan Hückelheim

Many-core adjoints

47

Speed of reorganised adjoint code (16 CPU threads)• Reorganisation slows down serial code, but scales better• Note: Serial times need recompilation without OpenMP

3

2.81

2.68

2.22
2
×2.44
0.93

ar

d

0.91

w
ck
ba

rw

ar

d

ic
om
at

ta

ng

en

al
im

Many-core adjoints

0.91

fo

0.7

pr


4.16

3.54

1

Jan Hückelheim

4.21

t

program runtime [s]4serialparallel48

Speed of reorganised adjoint code (240 MIC threads)• Overhead of atomics larger on many-core machine. Method pays off inthis example.

program runtime [s]

70
50

66.06
43.6

36.58

33.85
30

6

6

6

6

6

0.83

0.79

6

4.41
4
×5.31

Many-core adjoints

d
w
ar
ck

d

ic

ar
w
fo
r

om
at

ta

ng

en

t

0.37
pr
im
al



0.61

ba

2
0

Jan Hückelheim

78.15

serialparallel49

The good, the bad, the ugly
• The adjoint of a stencil computation can be a stencil computation.

The adjoint of GEMM can be GEMM (and look like one to thecompiler). This allows many compiler optimisations: blocking,polyhedral compilation, auto-vectorisation, ...
• Example: Polly (LLVM) can detect code that looks like GEMM,

achieve speedups of up to 9X• Both approaches shown here require certain symmetry conditions.

That rules out some ways of handling boundaries. Boundaries need tobe e.g. factored out into separate code part.
• Both approaches require high-level, user-given knowledge (e.g. throughpragmas)• None of this is available yet in an AD tool.
Our paper on this: Parallelisable adjoint stencil computations using transposed
forward-mode algorithmic differentiation, in review
Jan Hückelheim

Many-core adjoints

50Future work needed• New programming models emerge faster than AD can catch up• Adjoint MPI: Two decades of research.
• Adjoint OpenMP: Two PhD theses so far.
• Needed: Discussion with users to get priorities right. There are more

hard problems than we can solve.

Jan Hückelheim

Many-core adjoints

51

Thank you

Questions?

Jan Hückelheim

Many-core adjoints

52