diff --git a/Project.toml b/Project.toml index a75b616d..83292513 100644 --- a/Project.toml +++ b/Project.toml @@ -32,8 +32,9 @@ julia = "1.6" Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9" DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils"] +test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils", "JLArrays"] diff --git a/src/gradient.jl b/src/gradient.jl index f9711984..47651349 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -80,12 +80,12 @@ function extract_gradient!(::Type{T}, result::DiffResult, dual::Dual) where {T} end extract_gradient!(::Type{T}, result::AbstractArray, y::Real) where {T} = fill!(result, zero(y)) -extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}= copyto!(result, partials(T, dual)) +extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T} = + extract_gradient_chunk!(T, result, dual, 1, npartials(dual)) function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T} - offset = index - 1 - for i in 1:chunksize - result[i + offset] = partials(T, dual, i) + map!(view(Base.ReshapedArray(result, (length(result),), ()), index:index+chunksize-1), 1:chunksize) do i + @inbounds partials(T, dual, i) end return result end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 2a0075a7..8054b444 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -37,4 +37,26 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F end +@testset "Test extract_gradient! allocations" begin + T = Float64 + @testset "vector-mode size(result)=$size" for size in [(4,), (2,2)] + dual = ForwardDiff.Dual(0, (rand(T, size...)...,)) + y = Array{T}(undef, size) + alloc = @allocated ForwardDiff.extract_gradient!(Nothing, y, dual) + alloc = @allocated ForwardDiff.extract_gradient!(Nothing, y, dual) + @test alloc == 0 + end + @testset "chunk-mode size(result)=$size" for size in [(DEFAULT_CHUNK_THRESHOLD+1,), (DEFAULT_CHUNK_THRESHOLD+1, DEFAULT_CHUNK_THRESHOLD+1)] + Npartials = DEFAULT_CHUNK_THRESHOLD÷2 + dual = ForwardDiff.Dual(0, (rand(T, Npartials...)...,)) + y = Array{T}(undef, size) + alloc = @allocated ForwardDiff.extract_gradient_chunk!(Nothing, y, dual, 2, Npartials) + alloc = @allocated ForwardDiff.extract_gradient_chunk!(Nothing, y, dual, 2, Npartials) + @test alloc == 0 + alloc = @allocated ForwardDiff.extract_gradient_chunk!(Nothing, y, dual, 2, Npartials-1) + alloc = @allocated ForwardDiff.extract_gradient_chunk!(Nothing, y, dual, 2, Npartials-1) + @test alloc == 0 + end +end + end diff --git a/test/GradientTest.jl b/test/GradientTest.jl index a74f4304..7172f1b7 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -8,6 +8,8 @@ using ForwardDiff using ForwardDiff: Dual, Tag using StaticArrays using DiffTests +using JLArrays +JLArrays.allowscalar(false) include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -149,6 +151,26 @@ end @test isequal(ForwardDiff.gradient(t -> t[1]^t[2], [0.0, 1.5]), [0.0, 0.0]) end + +############################################## +# test GPUArray compatibility (via JLArrays) # +############################################## + +println(" ...testing GPUArray compatibility (via JLArrays)") + +@testset "size = $(size(x))" for x in JLArray.([ + rand(1), + rand(DEFAULT_CHUNK_THRESHOLD+1), + rand(1,1), + rand(DEFAULT_CHUNK_THRESHOLD+1,DEFAULT_CHUNK_THRESHOLD+1), + rand(1,1,1) +]) + + @test ForwardDiff.gradient(prod, x) isa typeof(x) + +end + + ############# # bug fixes # #############