Skip to content

Handle isbits wrapper eltypes in DiffCache #178

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ChrisRackauckas merged 1 commit into from
Jul 8, 2026
+78 −27
Merged

Handle isbits wrapper eltypes in DiffCache #178

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 41 additions & 27 deletions ext/PreallocationToolsForwardDiffExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
37 changes: 37 additions & 0 deletions test/core_dispatch.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Loading