Add GPUArraysCore update in place scalar (EnzymeAD#2220)

Author: wsmoses

Project.toml

@@ -21,13 +21,15 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [weakdeps]
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 EnzymeBFloat16sExt = "BFloat16s"
 EnzymeChainRulesCoreExt = "ChainRulesCore"
+EnzymeGPUArraysCoreExt = "GPUArraysCore"
 EnzymeLogExpFunctionsExt = "LogExpFunctions"
 EnzymeSpecialFunctionsExt = "SpecialFunctions"
 EnzymeStaticArraysExt = "StaticArrays"
@@ -52,6 +54,7 @@ julia = "1.10"
 [extras]
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

ext/EnzymeGPUArraysCore.jl

@@ -0,0 +1,26 @@
+module EnzymeGPUArraysCoreExt
+
+using GPUArraysCore
+using Enzyme
+
+@inline function Enzyme.onehot(x::AbstractGPUArray)
+    onehot_internal(zerosetfn, x, 0, length(x))
+end
+
+@inline function Enzyme.onehot(x::AbstractGPUArray, start::Int, endl::Int)
+    onehot_internal(zerosetfn, x, start-1, endl-start+1)
+end
+
+function Enzyme.zerosetfn(x::AbstractGPUArray, i::Int)
+    res = zero(x)
+    @allowscalar @inbounds res[i] = 1
+    return res
+end
+
+function Enzyme.zerosetfn!(x::AbstractGPUArray, i::Int, val)
+    @allowscalar @inbounds x[i] = += val
+    return
+end
+
+
+end # module

src/sugar.jl

@@ -7,6 +7,11 @@ function zerosetfn(x, i::Int)
     return res
 end
 
+function zerosetfn!(x, i::Int, val)
+    @inbounds x[i] += val
+    nothing
+end
+
 @generated function onehot_internal(fn::F, x::T, startv::Int, lengthv::Int) where {F, T<:Array}
     ir = GPUCompiler.JuliaContext() do ctx
         Base.@_inline_meta
@@ -927,7 +932,7 @@ this function will retun an AbstractArray of shape `size(output)` of values of t
                 dx = MD ? Ref(z) : z
                 res = primal(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx))
                 tape = res[1]
-                @inbounds res[3][i] += Compiler.default_adjoint(eltype(typeof(res[3])))
+                zerosetfn!(res[3], i, Compiler.default_adjoint(eltype(typeof(res[3]))))
                 adjoint(Const(f), MD ? MixedDuplicated(x, dx) : Duplicated(x, dx), tape)
                 return MD ? dx[] : dx, (i == 1 ? size(res[3]) : nothing)
             end
@@ -994,8 +999,8 @@ this function will retun an AbstractArray of shape `size(output)` of values of t
                 j = 0
                 for shadow in res[3]
                     j += 1
-                    @inbounds shadow[(i-1)*chunksize+j] +=
-                        Compiler.default_adjoint(eltype(typeof(shadow)))
+                    zerosetfn!(shadow, (i-1)*chunksize+j,
+                               Compiler.default_adjoint(eltype(typeof(shadow))))
                 end
                 (i == num ? adjoint2 : adjoint)(
                     Const(f),