Some highlights on Source-to-Source Adjoint AD
Laurent HascoëtUniversité Côte d’Azur, INRIA, France

NIPS 2017, Long Beach, California, USA.
Workshop:
Future of gradient-based ML software & techniques

Hascoët (INRIA)

ST-AD Highlights

NIPS 2017

1 / 27

Motivation
Adjoint derivatives by Algorithmic Differentiation (AD):compute gradients of numerical models,from the models source program,
more or less automatically,
at a cost independant of #inputs,...but facing serious challenges.
ML Back Propagation shares issues with adjoint AD.
⇒ Can we share a few solutions ?
Hascoët (INRIA)

ST-AD Highlights

NIPS 20172 / 27

Challenges of adjoint AD
Gradients are propagated backwards,using info from the (forward) primal code
⇒ Instruction flow reversal⇒ Data flow reversal
For the record, there are other challenges:
non-smoothness [Griewank et al.]stochastic or chaotic parts [Wang]
higher derivatives (cost, size...) [Walther, Wang, Pothen]...
Hascoët (INRIA)

ST-AD Highlights

NIPS 20173 / 27

AD models
Different AD models (⇒ different AD tools) explore
different strategies:
Instruction flow reversal⇒ either by storing runtime trace⇒ or by writing a new source
Data-Flow reversal⇒ either by storing values or partials⇒ or by recomputing them
Our directions:write a new source (Source-to-Source),
store intermediate values (Store-All) upon value kill
Hascoët (INRIA)

ST-AD Highlights

NIPS 20174 / 27

Source-to-Source adjoint AD, Store-Allx
x=x

op1

z=

op2

y=

x=

yx

x

y

op1

x+=
x

*

y

*

x+=

op2z
x=y+=

Hascoët (INRIA)

ST-AD Highlights

NIPS 2017

5 / 27

Source-to-Source adjoint AD, Store-Allx
x=x

op1

z=

op2

y=

x=yx

xy

op1

x+=x

*y*

x+=

op2z
x=y+=

Hascoët (INRIA)

ST-AD Highlights

NIPS 2017

5 / 27

The memory challenge
We are happy not to store the runtime instruction trace,− but we still need to store the intermediate values,− at a memory cost that grows linearly with runtime.
Can we master memory consumption ?
use every possible Data-Flow analysis→ can gain 40 to 70%... still linear memory cost
trade recomputation/storage (“Checkpointing”)→ achieves logarithmic growth
exploit profitable situations, (math or algorithm) e.g.
Linear solvers
Parallel loops
Fixed-Point iterations
Hascoët (INRIA)

ST-AD Highlights

NIPS 20176 / 27

Outline

1

Data-Flow Analysis

2

Checkpointing

3

Profitable Situations

Hascoët (INRIA)

ST-AD Highlights

NIPS 20177 / 27

Data-Flow Analysis
Naı̈ve application of the adjoint AD model would
execute all primal instructions
store every value before it is overwritten
execute the complete adjoint of each instruction
Forward constant propagation & backward slicing,
specialized for the particular structure of adjoint codes
Use static data-flow analysis (classic + and −),on the primal code,
then produce an optimized adjoint code
Hascoët (INRIA)

ST-AD Highlights

NIPS 20178 / 27

4 classic AD Data-Flow analyses
varied:[Fagan, Carle]
if current v depends on no “independent input”, then v is useless⇒ slice out computation of v
useful:if current v influences no “dependent output”, then v is zero⇒ propagate constant v and remove its initializationdiff-live:
if current v influences no useful derivative (may influence orig. result)⇒ slice out computation of v
TBR:[Naumann]
if current v not used in any derivative (e.g. only linear uses of v)⇒ slice out storage of v before it is overwritten
Hascoët (INRIA)

ST-AD Highlights

NIPS 20179 / 27

Discussion
These are just special cases of classic code optim.
Agressive compiler optim [Pearlmutter, Siskind] may be
more systematic (⇒ are we missing adjoint data-flow analyses?)... but there’s a limit to the window of code that thecompiler can examine, whereas fwd and bwd code arearbitrarily far apart
Adjoint data-flow analyses use structural knowledgeof adjoint codes, and run on the primal code. E.g.
TBR+ (I ) =Hascoët (INRIA)

(TBR− (I ) ∪ use(I 0 )) \ kill(I )TBR− (I ) ∪ use(I 0 )ST-AD Highlights

if I live
otherwise

NIPS 201710 / 27

Summary: good, but not sufficient
Adjoint data-flow analyses
are classical compiler analyses/optims specialized foradjoint codes.
bring substantial benefit
20% to 50% in runtime
40% to 70% in memory space

But memory still grows linearly with runtime⇒ we need something else.

Hascoët (INRIA)

ST-AD Highlights

NIPS 201711 / 27

Outline

1

Data-Flow Analysis

2

Checkpointing

3

Profitable Situations

Hascoët (INRIA)

ST-AD Highlights

NIPS 201712 / 27

Trading recomputation (time) for storage(memory)
Checkpointing: elementary stitch
D{C{{U

reduces peak storage
at the cost of duplicate execution
also costs a memory “Snapshot”, small enough:Snapshot ⊂ use(C) ∩ out(C) ∪ out(D)Hascoët (INRIA)

ST-AD Highlights

NIPS 201713 / 27

Combining Checkpointing and TBR
D{C{{UThe Snapshot may take care of TBR coming from U
The TBR sent to D can take care of the Snapshot
A range of “optimal” combinations exist.
E.g., given tbrU coming from U, “lazy” snapshot:Snapshot = out(C) ∩ use(C) ∪ tbrU
tbr to D = use(C) ∪ tbrU \ out(C)tbr to C = tbrU
Hascoët (INRIA)

ST-AD Highlights

NIPS 201714 / 27

Nesting checkpoints
Checkpoints must be (carefully) nested.
Optimal nesting (binomial) exists for time-stepping loops:
01630

38455156606257 58
52 53 54
46 47 48 49

39 40 41 42 43

31 32 33 34 35 36

1922201736924 25 26 27 28[Griewank, Walther]

10 11 12 13 1474
1

peak memory storage grows like log(runtime)
execution duplication grows like log(runtime)in real life, storage is fixed to q snapshots,
execution duplication grows like qth-root(runtime)
Hascoët (INRIA)

ST-AD Highlights

NIPS 201715 / 27

Checkpointing on calls
Nested checkpointing can be applied on procedure calls:
1111214122

22223151

32

32

32

32

41

41

51

515121213131

31

Sub-optimal(?), but still logarithmic if call tree isbalanced.
Applies also to code sections that could be procedures.
Hascoët (INRIA)

ST-AD Highlights

NIPS 201716 / 27

A few limitations

Checkpoints must respect code structure:
no checkpoint across proceduresno checkpoint across structured statements...well you could, but you need a flattened instruction tape

Checkpoints must contain both ends of system
resources lifespan:
read/write, alloc/free, send/recv, isend/wait. . .

Checkpointed code must be reentrant

Hascoët (INRIA)

ST-AD Highlights

NIPS 201717 / 27

Outline

1

Data-Flow Analysis

2

Checkpointing

3

Profitable Situations

Hascoët (INRIA)

ST-AD Highlights

NIPS 201718 / 27

Profitable Situations

Take advantage of algorithmic or mathematic knowledge
on parts of the code.
Just a selection:
Adjoint of Linear Solvers
Adjoint of Parallel Loops
Adjoint of Fixed-Point iterations

Hascoët (INRIA)

ST-AD Highlights

NIPS 201719 / 27

Adjoint of Linear Solvers
Avoid differentiation inside the source of linear solvers⇒ write their adjoint by hand, calling the solver itself!
SOLVE B(A,Ab,y,yb,b,bb) {At = TRANSPOSE(A)SOLVE(At,tmp,yb)bb[:] = bb[:] + tmp[:]
SOLVE(A,y,b)
for each i and each j {Ab[i,j] = Ab[i,j] - y[j]*tmp[i]}yb[:] = 0.0}
Hascoët (INRIA)

ST-AD Highlights[Giles]NIPS 201720 / 27

Data-Dependence Graph of Adjoints
Data-Dependence Graph is key to loop rescheduling.
Fewer arrows in the DDG ⇒ more rescheduling allowed.
(classical) No DDG arrow between successive reads of a variable.
No DDG arrow either between successive increments of a variable.
(assuming increments are atomic, or that memory is not shared)
The adjoint of a read(x) is an increment(x)The adjoint of an increment(x) is a read(x)The DDG of the backward sweep is a subset of the DDGof the primal code, only with arrows reversed
Therefore adjoint AD preserves most parallel properties!
Hascoët (INRIA)

ST-AD Highlights

NIPS 201721 / 27

Application to Parallel Loops// Parallel loop:for (i=0 ; i<=N ; ++i) {forward sweep iteration i}for (i=N ; i>=0 ; --i) {backward sweep iteration i}

Loop #2 is parallel: reverse iterations, fuse with loop #1:
for (i=0 ; i<=N ; ++i) {forward sweep iteration ibackward sweep iteration i}⇒ Reduces peak memory usage dramatically!
Hascoët (INRIA)

ST-AD Highlights

NIPS 201722 / 27

Adjoint of Fixed-Point iterations
until z converges:z = φ(z, x)

...

end

as many times:
x = x + z ∂φ/∂x ; z = z ∂φ/∂z

...

end

You should not do that!
all values from intermediate iterations are stored
poor convergence guarantees of the adjoint sweep
Hascoët (INRIA)

ST-AD Highlights

NIPS 201723 / 27

Two-Phases Adjoint
until z converges:z = φ(z, x)z = φ(z, x)

...
end

t=z
until z converges:z = z ∂φ/∂z + t[Christianson]

...

end
x = x + z ∂φ/∂x

Only the converged primal iteration is stored,
then is used several times.
The adjoint iteration has its own convergence control
Converges in one step if primal has quadratic
convergence
Hascoët (INRIA)

ST-AD Highlights

NIPS 201724 / 27

Loosing Warm-Start
Suppose Fixed-Point is included in another iteration:
FP iterations: 16, 16, 16, 16, . . .

Warm-Start uses previous converged z as next initial z.
⇒ convergence is reached earlier
FP iterations: 16, 9, 9, 9, . . .

Standard adjoint “inherits” this Warm-Start effectbut the Two-Phases adjoint doesn’t!
adjoint FP iterations: 46, 46, 46, 46, . . .

Hascoët (INRIA)

ST-AD Highlights

NIPS 201725 / 27

Retrieving Warm-Start?
Two-Phases adjoint modifies z between two adjoint Fixed-Points.
Notice that z has two uses:...
t=z
until z converges {z = z ∂φ/∂z + t}x = x + z ∂φ/∂xmodifications (reset!) on z...
(1) to set t, that influences converged value ⇒ don’t touch this!
(2) as the initial guess of the adjoint Fixed-Point ⇒ feel free to set it topreviously converged!
adjoint FP iterations: 46, 20, 21, 20, . . .
Is that correct? Can we automate it?
Hascoët (INRIA)

ST-AD Highlights

NIPS 201726 / 27

Thank you for your attention!

Hascoët (INRIA)

ST-AD Highlights

NIPS 201727 / 27