Seed dense arrays of isbits duals without scalar indexing (fixes GPU jacobians)

Project.toml

@@ -28,6 +28,7 @@ DiffRules = "1.4"
 DiffTests = "0.1"
 IrrationalConstants = "0.1, 0.2"
 JET = "0.9, 0.12"
+JLArrays = "0.1, 0.2"
 LogExpFunctions = "0.3, 1"
 NaNMath = "1"
 Preferences = "1"
@@ -41,9 +42,10 @@ DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 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", "IrrationalConstants", "JET", "JLArrays", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]

src/apiutils.jl

@@ -70,18 +70,45 @@ function structural_eachindex(x::Diagonal, y::AbstractArray)
     return diagind(x)
 end
 
+@inline function dense_seedable(duals, x, ::Type{V}) where {V}
+    return duals isa DenseArray && isbitstype(V) && !Base.has_offset_axes(duals, x)
+end
+
+struct SeededDual{D,S}
+    seeds::S
+    offset::Int
+end
+
+@inline (f::SeededDual{D})(x) where {D} = D(x, f.seeds)
+@inline (f::SeededDual{D})(x, i) where {D} = D(x, f.seeds[i - f.offset])
+
 # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is
 # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be
 # written before the target function reads them.
-seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} =
-    _seed_zero_partials!(duals, x, structural_eachindex(duals, x))
+function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N}
+    seed = zero(Partials{N,V})
+    if dense_seedable(duals, x, V) && axes(duals) == axes(x)
+        duals .= Dual{T,V,N}.(x, Ref(seed))
+        return duals
+    end
+    return _seed_zero_partials!(duals, x, structural_eachindex(duals, x))
+end
 
 # Zeroes the partials of `count` elements starting at structural position `index`. Chunk mode only
 # needs to clear the chunk it just seeded, so writing through to the end of the array would be O(n)
 # redundant work per chunk, i.e. O(n^2/N) per sweep. `count` mirrors the `chunksize` argument of
 # `seed!(duals, x, index, seeds, chunksize)`.
 function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                              count = N) where {T,V,N}
+    if dense_seedable(duals, x, V)
+        length(duals) == length(x) || throw(DimensionMismatch())
+        last_index = min(index + count - 1, length(duals))
+        dual_inds = index:last_index
+        seed = zero(Partials{N,V})
+        f = SeededDual{Dual{T,V,N},typeof(seed)}(seed, 0)
+        map!(f, view(duals, dual_inds), view(x, dual_inds))
+        return duals
+    end
     idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), count)
     return _seed_zero_partials!(duals, x, idxs)
 end
@@ -106,7 +133,12 @@ end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
                seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    if isbitstype(V)
+    if dense_seedable(duals, x, V)
+        length(duals) == length(x) || throw(DimensionMismatch())
+        dual_inds = 1:min(N, length(duals))
+        f = SeededDual{Dual{T,V,N},typeof(seeds)}(seeds, 0)
+        map!(f, view(duals, dual_inds), view(x, dual_inds), dual_inds)
+    elseif isbitstype(V)
         for (i, idx) in zip(1:N, structural_eachindex(duals, x))
             duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
         end
@@ -124,6 +156,14 @@ end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
+    if dense_seedable(duals, x, V)
+        length(duals) == length(x) || throw(DimensionMismatch())
+        shift = index - 1
+        dual_inds = (1 + shift):min(shift + chunksize, length(duals))
+        f = SeededDual{Dual{T,V,N},typeof(seeds)}(seeds, shift)
+        map!(f, view(duals, dual_inds), view(x, dual_inds), dual_inds)
+        return duals
+    end
     offset = index - 1
     idxs = Iterators.drop(structural_eachindex(duals, x), offset)
     if isbitstype(V)

test/GPUArraysTest.jl

@@ -0,0 +1,64 @@
+module GPUArraysTest
+
+using ForwardDiff, Test
+using JLArrays
+
+# Exercise GPU array semantics, including the scalar-indexing ban, without physical GPU hardware.
+JLArrays.allowscalar(false)
+
+@testset "ForwardDiff seeding on GPU arrays" begin
+    f(x) = x .^ 2 .+ 2 .* x
+
+    @testset "zero chunk tail" begin
+        values = collect(Float64, 1:20)
+        x = JLArray(values)
+        duals = JLArray([ForwardDiff.Dual{Nothing}(xi, 1.0) for xi in values])
+        ForwardDiff.seed_zero_partials!(duals, x, 5, 12)
+        result = Array(duals)
+        @test ForwardDiff.value.(result) == values
+        @test [ForwardDiff.partials(d)[1] for d in result] ==
+            [ones(4); zeros(12); ones(4)]
+    end
+
+    @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
+
+    @testset "jacobian of f! with GPU input and output" begin
+        f!(y, x) = (y .= x .^ 2 .+ 2 .* x; nothing)
+        x = collect(Float64, 1:8)
+        y = zeros(8)
+        J = ForwardDiff.jacobian(f!, JLArray(y), JLArray(x))
+        @test Array(J) == ForwardDiff.jacobian(f!, y, x)
+    end
+
+    @testset "jacobian with matrix input (chunk $c)" for c in (3, 6)
+        X = reshape(collect(Float64, 1:12), 4, 3)
+        g(x) = x .* sum(x)
+        cfg = ForwardDiff.JacobianConfig(g, JLArray(X), ForwardDiff.Chunk{c}())
+        @test Array(ForwardDiff.jacobian(g, JLArray(X), cfg)) ≈ ForwardDiff.jacobian(g, X)
+    end
+
+    @testset "jacobian with view input" begin
+        X = JLArray(reshape(collect(Float64, 1:18), 6, 3))
+        xv = view(X, :, 2)
+        @test Array(ForwardDiff.jacobian(f, xv)) == ForwardDiff.jacobian(f, Array(xv))
+    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 "GPUArrays" begin
+        println("##### Testing seeding on GPU arrays...")
+        t = @elapsed include("GPUArraysTest.jl")
+        println("##### done (took $t seconds).")
+    end
     @testset "Miscellaneous" begin
         println("##### Testing miscellaneous functionality...")
         t = @elapsed include("MiscTest.jl")