High-level atomic operations

[AntonOresten]

I'm interesting in having atomic operations that operate in tile space like ct.store but for e.g. atomically adding to an array instead of overwriting.

[AntonOresten]

Could be done with these functions:

@inline function element_indices(shape::NTuple{1, Int}, index::Integer)
    arange(shape, Int32) .+ Int32((index - One()) * shape[1])
end

@inline function element_indices(shape::NTuple{N, Int}, index::NTuple{N, Integer}) where {N}
    ntuple(Val(N)) do d
        bcast_shape = ntuple(i -> i == d ? shape[d] : 1, Val(N))
        base = Int32((index[d] - One()) * shape[d])
        reshape(arange((shape[d],), Int32), bcast_shape) .+ base
    end
end

So we have

ct.store(A, pid, tile)

and

ct.atomic_add(A, ct.element_indices(size(tile), pid), tile)

and then the store-like helper would be one that automatically calls element_indices with size(tile).

[AntonOresten]

@maleadt can decide whether this is worth adding to the package itself. It's simple enough that I could put it in a utils.jl when I'm writing kernels, but as I expressed in #97 (comment), this feels rather fundamental for tile-based programming DSL.

[maleadt]

Does cuTile Python support this, or are the changes from #96 sufficient to get feature parity?

For an API suggestion, it would be nice to have @atomic ct.store or @atomic tile .+= ... just like we do in CUDA.jl.

[AntonOresten]

#96 should bring parity with cuTile Python. This suggestion is more of a QoL, but I realized it just boils down to an indexing helper.

An @atomic macro would be nice, but since there is currently only ct.store which overwrites data, there's nothing to do an atomic version of, unless there was syntax for mutating a tile conditioned on it's current value in the array.

I considered something like A[(i, j)] .= a for store and ct.@atomic A[(i, j)] .+= a for atomic add, where the tuple is supposed to denote tile space I guess, but then we'd be straying from the more explicit API, and there's no neat getindex equivalent since we'd have to specify the tile size.