Add GPUArraysCore extension: broadcast-based seed! for GPU dual arrays

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "1.4.1"
+version = "1.5.0"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"
@@ -15,9 +15,11 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [weakdeps]
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
+ForwardDiffGPUArraysCoreExt = "GPUArraysCore"
 ForwardDiffStaticArraysExt = "StaticArrays"
 
 [compat]
@@ -26,8 +28,10 @@ CommonSubexpressions = "0.3"
 DiffResults = "1.1"
 DiffRules = "1.4"
 DiffTests = "0.1"
+GPUArraysCore = "0.1, 0.2"
 IrrationalConstants = "0.1, 0.2"
 JET = "0.9, 0.10, 0.11"
+JLArrays = "0.1, 0.2"
 LogExpFunctions = "0.3, 1"
 NaNMath = "1"
 Preferences = "1"
@@ -38,12 +42,14 @@ julia = "1.10"
 [extras]
 Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Calculus", "DiffTests", "IrrationalConstants", "JET", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]
+test = ["Calculus", "DiffTests", "GPUArraysCore", "IrrationalConstants", "JET", "JLArrays", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]

ext/ForwardDiffGPUArraysCoreExt.jl

@@ -0,0 +1,45 @@
+module ForwardDiffGPUArraysCoreExt
+
+using ForwardDiff: ForwardDiff, Dual, Partials
+using GPUArraysCore: AbstractGPUArray
+
+# ForwardDiff's default `seed!` methods (src/apiutils.jl) write each dual with a
+# scalar `setindex!` loop over `structural_eachindex`. On GPU arrays that
+# triggers a scalar-indexing error. GPU arrays are always dense, one-based, and
+# carry isbits element types, so the structural-index / unset-element handling of
+# the generic methods is unnecessary here and broadcast restores the
+# pre-1.0 GPU-compatible behavior.
+
+function ForwardDiff.seed!(
+        duals::AbstractGPUArray{Dual{T,V,N}}, x,
+        seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
+    duals .= Dual{T,V,N}.(x, Ref(seed))
+    return duals
+end
+
+function ForwardDiff.seed!(
+        duals::AbstractGPUArray{Dual{T,V,N}}, x,
+        seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
+    dual_inds = 1:N
+    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), seeds)
+    return duals
+end
+
+function ForwardDiff.seed!(
+        duals::AbstractGPUArray{Dual{T,V,N}}, x, index,
+        seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
+    dual_inds = index:length(duals)
+    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), Ref(seed))
+    return duals
+end
+
+function ForwardDiff.seed!(
+        duals::AbstractGPUArray{Dual{T,V,N}}, x, index,
+        seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
+    offset = index - 1
+    dual_inds = (1 + offset):(offset + chunksize)
+    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), seeds[1:chunksize])
+    return duals
+end
+
+end

test/GPUArraysCoreTest.jl

@@ -0,0 +1,34 @@
+module GPUArraysCoreTest
+
+using ForwardDiff, Test
+using JLArrays
+
+# JLArrays emulates GPU array semantics (including the scalar-indexing ban) on
+# the CPU, so the GPUArraysCore extension's broadcast-based `seed!` methods can
+# be exercised without a physical GPU.
+JLArrays.allowscalar(false)
+
+@testset "ForwardDiff seeding on GPU arrays" begin
+    f(x) = x .^ 2 .+ 2 .* x
+
+    @testset "jacobian, vector mode (length $n)" for n in (1, 4, 8)
+        x = collect(Float64, 1:n)
+        @test Array(ForwardDiff.jacobian(f, JLArray(x))) == ForwardDiff.jacobian(f, x)
+    end
+
+    # lengths above the chunk size exercise the chunked `seed!` methods
+    @testset "jacobian, chunk mode (length $n, chunk $c)" for n in (16, 20, 27), c in (4, 8)
+        x = collect(Float64, 1:n)
+        cfg = ForwardDiff.JacobianConfig(f, JLArray(x), ForwardDiff.Chunk{c}())
+        @test Array(ForwardDiff.jacobian(f, JLArray(x), cfg)) == ForwardDiff.jacobian(f, x)
+    end
+
+    @testset "jacobian! into a GPU array (length $n)" for n in (4, 16)
+        x = collect(Float64, 1:n)
+        out = JLArray(zeros(n, n))
+        ForwardDiff.jacobian!(out, f, JLArray(x))
+        @test Array(out) == ForwardDiff.jacobian(f, x)
+    end
+end
+
+end # module

test/runtests.jl

@@ -43,6 +43,11 @@ Random.seed!(SEED)
         t = @elapsed include("ConfusionTest.jl")
         println("##### done (took $t seconds).")
     end
+    @testset "GPUArraysCore" begin
+        println("##### Testing GPUArraysCore extension...")
+        t = @elapsed include("GPUArraysCoreTest.jl")
+        println("##### done (took $t seconds).")
+    end
     @testset "Miscellaneous" begin
         println("##### Testing miscellaneous functionality...")
         t = @elapsed include("MiscTest.jl")