Atomix.jl support

[jagoosw]

Atomix.jl implements atomic operations and prevents race conditions when kernels index into the same location (see https://juliagpu.github.io/KernelAbstractions.jl/stable/examples/atomix/).

In order for it to support Oceananigans fields so you can write:

Atomix.@atomic C[1, 2, 3]

where C is a field we need to add the following method:

using Oceananigans: Field

import Atomix: pointer

@inline function pointer(ref::Atomix.Internal.IndexableRef{<:Field, Tuple{Vararg{Int64, N}}} where {N, Indexable})
    i = LinearIndices(ref.data.data)[ref.indices...]
    return Base.pointer(ref.data.data, i)
end

Is this something we should add here or is the use case too limited?

I haven't checked that this works on GPU yet either.

[glwagner]

Can you give an example of a use case? Note that there are no Field on the GPU, we adapt to the underlying OffsetArray on the GPU.

[jagoosw]

My use case is having particles modifying a tracer field (e.g. releasing a tracer). Up to now I've been doing this serially to prevent a race condition (of different particles trying to add to the same cell at the same time), but now modifying that code for GPU I've had to make it all into kernels.

It was also not as straightforward as I though as Atomix doesn't seem to work interchangeably on GPU yet (it throws an error saying that the CPU address isn't available for the GPU object) so I had to implement what I was trying to do explicitly like:

@inline function pointer(ref::Atomix.Internal.IndexableRef{<:Field, Tuple{Vararg{Int64, N}}} where {N, Indexable})
    i = LinearIndices(ref.data.data)[ref.indices...]
    return Base.pointer(ref.data.data, i)
end

# Field on CPU
function atomic_add!(fld::Field, i, j, k, value)
    Atomix.@atomic fld[i, j, k] += value
end 

# Field on GPU which is adapted to field.data
function atomic_add!(fld::OffsetArray, i, j, k, value)
    linear_index = LinearIndices(fld)[i, j, k]
    CUDA.@atomic fld.parent[linear_index] += value
end

I may be missing something though.

[glwagner]

Interesting!

You could invert the order of the loop. Instead of looping over particles, loop over the grid indices. Then at each index, add the appropriate contribution from each particle. If the number of particles is not too huge, this might be simpler and the performance difference probably fairly negligible (still cheap compared to evaluating tendencies...)