diff --git a/Project.toml b/Project.toml index dedbfe5d..0f79a8ef 100644 --- a/Project.toml +++ b/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"] diff --git a/src/apiutils.jl b/src/apiutils.jl index 0615fdb3..e6338f76 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -70,11 +70,29 @@ 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) @@ -82,6 +100,15 @@ seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} = # `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) diff --git a/test/GPUArraysTest.jl b/test/GPUArraysTest.jl new file mode 100644 index 00000000..ab8f3ab2 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index e39f5e46..d03d5240 100644 --- a/test/runtests.jl +++ b/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")