Atomix example fails with complex array element type

Hi,

the Atomix example from the docs fails with Complex element types:

using Atomix, KernelAbstractions

function index_fun_fixed(arr; backend=get_backend(arr))
    out = similar(arr)
    fill!(out, 0)
    kernel! = my_kernel_fixed!(backend)
    kernel!(out, arr, ndrange=(size(arr, 1), size(arr, 2)))
    return out 
end

@kernel function my_kernel_fixed!(out, arr)
    i, j = @index(Global, NTuple)
    for k in 1:size(out, 1)
        Atomix.@atomic out[k, i] += arr[i, j]
    end 
end

Results in:

(DiffImageRotation) pkg> st
Project DiffImageRotation v0.2.1
Status `~/.julia/dev/DiffImageRotation.jl/Project.toml`
  [a9b6321e] Atomix v0.1.0
  [d360d2e6] ChainRulesCore v1.18.0
  [63c18a36] KernelAbstractions v0.9.14

julia> versioninfo
versioninfo (generic function with 2 methods)

julia> versioninfo()
Julia Version 1.9.4
Commit 8e5136fa297 (2023-11-14 08:46 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × Intel(R) Core(TM) i7-7600U CPU @ 2.80GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, skylake)
  Threads: 4 on 4 virtual cores
Environment:
  JULIA_NUM_THREADS = 4
  JULIA_IMAGE_THREADS = 1

julia> index_fun_fixed(randn((2,2)))
2×2 Matrix{Float64}:
 1.53179  -2.84534
 1.53179  -2.84534

julia> index_fun_fixed(randn(ComplexF32, (2,2)))
ERROR: TaskFailedException

    nested task error: MethodError: no method matching modify!(::Ptr{ComplexF32}, ::typeof(+), ::ComplexF32, ::UnsafeAtomics.Internal.LLVMOrdering{:seq_cst})

    Closest candidates are:
      modify!(::Core.LLVMPtr, ::OP, ::Any, ::UnsafeAtomics.Ordering) where OP
       @ UnsafeAtomicsLLVM ~/.julia/packages/UnsafeAtomicsLLVM/6HZfV/src/internal.jl:20
      modify!(::Ptr{T}, ::typeof(UnsafeAtomics.right), ::T, ::Any) where T
       @ UnsafeAtomics ~/.julia/packages/UnsafeAtomics/ugwrA/src/core.jl:197
      modify!(::Any, ::Any, ::Any)
       @ UnsafeAtomics ~/.julia/packages/UnsafeAtomics/ugwrA/src/core.jl:4
      ...

    Stacktrace:
     [1] modify!
       @ ~/.julia/packages/Atomix/F9VIX/src/core.jl:33 [inlined]
     [2] macro expansion
       @ ~/.julia/dev/DiffImageRotation.jl/examples/atomix.jl:14 [inlined]
     [3] cpu_my_kernel_fixed!
       @ ~/.julia/packages/KernelAbstractions/mS4kn/src/macros.jl:287 [inlined]
     [4] cpu_my_kernel_fixed!(__ctx__::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, CartesianIndex{2}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}}, out::Matrix{ComplexF32}, arr::Matrix{ComplexF32})
       @ Main ./none:0
     [5] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{CPU, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, typeof(cpu_my_kernel_fixed!)}, ndrange::Tuple{Int64, Int64}, iterspace::KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, args::Tuple{Matrix{ComplexF32}, Matrix{ComplexF32}}, dynamic::KernelAbstractions.NDIteration.DynamicCheck)
       @ KernelAbstractions ~/.julia/packages/KernelAbstractions/mS4kn/src/cpu.jl:115
     [6] (::KernelAbstractions.var"#18#21"{KernelAbstractions.Kernel{CPU, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, typeof(cpu_my_kernel_fixed!)}, Tuple{Int64, Int64}, KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, Tuple{Matrix{ComplexF32}, Matrix{ComplexF32}}, KernelAbstractions.NDIteration.DynamicCheck, Int64})()
       @ KernelAbstractions ./threadingconstructs.jl:416

...and 1 more exception.

Stacktrace:
 [1] sync_end(c::Channel{Any})
   @ Base ./task.jl:445
 [2] macro expansion
   @ ./task.jl:477 [inlined]
 [3] __run(obj::KernelAbstractions.Kernel{CPU, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, typeof(cpu_my_kernel_fixed!)}, ndrange::Tuple{Int64, Int64}, iterspace::KernelAbstractions.NDIteration.NDRange{2, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}}, args::Tuple{Matrix{ComplexF32}, Matrix{ComplexF32}}, dynamic::KernelAbstractions.NDIteration.DynamicCheck, static_threads::Bool)
   @ KernelAbstractions ~/.julia/packages/KernelAbstractions/mS4kn/src/cpu.jl:89
 [4] (::KernelAbstractions.Kernel{CPU, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, typeof(cpu_my_kernel_fixed!)})(::Matrix{ComplexF32}, ::Vararg{Matrix{ComplexF32}}; ndrange::Tuple{Int64, Int64}, workgroupsize::Nothing)
   @ KernelAbstractions ~/.julia/packages/KernelAbstractions/mS4kn/src/cpu.jl:44
 [5] Kernel
   @ ~/.julia/packages/KernelAbstractions/mS4kn/src/cpu.jl:37 [inlined]
 [6] index_fun_fixed(arr::Matrix{ComplexF32}; backend::CPU)
   @ Main ~/.julia/dev/DiffImageRotation.jl/examples/atomix.jl:7
 [7] index_fun_fixed(arr::Matrix{ComplexF32})
   @ Main ~/.julia/dev/DiffImageRotation.jl/examples/atomix.jl:3
 [8] top-level scope
   @ REPL[8]:1

Not sure if an Atomix or KernelAbstractions problem?

Best,

Felix

JuliaGPU / KernelAbstractions.jl

Atomix example fails with complex array element type #444