# Increment elements of array by index

**URL:** <https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694>\
**Category:** GPU\
**Tags:** tullio\
**Created:** [November 6, 2020, 6:00pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694 "2020-11-06T18:00:23Z")\
**Posts on this page:** 11\
**Page:** 1

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 6, 2020, 6:00pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/1 "2020-11-06T18:00:23Z")

</div>

Given an array `A`, indices `I` and values `v`, I want to add elements of `v` to `A` at `I`. Something like this:

```julia
A = zeros(4)
I = [1, 3]
v = [1, 1]

# partial solution
A[I] .+= v

```

Now the hard part:

1. Indices may repeat, e.g. `I = [1, 3, 1]`. In this case `A[1]` should be updated with values in both - `v[1]` and `v[3]`.
2. Code should work on CuArrays with `allowscalar(false)`.

Best ideas I came with all operate on iterate over indices and thus don’t work on CuArrays. Is there a workaround?

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [November 6, 2020, 6:05pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/2 "2020-11-06T18:05:38Z")

</div>

I am having a lot of trouble understanding what you mean here:

> [@dfdx](#):
>
> Indices may repeat, e.g. `I = [1, 3, 1]` . In this case `A[1]` should be updated with values in both - `v[1]` and `v[3]` .

Could you clarify?

Perhaps you can show your version that iterates over indices that will make it easier to help find a solution that is GPU friendly. Your current ‘partial solution’ errors, so I don’t really get what you’re trying to do.

Further questions:

- Does `I` need to be a mutable array? Can it be a `CartesianIndex`?

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 6, 2020, 6:22pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/3 "2020-11-06T18:22:44Z")

</div>

I’ve updated tge partial solution, now it should work. Iterative algorithm for 1D arrays would be:

```julia
for (i, x) in zip(I, v)
    A[i] += x
end

```

In other words, we iteratively take values from `v` and add them to elements of `A` specified by corresponding indices in `I`.

There are no restrictions on `I`.

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [November 6, 2020, 6:51pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/4 "2020-11-06T18:51:24Z")

</div>

Hm, I see now. When you say 'there are no restrictions on `I`, does that mean that `I` does not need to be a GPU array? How big do you anticipate `I` becoming?

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 6, 2020, 7:27pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/5 "2020-11-06T19:27:45Z")

</div>

Fully-GPU solution would be great, but in general I can convert `I` to CPU.

Let me give you some context. I’m trying to fix [this function](https://github.com/dfdx/Yota.jl/blob/master/src/helpers.jl#L20-L23) which is the derivative of `getindex(A, I...)`. It works fine in a typical scenario where elements of `I` are all unique, e.g. when you have in your code something like:

```julia
A[1:3]

```

the derivative of `A` in this case is:

```julia
dy = ...
dA = zeros(size(A))
dA[1:3] = dy

```

But the code breaks if indices repeat, e.g. in:

```julia
A[[1, 3, 1]]

```

Here the correct derivative would be:

```julia
dA = zeros(size(A))
dA[1] += dy[1]
dA[3] += dy[2]
dA[1] += dy[3]

```

Here `I = [1, 3, 1]`. In theory, it can be much larger, but if it’s not very optimal for, say, `length(I) > 10_000`, I think it’s fine 🙂

---

<div class="post-metadata">

**Author:** ![Mason](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mason/32/2423_2.png) [@Mason](https://discourse.julialang.org/u/Mason)\
**Post date:** [November 6, 2020, 8:43pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/6 "2020-11-06T20:43:08Z")

</div>

Aha, I see! This is a rather awkard thing to to on the GPU, though it should be possible to do if we can phrase it the right way. I’m not sure I have a good answer.

I had hoped Tullio.jl would handle this, but it seems to just silently assume that `I` does not have repeated indices:

```julia
#+begin_src julia
using CUDA, Tullio, KernelAbstractions

let A = CUDA.zeros(4), I = cu([1,3,1]), v = cu([1,2,3])
    @tullio A[I[i]] += v[i]
end
#+end_src

#+RESULTS:
: 4-element CuArray{Float32,1}:
: 1.0
: 0.0
: 2.0
: 0.0

```

@mcabbott Do you have any ideas here? Is this intended behaviour from Tullio?

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 6, 2020, 9:38pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/7 "2020-11-06T21:38:25Z")

</div>

That’s a bug, sorry. Or it was, hopefully fixed on master.

The macro does notice that the index `i` is (potentially) unsafe to break among different threads, and thus avoids that on the CPU. This wasn’t being passed to KernelAbstractions.

But I think the safe way will be slow, it’s got to work sequentially.

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 7, 2020, 12:01am UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/8 "2020-11-07T00:01:23Z")

</div>

Using Tullio#master I’m getting the following error from the code above:

```julia
julia> let A = CUDA.zeros(4), I = cu([1,3,1]), v = cu([1,2,3])
           @tullio A[I[i]] += v[i]
       end
ERROR: AssertionError: length(__workgroupsize) <= length(ndrange)
Stacktrace:
 [1] partition at /home/user/.julia/packages/KernelAbstractions/jAutM/src/nditeration.jl:103 [inlined]
 [2] partition(::KernelAbstractions.Kernel{CUDADevice,KernelAbstractions.NDIteration.StaticSize{(256,)},KernelAbstractions.NDIteration.DynamicSize,var"#gpu_##🇨🇺#253#3"}, ::Tuple{}, ::Nothing) at /home/user/.julia/packages/KernelAbstractions/jAutM/src/KernelAbstractions.jl:385
 [3] launch_config(::KernelAbstractions.Kernel{CUDADevice,KernelAbstractions.NDIteration.StaticSize{(256,)},KernelAbstractions.NDIteration.DynamicSize,var"#gpu_##🇨🇺#253#3"}, ::Tuple{}, ::Nothing) at /home/user/.julia/packages/KernelAbstractions/jAutM/src/backends/cuda.jl:156
 [4] (::KernelAbstractions.Kernel{CUDADevice,KernelAbstractions.NDIteration.StaticSize{(256,)},KernelAbstractions.NDIteration.DynamicSize,var"#gpu_##🇨🇺#253#3"})(::CuArray{Float32,1}, ::Vararg{Any,N} where N; ndrange::Tuple{}, dependencies::KernelAbstractions.CudaEvent, workgroupsize::Nothing, progress::Function) at /home/user/.julia/packages/KernelAbstractions/jAutM/src/backends/cuda.jl:163
 [5] 𝒜𝒸𝓉! at /home/user/.julia/packages/Tullio/RAkkV/src/macro.jl:1166 [inlined]
 [6] 𝒜𝒸𝓉! at /home/user/.julia/packages/Tullio/RAkkV/src/macro.jl:1163 [inlined]
 [7] threader(::var"#𝒜𝒸𝓉!#1", ::Type{CuArray{T,1} where T}, ::CuArray{Float32,1}, ::Tuple{CuArray{Int64,1},CuArray{Int64,1}}, ::Tuple{}, ::Tuple{Base.OneTo{Int64}}, ::Function, ::Int64, ::Bool) at /home/user/.julia/packages/Tullio/RAkkV/src/eval.jl:86
 [8] top-level scope at /home/user/.julia/packages/Tullio/RAkkV/src/macro.jl:1002
 [9] top-level scope at REPL[2]:2

```

But on CPU it indeed does the correct thing.

Another approach I thought about was grouping indices and summing corresponding values before adding them to `A`, but it seems to boil down to the same issue.

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [November 7, 2020, 9:14am UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/9 "2020-11-07T09:14:46Z")

</div>

That’s also no good, sorry. It’s trying to launch just one kernel (if I have the terminology straight) to do this serially, hence `length(ndrange) == 1`. Which KernelAbstractions allows on the CPU.

Any way to solve this seems quite serial. It’s possible that if you sorted the indices, you could then divide the work into safely disjoint sets?

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 7, 2020, 2:30pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/10 "2020-11-07T14:30:02Z")

</div>

It turns out Zygote has [the same issue](https://github.com/FluxML/Zygote.jl/issues/821) on CuArrays.

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [November 11, 2020, 9:51pm UTC](https://discourse.julialang.org/t/increment-elements-of-array-by-index/49694/11 "2020-11-11T21:51:52Z")

</div>

Eventually I adapted code from a nice library [ScatterNNlib.jl](https://github.com/yuehhua/ScatterNNlib.jl) which implements the operation in a pretty much [straightforward way](https://github.com/yuehhua/ScatterNNlib.jl/blob/74622bf40437f16e10468e8a35dae824527c83cf/src/cuda/cuarray.jl#L6-L17) and achieves roughly the same performance as PyTorch implementation. Semantics of the function `scatter_add!()` there is a bit confusing, so here’s my wrapper for the specific case in this post:

```julia
function scatter_add2!(A::AbstractArray, v::AbstractArray, I...)
    # replace (:) with actual index in A
    I = [i == (:) ? (1:size(A, d)) : i for (d, i) in enumerate(I)]
    II = collect(Iterators.product(I...))
    if A isa CuArray
        II = cu(II)
    end
    A_ = reshape(A, 1, size(A)...)
    v_ = reshape(v, 1, size(v)...)
    scatter_add!(A_, v_, II)
    return dropdims(A_, dims=1)
end

A = zeros(4);
I = [1, 3, 1];
v = [1, 1, 1];

scatter_add2!(A, v, I)
# ==> 4-element Array{Float64,1}:
# ==> 2.0
# ==> 0.0
# ==> 1.0
# ==> 0.0

```
