\textbf{D}^{3}_{c}[\tau]\rightarrow(\tau,\tau\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c))\\\
\mathrlap{\text{(Same as $\textrm{Deinterleave}^{2}$ in
\lx@cref{creftype~refnum}{fig:algo-monadic}, except with $\mathrm{id}$ and
$(\circ)$ in place of $0_{\text{Staged}}$ and
$(\mathbin{+_{\text{Staged}}})$)}}\end{array}$
$\displaystyle\begin{array}[]{@{}l}\textrm{Wrap}^{3}:(\sigma\rightarrow\tau)\rightsquigarrow(\sigma\rightarrow(\tau,\tau\multimap\sigma))\\\
\textrm{Wrap}^{3}[\lambda(x:\sigma).\ t]=\lambda(x:\sigma).\
\begin{array}[t]{@{}l@{}}\textbf{let}\
(x:\textbf{D}^{3}_{\sigma}[\sigma],i)=\textrm{Interleave}^{3}_{\sigma}\
(x,\textrm{SCotan})\ 0\\\ \textbf{in}\ \textbf{let}\
(y,d)=\textrm{Deinterleave}^{3}_{\tau}\
(\textrm{fst}(\textbf{D}^{3}_{\sigma}[t]\ i))\\\ \textbf{in}\
(y,\underline{\lambda}(z:\tau).\ \textrm{SResolve}\ (d\
z))\end{array}\end{array}$

Figure 9:  The Cayley-transformed code transformation, based on Fig. 8. Grey
parts are unchanged.

We show new code transformation in Fig. 9. Aside from the changes to types and
to the target monoid of the backpropagators, the only additional change is in
$\textrm{Interleave}^{3}$, which is adapted to accomodate the additional
Cayley transform on the $c$ argument of SCotan. Note that the backpropagators
in $\textrm{Interleave}^{3}$ do not create any $\underline{0}$ values for
untouched parts of the collected cotangent of type $c$, as promised, and that
the new type of SCotan has indeed eliminated all uses of $(+)$ on $c$, not
just moved them around.

###  7.2 Remaining Complexity Challenges

In Section 6.1, we pinpointed the three remaining complexity issues with the
reverse AD algorithm after function call staging: costly monoid operations on
$c$, costly one-hot backpropagators from Interleave, and logarithmic Map
operations in SResolve. The first issue has been solved by Cayley-transforming
$\operatorname*{\textrm{Staged}}c$: only $\underline{0}:c$ is still used, and
that only once (in the new SResolve). For the second issue, although
performance of the one-hot backpropagators has improved in most cases, it is
still unsatisfactory; for example, given the input type
$\sigma=((\mathbb{R},\mathbb{Z}),(\mathbb{R},\mathbb{R}))$, the backpropagator
for the second scalar looks as follows before and after the Cayley transform:

Before Cayley | After Cayley
---|---
$\underline{\lambda}(d:\underline{\mathbb{R}}).\ \textrm{SCotan}\ \biggl{(}$$\biggr{)}$$(,)$$(,)$$(,)$$0$$0$$d$$d$$0$ | $\underline{\lambda}(d:\underline{\mathbb{R}}).\ \textrm{SCotan}\ \biggl{(}$$\biggr{)}$$(,)$$\vphantom{()}\mathrm{id}$$(,)$$(+d)$$\mathrm{id}$

This yields complexity logarithmic in the size of the input if the input is
balanced, but can degrade to linear in the size of the input in the worst
case—which is no better than the previous version. We will make these
backpropagators properly logarithmic in the size of the input in Section 8,
essentially reducing issue two to issue three. Afterwards, in Section 9, we
finally properly resolve issue three (logarithmic overhead from Map
operations) by introducing mutable arrays. Doing so removes the final log-
factors from the algorithm’s complexity.

##  8 Keeping Just the Scalars: Efficient Gradient Updates

The codomain of the backpropagators is currently
$\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c$,
where $\operatorname*{\textrm{Staged}}c$ is defined as:

$\operatorname*{\textrm{Staged}}c=(c,\textrm{Map}\ \mathbb{Z}\
(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c),\underline{\mathbb{R}}))$

The final cotangent of the input to the program is collected in the first
component of the pair, of type $c$. This collector is initialised with
$\underline{0}:c$ in $0_{\text{Staged}}$, and added to by the one-hot input
backpropagators from Interleave, called in SResolve. The task of these input
backpropagators is to add the cotangent (of type $\underline{\mathbb{R}}$)
that they receive in their argument, to a particular scalar in the collector.

Hence, all we need of $c$ in $\operatorname*{\textrm{Staged}}c$ is really the
collection of its scalars: the rest simply came from $\underline{0}:c$ and is
never changed.191919If $c$ contains coproducts (sum types) as we discuss in
Section 11, this $\underline{0}:c$ becomes dependent on the actual input to
the program, copying the structure from there. Furthermore, the reason why the
one-hot input backpropagators currently do not finish in logarithmic time is
that $c$ may not be a balanced tree of its scalars. But if we are interested
only in the scalars anyway, we can _make_ the collector balanced—by replacing
it with $\textrm{Map}\ \mathbb{Z}\ \underline{\mathbb{R}}$:

$\operatorname*{\textrm{Staged}}c=({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\textrm{Map}\
\mathbb{Z}\ \underline{\mathbb{R}}},\textrm{Map}\ \mathbb{Z}\
(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c),\underline{\mathbb{R}}))$

To accomodate this, Interleave changes to number all the scalars in the input
with distinct IDs (conveniently with the same IDs as their corresponding input
backpropagators, but this is no fundamental requirement); the cotangent of the
input scalar with ID $i$ is stored in the Map at key $i$. The input
backpropagators can then modify the correct scalar in the collector (now of
type $\textrm{Map}\ \mathbb{Z}\ \underline{\mathbb{R}}$) in time logarithmic
in the size of the input. To be able to construct the final gradient from this
collection of just its scalars, $\textrm{Interleave}_{\tau}$ additionally
builds a _reconstruction_ function of type
$(\mathbb{Z}\rightarrow\underline{\mathbb{R}})\rightarrow\tau$, which we pass
a function that looks up the ID in the final collector Map to compute the
actual gradient value.

#### Complexity.

Now that we have fixed (in Section 7) the first complexity problem identified
in Section 6.1 (expensive monoid operations) and reduced the second (expensive
input backpropagators) to a logarithmic overhead over the original program, we
have reached the point where we satisfy the complexity requirement stated in
Sections 3 and 6.1 apart from log-factors. More precisely, $\textrm{Wrap}[P]$
computes the gradient of $P$ at $x$ not in time $O(\textrm{cost}(P\
x)+\textrm{size}(x))$ but in time $O((\textrm{cost}(P\
x)+\textrm{size}(x))\log(\textrm{cost}(P\ x)+\textrm{size}(x)))$. If we are
okay with logarithmic overhead, we could choose to stop here. However, if we
wish to strictly conform to the required complexity, or if we desire to lose
the non-negligible constant-factor overhead of dealing with a persistent tree-
map, we need to make the input backpropagators and Map operations in SResolve
constant-time; we do this using mutable arrays in Section 9.

##  9 Using Mutable Arrays to Shave Off Log Factors

The analysis in Section 6.1 showed that after the Cayley transform in Section
7, the strict complexity requirements are met if we make the input
backpropagators constant-time and make SResolve have only constant overhead
for each backpropagator that it calls. Luckily, in both cases the only
component that is not constant-time is the interaction with one of the Maps in
$\operatorname*{\textrm{Staged}}c$:

$\operatorname*{\textrm{Staged}}c=(\textrm{Map}\ \mathbb{Z}\
\underline{\mathbb{R}},\textrm{Map}\ \mathbb{Z}\
(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c),\underline{\mathbb{R}}))$

The input backpropagators perform (logarithmic-time) updates to the first Map
(the cotangent collector), and SResolve reads, deletes and updates entries in
the second Map (the staging map for recording delayed backpropagator calls).
Both of these Maps are keyed by increasing, consecutive integers starting from
0, and are thus ideal candidates to be replaced by an array:

$\operatorname*{\textrm{Staged}}c=(\mathrlap{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname*{\textrm{Array}}}}\hphantom{\textrm{Map}\
\mathbb{Z}}\
\underline{\mathbb{R}},\mathrlap{{\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname*{\textrm{Array}}}}\hphantom{\textrm{Map}\
\mathbb{Z}}\
(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c),\underline{\mathbb{R}}))$

To allocate an array, one must know how large it should be. Fortunately, at
the time when we allocate the initial $\operatorname*{\textrm{Staged}}c$ value
using $0_{\text{Staged}}$ in SResolve, the primal pass has already been
executed and we know (from the output ID of Interleave) how many input scalars
there are, and (from the output ID of the transformed program) how many
backpropagators there are. Hence, the size of these arrays is indeed known
when they are allocated; and while these arrays are large, the resulting space
complexity is equal to the worst case for reverse AD in general.

To get any complexity improvements from replacing a Map with an
$\operatorname*{\textrm{Array}}$ (indeed, to not pessimise the algorithm
completely!), the write operations to the arrays need to be done _mutably_.
These write operations occur in two places: in the updater functions (of type
$\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c$)
produced by backpropagators, and in SResolve. Hence, in these two places we
need an effectful function type that allows us to encapsulate the mutability
and ensure that it is not visible from the outside; options include a
resource-linear function type and a monad for local side-effects such as the
ST monad in Haskell. See Sections 9.1 and 9.2 below for more information about
how this works out.

#### Time complexity.

We now satisfy all the requirements of the analysis in Section 6.1, and hence
have the correct time complexity for reverse AD. In particular, let $I$ denote
the size of the input and $T$ the runtime of the original program. Let
$\textbf{D}^{4}$, $\textrm{Interleave}^{4}$, $\textrm{Deinterleave}^{4}$, etc.
be the definitions that use arrays as described above (see Section 9.2 for
details). Then we can observe the following:

  * •

The number of operations performed by $\textbf{D}^{4}_{c}[t]$ (with the
improvements from Sections 8 and 9) is only a constant factor times the number
of operations performed by $t$, and hence in $O(T)$. This was already observed
for $\textbf{D}^{2}_{c}[t]$ in Section 6.1, and still holds.

  * •

The number of backpropagators created while executing $\textbf{D}^{4}_{c}[t]$
is clearly also in $O(T)$.

  * •

The number of operations performed in any one backpropagator is constant. This
is new, and only true because $\mathrm{id}$ (replacing $0_{\text{Staged}}$),
$(\circ)$ (replacing $(\mathbin{+_{\text{Staged}}})$), SCotan (with a
constant-time mutable array updater as argument) and SCall are now all
constant-time.

  * •

Hence, because every backpropagator is invoked at most once thanks to our
staging, and because the overhead of SResolve is constant per invoked
backpropagator, the amount of work performed by calling the top-level input
backpropagator is again in $O(T)$.

  * •

Finally, the (non-constant-time) extra work performed in $\textrm{Wrap}^{4}$
is interleaving ($O(I)$), deinterleaving ($O(\text{size of output})$ and hence
$O(T+I)$), resolving ($O(T)$) and reconstructing the gradient from the scalars
in the $\operatorname*{\textrm{Array}}\underline{\mathbb{R}}$ in
$\operatorname*{\textrm{Staged}}c$ ($O(I)$); all this work is in $O(T+I)$.

Hence, calling $\textrm{Wrap}^{4}[t]$ with an argument and calling its
returned top-level derivative once takes time $O(T+I)$, i.e. at most
proportional to the runtime of calling $t$ with the same argument, plus the
size of the argument itself. This is indeed the correct time complexity for an
efficient reverse AD algorithm, as discussed in Section 3.

#### Space complexity.

The sizes of the allocated arrays are bounded by the total number of IDs that
we have generated, which equals the sum of the number of input scalars and the
number of primitive operations executed at runtime in the original
computation. This space complexity is unfortunate, but in the general case one
cannot do better; thus, absent optimisations for particular cases, all
programs get this space complexity.

###  9.1 Implementation Using Resource-Linear Types

For our algorithm, arguably the most elegant embedding of mutable arrays in a
purely functional setting uses resource-linear types. This subsection briefly
explains how this works; in Sections 9.2 and 12 we switch to a monadic
embedding to allow parallelism.

Resource-linear types capture the idea of passing around a (reference to an)
object in a _linear_ fashion: that is to say, without ever duplicating or
dropping it, as would be freely allowed in a normal type system. There are
various designs for including resource-linear types in a programming language;
broadly, there is a distinction between type systems based on _uniqueness
types_ (where the act of duplicating a “unique” value is disallowed without
(possibly implicitly) casting it to an immutable value first) and those based
on linear function types (where a “linear function” promises to handle its
input value linearly, but does not directly get any guarantees on the
uniqueness of its argument).

Using such a resource-linear type system, it is possible to write a safe
interface for a mutable array library: despite doing mutation, referential
transparency of code using the library remains intact. Let us take Linear
Haskell (Bernardy et al., 2018) with the mutable array library in linear-
base202020https://hackage.haskell.org/package/linear-base-0.2.0/docs/Data-
Array-Mutable-Linear.html as an example; this system is an extension of
traditional Haskell (available in GHC 9) that follows the second style using
“linear functions”. Write the (resource-)linear function type with
$\mathrel{\mathchoice{\overset{\smash{\raisebox{-4.09024pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-4.00412pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-3.06766pt}{\tiny{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-2.24966pt}{\scalebox{0.7}{\tiny{R}}}}}{\rightarrow}}}$.
The central trick that makes the interface of mutable arrays in Linear Haskell
work is the allocation function, which is higher-order:

$\mathrm{alloc}:(\textit{length}:\mathbb{Z})\rightarrow(\textit{initval}:\tau)\rightarrow(\operatorname*{\textrm{Array}}\tau\mathrel{\mathchoice{\overset{\smash{\raisebox{-4.09024pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-4.00412pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-3.06766pt}{\tiny{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-2.24966pt}{\scalebox{0.7}{\tiny{R}}}}}{\rightarrow}}}{!}\sigma)\mathrel{\mathchoice{\overset{\smash{\raisebox{-4.09024pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-4.00412pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-3.06766pt}{\tiny{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-2.24966pt}{\scalebox{0.7}{\tiny{R}}}}}{\rightarrow}}}{!}\sigma\\\
$

The ‘!’ syntax indicates an _unrestricted_ type (corresponding to the “of
course” modality of linear logic), which may be freely duplicated and dropped
even in a linear context, but the construction of which requires the right to
duplicate the value put in the unrestricted-box. The other functions in the
mutable array API, such as ‘$\mathrm{set}$’:

$\mathrm{set}:(\textit{index}:\mathbb{Z})\rightarrow(\textit{value}:\tau)\rightarrow\operatorname*{\textrm{Array}}\tau\mathrel{\mathchoice{\overset{\smash{\raisebox{-4.09024pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-4.00412pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-3.06766pt}{\tiny{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-2.24966pt}{\scalebox{0.7}{\tiny{R}}}}}{\rightarrow}}}\operatorname*{\textrm{Array}}\tau$

pass along the array linearly, meaning that the function passed to
‘$\mathrm{alloc}$’ must handle the passed array linearly, and cannot return
the array as part of its output; this allows ‘$\mathrm{alloc}$’ to deallocate
the array upon exit of the passed function. To support the use of other linear
resources within the function passed to ‘$\mathrm{alloc}$’, the final arrow in
$\mathrm{alloc}$’s type indicates that the passed function, and hence its
closure, is consumed linearly (called exactly once).

The effect of using this kind of linear types for our code transformation is
that the
$\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c$
functions become resource-linear:

$\operatorname*{\textrm{Staged}}c=(\operatorname*{\textrm{Array}}\underline{\mathbb{R}},\operatorname*{\textrm{Array}}\;(\underline{\mathbb{R}}\multimap({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname*{\textrm{Staged}}c\mathrel{\mathchoice{\overset{\smash{\raisebox{-4.09024pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-4.00412pt}{\\!\footnotesize{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-3.06766pt}{\tiny{R}}}}{\rightarrow}}{\overset{\smash{\raisebox{-2.24966pt}{\scalebox{0.7}{\tiny{R}}}}}{\rightarrow}}}\operatorname*{\textrm{Staged}}c}),\underline{\mathbb{R}}))$

The compositional code transformation itself remains unchanged except for its
type; updating the $\operatorname*{\textrm{Staged}}$ interface is not
difficult. A detailed exposition of the changes required and the resulting
definitions is given in an earlier preprint for this work (Smeding and Vákár,
2022, App. A).

###  9.2 Implementation Using Mutable References in a Monad

An alternative to resource-linear types, for using mutable arrays in a purely
functional language, is to use mutable references inside some monad. In
Haskell, the typical solution is to use the ST monad (Launchbury and Jones,
1994) together with a mutable array library that exposes an API using ST, such
as STVector in the
vector212121https://hackage.haskell.org/package/vector-0.13.1.0/docs/Data-
Vector-Mutable.html library. Because the ST monad is designed to be
deterministic, it has a pure handler:

$\mathrm{runST}:(\forall s.\ \texttt{ST}\ s\ \alpha)\rightarrow\alpha$

allowing the use of local mutability without compromising referential
transparency of the rest of the program.

However, precisely because of this design, ST does not support parallelism.
(Parallelism in the presence of mutable references trivially allows non-
deterministic behaviour.) For this reason we will write the definitions from
this point on in terms of IO, Haskell’s catch-all monad for impurity.
Fortunately, the only functionality we use from IO is parallelism and mutable
arrays and references, and furthermore the design of our algorithm is such
that the result is, in fact, deterministic even when the source program
includes parallelism.222222This is unfortunately not quite true, because
floating-point arithmetic is not associative or commutative. Given the nature
of the computations involved, however, we still think getting parallelism is
worth this caveat. Thus we can justify using
$\texttt{unsafePerformIO}::\texttt{IO}\ \alpha\rightarrow\alpha$ around the
differentiated program, making the interface to the differentiated program
pure again.

Letting the updater functions run in IO changes
$\operatorname*{\textrm{Staged}}c$ as follows:

$\operatorname*{\textrm{Staged}}c=(\operatorname*{\textrm{Array}}\underline{\mathbb{R}},\operatorname*{\textrm{Array}}{(\underline{\mathbb{R}}\multimap({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()}),\underline{\mathbb{R}})})$

for a suitable definition of $\operatorname*{\textrm{Array}}$, such as
IOVector from vector. To write this definition, we need to give a monoid
structure on $\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\ ()$;
fortunately, the only reasonable one ($f+g=\lambda x.\ f\
x\mathbin{>\\!\\!>}g\ x$) corresponds to the monoid structure on
$\operatorname*{\textrm{Staged}}c$ and is therefore precisely the one we want.

Note that this definition now no longer structurally depends on $c$! This is
to be expected, because the information about the structure of $c$ is now
contained in the _length_ of the first array in a
$\operatorname*{\textrm{Staged}}c$. For uniformity of notation, however, we
will continue to write the $c$ parameter to $\operatorname*{\textrm{Staged}}$.

#### Mutable arrays interface.

We assume an interface to mutable arrays that is similar to that for IOVector
in the Haskell vector library, cited earlier. In summary, we assume the
following functions:

$\begin{array}[]{@{}l@{\ }c@{\
}l@{}}\texttt{alloc}&:&\mathbb{Z}\rightarrow\alpha\rightarrow\texttt{IO}\
(\operatorname*{\textrm{Array}}\alpha)\\\
\texttt{get}&:&\mathbb{Z}\rightarrow\operatorname*{\textrm{Array}}\alpha\rightarrow\texttt{IO}\
\alpha\\\
\texttt{modify}&:&\mathbb{Z}\rightarrow(\alpha\rightarrow\alpha)\rightarrow\operatorname*{\textrm{Array}}\alpha\rightarrow\texttt{IO}\
()\\\
\texttt{freeze}&:&\operatorname*{\textrm{Array}}\alpha\rightarrow\texttt{IO}\
(\operatorname*{\textrm{IArray}}\alpha)\\\
(@)&:&\operatorname*{\textrm{IArray}}\alpha\rightarrow\mathbb{Z}\rightarrow\alpha\end{array}$

where $\operatorname*{\textrm{IArray}}$ is an immutable array type. alloc and
freeze are linear in the length of the array; get, modify and $(@)$ are
constant-time (in addition to calling the function once, of course, for
modify).

Let us see how the code transformation changes with mutable arrays.

#### Code transformation.

Using the new definition of $\operatorname*{\textrm{Staged}}$, we change the
code transformation once more, this time from Fig. 9 to the version given in
Fig. 10. The transformation on types and terms simply sees the type of scalar
backpropagators change to use effectful updating instead of functional
updating, so they do not materially change: we just give yet another
interpretation of $0_{\text{Staged}}$ and $(\mathbin{+_{\text{Staged}}})$,
using the monoid structure on
$\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\ ()$ defined above in
terms of $(\mathbin{>\\!\\!>})$. However, some important changes occur in the
$\operatorname*{\textrm{Staged}}c$ interface and the wrapper. Let us first
look at the algorithm from the top, by starting with $\textrm{Wrap}^{4}$;
after understanding the high-level idea, we explain how the other components
work.

On types:

$\begin{array}[]{l@{\qquad}l}\textbf{D}^{4}_{c}[\mathbb{R}]=(\mathbb{R},(\mathbb{Z},\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())))&{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\textbf{D}^{4}_{c}[\mathbb{Z}]=\mathbb{Z}}\qquad{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\textbf{D}^{4}_{c}[()]=()}\\\
{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\textbf{D}^{4}_{c}[\sigma\rightarrow\tau]=\textbf{D}^{4}_{c}[\sigma]\rightarrow\mathbb{Z}\rightarrow(\textbf{D}^{4}_{c}[\tau],\mathbb{Z})}&{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}\textbf{D}^{4}_{c}[(\sigma,\tau)]=(\textbf{D}^{4}_{c}[\sigma],\textbf{D}^{4}_{c}[\tau])}\\\
\end{array}$

On terms:

$\displaystyle{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>\Gamma\vdash t:\tau\ \text{then}\
\textbf{D}^{4}_{c}[\Gamma]\vdash\textbf{D}^{4}_{c}[t]:\mathbb{Z}\rightarrow(\textbf{D}^{4}_{c}[\tau],\mathbb{Z})}$
$\displaystyle\begin{array}[t]{@{}l@{}}\text{Same as $\textbf{D}^{2}_{c}$,
except with `$\lambda\\_.\ \mathbf{return}\ ()$' in place of
$0_{\text{Staged}}$ and `$\lambda f\ g\ x.\ f\ x\mathbin{>\\!\\!>}g\ x$'
in}\\\ \text{place of $(\mathbin{+_{\text{Staged}}})$.}\end{array}$

New $\operatorname*{\textrm{Staged}}$ interface:

$\displaystyle\begin{array}[]{@{}l@{}}\operatorname*{\textrm{Staged}}c=(\operatorname*{\textrm{Array}}\underline{\mathbb{R}},\operatorname*{\textrm{Array}}{}(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()),\underline{\mathbb{R}}))\\\\[5.69054pt]
\textrm{SAlloc}:\mathbb{Z}\rightarrow\mathbb{Z}\rightarrow\texttt{IO}\
(\operatorname*{\textrm{Staged}}c)\\\ \textrm{SAlloc}\ i_{\text{inp}}\
i_{\text{out}}=\texttt{alloc}\ i_{\text{inp}}\
0\mathbin{>\\!\\!>\\!\\!=}\lambda c.\ \texttt{alloc}\ i_{\text{out}}\
(\underline{\lambda}(z:\underline{\mathbb{R}}).\
\mathrm{id},0)\mathbin{>\\!\\!>\\!\\!=}\lambda m.\ \mathbf{return}\
(c,m)\\\\[5.69054pt]
\textrm{SCall}:(\mathbb{Z},\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()))\rightarrow\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())\\\ \textrm{SCall}\ (i,f)\ a\ (c,m)=\texttt{modify}\ i\
(\lambda(\mathunderscore,a^{\prime}).\ (f,a+a^{\prime}))\ m\\\\[5.69054pt]
\textrm{SOneHot}:\mathbb{Z}\rightarrow\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())\\\ \textrm{SOneHot}\ i\ a\ (c,m)=\texttt{modify}\ i\
(\lambda(a^{\prime}:\mathbb{R}).\ a+a^{\prime})\ c\\\\[5.69054pt]
\textrm{SResolve}:\mathbb{Z}\rightarrow\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
(\operatorname*{\textrm{IArray}}\underline{\mathbb{R}})\\\ \textrm{SResolve}\
i_{\text{out}}\ (c,m)=\textit{loop}\
(i_{\text{out}}-1)\mathbin{>\\!\\!>}\texttt{freeze}\ c\\\ \quad\text{where}\
\begin{array}[t]{@{}l@{\ }c@{\ }l@{}}\textit{loop}\ (-1)&=&\mathbf{return}\
()\\\ \textit{loop}\ i&=&\texttt{get}\ i\
m\mathbin{>\\!\\!>\\!\\!=}\lambda(f,a).\ f\ a\
(c,m)\mathbin{>\\!\\!>}\textit{loop}\ (i-1)\end{array}\end{array}$

Wrapper:

$\displaystyle\begin{array}[]{@{}l@{\;}c@{\;}l@{}}\textrm{Interleave}^{4}_{\tau}&:&\tau\rightarrow\mathbb{Z}\rightarrow((\textbf{D}^{4}_{c}[\tau],\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}\rightarrow\tau),\mathbb{Z})\\\
\textrm{Interleave}^{4}_{\mathbb{R}}&=&\lambda x.\ \lambda i.\
(((x,(i,\textrm{SOneHot}\ i)),\lambda a.\ a\mathbin{@}i),i+1)\\\
\textrm{Interleave}^{4}_{()}&=&\lambda().\ \lambda i.\ (((),\lambda a.\
()),i)\\\ \textrm{Interleave}^{4}_{(\sigma,\tau)}&=&\lambda(x,y).\ \lambda i.\
\begin{array}[t]{@{}l@{}}\textbf{let}\
((x^{\prime},f_{1}),i^{\prime})=\textrm{Interleave}^{4}_{\sigma}\ x\ i\\\
\textbf{in}\ \textbf{let}\
((y^{\prime},f_{2}),i^{\prime\prime})=\textrm{Interleave}^{4}_{\tau}\ y\
i^{\prime}\\\ \textbf{in}\ (((x^{\prime},y^{\prime}),\lambda a.\ (f_{1}\
a,f_{2}\ a)),i^{\prime\prime})\end{array}\\\
\textrm{Interleave}^{4}_{\mathbb{Z}}&=&\lambda n.\ \lambda i.\ ((n,\lambda a.\
n),i)\end{array}$
$\displaystyle\begin{array}[]{@{}l}\textrm{Deinterleave}^{4}_{\tau}:\textbf{D}^{4}_{c}[\tau]\rightarrow(\tau,\tau\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()))\\\ \text{(Same as $\textrm{Deinterleave}^{2}$, except with the same
monoid changes as $\textbf{D}^{4}_{c}$ above)}\end{array}$
$\displaystyle\begin{array}[]{@{}l}\textrm{Wrap}^{4}:(\sigma\rightarrow\tau)\rightsquigarrow(\sigma\rightarrow(\tau,\tau\multimap\sigma))\\\
\textrm{Wrap}^{4}[\lambda(x:\sigma).\ t]=\lambda(x:\sigma).\\\
\qquad\textbf{let}\
((x:\textbf{D}^{4}_{\sigma}[\sigma],\textit{rebuild}:\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}\rightarrow\sigma),i)=\textrm{Interleave}^{4}_{\sigma}\
x\ 0\\\ \qquad\textbf{in}\ \textbf{let}\
(y^{\prime},i^{\prime})=\textbf{D}^{4}_{\sigma}[t]\ i\\\ \qquad\textbf{in}\
\textbf{let}\
(y,d:\tau\multimap(\operatorname*{\textrm{Staged}}\sigma\rightarrow\texttt{IO}\
()))=\textrm{Deinterleave}^{4}_{\tau}\ y^{\prime}\\\ \qquad\textbf{in}\
(y,\underline{\lambda}(z:\tau).\ \textit{rebuild}\
(\texttt{unsafePerformIO}\\\ \hphantom{\qquad\textbf{in}\
(y,\underline{\lambda}(z:\tau).\ \textit{rebuild}\ (}\qquad(\textrm{SAlloc}\
i\ i^{\prime}\mathbin{>\\!\\!>\\!\\!=}\lambda s.\ d\ z\
s\mathbin{>\\!\\!>}\textrm{SResolve}\ i^{\prime}\ s)))\end{array}$

Figure 10:  Code transformation plus wrapper using mutable arrays, modified
from Fig. 9. Grey parts are unchanged.

In basis, $\textrm{Wrap}^{4}$ does the same as $\textrm{Wrap}^{3}$ from Fig.
9: interleave injector backpropagators with the input of type $\sigma$,
execute the transformed function body using the interleaved input, and then
deinterleave the result. However, because we (since Section 8) represent the
final cotangent not directly as a value of type $\sigma$ in a
$\operatorname*{\textrm{Staged}}\sigma$ but instead as an array of only the
embedded scalars ($\operatorname*{\textrm{Array}}\underline{\mathbb{R}}$),
some more work needs to be done.

Firstly, $\textrm{Interleave}^{4}_{\sigma}$ (monadically, in the ID generation
monad that we are writing out explicitly) produces, in addition to the
interleaved input, also a _reconstruction_ function232323Implementing the
$(\mathbb{Z}\rightarrow\underline{\mathbb{R}})$ in
$(\mathbb{Z}\rightarrow\underline{\mathbb{R}})\rightarrow\tau$ from Section 8
as $\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}$. of type
$\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}\rightarrow\sigma$. This
rebuilder takes an array with precisely as many scalars as were in the input,
and produces a value of type $\sigma$ with the structure (and discrete-typed
values) of the input, but the scalars from the array. The mapping between
locations in $\sigma$ and indices in the array is the same as the numbering
performed by $\textrm{Interleave}^{4}$.

Having $x$, rebuild and $i$ (the next available ID), we execute the
transformed term $\textbf{D}^{4}_{\sigma}[t]$ monadically (with $x$ in scope),
resulting in an output $y^{\prime}:\textbf{D}^{4}_{\sigma}[\tau]$. This output
we deinterleave to $y:\tau$ and
$d:\tau\multimap(\operatorname*{\textrm{Staged}}\sigma\rightarrow\texttt{IO}\
())$.

The final result then consists of the regular function result ($y$) as well as
the top-level derivative function of type $\tau\multimap\sigma$. In the
derivative function, we allocate two arrays to initialise an empty
$\operatorname*{\textrm{Staged}}\sigma$ (note that the given sizes are indeed
precisely large enough), and apply $d$ to the incoming $\tau$ cotangent. This
gives us an updater function that (because of how $\textrm{Deinterleave}^{4}$
works) stages calls to the top-level backpropagators contained in $y^{\prime}$
in the $\operatorname*{\textrm{Staged}}$ arrays, and we apply this function to
the just-allocated $\operatorname*{\textrm{Staged}}$ object. Then we use the
new SResolve to propagate the cotangent contributions backwards, by invoking
each backpropagator in turn in descending order of IDs. Like before in the
Cayley-transformed version of our AD technique, those backpropagators update
the state (now mutably) to record their own contributions to (i.e. invocations
of) other backpropagators. At the end of SResolve, the backpropagator staging
array is dropped and the cotangent collection array is frozen and returned as
an $\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}$ (corresponding to
the $c$ value in a $\operatorname*{\textrm{Staged}}c$ for the Cayley-
transformed version in Fig. 9).

This whole derivative computation is, in the end, pure (in that it is
deterministic and has no side-effects), so we can safely evaluate the IO using
unsafePerformIO to get a pure
$\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}$, which contains the
scalar cotangents that rebuild from $\textrm{Interleave}^{4}$ needs to put in
the correct locations in the input, thus constructing the final gradient.

#### Implementation of the components.

Having discussed the high-level sequence of operations, let us briefly discuss
the implementation of the $\operatorname*{\textrm{Staged}}c$ interface and
(de)interleaving. In $\textrm{Interleave}^{4}$, instead of passing structure
information down in the form of a setter
($(\tau\rightarrow\tau)\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\operatorname*{\textrm{Staged}}c)$)
like we did in $\textrm{Interleave}^{3}$ in Fig. 9, we build structure
information up in the form of a getter
($\operatorname*{\textrm{IArray}}\underline{\mathbb{R}}\rightarrow\tau$). This
results in a somewhat more compact presentation, but in some sense the same
information is still communicated.

The program text of $\textrm{Deinterleave}^{4}$ is again unchanged, because it
is agnostic about the codomain of the backpropagators, as long as it is a
monoid, which it remains.

On the $\operatorname*{\textrm{Staged}}$ interface the transition to mutable
arrays had a significant effect. The $0_{\text{Staged}}$ created by SResolve
in Section 7 is now created in SAlloc, which uses alloc to allocate a zero-
filled cotangent collection array of size $i_{\text{inp}}$, and the
backpropagator staging array of size $i_{\text{out}}$ filled with zero-
backpropagators with an accumulated argument of zero.

SCall has essentially the same type, but its implementation differs because it
now performs a constant-time mutable update on the backpropagator staging
array instead of a logarithmic-complexity immutable Map update. Note that,
unlike in Section 7, there is no special case if $i$ is not yet in the array,
because unused positions are already filled with zeros.

SOneHot takes the place of SCotan, with the difference that we have
specialised it using the knowledge that all relevant $c\rightarrow c$
functions add a particular scalar to a particular index in the input, and that
these functions can hence be defunctionalised to a pair
$(\mathbb{Z},\underline{\mathbb{R}})$. The monoid-linearity here is in the
real scalar, as it was before, hence the placement of the $\multimap$-arrow.

Finally, SResolve takes an additional $\mathbb{Z}$ argument that should
contain the output ID of $\textbf{D}^{4}_{c}[t]$, i.e. one more than the
largest ID generated. loop then does what the original SResolve did directly,
iterating over all IDs in descending order and applying the state updaters in
the backpropagator staging array one-by-one to the state. After the loop is
complete, we freeze and return just the cotangent collection array, because we
have no need for the staging array any more. This frozen collection array will
then be used to build the final gradient in $\textrm{Wrap}^{4}$.

##  10 Was It Taping All Along?

In this section we first apply one more optimisation to our algorithm to make
it slightly more efficient (Section 10.1). Afterwards, we show that
defunctionalising the backpropagators (Section 10.2) essentially reduces the
technique to classical taping approaches (Section 10.3).

###  10.1 Dropping the Cotangent Collection Array

Recall that the final transformation of Section 9 used two mutable arrays
threaded through the backpropagators in the $\operatorname*{\textrm{Staged}}c$
pair: a cotangent collection array of type
$\operatorname*{\textrm{Array}}\underline{\mathbb{R}}$ and a backpropagator
call staging array of type
$\operatorname*{\textrm{Array}}{}(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()),\underline{\mathbb{R}})$, using the monad-based implementation from
Section 9.2. The first array is modified by $\textrm{Interleave}_{\mathbb{R}}$
and the second by SCall. No other functions modify these arrays.

Looking at the function of $\textrm{Interleave}_{\mathbb{R}}$ in the
algorithm, all it does is produce input backpropagators with some ID $i$,
which act by adding their argument to index $i$ in the cotangent collection
array. This means that we have $c[i]=\textrm{snd}(m[i])$ for all $i$ for which
$c[i]$ is defined, if $(c,m)$ is the input to SResolve for which the recursion
terminates. Therefore, the cotangent collection array is actually unnecessary:
its information can be read off directly from the backpropagator staging
array.

With this knowledge, we can instead use
$\operatorname*{\textrm{Staged}}c=\operatorname*{\textrm{Array}}{(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()),\underline{\mathbb{R}})}$ as our definition. The reconstruction functions
of Section 8 simply take the second projection of the correponding array
element.

###  10.2 Defunctionalisation of Backpropagators

In the core code transformation ($\textbf{D}_{c}$, excluding the wrapper), all
backpropagators are (now) of type
$\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())$, and, as observed earlier in Section 6, these backpropagators come in
only a limited number of forms:

  1. 1.

The input backpropagators created in $\textrm{Interleave}_{\mathbb{R}}$,
reduced to $(\underline{\lambda}(z:\underline{\mathbb{R}}).\ \mathbf{return}\
())$ in Section 10.1;

  2. 2.

$(\underline{\lambda}(z:\underline{\mathbb{R}}).\ \mathbf{return}\ ())$
created in $\textbf{D}^{4}_{c}[r]$ for scalar constants $r$;

  3. 3.

$(\underline{\lambda}(z:\underline{\mathbb{R}}).\ \textrm{SCall}\ d_{1}\
(\partial_{1}\textit{op}(x_{1},\ldots,x_{n})(z))\circ\cdots\circ\textrm{SCall}\
d_{n}\ (\partial_{n}\textit{op}(x_{1},\ldots,x_{n})(z)))$ created in
$\textbf{D}_{c}[\textit{op}(x_{1},\ldots,x_{n})]$ for primitive operations op.

Furthermore, the information contained in an operator backpropagator of form
(3) can actually be described without reference to the value of its argument
$z$: because our operators return a single scalar (as opposed to e.g. a
vector), we have

$\partial_{i}\textit{op}(x_{1},\ldots,x_{n})(z)=z\cdot\partial_{i}\textit{op}(x_{1},\ldots,x_{n})(1)$

Hence, we can defunctionalise (Reynolds, 1998) and change all occurrences of
the type
$\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())$ to Contrib, where
$\textrm{Contrib}=[(\underline{\mathbb{R}},(\mathbb{Z},\textrm{Contrib}))]$: a
list of triples of a scalar, an integer ID, and a recursive Contrib structure.
The ID is the ID of the Contrib structure (i.e. the backpropagator) that it is
adjacent to. (As before, these IDs make sharing observable.) In this
representation, we think of
$[(a_{1},(i_{1},\textit{cb}_{1})),\ldots,(a_{n},(i_{n},\textit{cb}_{n}))]$ of
type Contrib as the backpropagator

$\underline{\lambda}(z:\underline{\mathbb{R}}).\ \textrm{SCall}\
(i_{1},\textit{cb}_{1})\ (z\cdot a_{1})\circ\cdots\circ\textrm{SCall}\
(i_{n},\textit{cb}_{n})\ (z\cdot a_{n})$

For example, suppose we differentiate the following simple program:

$\displaystyle\lambda(x,y).\ \textbf{let}\ z=x+y\ \textbf{in}\ x\cdot z$

using the final algorithm of Section 9.2. The return value from the
$\textbf{D}_{(\mathbb{R},\mathbb{R})}$-transformed code (when applied to the
output from $\textrm{Interleave}_{(\mathbb{R},\mathbb{R})}$) has the sharing
structure shown in Fig. 11(a). This shows how the backpropagators refer to
each other in their closures.

If we perform the type replacement and the defunctionalisation, Interleave
simplifies and SCall disappears; backpropagators of forms (1) and (2) become
$[]$ (the empty list) and those of form (3) become:

$[(\partial_{1}\textit{op}(x_{1},\ldots,x_{n})(1),d_{1}),\ldots,(\partial_{n}\textit{op}(x_{1},\ldots,x_{n})(1),d_{n})]$

SResolve then interprets a list of such $(a,(i,\textit{cb}))$ by iterating
over the list and for each such triple, replacing
$(\textit{cb}^{\prime},a^{\prime})$ at index $i$ in the staging array with
$(\textit{cb},a^{\prime}+a\cdot d)$, where $d$ is the cotangent recorded in
the array cell where the list was found.

${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(x,(1,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\textstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptscriptstyle\underline{\lambda}$}}d.\
...{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(y,(2,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\textstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptscriptstyle\underline{\lambda}$}}d.\
...{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(z,(3,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\textstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptstyle\underline{\lambda}$}}{{}{}{{}}\hbox{$\scriptscriptstyle\underline{\lambda}$}}d.\
\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}\
d\circ\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}\
d{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$$(x\cdot
z,(4,\underline{\lambda}d.\
\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}\
(z\cdot
d)\circ\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}\
(x\cdot d)))$ (a) Before defunctionalisation

${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(x,(1,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle[]$}}{{}{}{{}}\hbox{$\textstyle[]$}}{{}{}{{}}\hbox{$\scriptstyle[]$}}{{}{}{{}}\hbox{$\scriptscriptstyle[]$}}{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(y,(2,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle[]$}}{{}{}{{}}\hbox{$\textstyle[]$}}{{}{}{{}}\hbox{$\scriptstyle[]$}}{{}{}{{}}\hbox{$\scriptscriptstyle[]$}}{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$${\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}(z,(3,\,}\mathchoice{{}{}{{}}\hbox{$\displaystyle[$}}{{}{}{{}}\hbox{$\textstyle[$}}{{}{}{{}}\hbox{$\scriptstyle[$}}{{}{}{{}}\hbox{$\scriptscriptstyle[$}}(1.0,(1,\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}})),(1.0,(2,\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}))]{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}))}$$(x\cdot
z,(4,[(z,(1,\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}})),(x,(3,\mathchoice{{}{}{{}}\hbox{$\displaystyle\Box$}}{{}{}{{}}\hbox{$\textstyle\Box$}}{{}{}{{}}\hbox{$\scriptstyle\Box$}}{{}{}{{}}\hbox{$\scriptscriptstyle\Box$}}))]))$
(b) After defunctionalisation

Figure 11:  The sharing structure before and after defunctionalisation. SCall
is elided here; in Fig. 11(a), the backpropagator calls are depicted as if
they are still normal calls. Boxes ($\Box$) are the same in-memory value as
the value their arrow points to; two boxes pointing to the same value
indicates that this value is _shared_ : referenced in two places.

###  10.3 Was It Taping All Along?

After the improvements from Sections 10.1 and 10.2, what previously was a tree
of (staged) calls to backpropagator functions is now a tree of Contrib values
with attached IDs242424Note that we now have
$\textbf{D}[\mathbb{R}]=(\mathbb{R},(\mathbb{Z},\textrm{Contrib}))$, the
integer being the ID of the Contrib value. that are interpreted by SResolve.
This interpretation (eventually) writes the Contrib value with ID $i$ to index
$i$ in the staging array (possibly multiple times), and furthermore
accumulates argument cotangents in the second component of the pairs in the
staging array. While the argument cotangents must be accumulated in reverse
order of program execution (indeed, that is the whole point of _reverse_ AD),
the mapping from ID to Contrib value can be fully known in the forward pass:
the partial derivatives of operators,
$\partial_{i}\textit{op}(x_{1},\ldots,x_{n})(1)$, can be computed in the
forward pass already.

This means that we can already compute the Contrib lists and write them to the
array in the forward pass, if we change the ID generation monad that the
differentiated code already lives in (which is a state monad with a single
$\mathbb{Z}$ as state) to additionally carry the staging array, and
furthermore change the monad to thread its state in a way that allows
mutation, again using the techniques from Section 9, but now in the forward
pass too. All that SResolve then has to do is loop over the array in reverse
order (as it already does) and add cotangent contributions to the correct
positions in the array according to the Contrib lists that it finds there.

At this point, there is no meaningful difference any more between this
algorithm and what is classically known as taping: we have a tape (the staging
array) that we write the performed operations to in the forward pass
(automatically growing the array as necessary)—although the tape entries are
the already-differentiated operations in this case, and not the original ones.
In this way, we have related the naive version of dual-numbers reverse AD,
which admits neat correctness proofs, to the classical, very imperative
approach to reverse AD based on taping, which is used in industry-standard
implementations of reverse AD (e.g. PyTorch (Paszke et al., 2017)).

##  11 Extending the Source Language

The source language (Fig. 4) that the algorithm discussed so far works on, is
a higher-order functional language including product types and primitive
operations on scalars. However, dual-numbers reverse AD generalises to much
richer languages in a very natural way, because most of the interesting work
happens in the scalar primitive operations. The correctness proof for the
algorithm can be extended to many expressive language constructs in the source
language, such as coproducts and recursive types by using suitable logical
relations arguments (Nunes and Vákár, 2023). Further, the efficiency of the
algorithm is independent of the language constructs in the source language.
Indeed, in the forward pass, the code transformation is fully structure-
preserving outside of the scalar constant and primitive operation cases; and
in the reverse pass (in SResolve), all program structure is forgotten anyway,
because the computation is flattened to the (reverse of the) linear sequence
of primitive operations on scalars that was performed in the particular
execution of the forward pass.

#### (Mutual) recursion.

For example, we can allow recursive functions in our source language by adding
recursive let-bindings with the following typing rule:

$\frac{\Gamma,f:\sigma\rightarrow\tau,x:\sigma\vdash
s:\tau\qquad\Gamma,f:\sigma\rightarrow\tau\vdash t:\rho}{\textbf{letrec}\ f\
x=s\ \textbf{in}\ f\ t:\rho}$

The code transformation $\textbf{D}^{i}$ for all $i$ then treats letrec
exactly the same as let—note that the only syntactic difference between letrec
and let is the scoping of $f$—and the algorithm remains both correct and
efficient. Note that due to the assumed call-by-value semantics, recursive
non-function definitions would make little sense.

#### Coproducts.

To support dynamic control flow (necessary to make recursion useful), we can
easily add coproducts to the source language. First add coproducts to the
syntax for types
($\sigma,\tau\mathbin{\raisebox{0.5pt}{::}{=}}\cdots\mid\sigma\sqcup\tau$)
both in the source language and in the target language, and add constructors
and eliminators to all term languages (both linear and non-linear):

$\displaystyle
s,t\mathbin{\raisebox{0.5pt}{::}{=}}\cdots\mid\textrm{inl}(t)\mid\textrm{inr}(t)\mid\textbf{case}\
{s}\ \textbf{of}\ \\{\textrm{inl}(x)\rightarrow
t_{1};\textrm{inr}(y)\rightarrow t_{2}\\}$

where $x$ and $y$ are in scope in $t_{1}$ and $t_{2}$. Then the type and code
transformations extend in the unique structure-preserving manner:

$\displaystyle\textbf{D}^{1}_{c}[\sigma\sqcup\tau]=\textbf{D}^{1}_{c}[\sigma]\sqcup\textbf{D}^{1}_{c}[\tau]$
$\displaystyle\textbf{D}^{1}_{c}[\textrm{inl}(t)]=\textrm{inl}(\textbf{D}^{1}_{c}[t])\qquad\textbf{D}^{1}_{c}[\textrm{inr}(t)]=\textrm{inr}(\textbf{D}^{1}_{c}[t])$
$\displaystyle\textbf{D}^{1}_{c}[\textbf{case}\ {s}\ \textbf{of}\
\\{\textrm{inl}(x)\rightarrow t_{1};\textrm{inr}(x)\rightarrow
t_{2}\\}]=\textbf{case}\ {\textbf{D}^{1}_{c}[s]}\ \textbf{of}\
\\{\begin{array}[t]{@{}l@{\;\ra\;}l@{}}\textrm{inl}(x)\hfil\;\rightarrow&\textbf{D}^{1}_{c}[t_{1}];\\\
\textrm{inr}(x)\hfil\;\rightarrow&\textbf{D}^{1}_{c}[t_{2}]\\}\end{array}$

To create an interesting interaction between control flow and differentiation,
we can add a construct ‘sign’ with the unsurprising typing rule

$\frac{\Gamma\vdash t:\mathbb{R}}{\Gamma\vdash\textrm{sign}(t):\mathrm{Bool}}$

where $\mathrm{Bool}=()\sqcup()$, which allows us to perform a case
distinction on the sign of a real number. For differentiation of this
construct it suffices to define
$\textbf{D}^{1}_{c}[\textrm{sign}(t)]=\textrm{sign}(\textrm{fst}(\textbf{D}^{1}_{c}[t]))$;
if one wants to give up the structure-preserving aspect of the transformation,
it is also possible to prevent redundant differentiation of $t$ and just make
clever substitutions in $t$’s free variables to convert back from dual-numbers
form to the plain data representation.

The type transformation stays unchanged when moving to $\textbf{D}^{4}_{c}$,
and the only change for the term definitions is to transition to monadic code
in $\textbf{D}^{2}_{c}$. Lifting a computation to monadic code is a well-
understood process. The corresponding cases in Interleave and Deinterleave are
the only reasonable definitions that type-check.

#### Polymorphic and (mutually) recursive types.

In Haskell one can define (mutually) recursive data types e.g. as follows:

$\displaystyle\textbf{data}\ T_{1}\ \alpha=C_{1}\ \alpha\ (T_{2}\ \alpha)\mid
C_{2}\ \mathbb{R}$ $\displaystyle\textbf{data}\ T_{2}\ \alpha=C_{3}\
\mathbb{Z}\ (T_{1}\ \alpha)\ (T_{2}\ \alpha)$

If the user has defined some data types, then we can allow these data types in
the code transformation. We generate new data type declarations that simply
apply $\textbf{D}^{1}_{c}[-]$ to all parameter types of all constructors:

$\displaystyle\textbf{data}\ DT_{1}\ \alpha=DC_{1}\ \alpha\ (DT_{2}\
\alpha)\mid DC_{2}\ (\mathbb{R},\underline{\mathbb{R}}\multimap c)$
$\displaystyle\textbf{data}\ DT_{2}\ \alpha=DC_{3}\ \mathbb{Z}\ (DT_{1}\
\alpha)\ (DT_{2}\ \alpha)$

and we add one rule for each data type that simply maps:

$\displaystyle\textbf{D}^{1}_{c}[T_{1}\ \tau]=DT_{1}\
\textbf{D}^{1}_{c}[\tau]\qquad\textbf{D}^{1}_{c}[T_{2}\ \tau]=DT_{2}\
\textbf{D}^{1}_{c}[\tau]$

Furthermore, for plain type variables, we set
$\textbf{D}^{1}_{c}[\alpha]=\alpha$.252525As declaring new data types is
inconvenient in Template Haskell, our current implementation only handles
recursive data types that do not contain explicit scalar values. As we can
pass all required scalar types by instantiating their type parameters with a
type containing $\mathbb{R}$, this is not a real restriction.

The code transformation on terms is completely analogous to a combination of
coproducts (given above in this section, where we take care to match up
constructors as one would expect: $C_{i}$ gets sent to $DC_{i}$) and products
(given already in Fig. 6). The wrapper also changes analogously: Interleave
and Deinterleave get clauses for $\textrm{Interleave}_{(T_{i}\,\tau)}$ and
$\textrm{Deinterleave}_{(T_{i}\,\tau)}$.

Finally, we note that with the mentioned additional rule that
$\textbf{D}^{1}_{c}[\alpha]=\alpha$, polymorphic functions can also be
differentiated transparently, similarly to how the above handles polymorphic
data types.

##  12 Parallelism

So far, we have assigned sequentially increasing IDs to backpropagators in the
forward pass and resolved them in their linear order from top to bottom during
the reverse pass. As long as the source program is executed sequentially, such
ID generation is appropriate. However, if the source program uses parallelism
in its execution, such linear ID assignment discards this parallelism
structure and prevents us from exploiting it for computing the derivative in
parallel.

In this section, we explore how to perform dual numbers reverse AD on source
programs that contain fork-join task parallelism using a simple, but
representative, parallel combination construct $(\star)$ that has the
semantics that $s\star t$ computes $s$ and $t$ in parallel and returns the
pair $(s,t)$. We discuss a different ID assignment scheme for backpropagators
that takes parallelism into account, and we show that we can take advantage of
these new IDs to resolve backpropagators in parallel during the reverse pass.

###  12.1 Fork-Join Parallelism

We work with a parallel operational model in which computations execute in a
_job_. A job is sequential: forking terminates a job, and if there is more
work to do after the corresponding join, a new job must be started to run the
continuation. Each new job gets a fresh job ID that is its unique identifier.
The idea is that independent jobs can execute in parallel in different
threads. A typical runtime for this model is a thread pool together with a job
queue to distribute work over the operating system threads: new jobs are
submitted to the queue when they are created, and threads from the pool pick
up waiting jobs from the queue when idle.

We extend our source language syntax with a parallel pairing construct
$(\star)$ with the following typing rule:

$\frac{\Gamma\vdash s:\sigma\qquad\Gamma\vdash t:\tau}{\Gamma\vdash s\star
t:(\sigma,\tau)}$

Operationally, when we encounter the term $t_{1}\star t_{2}$ while evaluating
some term $t$ in job $\alpha$, two new jobs are created with fresh job IDs
$\beta$ and $\gamma$. The term $t_{1}$ starts evaluating in job $\beta$,
potentially does some forks and joins, and finishes in a (potentially)
different job $\beta^{\prime}$, returning a result $v$; $t_{2}$ starts
evaluating in job $\gamma$ and finishes in $\gamma^{\prime}$, returning a
result $w$. When $\beta^{\prime}$ and $\gamma^{\prime}$ terminate, evaluation
of $t$ continues in a new job $\delta$ with the result $(v,w)$ for $t_{1}\star
t_{2}$. We say that $\alpha$ _forks_ into $\beta$ and $\gamma$ and that
$\beta^{\prime}$ and $\gamma^{\prime}$ _join_ into $\delta$. The two parallel
subgraphs, one from $\beta$ to $\beta^{\prime}$ and one from $\gamma$ to
$\gamma^{\prime}$, we call _tasks_. The operational model does not know about
tasks (it only knows about jobs), but we will use the concept of tasks in
Section 12.3 as a compositional building block for the job graph. Figure 12
contains a diagrammatic representation of the preceding paragraph.

$\alpha$$\beta$$\gamma$$\beta\mathrlap{\smash{{}^{\prime}}}$$\gamma\mathrlap{\smash{{}^{\prime}}}$$\delta\hskip
1.0pt$jobstask Figure 12:  Schematic view of the operational model underlying
$(\star)$.

The algorithm in this section extends readily to $n$-ary parallel tupling
constructs:

$\frac{\Gamma\vdash t_{1}:\tau_{1}\quad\cdots\quad\Gamma\vdash
t_{n}:\tau_{n}}{\Gamma\vdash(t_{1},\ldots,t_{n})^{\star}:(\tau_{1},\ldots,\tau_{n})}$

or parallel combination constructs of dynamic arity:

$\frac{\Gamma\vdash t:\text{list}(()\to\tau)}{\Gamma\vdash\star
t:\text{list}(\tau)}$

but for simplicity of presentation, we restrict ourselves here to binary
forking.

###  12.2 Parallel IDs and Their Partial Order

To avoid discarding the (explicit, using $(\star)$) parallelism structure in
the source program, we have to somehow record the dependency graph of the
backpropagators (the same graph as the computation graph of scalars in the
source program) in a way that is more precise than the chronological linear
order used for the sequential algorithm. Specifically, we want backpropagators
that were created in parallel jobs in the forward pass to not depend on each
other in the dependency graph, not even transitively. In other words, they
should be incomparable in the partial order that informs the reverse pass
(SResolve) what backpropagator to resolve next.

To support the recording of this additional dependency information, we switch
to _compound IDs_ , consisting of two integers instead of one:

  * •

The _job ID_ that uniquely identifies the job a backpropagator is created
from. This requires that we have a way to generate a unique ID in parallel
every time a job forks or two jobs join. Job IDs do _not_ carry a special
partial order; see below.

  * •

The _ID within the job_ , which we assign sequentially (starting from 0) to
all backpropagators created in a job. The operations within one forward-pass
job are sequential (because of our fork-join model where a fork ends a job;
see the previous subsection); this ID-within-job simply witnesses this.

Compound IDs have lexicographic order where
$(\alpha,i)\leq_{\mathrm{c}}(\beta,i^{\prime})$ iff
$\alpha\leq_{\mathrm{j}}\beta$ and $i\leq_{\mathbb{Z}}i^{\prime}$: the order
on sequential IDs (within a job) is simply the standard linear order
$\leq_{\mathbb{Z}}$. The partial order on job IDs, on the other hand, is
different and is defined as the transitive closure of the following three
cases:

  1. 1.

$\alpha\leq_{\mathrm{j}}\alpha$;

  2. 2.

If $\alpha$ forks into $\beta$ and $\gamma$, then
$\alpha\leq_{\mathrm{j}}\beta$ and $\alpha\leq_{\mathrm{j}}\gamma$;

  3. 3.

If $\alpha$ and $\beta$ join into $\gamma$, then
$\alpha\leq_{\mathrm{j}}\gamma$ and $\beta\leq_{\mathrm{j}}\gamma$.

In Fig. 13, we give an example term together with graphs showing the
(generators of the) partial orders $\leq_{\mathrm{j}}$ on job IDs and
$\leq_{\mathrm{c}}$ on compound IDs. Both graphs have arrows pointing to the
successors of each node: $n_{2}$ is a successor of $n_{1}$ (and $n_{1}$ a
predecessor of $n_{2}$) if $n_{1}<n_{2}$ (i.e. $n_{1}\leq n_{2}$ and
$n_{1}\not=n_{2}$) and there is no $m$ such that $n_{1}<m<n_{2}$. The graphs
generate their respective partial orders if one takes the transitive closure
and includes trivial self-loops. For Fig. 13(b) ($\leq_{\mathrm{j}}$), the
arrows thus show the fork/join relationships; for Fig. 13(c)
($\leq_{\mathrm{c}}$), this is refined with the linear order on sequential IDs
within each job.

$\mathbf{let}\ (z_{1},z_{2})=\\\ \phantom{....}\Big{(}\mathbf{let}\
x=a*b+\sin(b)\\\ \phantom{....\Big{(}}\mathbf{in}\ \mathbf{let}\
(y_{1},y_{2})=(x*a+a)\star\cos(x)\\\ \phantom{....\Big{(}}\mathbf{in}\
x+y_{1}*y_{2}\Big{)}\star\bigl{(}\exp(a)+b\bigr{)}\\\ \mathbf{in}\
z_{1}*z_{2}$

(a) Example program

$\alpha$$\beta$$\delta$$\varepsilon$$\zeta$$\gamma$$\eta$ (b) Job IDs
($\leq_{\mathrm{j}}$)

$(\beta,0)$$(\beta,1)$$(\beta,2)$$(\delta,0)$$(\delta,1)$$(\varepsilon,0)$$(\zeta,0)$$(\zeta,1)$$(\gamma,0)$$(\gamma,1)$$(\eta,0)$
(c) Compound IDs ($\leq_{\mathrm{c}}$)

Figure 13:  An example program. Note that the program starts by forking,
before performing any primitive operations, hence job $\alpha$ is empty and
the partial order on compound IDs happens to have multiple minimal elements.

The reverse pass (SResolve) needs to traverse the dependency graph on compound
IDs (Fig. 13(c)) in reverse dependency order (taking advantage of task
parallelism), but it is actually unnecessary to construct the full graph at
runtime. It is sufficient to construct the dependency graph on _job IDs_ (Fig.
13(b)) together with, for each job $\alpha$, the number $n_{\alpha}$ of
sequential IDs generated in that job (note that this number may be zero if no
scalar computation was done while running in that job). With this information,
SResolve can walk the job graph in reverse topological order, resolving
parallel tasks in parallel, and for each job $\alpha$ sequentially resolve the
individual backpropagators from $n_{\alpha}-1$ to $0$. We will collect this
additional information during the forward pass by extending the monad that the
forward pass runs in.

###  12.3 Extending the Monad

In the first optimisation that we applied to the (sequential) algorithm,
namely linear factoring via staging of backpropagators (Section 6, Fig. 8), we
modified the code transformation to produce code that runs in an state monad
$\operatorname*{\mathcal{M}}\tau=\textrm{ID}\rightarrow(\tau,\textrm{ID})$,
for $\textrm{ID}=\mathbb{Z}$:

$\text{If}\ \Gamma\vdash t:\tau\ \text{then}\
\textbf{D}^{2}_{c}[\Gamma]\vdash\textbf{D}^{2}_{c}[t]:\operatorname*{\mathcal{M}}\textbf{D}^{2}_{c}[\tau]$

In fact, this state monad was simply the natural implementation of an _ID
generation monad_ with one method:

$\begin{array}[]{@{}l@{\ }c@{\
}l@{}}\mathrm{genID}&:&\operatorname*{\mathcal{M}}\textrm{ID}\\\
\mathrm{genID}&=&\lambda(i:\mathbb{Z}).\ (i+1,i)\end{array}$

We saw above in Section 12.2 that we need to extend this monad to be able to
do two things: (1) generate compound IDs, not just sequential IDs, and (2)
record the job graph resulting from parallel execution using $(\star)$. Write
JID for the type of job IDs (in an implementation we can simply set
$\textrm{JID}\coloneqq\mathbb{Z}$) and write
$\textrm{CID}\coloneqq(\textrm{JID},\mathbb{Z})$ for the type of compound IDs.
The extended monad needs two methods:

$\begin{array}[]{@{}r@{\ }c@{\
}l@{}}\mathrm{genID}&:&\operatorname*{\mathcal{M}}\textrm{CID}\\\
(\star^{\prime})&:&\operatorname*{\mathcal{M}}\sigma\to\operatorname*{\mathcal{M}}\tau\to\operatorname*{\mathcal{M}}{(\sigma,\tau)}\end{array}$

The updated $\mathrm{genID}$ reads the current job ID from an atomic mutable
cell inside the monad and pairs that with a sequential ID generated in the
standard fashion with monadic state. The monadic parallel combination method,
$(\star^{\prime})$, generates some fresh job IDs by incrementing the mentioned
atomic mutable cell, runs the two monadic computations _in parallel_ by
spawning jobs as described in Section 12.1, and records the structure of the
job graph thus created in some state inside the monad.

In an implementation, one can take the definitions in Fig. 14. Working in
Haskell, we write IORef for a mutable cell and use IO as the base monad in
which we can access that mutable cell as well as spawn and join parallel
threads. (We do not use IO in any other way in $\operatorname*{\mathcal{M}}$,
although SResolve, which runs after the forward pass, will also use mutable
arrays as in Section 9.2.) The given implementation of
$\operatorname*{\mathcal{M}}$ has the atomic mutable cell in a reader context,
and is a state monad in ‘$\mathrm{JobDescr}$’: a description of a job that
contains the _history_ of a job together with its job ID and the number of
sequential IDs generated in that job (numbered $0,\ldots,n-1$). The history of
a job $\alpha$ is the subgraph of the job graph given by all jobs $\beta$
satisfying $\beta<\alpha$ in the smallest _task_ (recall Fig. 12) containing
the job. For the special case of the (unique) last job of a task, its history
is precisely the whole task excluding itself. This definition of a “history”
ensures that $(\star^{\prime})$ has precisely the parts of the job graph that
it needs to build up the job graph of the task that it itself is running in,
which makes everything compose.

$\begin{array}[]{@{}c@{}}\operatorname*{\mathcal{M}}\alpha=\texttt{IORef}\
\textrm{JID}\rightarrow\mathrm{JobDescr}\rightarrow\texttt{IO}\
(\mathrm{JobDescr},\alpha)\\\\[5.0pt]
\begin{array}[]{@{}l@{\qquad}l@{}}\begin{array}[]{@{}c@{}}\mathrm{JobDescr}=(\mathrm{History},\textrm{JID},\mathbb{Z})\\\\[5.0pt]
\mathbf{data}\
\mathrm{History}\begin{array}[t]{@{\;}c@{\;}l@{}}=&\mathrm{Start}\\\
\mid&\mathrm{Fork}\
\mathrlap{\smash{\overset{(\text{A})}{\mathrm{JobDescr}}}}\hphantom{\mathrm{JobDescr}}\
\mathrlap{\smash{\overset{(\text{B})}{\mathrm{JobDescr}}}}\hphantom{\mathrm{JobDescr}}\
\mathrlap{\smash{\overset{(\text{C})}{\mathrm{JobDescr}}}}\hphantom{\mathrm{JobDescr}}\end{array}\end{array}&\mathrlap{\raisebox{-17.64056pt}{
\leavevmode\hbox to38.66pt{\vbox
to35.88pt{\pgfpicture\makeatletter\hbox{\hskip 19.32782pt\lower-3.55556pt\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{
}\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{
}\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox
to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{{}}
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-3.0pt}{25.71944pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{\footnotesize{A}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-9.94653pt}{11.49306pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{\footnotesize{B}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{4.2243pt}{11.49306pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{\footnotesize{C}}
}}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
{{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{}{{ {}{}}}{ {}{}}
{{}{{}}}{{}{}}{}{{}{}} { }{{{{}}\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-19.32782pt}{-2.0pt}\pgfsys@invoke{
}\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{
}\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{\footnotesize{current
job}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}}
\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} { {}{}{}}{}{ {}{}{}}
{}{{{{{}}{
{}{}}{}{}{{}{}}}}}{{}}{}{}{}{}{}{}{{}}\pgfsys@moveto{-2.8516pt}{24.38133pt}\pgfsys@lineto{-5.61258pt}{20.43994pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{-0.57375}{-0.81903}{0.81903}{-0.57375}{-4.82774pt}{21.56032pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} { {}{}{}}{}{ {}{}{}} {}{{{{{}}{
{}{}}{}{}{{}{}}}}}{{}}{}{}{}{}{}{}{{}}\pgfsys@moveto{2.8516pt}{24.38133pt}\pgfsys@lineto{5.61258pt}{20.43994pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{0.57375}{-0.81903}{0.81903}{0.57375}{4.82774pt}{21.56032pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} { {}{}{}}{}{{}}{}{ {}{}{}} {{{{{}}{
{}{}}{}{}{{}{}}}}}{}{}{}{}{}{}{}{{}}\pgfsys@moveto{-7.11319pt}{10.15495pt}\pgfsys@lineto{-4.13097pt}{5.89774pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{0.57375}{-0.81905}{0.81905}{0.57375}{-4.9158pt}{7.01813pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} { {}{}{}}{}{{}}{}{ {}{}{}} {{{{{}}{
{}{}}{}{}{{}{}}}}}{}{}{}{}{}{}{}{{}}\pgfsys@moveto{7.11319pt}{10.15495pt}\pgfsys@lineto{4.13097pt}{5.89774pt}\pgfsys@stroke\pgfsys@invoke{
}{{}{{}}{}{}{{}}{{{}}{{{}}{\pgfsys@beginscope\pgfsys@invoke{
}\pgfsys@transformcm{-0.57375}{-0.81905}{0.81905}{-0.57375}{4.9158pt}{7.01813pt}\pgfsys@invoke{
}\pgfsys@invoke{ \lxSVG@closescope }\pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope}}{{}}}} \pgfsys@invoke{\lxSVG@closescope
}\pgfsys@endscope{}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope
}\pgfsys<EMAIL_ADDRESS>\begin{array}[]{@{}l@{}}f\star^{\prime}g=\lambda\mathit{ref}\
\mathit{jd}_{0}.\ \mathbf{do}\\\
\qquad\begin{array}[t]{@{}l@{}}(j_{1},j_{2},j_{3})\leftarrow\texttt{atomicModifyIORef'}\
\mathit{ref}\ (\lambda j.\ (j+3,(j,j+1,j+2)))\\\
((\mathit{jd}_{1},x),(\mathit{jd}_{2},y))\leftarrow\texttt{inParallel}\ (f\
\mathit{ref}\ (\mathrm{Start},j_{1},0))\ (\mathrlap{g}\hphantom{f}\
\mathit{ref}\ (\mathrm{Start},j_{2},0))\\\ \mathbf{return}\ ((\mathrm{Fork}\
\mathit{jd}_{0}\ \mathit{jd}_{1}\
\mathit{jd}_{2},j_{3},0),(x,y))\end{array}\\\ \quad\text{where:}\
\begin{array}[t]{@{}l@{}}\texttt{atomicModifyIORef'}::\texttt{IORef}\
\alpha\rightarrow(\alpha\rightarrow(\alpha,\beta))\rightarrow\texttt{IO}\
\beta\\\ \texttt{inParallel}::\texttt{IO}\ \alpha\rightarrow\texttt{IO}\
\beta\rightarrow\texttt{IO}\ (\alpha,\beta)\end{array}\end{array}\end{array}$
Figure 14:  Sketch of the implementation of the monad
$\operatorname*{\mathcal{M}}$. The diagram shows the meaning of the job
descriptions in ‘$\mathrm{Fork}$’: the first field (labeled ‘$A$’) contains
the history up to the last fork in this task (excluding subtasks), and the
fields labeled $B$ and $C$ describe the subtasks spawned by that fork. The
first job in a task has no history, indicated with ‘$\mathrm{Start}$’.

#### Differentiation.

We keep the differentiation rules for all existing language constructs the
same, except for changing the type of IDs to CID and using the monad
$\operatorname*{\mathcal{M}}$ instead of doing manual state passing of the
next ID to generate. A representative rule showing how this looks is: (compare
Fig. 8)

$\textbf{D}^{5}_{c}[(s,t)]=\textbf{D}^{5}_{c}[s]\mathbin{>\\!\\!>\\!\\!=}\lambda
x.\ \textbf{D}^{5}_{c}[t]\mathbin{>\\!\\!>\\!\\!=}\lambda y.\ \mathbf{return}\
(x,y)$

Because we now generate fresh compound IDs rather than plain integer IDs when
executing primitive operations, we change $\mathbb{Z}$ to CID in
$\textbf{D}^{4}_{c}[\mathbb{R}]$:262626Here we work from the transformation
$\textbf{D}^{4}$ as described in Section 9.2 on monadic mutable arrays. With
the defunctionalisation described in Section 10.2, the backpropagator would
simply read ‘Contrib’ instead.

$\textbf{D}^{5}_{c}[\mathbb{R}]=(\mathbb{R},({\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}\textrm{CID}},\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
())))$

The new rule for the parallel pairing construct is simply:

$\textbf{D}^{5}_{c}[t\star
s]=\textbf{D}^{5}_{c}[t]\star^{\prime}\textbf{D}^{5}_{c}[s]$

This is the only place where we use the operation $(\star^{\prime})$, and thus
the only place in the forward pass where we directly use the extended
functionality of our monad $\operatorname*{\mathcal{M}}$.

###  12.4 Updating the Wrapper

The wrapper is there to glue the various components together and to provide
the (now parallel) definition of SResolve. We discuss the main ideas behind
the parallel wrapper implementation in this section, skipping over some
implementation details.

In Section 9, the backpropagators staged their backpropagator calls in a
single array, indexed by their ID. With compound IDs, we still need one array
slot for each backpropagator, meaning that we need a _nested_ staging array:

$\operatorname*{\textrm{Staged}}c=(\operatorname*{\textrm{Array}}\underline{\mathbb{R}},\operatorname*{\textrm{Array}}{(\operatorname*{\textrm{Array}}{(\underline{\mathbb{R}}\multimap(\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
()),\underline{\mathbb{R}})})})$

where the outer array is indexed by the job ID and the inner by the sequential
ID within that job. Because jobs have differing lengths, the nested arrays
also have differing lengths and we cannot use a rectangular two-dimensional
array. The implementation of SOneHot must be changed to update the cotangent
collection array atomically, because backpropagators, and hence SOneHot, will
now be called from multiple threads. SCall needs a simple modification to
(atomically) modify the correct element in the now-nested staging array.

To construct the initial $\operatorname*{\textrm{Staged}}$ object, SAlloc
needs to know the correct length of all the nested staging arrays; it can get
this information from the job graph (in the $\mathrm{History}$ structure) that
Wrap receives from the monadic evaluation of the forward pass.

This leaves the parallel implementation of SResolve. The idea here is to have
two mutually recursive functions: one that resolves a _task_ (this one is new
and handles all parallelism), and one that resolves a _job_ (and is
essentially identical to the last version of SResolve, from Fig. 10). An
example implementation is given in Fig. 15, where the first function is called
$\mathit{resolveTask}$ and the second $\mathit{resolveJob}$. In this way,
SResolve traverses the job graph (Fig. 13(b)) from the terminal job backwards
to the initial job, doing the usual sequential resolving process for each
sequential job in the graph.

$\begin{array}[]{@{}l@{}}\textrm{SResolve}:\mathrm{JobDescr}\rightarrow\operatorname*{\textrm{Staged}}c\rightarrow\texttt{IO}\
(\operatorname*{\textrm{IArray}}\underline{\mathbb{R}})\\\ \textrm{SResolve}\
\mathit{jd}\ (c,m)=\mathit{resolveTask}\ \mathit{jd}\\\ \quad\mathrm{where}\
\begin{array}[t]{@{}l@{}}\mathit{resolveTask}\
(\mathit{history},\mathit{jid},i)=\mathbf{do}\\\
\quad\begin{array}[t]{@{}l@{}}\mathit{jobarr}\leftarrow\texttt{get}\
\mathit{jid}\ m\\\ \mathit{resolveJob}\ (i-1)\ \mathit{jobarr}\\\
\mathbf{case}\ \mathit{history}\ \mathbf{of}\
\begin{array}[t]{@{}l@{}}\mathrm{Start}\rightarrow\mathbf{return}\ ()\\\
\mathrm{Fork}\ \mathit{jd}_{0}\ \mathit{jd}_{1}\
\mathit{jd}_{2}\rightarrow\mathbf{do}\\\ \quad\texttt{inParallel}\
(\mathit{resolveTask}\ \mathit{jd}_{1})\ (\mathit{resolveTask}\
\mathit{jd}_{2})\\\ \quad\mathit{resolveTask}\ \mathit{jd}_{0}\end{array}\\\
\end{array}\\\ {}\\\\[-8.00003pt] \mathit{resolveJob}\ (-1)\
\mathit{arr}=\mathbf{return}\ ()\\\ \mathit{resolveJob}\ i\
\mathit{arr}=\mathbf{do}\\\
\quad\begin{array}[t]{@{}l@{}}(f,a)\leftarrow\texttt{get}\ i\ \mathit{arr}\\\
f\ a\ (c,m)\\\ \mathit{resolveJob}\ (i-1)\
\mathit{arr}\end{array}\end{array}\end{array}$ Figure 15:  Implementation of
SResolve for the parallel-ready dual-numbers reverse AD algorithm. The
inParallel function is as in Fig. 14.

#### Duality.

The usual mantra in reverse-mode AD is that “sharing in the primal becomes
addition in the dual”. When parallelism comes into play, not only do we have
this duality in the data flow graph, we also get an interesting duality in the
control-flow graph: SResolve forks where the primal joined, and joins where
the primal forked. Perhaps we can add a mantra for task-parallel reverse AD:
“forks in the primal become joins in the dual”.

##  13 Implementation

To show the practicality of our method, we provide a prototype
implementation272727The code is available at https://github.com/tomsmeding/ad-
dualrev-th. of the parallel algorithm of Section 12, together with the
improvements from Sections 10.1 and 10.2, that differentiates a sizeable
fragment of Haskell98 including recursive types (reinterpreted as a strict,
call-by-value language) using Template Haskell. As described in Section 9.2,
we realise the mutable arrays using IOVectors. The implementation does not
incorporate the changes given in Section 10.3 that transform the algorithm
into classical taping (because implementations of taping already abound), but
it does include support for recursive functions, coproduct types, and user-
defined data types as described in Section 11.

Template Haskell (Sheard and Jones, 2002) is a built-in metaprogramming
facility in GHC Haskell that (roughly) allows the programmer to write a
Haskell function that takes a block of user-written Haskell code, do whatever
it wants with the AST of that code, and finally splice the result back into
the user’s program. The resulting code is still type-checked as usual. The AST
transformation that we implement is, of course, differentiation.

#### Benchmarks.

To check that our implementation has reasonable performance in practice, we
benchmark (in bench/Main.hs) against Kmett’s Haskell ad library (Kmett and
contributors, 2021) (version 4.5.6) on a few basic functions. These functions
are the followign (abbreviating Double as ‘D’):

  * •

A single scalar multiplication of type (D, D) -> D;

  * •

Dot product of type ([D], [D]) -> D;

  * •

Matrix–vector multiplication, then sum: of type ([[D]], [D]) -> D;

  * •

The rotate_vec_by_quat example from (Krawiec et al., 2022) of type (Vec3 D,
Quaternion D) -> Vec3 D, with data Vec3 s = Vec3 s s s and data Quaternion s =
Quaternion s s s s;

  * •

Simulation of 4 particles in a simple force field with friction for 1000 time
steps; this example has type [((D, D), (D, D))] -> D. The input is a list (of
length 4) of initial positions and velocities; the output is
$\sum_{(x,y)}x\cdot y$ ranging over the 4 result positions $p$, to ensure that
the computation has a single scalar as output. The four particles are
simulated in parallel using the $(\star)$ combinator from Section 12.

The fourth test case, rotate_vec_by_quat, has a non-trivial return type; the
benchmark executes its reverse pass three times (‘3’ being the number of
scalars in the function output) to get the full Jacobian. The fifth test case,
‘particles’, is run on 1, 2 and 4 threads, where the ideal result would be
perfect scaling due to the four paricles being independent.

Table 1:  Benchmark results of Section 12 \+ Sections 10.1 and 10.2 versus ad-4.5.6. The ‘TH’ and ‘ad’ columns indicate runtimes on one machine for our implementation and the ad library, respectively. The last column shows the ratio between the previous two columns. We give the size of the largest side of criterion’s 95% bootstrapping confidence interval, rounded to 2 decimal digits. Setup: GHC 9.6.4 on Linux, Intel i9-10900K CPU, with Intel Turbo Boost disabled (i.e. running at a consistent 3.7 GHz).  | TH | ad | TH / ad
---|---|---|---
scalar mult. | 0 | .34 |  $\mu$s $\pm 0.00$ | 0 | .90 |  $\mu$s $\pm 0.00$ |  $\approx$0.4
dot product | 0 | .91 |  $\mu$s $\pm 0.03$ | 2 | .41 |  $\mu$s $\pm 0.08$ |  $\approx$0.4
sum-mat-vec | 0 | .59 |  $\mu$s $\pm 0.02$ | 1 | .56 |  $\mu$s $\pm 0.07$ |  $\approx$0.4
rotate_vec_by_quat | 5 | .69 |  $\mu$s $\pm 0.00$ | 8 | .11 |  $\mu$s $\pm 0.01$ |  $\approx$0.7
particles (1 thr.) | 8 | .3 | ms $\pm 0.20$ | 15 | .41 | ms $\pm 0.16$ |  $\approx$0.5
particles (2 thr.) | 5 | .63 | ms $\pm 0.06$ | — | —
particles (4 thr.) | 2 | .66 | ms $\pm 0.18$ | — | —

The benchmark results are summarised in Table 1 and shown in detail (including
evidence of the desired linear complexity scaling) in LABEL:sec:criterion-
benchmark-details. The benchmarks are measured using the criterion282828By
Bryan O’Sullivan: https://hackage.haskell.org/package/criterion library. To
get statistically significant results, we measure how the timings scale with
increasing $n$:

  * •

Scalar multiplication, rotate_vec_by_quat and the particle simulation are
simply differentiated $n$ times;

  * •

Dot product is performed on lists of length $n$;

  * •

Matrix multiplication is done for a matrix and vector of size $\sqrt{n}$, to
get linear scaling in $n$.

The reported time is the deduced actual runtime for $n=1$.

By the results in Table 1, we see that on these simple benchmark programs, our
implementation is faster than the existing ad library. While this is
encouraging, it is not overly surprising: because our algorithm is implemented
as a compile-time code transformation, the compiler is able to optimise the
derivative code before it gets executed. And while the programs are
monomorphic in our benchmark, meaning that the type-class-driven approach of
ad does not have the disadvantage of runtime polymorphism, its implementation
is nevertheless generic and harder to optimise for the compiler.

Of course, performance results may well be different for other programs, and
AD implementations that have native support for array operations can handle
some of these programs much more efficiently. Furthermore, there are various
implementation choices for the code transformation that may seem relatively
innocuous but have a large impact on performance (we give some more details
below in Section 13.1).

For these reasons, our goal here is merely to substantiate the claim that the
implementation exhibits constant-factor performance in the right ballpark (in
addition to it having the right asymptotic complexity, as we have argued).
Nevertheless, our benchmarks include key components of many AD applications,
such as neural nets, and seeing that we have not at all special-cased their
implementation (the implementation knows only primitive arithmetic operations
such as scalar addition, multiplication, etc.), we believe that they suffice
to demonstrate our limited claim.

###  13.1 Considerations for Implementation Performance

The target language of the code transformation in our implementation is
Haskell, which is a lazy, garbage-collected language. This has various effects
on the performance characteristics of our implementation.

#### Garbage collection.

The graph of backpropagators (a normal data structure, ‘Contrib’, since
Section 10.2) is a big data structure of size proportional to the number of
scalar operations performed in the source program. While this data structure
grows during the forward pass, the nursery (zeroth generation) of GHC’s
generational garbage collector (GC) repeatedly fills up, triggering garbage
collection of the heap. Because a GC pass takes time on the order of the
amount of live data on the heap, these passes end up very expensive: a naive
taping reverse AD algorithm becomes _quadratic_ in the source program runtime.
Using a GHC runtime system flag (e.g. -A500m) to set the nursery size of GHC’s
GC sufficiently large to prevent the GC from running during a benchmark
(criterion runs the GC explicitly between each benchmark invocation), timings
on some benchmarks above decrease significantly: on the largest benchmark
(particles), this saves more than 35% off ad’s runtime and more than 25% off
our runtime. The times reported in Table 1 are with GHC’s default GC settings.

While this is technically a complexity problem in our implementation, we gloss
over this because it is not fundamental to the algorithm: the backpropagator
graph does not contain cycles, so it could be tracked with reference-counting
GC to immediately eliminate the quadratic blowup. Using a custom, manual
allocator, one could also eliminate all tracking of the liveness of the graph
because we know from the structure of the algorithm exactly when we can free a
node in the graph (namely when we have visited it in the reverse pass). Our
reference implementation does not do these things to be able to focus on the
workings of the algorithm.

#### Laziness.

Because data types are lazy by default in Haskell, a naive encoding of the
Contrib data type from Section 10.2 would make the whole graph lazily
evaluated (because it is never demanded during the forward pass). This results
in a significant constant-factor overhead (more than $2\times$), and
furthermore means that part of the work of the forward pass happens when the
reverse pass first touches the top-level backpropagator; this work then
happens sequentially, even if the forward pass was meant to be parallel. To
achieve proper parallel scaling, it was necessary to make the Contrib graph
strict, and furthermore to make the
$\textbf{D}_{c}[\mathbb{R}]=(\mathbb{R},(\textrm{CID},\textrm{Contrib}))$
triples strict, to ensure that the graph is fully evaluated as it is being
_constructed_ , not when it is demanded in the reverse pass.

Using some well-chosen `{-# UNPACK #-}` pragmas on some of these fields also
had a significant positive effect on performance.

#### Thread pool.

In Section 12 we used an abstract inParallel operation for running two jobs in
parallel, assuming some underlying thread pool for efficient evaluation of the
resulting parallel job graph. In a standard thread pool implementation,
spawning tasks from within tasks can result in deadlocks. Because nested tasks
are essential to our model of task parallelism, the implementer has to take
care to further augment $\operatorname*{\mathcal{M}}$ from Section 12 to be a
continuation monad as well: this allows the continuation of the inParallel
operation to be captured and scheduled separately in the thread pool, so that
thread pool jobs are indeed individual _jobs_ as defined in Section 12.

The GHC runtime system has a thread scheduler that should handle this
completely transparently, but in our tests it was not eager enough in
assigning virtual Haskell threads to actual separate kernel threads, resulting
in a sequential benchmark. This motivated a (small) custom thread pool
implementation that sufficed for our benchmarks, but has a significant amount
of overhead that can be optimised with more engineering effort.

#### Imperfect scaling.

Despite efforts to the contrary, it is evident from Table 1 that even on an
embarrassingly parallel task like ‘particles’, the implementation does not
scale perfectly. From inspection of more granular timing of our
implementation, we suspect that this is a combination of thread pool overhead
and the fact that the forward pass simply allocates too quickly, exhausting
the memory bandwidth of our system when run sufficiently parallel. However,
more research is needed here to uncover the true bottlenecks.

##  14 Conclusions

One may ask: if the final algorithm from Section 9 can be argued to be “just
taping” (Section 10.3), which is already widely used in practice—and the
parallel extension is _still_ just taping, except on a non-linear tape—what
was the point? The point is the observation that optimisations are key to
implementing efficient AD and that multiple kinds of reverse AD algorithms
proposed by the programming languages community (in particular the one from
Fig. 6, studied by Brunel et al. (2020) and (Huot et al., 2020, Section 6)—for
further examples, see below in Section 15.2) tend to all reduce to taping
after optimisation. We hope to have demonstrated that these techniques are
quite flexible, allowing the differentiation of rich source languages with
dynamic control flow, recursion and parallelism, and that the resulting
algorithm can be relatively straightforward by starting from a naive
differentiation algorithm and next optimising it to achieve the desired
complexity.

The first of our optimisations (linear factoring) is quite specific to
starting AD algorithms that need some kind of distributive law to become
efficient (e.g. also (Krawiec et al., 2022)). However, we think that the other
optimisations are more widely applicable (and are, for example, also related
to the optimisations that fix the time complexity of CHAD (Smeding and Vákár,
2024)): sparse vectors are probably needed in most functional reverse AD
algorithms to efficiently represent the one-hot vectors resulting from
projections (fst/snd as well as random access into arrays, through indexing),
and mutable arrays are a standard solution to remove the ubiquitous log-factor
in the complexity of purely functional algorithms.

If one desires to take the techniques in this paper further to an algorithm
that is useful in practice, it will be necessary to add _arrays_ of scalars as
a primitive type in the transformation. This would allow us to significantly
reduce the size of the allocated tape, and reuse much more of the structure of
the original program in the reverse pass. Since many AD applications tend to
involve very large arrays of scalars, we expect to be able to gain significant
constant factor performance by replacing an array of many scalar
backpropagators with a single vector backpropagator. We are currently
exploring this direction of research.

##  15 Origins of Dual-Numbers Reverse AD, Relationship With Vectorised
Forward AD and Other Related Work

The literature about automatic differentiation spans many decades and academic
subcommunities (scientific computing, machine learning and—most
recently—programming languages). Important early references are (Wengert,
1964; Linnainmaa, 1970; Speelpenning, 1980). Good surveys can be found in
(Baydin et al., 2017; Margossian, 2019). In the rest of this section, we focus
on the more recent literature that studies AD from a programming languages
(PL) point of view, to extend the scope of our discussion in Section 10.

###  15.1 Theoretical Foundations for Our Algorithm

The first mention that we know of the naive dual-numbers reverse mode AD
algorithm that we analyse in this paper is (Pearlmutter and Siskind, 2008,
page 12), where it is quickly dismissed before a different technique is
pursued. The algorithm is first thoroughly studied by Brunel et al. (2020)
using operational semantics and in (Huot et al., 2020, Section 6) using
denotational semantics. Brunel et al. (2020) introduce the key idea that
underlies the optimisations in our paper: the linear factoring rule, stating
that a term $f\ x+f\ y$, with $f$ a linear function, may be reduced to $f\
(x+y)$. We build on their use of this rule as a tool in a complexity proof to
make it a suitable basis for a performant implementation.

Mazza and Pagani (2021) extend the work of Brunel et al. (2020) to apply to a
language with term recursion, showing that dual-numbers reverse AD on PCF is
almost everywhere correct. Similarly, Nunes and Vákár (2022, 2023) extend the
work of Huot et al. (2020) to apply to partial programs involving iteration,
recursion and recursive types, thus giving a correctness proof for the initial
dual-numbers reverse AD transformation of Fig. 6 applied to, essentially,
idealised Haskell98.

###  15.2 Vectorised Forward AD

There are strong parallels between our optimisations to the sequential
algorithm in Sections 5, 6, 7, 8 and 9 and the derivation by Krawiec et al.
(2022). Like the present paper, they give a sequence of steps that optimise a
simple algorithm to an efficient implementation—but the starting algorithm is
vectorised forward AD (vfad) instead of backpropagator-based dual-numbers
reverse AD (dnrad). In our notation, their initial type transformation does
not have
$\textbf{D}^{1}_{c}[\mathbb{R}]=(\mathbb{R},\underline{\mathbb{R}}\multimap
c)$, but instead $\textbf{D}^{1}_{c}[\mathbb{R}]=(\mathbb{R},c)$. (As befits a
dual-numbers algorithm, the rest of the type transformation is simply
structurally recursive.)

Linear algebra tells us that the vector spaces
$\underline{\mathbb{R}}\multimap c$ and $c$ are isomorphic, and indeed
inspection of the term transformations shows that both naive algorithms
compute the same thing. Their operational behaviour, on the other hand, is
very different: the complexity problem with dnrad is exponential blowup in the
presence of sharing, whereas vfad is “simply” $n$ times too slow, where $n$ is
the number of scalars in the input.

But the first optimisation on vfad, which defunctionalises the zero, one-hot,
addition and scaling operations on the $c$ tangent vector, introduces the same
sharing-induced complexity problem as we have in naive dnrad as payment for
fixing the factor-$n$ overhead. The two algorithms are now on equal footing:
we could defunctionalise the backpropagators in dnrad just as easily (and
indeed we do so, albeit later in Section 10.2).

Afterwards, vfad is lifted to a combination (stack) of an ID generation monad
and a Writer monad. Each scalar result of a primitive operation gets an ID,
and the Writer monad records for each ID (hence, scalar in the program) its
defunctionalised tangent vector (i.e. an expression) in terms of other
already-computed tangent vectors from the Writer record. These expressions
correspond to our primitive operation backpropagators with calls replaced with
SCall: where we replace calls with ID-tagged pairs of function and argument,
vfad replaces the usage of already-computed tangent vectors with scaled
references to the IDs of those vectors. The choice in our SResolve of
evaluation order from highest ID to lowest ID (Section 6) is encoded in vfad’s
definitions of runDelta and eval, which process the record back-to-front.

Finally, our Cayley transform is encoded in the type of vfad’s eval function,
which interprets the defunctionalised operations on tangent vectors (including
explicit sharing using the Writer log) into an actual tangent vector—the final
gradient: its gradient return type is Cayley-transformed. Our final
optimisation to mutable arrays to eliminate log-factors in the complexity is
also mirrored in vfad.

#### Distributive law.

Under the isomorphism $\underline{\mathbb{R}}\multimap c\cong c$, the type
$\operatorname*{\textrm{Staged}}c$ can be thought of as a type
$\operatorname*{\textrm{Expr}}c$ of ASTs of expressions (with sharing) of type
$c$.292929In (Krawiec et al., 2022), $\operatorname*{\textrm{Expr}}c$ is
called Delta. The linear factoring rule $f\ x+f\ y\leadsto f\ (x+y)$ for a
linear function $f:\underline{\mathbb{R}}\multimap c$ that rescales a vector
$v:c$ with a scalar then corresponds to the distributive law $v\cdot x+v\cdot
y\leadsto v\cdot(x+y)$. This highlights the relationship between our work and
that of Shaikhha et al. (2019), who try to (statically) optimise vectorised
forward AD to reverse AD using precisely this distributive law. A key
distinction is that we apply this law (in the form of the linear factoring
rule) at runtime rather than compile time, allowing us to always achieve the
complexity of reverse AD, rather than merely on some specific programs with
straightforward control and data flow. The price we pay for this generality is
a runtime overhead, similar to the usual overhead of taping.

###  15.3 Other PL Literature About AD

#### CHAD and category theory inspired AD.

Rather than interleaving backpropagators by pairing them with scalars in a
type, we can also try to directly implement reverse AD as a structurally
recursive code transformation that does not need a (de)interleaving wrapper:
$\mathcal{R}_{2}$ from Section 4. This is the approach taken by Elliott
(2018). It pairs vectors (and values of other composite types) with a single
composite backpropagator, rather than decomposing to the point where each
scalar is paired with a mini-backpropagator like in our dual-numbers approach.
The resulting algorithm is extended to source languages with function types in
(Vákár, 2021; Vákár and Smeding, 2022; Vytiniotis et al., 2019) and to sum and
(co)inductive types in (Nunes and Vákár, 2023). Smeding and Vákár (2024) give
a mechanised proof that the resulting algorithm attains the correct
computational complexity, after a series of optimisations that are similar to
the ones considered in this paper. Like our dual-numbers reverse AD approach,
the algorithm arises as a canonical structure-preserving functor on the syntax
of a programming language. However, due to a different choice in target
category (a Grothendieck construction of a linear $\lambda$-calculus for CHAD
rather than the syntax of a plain $\lambda$-calculus for dual-numbers AD), the
resulting algorithm looks very different.

#### Approaches utilising non-local control flow.

Another family of approaches to AD recently taken by the PL community contains
those that rely on forms of non-local control flow such as delimited
continuations (Wang and Rompf, 2018) or effect handlers (Sigal, 2021; de
Vilhena and Pottier, 2023). These techniques are different in the sense that
they generate code that is not purely functional. This use of non-local
control flow makes it possible to achieve an efficient implementation of
reverse AD that looks strikingly simple compared to alternative approaches.
Where the CHAD approaches and our dual-numbers reverse AD approach both have
to manually invert control flow at compile time by making use of continuations
that get passed around, combined with smart staging of execution of those
continuations in this paper, this inversion of control can be deferred to run
time by clever use of delimited control operators or effect handlers. de
Vilhena and Pottier (2023) give a mechanised proof that the resulting (rather
subtle) code transformation is correct.

#### AD of parallel code.

Bischof et al. (1991); Bischof (1991) present some of the first work in
parallel AD. Rather than starting with a source program (with potential
dynamic control flow) that has (fork-join) parallelism like we do and
mirroring that in the reverse pass of the algorithm, they focus on code
without dynamic control flow and analyse the dependency structure of the
reverse pass at compile time to automatically parallelise it. This approach is
developed further by Bucker et al. (2002).

Building on the classic imperative AD tool Tapenade (Hascoët and Pascual,
2013), Hückelheim and Hascoët (2022) discuss a method for differentiating
parallel for-loops (with shared memory) in a parallelism-preserving way.
Industrial machine learning frameworks such as TensorFlow (Abadi et al.,
2016), JAX (Bradbury et al., 2018) and PyTorch (Paszke et al., 2017) focus on
data parallelism through parallel array operations—typically first-order ones.

Kaler et al. (2021) focus on AD of programs with similar forms of fork-join
parallelism as we consider in this paper. Their implementation builds on the
Adept C++ library, which implements an AD algorithm that is very different
from the one discussed in this paper. A unique feature of their work is that
they give a formal analysis of the complexity of the technique and give bounds
for both the work and the span of the resulting derivative code; it is
possible that ideas from the cited work can be used to improve and bound the
span of our parallel algorithm.

Other recent work has focussed on developing data-parallel derivatives for
data-parallel array programs (Paszke et al., 2021b; Schenck et al., 2022;
Paszke et al., 2021a). These methods are orthogonal to the ideas focussed on
task parallelism that we present in the present paper.

In recent work, we have shown that the an optimise version of CHAD seems to
preserve task and data parallelism (Smeding and Vákár, 2024). We are working
on giving formal guarantees about the extent (work and span) to which
different forms of parallelism is preserved. Compared to the dual numbers
reverse AD method of the present paper, CHAD seems to accommodate higher order
array combinators more easily, including their data parallel versions. This
makes us hopeful for the use of CHAD as a candidate algorithm for parallel AD.

#### Conflicts of interest.

None.

## Appendix

On the following page, we show benchmark details from our implementation on
some sequential programs, as determined by the Haskell library criterion. See
Section 13.

See criterion-report.pdf

See pages 2- of criterion-report.pdf

## References

  * Abadi et al. (2016) Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., Kudlur, M., Levenberg, J., Monga, R., Moore, S., Murray, D. G., Steiner, B., Tucker, P. A., Vasudevan, V., Warden, P., Wicke, M., Yu, Y. & Zheng, X. (2016) Tensorflow: A system for large-scale machine learning. 12th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2016, Savannah, GA, USA, November 2-4, 2016\. USENIX Association. pp. 265–283.
  * Abadi and Plotkin (2020) Abadi, M. & Plotkin, G. D. (2020) A simple differentiable programming language. Proc. ACM Program. Lang. 4(POPL), 38:1–38:28.
  * Baydin et al. (2017) Baydin, A. G., Pearlmutter, B. A., Radul, A. A. & Siskind, J. M. (2017) Automatic differentiation in machine learning: a survey. J. Mach. Learn. Res. 18, 153:1–153:43.
  * Bernardy et al. (2018) Bernardy, J., Boespflug, M., Newton, R. R., Jones, S. P. & Spiwack, A. (2018) Linear haskell: practical linearity in a higher-order polymorphic language. Proc. ACM Program. Lang. 2(POPL), 5:1–5:29.
  * Bischof (1991) Bischof, C. (1991) Issues in parallel automatic differentiation. Argonne National Lab.
  * Bischof et al. (1991) Bischof, C., Griewank, A. & Juedes, D. (1991) Exploiting parallelism in automatic differentiation. Proceedings of the 5th international conference on Supercomputing. pp. 146–153.
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S. & Zhang, Q. (2018) JAX: composable transformations of Python+NumPy programs .
  * Brunel et al. (2020) Brunel, A., Mazza, D. & Pagani, M. (2020) Backpropagation in the simply typed lambda-calculus with linear negation. Proc. ACM Program. Lang. 4(POPL), 64:1–64:27.
  * Bucker et al. (2002) Bucker, H., Lang, B., Rasch, A., Bischof, C. H. & an Mey, D. (2002) Explicit loop scheduling in openmp for parallel automatic differentiation. Proceedings 16th Annual International Symposium on High Performance Computing Systems and Applications. IEEE. pp. 121–126.
  * de Vilhena and Pottier (2023) de Vilhena, P. E. & Pottier, F. (2023) Verifying an effect-handler-based define-by-run reverse-mode ad library. Logical Methods in Computer Science. 19.
  * Elliott (2018) Elliott, C. (2018) The simple essence of automatic differentiation. Proc. ACM Program. Lang. 2(ICFP), 70:1–70:29.
  * Griewank and Walther (2008) Griewank, A. & Walther, A. (2008) Evaluating derivatives - principles and techniques of algorithmic differentiation, Second Edition. SIAM.
  * Hascoët et al. (2005) Hascoët, L., Naumann, U. & Pascual, V. (2005) ”to be recorded” analysis in reverse-mode automatic differentiation. Future Gener. Comput. Syst. 21(8), 1401–1417.
  * Hascoët and Pascual (2013) Hascoët, L. & Pascual, V. (2013) The Tapenade automatic differentiation tool: Principles, model, and specification. ACM Trans. Math. Softw. 39(3), 20:1–20:43.
  * Hückelheim and Hascoët (2022) Hückelheim, J. & Hascoët, L. (2022) Source-to-source automatic differentiation of openmp parallel loops. ACM Transactions on Mathematical Software (TOMS). 48(1), 1–32.
  * Hughes (1986) Hughes, R. J. M. (1986) A novel representation of lists and its application to the function ”reverse”. Inf. Process. Lett. 22(3), 141–144.
  * Huot et al. (2020) Huot, M., Staton, S. & Vákár, M. (2020) Correctness of automatic differentiation via diffeologies and categorical gluing. Foundations of Software Science and Computation Structures - 23rd International Conference, FOSSACS 2020, Held as Part of the European Joint Conferences on Theory and Practice of Software, ETAPS 2020, Dublin, Ireland, April 25-30, 2020, Proceedings. Springer. pp. 319–338.
  * Kaler et al. (2021) Kaler, T., Schardl, T. B., Xie, B., Leiserson, C. E., Chen, J., Pareja, A. & Kollias, G. (2021) Parad: A work-efficient parallel algorithm for reverse-mode automatic differentiation. Symposium on Algorithmic Principles of Computer Systems (APOCS). SIAM. pp. 144–158.
  * Kmett and contributors (2021) Kmett, E. & contributors. (2021) ad: Automatic differentiation (Haskell library on Hackage).
  * Krawiec et al. (2022) Krawiec, F., Jones, S. P., Krishnaswami, N., Ellis, T., Eisenberg, R. A. & Fitzgibbon, A. W. (2022) Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation. Proc. ACM Program. Lang. 6(POPL), 1–30.
  * Launchbury and Jones (1994) Launchbury, J. & Jones, S. L. P. (1994) Lazy functional state threads. Proceedings of the ACM SIGPLAN’94 Conference on Programming Language Design and Implementation (PLDI), Orlando, Florida, USA, June 20-24, 1994. ACM. pp. 24–35.
  * Linnainmaa (1970) Linnainmaa, S. (1970) The representation of the cumulative rounding error of an algorithm as a taylor expansion of the local rounding errors. Master’s Thesis (in Finnish), Univ. Helsinki.
  * Margossian (2019) Margossian, C. C. (2019) A review of automatic differentiation and its efficient implementation. Wiley Interdiscip. Rev. Data Min. Knowl. Discov. 9(4).
  * Mazza and Pagani (2021) Mazza, D. & Pagani, M. (2021) Automatic differentiation in PCF. Proc. ACM Program. Lang. 5(POPL), 1–27.
  * Nunes and Vákár (2022) Nunes, F. L. & Vákár, M. (2022) Logical Relations for Partial Features and Automatic Differentiation Correctness. CoRR. abs/2210.08530.
  * Nunes and Vákár (2023) Nunes, F. L. & Vákár, M. (2023) Automatic Differentiation for ML-family languages: correctness via logical relations. CoRR. abs/2210.07724.
  * Nunes and Vákár (2023) Nunes, F. L. & Vákár, M. (2023) CHAD for expressive total languages. Math. Struct. Comput. Sci. 33(4-5), 311–426.
  * Paszke et al. (2017) Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L. & Lerer, A. (2017) Automatic differentiation in PyTorch. NIPS 2017 Autodiff Workshop: The future of gradient-based machine learning software and techniques. Red Hook, NY, USA. Curran Associates, Inc.
  * Paszke et al. (2021a) Paszke, A., Johnson, D. D., Duvenaud, D., Vytiniotis, D., Radul, A., Johnson, M. J., Ragan-Kelley, J. & Maclaurin, D. (2021a) Getting to the point: index sets and parallelism-preserving autodiff for pointful array programming. Proc. ACM Program. Lang. 5(ICFP), 1–29.
  * Paszke et al. (2021b) Paszke, A., Johnson, M. J., Frostig, R. & Maclaurin, D. (2021b) Parallelism-preserving automatic differentiation for second-order array languages. Proceedings of the 9th ACM SIGPLAN International Workshop on Functional High-Performance and Numerical Computing. pp. 13–23.
  * Pearlmutter and Siskind (2008) Pearlmutter, B. A. & Siskind, J. M. (2008) Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator. ACM Trans. Program. Lang. Syst. 30(2), 7:1–7:36.
  * Reynolds (1998) Reynolds, J. C. (1998) Definitional interpreters for higher-order programming languages. High. Order Symb. Comput. 11(4), 363–397.
  * Schenck et al. (2022) Schenck, R., Rønning, O., Henriksen, T. & Oancea, C. E. (2022) AD for an array language with nested parallelism. SC22: International Conference for High Performance Computing, Networking, Storage and Analysis, Dallas, TX, USA, November 13-18, 2022. IEEE. pp. 58:1–58:15.
  * Shaikhha et al. (2019) Shaikhha, A., Fitzgibbon, A., Vytiniotis, D. & Jones, S. P. (2019) Efficient differentiable programming in a functional array-processing language. Proc. ACM Program. Lang. 3(ICFP), 97:1–97:30.
  * Sheard and Jones (2002) Sheard, T. & Jones, S. L. P. (2002) Template meta-programming for haskell. ACM SIGPLAN Notices. 37(12), 60–75.
  * Sigal (2021) Sigal, J. (2021) Automatic differentiation via effects and handlers: An implementation in frank. arXiv preprint arXiv:2101.08095.
  * Siskind and Pearlmutter (2018) Siskind, J. M. & Pearlmutter, B. A. (2018) Divide-and-conquer checkpointing for arbitrary programs with no user annotation. Optimization Methods and Software. 33(4-6), 1288–1330.
  * Smeding and Vákár (2022) Smeding, T. & Vákár, M. (2022) Efficient dual-numbers reverse AD via well-known program transformations. CoRR. abs/2207.03418v2.
  * Smeding and Vákár (2023) Smeding, T. & Vákár, M. (2023) Efficient dual-numbers reverse AD via well-known program transformations. Proc. ACM Program. Lang. 7(POPL), 1573–1600.
  * Smeding and Vákár (2024) Smeding, T. & Vákár, M. (2024) Efficient CHAD. Proc. ACM Program. Lang. 8(POPL), 1060–1088.
  * Speelpenning (1980) Speelpenning, B. (1980) Compiling fast partial derivatives of functions given by algorithms. Technical report. Illinois University.
  * Vákár (2021) Vákár, M. (2021) Reverse AD at higher types: Pure, principled and denotationally correct. Programming Languages and Systems. Springer. pp. 607–634.
  * Vákár and Smeding (2022) Vákár, M. & Smeding, T. (2022) CHAD: combinatory homomorphic automatic differentiation. pp. 20:1–20:49.
  * Vytiniotis et al. (2019) Vytiniotis, D., Belov, D., Wei, R., Plotkin, G. & Abadi, M. (2019) The differentiable curry. NeurIPS Workshop on Program Transformations.
  * Wang and Rompf (2018) Wang, F. & Rompf, T. (2018) A language and compiler view on differentiable programming. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Workshop Track Proceedings. OpenReview.net.
  * Wengert (1964) Wengert, R. E. (1964) A simple automatic derivative evaluation program. Commun. ACM. 7(8), 463–464.
  * Westrick et al. (2024) Westrick, S., Fluet, M., Rainey, M. & Acar, U. A. (2024) Automatic parallelism management. Proc. ACM Program. Lang. 8(POPL), 1118–1149.