Bring back store!, accumulate, and accumulate!?

[nickrobinson251]

We removed store! in #110, but one day we might want something like that back.

Paraphrasing @willtebbutt's comment

Just as it's useful to have accumulate!(::Array, ::Array) (i.e. add!) to avoid accumulating differentials into a third array, it maybe useful to have store! for storing the result into a pre-allocated array. E.g. If you're a reverse-mode AD like ReverseDiff, and store! would allow you the option to pre-allocate all of the memory required for the forwards- and reverse- passes over a programme.

---

for ease of reference, the old code was

"""
    store!(Δ, ∂)
Stores `∂`, in `Δ`, overwriting what ever was in `Δ` before.
potentially avoiding intermediate temporary allocations that might be
necessary for alternative approaches  (e.g. `copyto!(Δ, extern(∂))`)
Like [`accumulate`](@ref) and [`accumulate!`](@ref), this function is intended
to be customizable for specific rules/input types.
See also: [`accumulate`](@ref), [`accumulate!`](@ref)
"""
store!(Δ, ∂) = materialize!(Δ, broadcastable(∂))

[nickrobinson251]

actually #110 ends up removing accumulate!(x, y) in favour of x .+= y, but the above point may still stand

[willtebbutt]

Yeah. x .+= y doesn't seem quite as general as an add! / accumulate! function -- we shouldn't need to have a notion of broadcasting in general I don't think. (Obviously fine in your case, just not a substitute in general)

[oxinabox]

I supect we want:

• accumulate remove we have +
• accumulate! replace with add! or something but keep behavour broadly the same
• store! replace with copyto! but keep behavour broadly the same

Also we should wait for mutation to be happening as that affects this signficantly.

I am happy enough for now to let the AD solve this.
(They have the InplaceThunk still)

And then if we find common patterns look at abstracting them

[nickrobinson251]

xref #62 #24 #11

[oxinabox]

Have thought more about it

for accucumulate! we want to support maybe in-place maybe not.
The core semeantics we want is

accumulate!(x, y) = x + y
accumulate!(x::AbstractZero, y::InplaceThunk) = unthunk(y)
accumulate!(x, y::InplaceThunk) = (y.add!(x); x)
accumulate!(x::InplaceThunk, y) = accumulate!(y, x)
accumulate!(x::AbstractThunk, y::InplaceThunk) = accumulate!(unthunk(x), y)

the extension to make sure it does the right thing, and is fast for multiple arguments is:

@generated function accumulate!(args...)
    # need to make sure optimized multiarg + gets hit for arrays etc.
    # so need to group operations so we get `add!(add!(add!(n1+n2+n3, i1), i2), i3)`
    out_of_place_elements = Expr[]
    out_of_place_op = Expr(:call)
    inplace_op = out_of_place_op  # initially we will just have the out of place sum
    for (ind, argT) in enumerate(args)
        element = :(@inbounds arg[ind])
        if argT == InplaceThunk
            inplace_op = :(accumulate!($inplace_op, $element))
        else
            push!(out_of_place_elements, element)
        end
    end
    
    if isempty(out_of_place_args)
        #it is empty, so nothing to accumulate! against
        out_of_place_op.args = [:Zero]
    else # not empty so we have things to sum
        out_of_place_op.args = [:+; out_of_place_elements]
    end
    return inplace_op
end

basically we need to first do all the out of place accumulation in a single call to + because that will hit efficient methods for things like arrays, that will not be hit if we do +(a, +(b, c)), I checked:

julia> const z = rand(1024, 1024);

julia> @btime z + z + z + z;
  1.305 ms (2 allocations: 8.00 MiB)

julia> @btime +(+(+(z, z), z), z);
  2.928 ms (6 allocations: 24.00 MiB)

The fact that this is already fast is why we don't need to do special things with broadcast ourself for handling making sure multiple arrays are fast.
Then once we have summed up all the out-of-place elements we need to add! in all the inplace elements.
This we can do with nested calls to accumulate!.
We will handle pure out of place naturally by not wrapping it in any accumulate!s,
and we handle pure inplace by starting from Zero() if there are no out of place elements.

Above code is not tested.

We want this for Zygote to start using when it starts using JuliaDiff/ChainRules.jl#240

[oxinabox]

maybe accumulate! could also be overloading for

accumulate!(x::Array, y::Array) = x .+= y

That seems like it might start to get complicated particularly with the multi-arg versions
Not for the initial PR that adds this, i guess