FP atomic add not lowering to CAS on Metal.jl

[Moelf]

Metal does not support floating point atomic addition on shared memory (block-local), but right now KernalAbstractions.jl would lower to invalid instructions, which prevents kernel from running.

I'm happy to make a PR, but is there a rope I can follow in terms to adding a different lowering for this?

[Moelf]

diff --git a/ext/AtomixMetalExt.jl b/ext/AtomixMetalExt.jl
index 7df0bb6..3622b20 100644
--- a/ext/AtomixMetalExt.jl
+++ b/ext/AtomixMetalExt.jl
@@ -35,7 +35,8 @@ end
     ptr = Atomix.pointer(ref)
     begin
         old = if op === (+)
-            Metal.atomic_fetch_add_explicit(ptr, x)
+            # Metal.atomic_fetch_add_explicit(ptr, x)
+            Metal.atomic_fetch_op_explicit(ptr, +, x)
         elseif op === (-)
             Metal.atomic_fetch_sub_explicit(ptr, x)
         elseif op === (&)

this doesn't work somehow because of error:

ERROR: InvalidIRError: compiling MethodInstance for gpu_fp_sharemem(::KernelAbstractions.CompilerMetadata{…}, ::MtlDeviceVector{…}, ::Val{…}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to julia.get_pgcstack)
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erroneous code with Cthulhu.jl
Stacktrace:
  [1] check_ir(job::GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, args::LLVM.Module)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/validation.jl:167
  [2] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:385 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/Tracy/slmNc/src/tracepoint.jl:163 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:384 [inlined]
  [5] emit_llvm(job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/utils.jl:116
  [6] emit_llvm(job::GPUCompiler.CompilerJob)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/utils.jl:114
  [7] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:95
  [8] compile_unhooked
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:80 [inlined]
  [9] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:67
 [10] compile
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:55 [inlined]
 [11] (::Metal.var"#159#167"{GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}})(ctx::LLVM.Context)
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:108
 [12] JuliaContext(f::Metal.var"#159#167"{GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}}; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:34
 [13] JuliaContext(f::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:25
 [14] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:107 [inlined]
 [15] macro expansion
    @ ~/.julia/packages/ObjectiveC/UNTzb/src/os.jl:264 [inlined]
 [16] compile(job::GPUCompiler.CompilerJob)
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:105
 [17] actual_compilation(cache::Dict{…}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{…}, compiler::typeof(Metal.compile), linker::typeof(Metal.link))
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/execution.jl:245
 [18] cached_compilation(cache::Dict{Any, Any}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, compiler::Function, linker::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/execution.jl:159
 [19] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:189 [inlined]
 [20] macro expansion
    @ ./lock.jl:273 [inlined]
 [21] mtlfunction(f::typeof(gpu_fp_sharemem), tt::Type{Tuple{KernelAbstractions.CompilerMetadata{…}, MtlDeviceVector{…}, Val{…}}}; name::Nothing, kwargs::@Kwargs{})
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:184
 [22] mtlfunction(f::typeof(gpu_fp_sharemem), tt::Type{Tuple{KernelAbstractions.CompilerMetadata{…}, MtlDeviceVector{…}, Val{…}}})
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:182
 [23] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:85 [inlined]
 [24] (::KernelAbstractions.Kernel{…})(::MtlVector{…}, ::Vararg{…}; ndrange::Tuple{…}, workgroupsize::Nothing)
    @ Metal.MetalKernels ~/.julia/packages/Metal/uzlNf/src/MetalKernels.jl:113
 [25] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/reflection.jl:257 [inlined]
 [26] top-level scope
    @ REPL[3]:1
Some type information was truncated. Use `show(err)` to see complete types.

[Moelf]

diff --git a/ext/AtomixMetalExt.jl b/ext/AtomixMetalExt.jl
index 7df0bb6..3622b20 100644
--- a/ext/AtomixMetalExt.jl
+++ b/ext/AtomixMetalExt.jl
@@ -35,7 +35,8 @@ end
     ptr = Atomix.pointer(ref)
     begin
         old = if op === (+)
-            Metal.atomic_fetch_add_explicit(ptr, x)
+            # Metal.atomic_fetch_add_explicit(ptr, x)
+            Metal.atomic_fetch_op_explicit(ptr, +, x)
         elseif op === (-)
             Metal.atomic_fetch_sub_explicit(ptr, x)
         elseif op === (&)

this doesn't work somehow because of error:

ERROR: InvalidIRError: compiling MethodInstance for gpu_fp_sharemem(::KernelAbstractions.CompilerMetadata{…}, ::MtlDeviceVector{…}, ::Val{…}) resulted in invalid LLVM IR
Reason: unsupported call to an unknown function (call to julia.get_pgcstack)
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erroneous code with Cthulhu.jl
Stacktrace:
  [1] check_ir(job::GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, args::LLVM.Module)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/validation.jl:167
  [2] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:385 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/Tracy/slmNc/src/tracepoint.jl:163 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:384 [inlined]
  [5] emit_llvm(job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/utils.jl:116
  [6] emit_llvm(job::GPUCompiler.CompilerJob)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/utils.jl:114
  [7] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:95
  [8] compile_unhooked
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:80 [inlined]
  [9] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:67
 [10] compile
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:55 [inlined]
 [11] (::Metal.var"#159#167"{GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}})(ctx::LLVM.Context)
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:108
 [12] JuliaContext(f::Metal.var"#159#167"{GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}}; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:34
 [13] JuliaContext(f::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:25
 [14] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:107 [inlined]
 [15] macro expansion
    @ ~/.julia/packages/ObjectiveC/UNTzb/src/os.jl:264 [inlined]
 [16] compile(job::GPUCompiler.CompilerJob)
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/compilation.jl:105
 [17] actual_compilation(cache::Dict{…}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{…}, compiler::typeof(Metal.compile), linker::typeof(Metal.link))
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/execution.jl:245
 [18] cached_compilation(cache::Dict{Any, Any}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}, compiler::Function, linker::Function)
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/execution.jl:159
 [19] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:189 [inlined]
 [20] macro expansion
    @ ./lock.jl:273 [inlined]
 [21] mtlfunction(f::typeof(gpu_fp_sharemem), tt::Type{Tuple{KernelAbstractions.CompilerMetadata{…}, MtlDeviceVector{…}, Val{…}}}; name::Nothing, kwargs::@Kwargs{})
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:184
 [22] mtlfunction(f::typeof(gpu_fp_sharemem), tt::Type{Tuple{KernelAbstractions.CompilerMetadata{…}, MtlDeviceVector{…}, Val{…}}})
    @ Metal ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:182
 [23] macro expansion
    @ ~/.julia/packages/Metal/uzlNf/src/compiler/execution.jl:85 [inlined]
 [24] (::KernelAbstractions.Kernel{…})(::MtlVector{…}, ::Vararg{…}; ndrange::Tuple{…}, workgroupsize::Nothing)
    @ Metal.MetalKernels ~/.julia/packages/Metal/uzlNf/src/MetalKernels.jl:113
 [25] macro expansion
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/reflection.jl:257 [inlined]
 [26] top-level scope
    @ REPL[3]:1
Some type information was truncated. Use `show(err)` to see complete types.