Avoid device-to-host copy in ∇getindex!

[pxl-th]

• Implement kernel for accumulation in ∇getindex! for generic index types.
  Fallback to old implementation if KernelAbstractions backend does not support atomics (currently this is Metal.jl).
  This is implemented as an extension which is triggered by Atomix, KernelAbstractions, GPUArrays (which every GPU backend has as dependencies). This way we won't need to manually specify every GPU backend.

Closes #800.

Benchmarking:

using AMDGPU
using BenchmarkTools
using Zygote
using Random

function main()
    x = ROCArray(zeros(Float32, 8192, 1, 32))
    ids = randperm(8192)
    Δ = ROCArray(ones(Float32, 1))

    y, back = Zygote.pullback(x) do x
        xd = x[ids, :, :]
        sum(xd; dims=(1:ndims(xd)...,))
    end
    back(Δ)
    AMDGPU.synchronize()

    @btime AMDGPU.@sync $back($Δ)
    return
end

Before:

julia> main()
  1.528 ms (69 allocations: 2.00 MiB)

Now:

julia> main()
  102.913 μs (223 allocations: 7.84 KiB)

[oxinabox]

Can this package extension live in GPUArrays or one of the other packages this is a extension on?

It;s out general policy not to have rules that are applicable to specific packages in ChainRules.jl but rather just have the rules for Base and StdLibs.
Its beyond our capacity to maintain rules for the whole ecosystem.
Generally the authors of the packages themselves are best placed to do that, since the rules often depend on the fine details of the packages in question

This case is a little less clear since it overloading ChainRules.∇getindex! rather than ChainRulesCore.rrule(::typeof(getindex), ...)
but perhaps it can be written in as a normal rrule or two.

[pxl-th]

Would it be fine to remove ∇getindex!(dx::AbstractGPUArray, dy, inds...) from ChainRules then?
Otherwise it'd break precompilation I think, if another package overwrites it.

[pxl-th]

Is the problem that it uses other packages and adds them as weak dependencies?
Because the ∇getindex!(dx::AbstractGPUArray, dy, inds...) method already existed in ChainRules.jl, but it was not optimal w.r.t. performance.
This PR just adds a more efficient path.

[oxinabox]

That is a fair point, we did apparently already make an exception for AbstractGPUArrays for some reason I can not recall

[ToucheSir]

Could this Gordion Knot be sliced by creating an API in CRC which both ChainRules and GPU array packages could use and overload respectively? Maybe something adjacent to the current accumulation functionality (add!! and co). I feel like array indexing is sufficiently fundamental that it wouldn't be too outlandish of an API to expose.

[pxl-th]

Could this Gordion Knot be sliced by creating an API in CRC which both ChainRules and GPU array packages could use and overload respectively? Maybe something adjacent to the current accumulation functionality (add!! and co).

Since this is basically scatter with +, we can add something like this in CRC which is defined only for AbstractArray.
Then in an extension for GPUArrays or KernelAbstractions package (which every GPU backend has as a dependency) we'd have GPU-specific version for AbstractGPUArray. It does not make sense to define this separately in CUDA.jl, AMDGPU.jl, etc.

BTW, I'm not entirely sure if it is OK that an extension for KernelAbstractions extends methods for CRC?