diff --git a/ext/PreallocationToolsForwardDiffExt.jl b/ext/PreallocationToolsForwardDiffExt.jl index b4a3fba..4bd46ca 100644 --- a/ext/PreallocationToolsForwardDiffExt.jl +++ b/ext/PreallocationToolsForwardDiffExt.jl @@ -24,6 +24,44 @@ PreallocationTools.forwarddiff_compat_chunk_size(x::Int) = ForwardDiff.pickchunk # Define chunksize for ForwardDiff.Dual types PreallocationTools.chunksize(::Type{ForwardDiff.Dual{T, V, N}}) where {T, V, N} = N +function dual_eltype(::Type{CacheT}, ::Type{DualT}) where { + CacheT, Tag, V, N, DualT <: ForwardDiff.Dual{Tag, V, N}, + } + dual_cache_t = replace_type_parameter(CacheT, V, DualT) + return dual_cache_t === CacheT ? DualT : dual_cache_t +end + +function replace_type_parameter(::Type{T}, ::Type{From}, ::Type{To}) where {T, From, To} + T === From && return To + T <: Number && From <: Number && promote_type(T, From) === From && return To + + typename = Base.typename(T) + wrapper = typename.wrapper + parameters = T.parameters + isempty(parameters) && return T + + new_parameters = map(parameters) do parameter + parameter isa Type ? replace_type_parameter(parameter, From, To) : parameter + end + + return new_parameters == parameters ? T : Core.apply_type(wrapper, new_parameters...) +end + +function diffcache_dual_tmp(dc::PreallocationTools.DiffCache, ::Type{T}) where {T <: ForwardDiff.Dual} + cache_eltype = dual_eltype(eltype(dc.du), T) + return if isbitstype(cache_eltype) + nelem = div(sizeof(cache_eltype), sizeof(eltype(dc.dual_du))) * length(dc.du) + if nelem > length(dc.dual_du) + PreallocationTools.enlargediffcache!(dc, nelem) + end + PreallocationTools._restructure( + dc.du, reinterpret(cache_eltype, view(dc.dual_du, 1:nelem)) + ) + else + PreallocationTools._restructure(dc.du, zeros(cache_eltype, size(dc.du))) + end +end + # Define get_tmp methods for ForwardDiff.Dual types function PreallocationTools.get_tmp(dc::PreallocationTools.FixedSizeDiffCache, u::T) where {T <: ForwardDiff.Dual} x = reinterpret(T, dc.dual_du) @@ -53,39 +91,15 @@ function PreallocationTools.get_tmp(dc::PreallocationTools.FixedSizeDiffCache, u end function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, u::T) where {T <: ForwardDiff.Dual} - return if isbitstype(T) - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) - if nelem > length(dc.dual_du) - PreallocationTools.enlargediffcache!(dc, nelem) - end - PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) - else - PreallocationTools._restructure(dc.du, zeros(T, size(dc.du))) - end + return diffcache_dual_tmp(dc, T) end function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, ::Type{T}) where {T <: ForwardDiff.Dual} - return if isbitstype(T) - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) - if nelem > length(dc.dual_du) - PreallocationTools.enlargediffcache!(dc, nelem) - end - PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) - else - PreallocationTools._restructure(dc.du, zeros(T, size(dc.du))) - end + return diffcache_dual_tmp(dc, T) end function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual} - return if isbitstype(T) - nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du) - if nelem > length(dc.dual_du) - PreallocationTools.enlargediffcache!(dc, nelem) - end - PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem))) - else - PreallocationTools._restructure(dc.du, zeros(T, size(dc.du))) - end + return diffcache_dual_tmp(dc, T) end @setup_workload begin diff --git a/test/core_dispatch.jl b/test/core_dispatch.jl index 0ad8cfa..ce8528b 100644 --- a/test/core_dispatch.jl +++ b/test/core_dispatch.jl @@ -26,6 +26,13 @@ function structequal(struct1, struct2) end end +struct IsbitsPair{T} + x::T + y::T +end + +Base.zero(::Type{IsbitsPair{T}}) where {T} = IsbitsPair(zero(T), zero(T)) + #Setup Base Array tests chunk_size = 5 u0 = ones(5, 5) @@ -195,3 +202,33 @@ dual_cache = FixedSizeDiffCache([0.0, 0.0, 0.0], 3) dl = zeros(ForwardDiff.Dual{Nothing, Float64, 3}, 3) @test !(get_tmp(dual_cache, dl) isa Base.ReinterpretArray) + +@testset "DiffCache preserves isbits wrappers around duals" begin + z = zeros(ComplexF64, 20) + zd = DiffCache(z) + + function sum_cis(θ) + ztmp = get_tmp(zd, θ) + @test eltype(ztmp) <: Complex + for i in eachindex(ztmp) + ztmp[i] = cis(i * θ) + end + abs(sum(ztmp)) + end + + @test ForwardDiff.derivative(sum_cis, 1.1) isa Float64 + + DualT = ForwardDiff.Dual{Nothing, Float64, 2} + pair_cache = DiffCache(zeros(IsbitsPair{Float64}, 3), 2) + pair_tmp = get_tmp(pair_cache, DualT) + + @test size(pair_tmp) == (3,) + @test eltype(pair_tmp) == IsbitsPair{DualT} + + pair_tmp[1] = IsbitsPair(zero(DualT), zero(DualT)) + @test pair_tmp[1] == IsbitsPair(zero(DualT), zero(DualT)) + + complex32_cache = DiffCache(zeros(ComplexF32, 3), 5) + complex32_tmp = get_tmp(complex32_cache, DualT) + @test eltype(complex32_tmp) == Complex{DualT} +end