Exploit symmetry in the Hessian, on top of #840

ext/ForwardDiffStaticArraysExt.jl

@@ -7,7 +7,7 @@ using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, Jacobian
                    gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
                    extract_gradient!, extract_jacobian!, extract_value!, structural_indices,
                    vector_mode_gradient, vector_mode_gradient!,
-                   vector_mode_jacobian, vector_mode_jacobian!, valtype, value
+                   vector_mode_jacobian, vector_mode_jacobian!, HESSIAN_ERROR, valtype, value
 using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
 
 @generated function dualize(::Type{T}, x::StaticArray) where T
@@ -107,11 +107,36 @@ end
 end
 
 # Hessian
-ForwardDiff.hessian(f::F, x::StaticArray) where {F} = jacobian(Base.Fix1(gradient, f), x)
+@inline function extract_hessian(::Type{T}, ydual::Partials, x::StaticArray) where {T}
+    H = extract_jacobian(T, ydual, x)
+    return typeof(H)(Symmetric(H, :U))
+end
+
+# An `f` ignoring its argument returns no partials at all, not `length(x)` zero ones, so the method
+# above would build a result with no rows. Reached for an empty `x` too.
+@inline function extract_hessian(::Type{T}, ydual::Partials{0}, x::S) where {T,S<:StaticArray}
+    R = StaticArrays.similar_type(S, valtype(T, eltype(ydual)), Size(length(x), length(x)))
+    return zero(R)
+end
+
+@inline function ForwardDiff.hessian(f::F, x::StaticArray) where {F}
+    T = typeof(Tag(f, eltype(x)))
+    ydual = f(dualize(T, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    return extract_hessian(T, partials(T, ydual), x)
+end
+
 ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig) where {F} = hessian(f, x)
 ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} = hessian(f, x)
 
-ForwardDiff.hessian!(result::AbstractArray, f::F, x::StaticArray) where {F} = jacobian!(result, Base.Fix1(gradient, f), x)
+@inline function ForwardDiff.hessian!(result::AbstractArray, f::F, x::StaticArray) where {F}
+    T = typeof(Tag(f, eltype(x)))
+    ydual = f(dualize(T, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    H = ForwardDiff.reshape_hessian(result, x)
+    ForwardDiff.extract_hessian_chunk!(T, H, ydual, structural_indices(x), 0, 0, length(x), length(x))
+    return result
+end
 
 ForwardDiff.hessian!(result::MutableDiffResult, f::F, x::StaticArray) where {F} = hessian!(result, f, x, HessianConfig(f, result, x))
 
@@ -123,9 +148,10 @@ function ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray)
     d1 = dualize(T, x)
     d2 = dualize(T, d1)
     fd2 = f(d2)
+    fd2 isa Real || throw(HESSIAN_ERROR)
     val = value(T,value(T,fd2))
     grad = extract_gradient(T,value(T,fd2), x)
-    hess = extract_jacobian(T,partials(T,fd2), x)
+    hess = extract_hessian(T,partials(T,fd2), x)
     result = DiffResults.hessian!(result, hess)
     result = DiffResults.gradient!(result, grad)
     result = DiffResults.value!(result, val)

src/apiutils.jl

@@ -174,3 +174,17 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index,
         Dual{T,V,N}(value, seeds[i])
     end
 end
+
+# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer;
+# `seed_zero_partials!` cannot, as it would pass the primal where a nested `Dual` is wanted.
+function seed_hessian_chunk!(duals::AbstractArray{Dual{T,Dual{T,V,N},N}}, x, indices, index,
+                             iseeds::Union{Nothing,NTuple{N,Partials{N,V}}},
+                             oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}},
+                             chunksize = N) where {T,V,N}
+    izero = iseeds === nothing ? zero(Partials{N,V}) : nothing
+    ozero = oseeds === nothing ? zero(Partials{N,Dual{T,V,N}}) : nothing
+    return _seed!(duals, x, structural_chunk(indices, index, chunksize)) do value, i
+        inner = Dual{T,V,N}(value, iseeds === nothing ? izero : iseeds[i])
+        Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
+    end
+end

src/config.jl

@@ -206,9 +206,11 @@ Base.eltype(::Type{JacobianConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N}
 # HessianConfig #
 #################
 
-struct HessianConfig{T,V,N,DG,DJ,IG,IJ} <: AbstractConfig{N}
-    jacobian_config::JacobianConfig{T,V,N,DJ,IJ}
-    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG,IG}
+struct HessianConfig{T,V,N,D,I} <: AbstractConfig{N}
+    iseeds::NTuple{N,Partials{N,V}}
+    oseeds::NTuple{N,Partials{N,Dual{T,V,N}}}
+    duals::D
+    indices::I
 end
 
 """
@@ -218,10 +220,9 @@ Return a `HessianConfig` instance based on the type of `f` and type/shape of the
 vector `x`.
 
 The returned `HessianConfig` instance contains all the work buffers required by
-`ForwardDiff.hessian` and `ForwardDiff.hessian!`. For the latter, the buffers are
-configured for the case where the `result` argument is an `AbstractArray`. If
-it is a `DiffResult`, the `HessianConfig` should instead be constructed via
-`ForwardDiff.HessianConfig(f, result, x, chunk)`.
+`ForwardDiff.hessian` and `ForwardDiff.hessian!`, including when the latter stores into a
+`DiffResult`. The `ForwardDiff.HessianConfig(f, result, x, chunk)` constructor may also be
+used with any of these methods.
 
 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
@@ -231,11 +232,13 @@ This constructor does not store/modify `x`.
 """
 function HessianConfig(f::F,
                        x::AbstractArray{V},
-                       chunk::Chunk = Chunk(x),
-                       tag = Tag(f, V)) 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)
+                       ::Chunk{N} = Chunk(x),
+                       ::T = Tag(f, V)) where {F,V,N,T}
+    iseeds = construct_seeds(Partials{N,V})
+    oseeds = construct_seeds(Partials{N,Dual{T,V,N}})
+    duals = similar(x, Dual{T,Dual{T,V,N},N})
+    indices = structural_indices(duals)
+    return HessianConfig{T,V,N,typeof(duals),typeof(indices)}(iseeds, oseeds, duals, indices)
 end
 
 """
@@ -244,25 +247,20 @@ end
 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`.
+Equivalent to `ForwardDiff.HessianConfig(f, x, chunk)`: the work buffers do not depend on
+`result`. The result-aware form is retained for compatibility.
 
 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).
 
-This constructor does not store/modify `x`.
+This constructor does not store/modify `result` or `x`.
 """
-function HessianConfig(f::F,
-                       result::DiffResult,
-                       x::AbstractArray{V},
-                       chunk::Chunk = Chunk(x),
-                       tag = Tag(f, V)) 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)
-end
+HessianConfig(f::F,
+              ::DiffResult,
+              x::AbstractArray{V},
+              chunk::Chunk = Chunk(x),
+              tag = Tag(f, V)) where {F,V} = HessianConfig(f, x, chunk, tag)
 
 checktag(::HessianConfig{T},f,x) where {T} = checktag(T,f,x)
-Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ,IG,IJ}}) where {T,V,N,DG,DJ,IG,IJ} =
-    Dual{T,Dual{T,V,N},N}
+Base.eltype(::Type{HessianConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,Dual{T,V,N},N}

src/hessian.jl

@@ -5,7 +5,12 @@
 """
     ForwardDiff.hessian(f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x), check=Val{true}())
 
-Return `H(f)` (i.e. `J(∇(f))`) evaluated at `x`, assuming `f` is called as `f(x)`.
+Return `H(f)` evaluated at `x`, assuming `f` is called as `f(x)`.
+Multidimensional arrays are flattened in iteration order: the array
+`H(f)` has shape `length(x) × length(x)`, and its elements are
+`H(f)[j,k] = ∂²f(x)/∂x[j]∂x[k]`.
+The returned Hessian is exactly symmetric: its two triangles are filled from the same
+derivative values.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -14,15 +19,17 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function hessian(f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T,CHK}
     require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
-    return jacobian(∇f, x, cfg.jacobian_config, Val{false}())
+    checkstructure(cfg, x)
+    H, _ = symmetric_hessian(f, x, cfg, nothing)
+    return H
 end
 
 """
     ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x), check=Val{true}())
 
-Compute `H(f)` (i.e. `J(∇(f))`) evaluated at `x` and store the result(s) in `result`,
-assuming `f` is called as `f(x)`.
+Compute `H(f)` evaluated at `x` and store the result(s) in `result`, assuming `f` is
+called as `f(x)`. The stored Hessian is exactly symmetric: its two triangles are filled
+from the same derivative values.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -31,41 +38,131 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
     require_one_based_indexing(result, x)
     CHK && checktag(T, f, x)
-    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
-    jacobian!(result, ∇f, x, cfg.jacobian_config, Val{false}())
+    checkstructure(cfg, x)
+    symmetric_hessian!(reshape_hessian(result, x), f, x, cfg, nothing)
     return result
 end
 
-
-# We use this struct below instead of an
-# equivalent closure in order to avoid
-# JuliaLang/julia#15276-related performance
-# issues. See #316.
-mutable struct InnerGradientForHess{R,C,F}
-    result::R
-    cfg::C
-    f::F
-end
-
-function (g::InnerGradientForHess)(y, z)
-    inner_result = DiffResult(zero(eltype(y)), y)
-    gradient!(inner_result, g.f, z, g.cfg.gradient_config, Val{false}())
-    g.result = DiffResults.value!(g.result, value(DiffResults.value(inner_result)))
-    return y
-end
-
 """
     ForwardDiff.hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, result, x), check=Val{true}())
 
-Exactly like `ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig)`, but
-because `isa(result, DiffResult)`, `cfg` is constructed as `HessianConfig(f, result, x)` instead of
-`HessianConfig(f, x)`.
+Exactly like `ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig)`,
+but also stores the value and gradient in `result`. The default `cfg` is constructed as
+`HessianConfig(f, result, x)`, though a config constructed as `HessianConfig(f, x)` may also
+be used.
 
 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
 function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
+    require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    ∇f! = InnerGradientForHess(result, cfg, f)
-    jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}())
-    return ∇f!.result
+    checkstructure(cfg, x)
+    _, ydual = symmetric_hessian!(reshape_hessian(result, x), f, x, cfg,
+                                  DiffResults.gradient(result))
+    result = DiffResults.value!(result, value(T, value(T, ydual)))
+    return result
+end
+
+############################
+# symmetric Hessian kernel #
+############################
+
+const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?")
+
+function reshape_hessian(result::AbstractMatrix, x)
+    require_one_based_indexing(result)
+    size(result) == (length(x), length(x)) || throw(DimensionMismatch(
+        lazy"cannot store the $(length(x))×$(length(x)) Hessian in a result of size $(size(result))"))
+    return result
+end
+reshape_hessian(result::AbstractArray, x) = reshape(result, length(x), length(x))
+reshape_hessian(result::DiffResult, x) = reshape_hessian(DiffResults.hessian(result), x)
+
+# Copy a block from the nested partials and fill its transpose. On diagonal blocks, read
+# only the upper triangle so the result is exactly symmetric. Both axes are indexed by the linear
+# indices of `x`, as the columns of a Jacobian are, so `indices` gives the row and column of a block.
+function extract_hessian_chunk!(::Type{T}, H, ydual, indices, roffset, coffset, rsize, csize) where {T}
+    rows = structural_chunk(indices, roffset + 1, rsize)
+    cols = structural_chunk(indices, coffset + 1, csize)
+    for r in 1:rsize
+        drow = partials(T, ydual, r)
+        i = rows[r]
+        for c in (roffset == coffset ? r : 1):csize
+            h = partials(T, drow, c)
+            j = cols[c]
+            H[i, j] = h
+            H[j, i] = h
+        end
+    end
+    return H
+end
+
+# The inner partials of a diagonal block contain the corresponding gradient chunk.
+extract_hessian_gradient_chunk!(::Type{T}, ::Nothing, ydual, indices, index, chunksize) where {T} = nothing
+extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, indices, index, chunksize) where {T} =
+    extract_gradient_chunk!(T, grad, value(T, ydual), indices, index, chunksize)
+
+# Evaluate one pair of chunks at a time using nested duals. Only one triangle of block
+# pairs is evaluated; the other is filled by symmetry (see #836).
+function symmetric_hessian_expr(result_definition::Expr)
+    return quote
+        xlen = structural_length(x)
+        if xlen < N
+            throw(ArgumentError(lazy"chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))"))
+        end
+
+        # `N == 0` only for empty inputs, which still need one evaluation to determine the
+        # output type and value.
+        nblocks = xlen == 0 ? 1 : cld(xlen, N)
+
+        xdual = cfg.duals
+        indices = cfg.indices
+        iseeds = cfg.iseeds
+        oseeds = cfg.oseeds
+
+        # The first evaluation determines the output type. Seeding the first block and clearing
+        # the untouched tail partitions the fresh buffer, so every element is initialized once.
+        seed_hessian_chunk!(xdual, x, indices, 1, iseeds, oseeds)
+        seed_hessian_chunk!(xdual, x, indices, N + 1, nothing, nothing, xlen - N)
+        ydual1 = f(xdual)
+        ydual1 isa Real || throw(HESSIAN_ERROR)
+        $(result_definition)
+        # unwrapped once so that the helpers' `valtype` reaches the value type
+        zero_unseeded_columns!(T, H, value(T, ydual1), x)
+        grad === nothing || zero_unseeded!(T, grad, value(T, ydual1), x)
+        extract_hessian_chunk!(T, H, ydual1, indices, 0, 0, N, N)
+        extract_hessian_gradient_chunk!(T, grad, ydual1, indices, 1, N)
+        nblocks > 1 && seed_hessian_chunk!(xdual, x, indices, 1, nothing, nothing)
+
+        for q in 2:nblocks
+            qoffset = (q - 1) * N
+            qsize = min(N, xlen - qoffset)
+            # Outer-i inner-j and outer-j inner-i round differently, so the outer layer always
+            # takes the earlier position -- else the result would depend on the chunk size.
+            seed_hessian_chunk!(xdual, x, indices, qoffset + 1, iseeds, nothing, qsize)
+            for p in 1:(q - 1)
+                poffset = (p - 1) * N
+                seed_hessian_chunk!(xdual, x, indices, poffset + 1, nothing, oseeds)
+                ydual = f(xdual)
+                extract_hessian_chunk!(T, H, ydual, indices, poffset, qoffset, N, qsize)
+                seed_hessian_chunk!(xdual, x, indices, poffset + 1, nothing, nothing)
+            end
+            # The diagonal block adds q's outer seeds while retaining its inner seeds.
+            seed_hessian_chunk!(xdual, x, indices, qoffset + 1, iseeds, oseeds, qsize)
+            ydual = f(xdual)
+            extract_hessian_chunk!(T, H, ydual, indices, qoffset, qoffset, qsize, qsize)
+            extract_hessian_gradient_chunk!(T, grad, ydual, indices, qoffset + 1, qsize)
+            seed_hessian_chunk!(xdual, x, indices, qoffset + 1, nothing, nothing, qsize)
+        end
+
+        return H, ydual1
+    end
+end
+
+@eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N}
+    $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), length(x), length(x)))))
+end
+
+@eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N}
+    $(symmetric_hessian_expr(:()))
 end

test/AllocationsTest.jl

@@ -11,9 +11,7 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
 @testset "Test seed!/seed_zero_partials! allocations" begin
     x = rand(1000)
     cfg = ForwardDiff.GradientConfig(nothing, x)
-    duals = cfg.duals
-    seeds = cfg.seeds
-    indices = cfg.indices
+    (; duals, seeds, indices) = cfg
 
     allocs_seed!(args...) = @allocated ForwardDiff.seed!(args...)
     allocs_seed!(duals, x, indices, seeds)
@@ -34,6 +32,33 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
     @test iszero(allocs_convert_test_574())
 end
 
+@testset "Test seed_hessian_chunk! allocations" begin
+    x = rand(1000)
+    cfg = ForwardDiff.HessianConfig(nothing, x)
+    (; duals, indices, iseeds, oseeds) = cfg
+
+    allocs_hseed!(args...) = @allocated ForwardDiff.seed_hessian_chunk!(args...)
+    # all four seed combinations, the mixed ones being what the off-diagonal blocks use
+    @testset "iseeds=$(i === nothing) oseeds=$(o === nothing)" for (i, o) in
+                                                                  ((iseeds, oseeds),
+                                                                   (iseeds, nothing),
+                                                                   (nothing, oseeds),
+                                                                   (nothing, nothing))
+        allocs_hseed!(duals, x, indices, 1, i, o)
+        @test iszero(allocs_hseed!(duals, x, indices, 1, i, o))
+        allocs_hseed!(duals, x, indices, 1, i, o, 4)
+        @test iszero(allocs_hseed!(duals, x, indices, 1, i, o, 4))
+    end
+
+    # a zero of a non-isbits value type does allocate, so supplying both seeds must not build one
+    @testset "BigFloat" begin
+        y = BigFloat[1, 2, 3]
+        cfg = ForwardDiff.HessianConfig(nothing, y, ForwardDiff.Chunk{3}())
+        allocs_hseed!(cfg.duals, y, cfg.indices, 1, cfg.iseeds, cfg.oseeds)
+        @test iszero(allocs_hseed!(cfg.duals, y, cfg.indices, 1, cfg.iseeds, cfg.oseeds))
+    end
+end
+
 @testset "Test jacobian! allocations" begin
     # jacobian! should not allocate when called with a pre-allocated result Matrix.
     # Previously, reshape() inside extract_jacobian! allocated a wrapper