Skip to content

Add SmallTag type for more compact Dual types #748

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

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Add SmallTag type for more compact Dual types #748

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a515f5a
Add `SmallTag` type
topolarity Apr 23, 2025
438fe5f
Support `tag = :small` in `Config(...)` constructors
topolarity Apr 23, 2025
73e8b14
Mark allocation tests as fixed on 1.10
topolarity Apr 23, 2025
d88bf58
Fix pre-compilation on Julia 1.6
topolarity Apr 23, 2025
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
228 changes: 194 additions & 34 deletions src/config.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,46 @@

Tag(::Nothing, ::Type{V}) where {V} = nothing


@inline function ≺(::Type{Tag{F1,V1}}, ::Type{Tag{F2,V2}}) where {F1,V1,F2,V2}
tagcount(Tag{F1,V1}) < tagcount(Tag{F2,V2})
end

# SmallTag is similar to a Tag, but carries just a small UInt64 hash, instead
# of the full type, which makes stacktraces / types easier to read while still
# providing good resilience to perturbation confusion.
struct SmallTag{H}
end

@generated function tagcount(::Type{SmallTag{H}}) where {H}
:($(Threads.atomic_add!(TAGCOUNT, UInt(1))))
end

function SmallTag(f::F, ::Type{V}) where {F,V}
H = if F <: Tuple
# no easy way to check Jacobian tag used with Hessians as multiple functions may be used
# see checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Tuple,VT,F,V}
nothing
else
hash(F) ⊻ hash(V)
end
tagcount(SmallTag{H}) # trigger generated function
SmallTag{H}()
end

SmallTag(::Nothing, ::Type{V}) where {V} = nothing

Check warning on line 49 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L49 

Added line #L49 was not covered by tests

@inline function ≺(::Type{SmallTag{H1}}, ::Type{Tag{F2,V2}}) where {H1,F2,V2}
tagcount(SmallTag{H1}) < tagcount(Tag{F2,V2})

Check warning on line 52 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L51-L52 

Added lines #L51 - L52 were not covered by tests
end

@inline function ≺(::Type{Tag{F1,V1}}, ::Type{SmallTag{H2}}) where {F1,V1,H2}
tagcount(Tag{F1,V1}) < tagcount(SmallTag{H2})

Check warning on line 56 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L55-L56 

Added lines #L55 - L56 were not covered by tests
end

@inline function ≺(::Type{SmallTag{H1}}, ::Type{SmallTag{H2}}) where {H1,H2}
tagcount(SmallTag{H1}) < tagcount(SmallTag{H2})

Check warning on line 60 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L59-L60 

Added lines #L59 - L60 were not covered by tests
end

struct InvalidTagException{E,O} <: Exception
end

Expand All @@ -36,13 +71,22 @@

checktag(::Type{Tag{F,V}}, f::F, x::AbstractArray{V}) where {F,V} = true

# SmallTag is a smaller tag, that only confirms the hash
function checktag(::Type{SmallTag{HT}}, f::F, x::AbstractArray{V}) where {HT,F,V}
H = hash(F) ⊻ hash(V)
if HT == H || HT === nothing
true
else
throw(InvalidTagException{SmallTag{H},SmallTag{HT}}())

Check warning on line 80 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L80 

Added line #L80 was not covered by tests
end
end

# no easy way to check Jacobian tag used with Hessians as multiple functions may be used
checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Tuple,VT,F,V} = true

# custom tag: you're on your own.
checktag(z, f, x) = true


##################
# AbstractConfig #
##################
Expand All @@ -55,6 +99,18 @@

@inline (chunksize(::AbstractConfig{N})::Int) where {N} = N

function maketag(kind::Union{Symbol,Nothing}, f, X)
if kind === :default
return Tag(f, X)
elseif kind === :small
return SmallTag(f, X)
elseif kind === nothing
return nothing
else
throw(ArgumentError("tag may be :default, :small, or nothing"))

Check warning on line 110 in src/config.jl

View check run for this annotation

Codecov / codecov/patch

src/config.jl#L110 

Added line #L110 was not covered by tests
end
end

####################
# DerivativeConfig #
####################
Expand All @@ -64,7 +120,7 @@
end

"""
ForwardDiff.DerivativeConfig(f!, y::AbstractArray, x::Real)
ForwardDiff.DerivativeConfig(f!, y::AbstractArray, x::Real; tag::Union{Symbol,Nothing} = :default)

Return a `DerivativeConfig` instance based on the type of `f!`, and the types/shapes of the
output vector `y` and the input value `x`.
Expand All @@ -77,12 +133,29 @@
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `y` or `x`.
"""
@inline function DerivativeConfig(f::F,
y::AbstractArray{Y},
x::X;
tag::Union{Symbol,Nothing} = :default) where {F,X<:Real,Y<:Real}
# @inline ensures that, e.g., DerivativeConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, X)
return @noinline DerivativeConfig(f,y,x,T)
else
T = maketag(tag, f, X)
return DerivativeConfig(f,y,x,T)
end
end

function DerivativeConfig(f::F,
y::AbstractArray{Y},
x::X,
tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T}
tag::T) where {F,X<:Real,Y<:Real,T}
duals = similar(y, Dual{T,Y,1})
return DerivativeConfig{T,typeof(duals)}(duals)
end
Expand All @@ -100,24 +173,41 @@
end

"""
ForwardDiff.GradientConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
ForwardDiff.GradientConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)

Return a `GradientConfig` instance based on the type of `f` and type/shape of the input
vector `x`.

The returned `GradientConfig` instance contains all the work buffers required by
`ForwardDiff.gradient` and `ForwardDiff.gradient!`.

If `f` is `nothing` instead of the actual target function, then the returned instance can
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `x`.
"""
@inline function GradientConfig(f::F,
x::AbstractArray{V},
c::Chunk{N} = Chunk(x);
tag::Union{Symbol,Nothing} = :default) where {F,V,N}
# @inline ensures that, e.g., GradientConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, V)
return @noinline GradientConfig(f,x,c,T)
else
T = maketag(tag, f, V)
return GradientConfig(f,x,c,T)
end
end

function GradientConfig(f::F,
x::AbstractArray{V},
::Chunk{N} = Chunk(x),
::T = Tag(f, V)) where {F,V,N,T}
::Chunk{N},
::T) where {F,V,N,T}
seeds = construct_seeds(Partials{N,V})
duals = similar(x, Dual{T,V,N})
return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
Expand All @@ -136,7 +226,7 @@
end

"""
ForwardDiff.JacobianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
ForwardDiff.JacobianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)

Return a `JacobianConfig` instance based on the type of `f` and type/shape of the input
vector `x`.
Expand All @@ -145,23 +235,40 @@
`ForwardDiff.jacobian` and `ForwardDiff.jacobian!` when the target function takes the form
`f(x)`.

If `f` is `nothing` instead of the actual target function, then the returned instance can
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `x`.
"""
@inline function JacobianConfig(f::F,
x::AbstractArray{V},
c::Chunk{N} = Chunk(x);
tag::Union{Symbol,Nothing} = :default) where {F,V,N}
# @inline ensures that, e.g., JacobianConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, V)
return @noinline JacobianConfig(f,x,c,T)
else
T = maketag(tag, f, V)
return JacobianConfig(f,x,c,T)
end
end

function JacobianConfig(f::F,
x::AbstractArray{V},
::Chunk{N} = Chunk(x),
::T = Tag(f, V)) where {F,V,N,T}
::Chunk{N},
::T) where {F,V,N,T}
seeds = construct_seeds(Partials{N,V})
duals = similar(x, Dual{T,V,N})
return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals)
end

"""
ForwardDiff.JacobianConfig(f!, y::AbstractArray, x::AbstractArray, chunk::Chunk = Chunk(x))
ForwardDiff.JacobianConfig(f!, y::AbstractArray, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)

Return a `JacobianConfig` instance based on the type of `f!`, and the types/shapes of the
output vector `y` and the input vector `x`.
Expand All @@ -170,17 +277,35 @@
`ForwardDiff.jacobian` and `ForwardDiff.jacobian!` when the target function takes the form
`f!(y, x)`.

If `f!` is `nothing` instead of the actual target function, then the returned instance can
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
If `f!` or `tag` is `nothing`, then the returned instance can be used with any target function.
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `y` or `x`.
"""
@inline function JacobianConfig(f::F,
y::AbstractArray{Y},
x::AbstractArray{X},
c::Chunk{N} = Chunk(x);
tag::Union{Symbol,Nothing} = :default) where {F,Y,X,N}
# @inline ensures that, e.g., JacobianConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, X)
return @noinline JacobianConfig(f,y,x,c,T)
else
T = maketag(tag, f, X)
return JacobianConfig(f,y,x,c,T)
end
end

function JacobianConfig(f::F,
y::AbstractArray{Y},
x::AbstractArray{X},
::Chunk{N} = Chunk(x),
::T = Tag(f, X)) where {F,Y,X,N,T}
::Chunk{N},
::T) where {F,Y,X,N,T}
seeds = construct_seeds(Partials{N,X})
yduals = similar(y, Dual{T,Y,N})
xduals = similar(x, Dual{T,X,N})
Expand All @@ -201,7 +326,7 @@
end

"""
ForwardDiff.HessianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
ForwardDiff.HessianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)

Return a `HessianConfig` instance based on the type of `f` and type/shape of the input
vector `x`.
Expand All @@ -212,41 +337,76 @@
it is a `DiffResult`, the `HessianConfig` should instead be constructed via
`ForwardDiff.HessianConfig(f, result, x, chunk)`.

If `f` is `nothing` instead of the actual target function, then the returned instance can
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `x`.
"""
@inline function HessianConfig(f::F,
x::AbstractArray{V},
chunk::Chunk = Chunk(x);
tag::Union{Symbol,Nothing} = :default) where {F,V}
# @inline ensures that, e.g., HessianConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, V)
return @noinline HessianConfig(f, x, chunk, T)
else
T = maketag(tag, f, V)
return HessianConfig(f, x, chunk, T)
end
end

function HessianConfig(f::F,
x::AbstractArray{V},
chunk::Chunk = Chunk(x),
tag = Tag(f, V)) where {F,V}
chunk::Chunk,
tag) where {F,V}
jacobian_config = JacobianConfig(f, x, chunk, tag)
gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tag)
return HessianConfig(jacobian_config, gradient_config)
end

"""
ForwardDiff.HessianConfig(f, result::DiffResult, x::AbstractArray, chunk::Chunk = Chunk(x))
ForwardDiff.HessianConfig(f, result::DiffResult, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)

Return a `HessianConfig` instance based on the type of `f`, types/storage in `result`, and
type/shape of the input vector `x`.

The returned `HessianConfig` instance contains all the work buffers required by
`ForwardDiff.hessian!` for the case where the `result` argument is an `DiffResult`.

If `f` is `nothing` instead of the actual target function, then the returned instance can
be used with any target function. However, this will reduce ForwardDiff's ability to catch
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).

If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
with similar accuracy, but is much smaller when printing types.

This constructor does not store/modify `x`.
"""
@inline function HessianConfig(f::F,
result::DiffResult,
x::AbstractArray{V},
chunk::Chunk = Chunk(x);
tag::Union{Symbol,Nothing} = :default) where {F,V}
# @inline ensures that, e.g., HessianConfig(...; tag = :small) will be well-inferred
@static if VERSION ≥ v"1.8"
T = @inline maketag(tag, f, V)
return @noinline HessianConfig(f, result, x, chunk, T)
else
T = maketag(tag, f, V)
return HessianConfig(f, result, x, chunk, T)
end
end

function HessianConfig(f::F,
result::DiffResult,
x::AbstractArray{V},
chunk::Chunk = Chunk(x),
tag = Tag(f, V)) where {F,V}
chunk::Chunk,
tag) where {F,V}
jacobian_config = JacobianConfig((f,gradient), DiffResults.gradient(result), x, chunk, tag)
gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, tag)
return HessianConfig(jacobian_config, gradient_config)
Expand Down
12 changes: 2 additions & 10 deletions test/AllocationsTest.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,20 +24,12 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
index = 1
alloc = @allocated ForwardDiff.seed!(duals, x, index, seeds)
alloc = @allocated ForwardDiff.seed!(duals, x, index, seeds)
if VERSION < v"1.9" || VERSION >= v"1.11"
@test alloc == 0
else
@test_broken alloc == 0
end
@test alloc == 0

index = 1
alloc = @allocated ForwardDiff.seed!(duals, x, index, seed)
alloc = @allocated ForwardDiff.seed!(duals, x, index, seed)
if VERSION < v"1.9" || VERSION >= v"1.11"
@test alloc == 0
else
@test_broken alloc == 0
end
@test alloc == 0

alloc = @allocated convert_test_574()
alloc = @allocated convert_test_574()
Expand Down
Loading
Loading