[ReverseAD] Add support for user-defined hessians

Author: odow

src/Nonlinear/ReverseAD/forward_over_reverse.jl

@@ -126,6 +126,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T}
     for i in ex.dependent_subexpressions
         subexpr = d.subexpressions[i]
         subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(
+            d,
             subexpr,
             _reinterpret_unsafe(T, subexpr.forward_storage_ϵ),
             _reinterpret_unsafe(T, subexpr.partials_storage_ϵ),
@@ -135,6 +136,7 @@ function _hessian_slice_inner(d, ex, input_ϵ, output_ϵ, ::Type{T}) where {T}
         )
     end
     _forward_eval_ϵ(
+        d,
         ex,
         _reinterpret_unsafe(T, d.forward_storage_ϵ),
         _reinterpret_unsafe(T, d.partials_storage_ϵ),
@@ -178,6 +180,7 @@ end
 
 """
     _forward_eval_ϵ(
+        d::NLPEvaluator,
         ex::Union{_FunctionStorage,_SubexpressionStorage},
         storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
         partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
@@ -195,6 +198,7 @@ components separate so that we don't need to recompute the real components.
 This assumes that `_reverse_model(d, x)` has already been called.
 """
 function _forward_eval_ϵ(
+    d::NLPEvaluator,
     ex::Union{_FunctionStorage,_SubexpressionStorage},
     storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
     partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}},
@@ -328,10 +332,35 @@ function _forward_eval_ϵ(
                         recip_denominator,
                     )
                 elseif op > 6
-                    error(
-                        "User-defined operators not supported for hessian " *
-                        "computations",
+                    f_input = _UnsafeVectorView(d.jac_storage, n_children)
+                    for (i, c) in enumerate(children_idx)
+                        f_input[i] = ex.forward_storage[children_arr[c]]
+                    end
+                    H = _UnsafeHessianView(d.user_output_buffer, n_children)
+                    has_hessian = Nonlinear.eval_multivariate_hessian(
+                        user_operators,
+                        user_operators.multivariate_operators[node.index],
+                        H,
+                        f_input,
                     )
+                    if !has_hessian
+                        continue
+                    end
+                    for col in 1:n_children
+                        dual = zero(ForwardDiff.Partials{N,T})
+                        for row in 1:n_children
+                            # Make sure we get the lower-triangular component.
+                            h = row >= col ? H[row, col] : H[col, row]
+                            # Performance optimization: hessians can be quite
+                            # sparse
+                            if !iszero(h)
+                                i = children_arr[children_idx[row]]
+                                dual += h * storage_ϵ[i]
+                            end
+                        end
+                        i = children_arr[children_idx[col]]
+                        partials_storage_ϵ[i] = dual
+                    end
                 end
             elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE
                 @inbounds child_idx = children_arr[ex.adj.colptr[k]]

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -7,8 +7,10 @@
 function MOI.features_available(d::NLPEvaluator)
     # Check if we have any user-defined multivariate operators, in which case we
     # need to disable hessians. The result of features_available depends on this.
-    d.disable_2ndorder =
-        length(d.data.operators.registered_multivariate_operators) > 0
+    d.disable_2ndorder = any(
+        op -> op.∇²f === nothing,
+        d.data.operators.registered_multivariate_operators,
+    )
     if d.disable_2ndorder
         return [:Grad, :Jac, :JacVec]
     end
@@ -286,6 +288,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
     for i in d.subexpression_order
         subexpr = d.subexpressions[i]
         subexpr_forward_values_ϵ[i] = _forward_eval_ϵ(
+            d,
             subexpr,
             reinterpret(T, subexpr.forward_storage_ϵ),
             reinterpret(T, subexpr.partials_storage_ϵ),
@@ -302,6 +305,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
     fill!(output_ϵ, zero(T))
     if d.objective !== nothing
         _forward_eval_ϵ(
+            d,
             d.objective,
             reinterpret(T, d.forward_storage_ϵ),
             reinterpret(T, d.partials_storage_ϵ),
@@ -322,6 +326,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ)
     end
     for (i, con) in enumerate(d.constraints)
         _forward_eval_ϵ(
+            d,
             con,
             reinterpret(T, d.forward_storage_ϵ),
             reinterpret(T, d.partials_storage_ϵ),

src/Nonlinear/ReverseAD/utils.jl

@@ -60,6 +60,68 @@ function _UnsafeVectorView(x::Vector, N::Int)
     return _UnsafeVectorView(0, N, pointer(x))
 end
 
+"""
+    _UnsafeHessianView(x, N)
+
+Lightweight unsafe view that converts a vector `x` into the lower-triangular
+component of a symmetric `N`-by-`N` matrix.
+
+## Motivation
+
+`_UnsafeHessianView` is needed as an allocation-free equivalent of `view`. Other
+alternatives, like `reshape(view(x, 1:N^2), N, N)` or a struct like
+```julia
+struct _SafeView{T}
+    x::Vector{T}
+    len::Int
+end
+```
+will allocate so that `x` can be tracked by Julia's GC.
+`_UnsafeHessianView` relies on the fact that the use-cases of
+`_UnsafeHessianView` only temporarily wrap a long-lived vector like
+`d.jac_storage` so that we don't have to worry about the GC removing
+`d.jac_storage` while `_UnsafeHessianView` exists. This lets us use a `Ptr{T}`
+and create a struct that is `isbitstype` and therefore does not allocate.
+
+## Unsafe behavior
+
+`_UnsafeHessianView` is unsafe because it assumes that the vector `x` remains
+valid during the usage of `_UnsafeHessianView`.
+"""
+struct _UnsafeHessianView <: AbstractMatrix{Float64}
+    N::Int
+    ptr::Ptr{Float64}
+end
+
+Base.size(x::_UnsafeHessianView) = (x.N, x.N)
+
+function _linear_index(row, col)
+    if row < col
+        error("Unable to access upper-triangular component: ($row, $col)")
+    end
+    return div((row - 1) * row, 2) + col
+end
+
+function Base.getindex(x::_UnsafeHessianView, i, j)
+    return unsafe_load(x.ptr, _linear_index(i, j))
+end
+
+function Base.setindex!(x::_UnsafeHessianView, value, i, j)
+    unsafe_store!(x.ptr, value, _linear_index(i, j))
+    return value
+end
+
+function _UnsafeHessianView(x::Vector, N::Int)
+    z = div(N * (N + 1), 2)
+    if length(x) < z
+        resize!(x, z)
+    end
+    for i in 1:z
+        x[i] = 0.0
+    end
+    return _UnsafeHessianView(N, pointer(x))
+end
+
 function _reinterpret_unsafe(::Type{T}, x::Vector{R}) where {T,R}
     # how many T's fit into x?
     @assert isbitstype(T) && isbitstype(R)

test/Nonlinear/ReverseAD.jl

@@ -287,14 +287,13 @@ function test_hessian_sparsity_registered_function()
     Nonlinear.set_objective(model, :(f($x, $z) + $y^2))
     evaluator =
         Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y, z])
-    @test_broken :Hess in MOI.features_available(evaluator)
-    # TODO(odow): re-enable these tests when user-defined hessians are supported
-    # MOI.initialize(evaluator, [:Grad, :Jac, :Hess])
-    # @test MOI.hessian_lagrangian_structure(evaluator) ==
-    #       [(1, 1), (2, 2), (3, 3), (3, 1)]
-    # H = fill(NaN, 4)
-    # MOI.eval_hessian_lagrangian(evaluator, H, rand(3), 1.5, Float64[])
-    # @test H == 1.5 .* [2.0, 2.0, 2.0, 0.0]
+    @test :Hess in MOI.features_available(evaluator)
+    MOI.initialize(evaluator, [:Grad, :Jac, :Hess])
+    @test MOI.hessian_lagrangian_structure(evaluator) ==
+          [(1, 1), (2, 2), (3, 3), (3, 1)]
+    H = fill(NaN, 4)
+    MOI.eval_hessian_lagrangian(evaluator, H, rand(3), 1.5, Float64[])
+    @test H == 1.5 .* [2.0, 2.0, 2.0, 0.0]
     return
 end
 
@@ -318,13 +317,12 @@ function test_hessian_sparsity_registered_rosenbrock()
     Nonlinear.set_objective(model, :(rosenbrock($x, $y)))
     evaluator =
         Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y])
-    @test_broken :Hess in MOI.features_available(evaluator)
-    # TODO(odow): re-enable these tests when user-defined hessians are supported
-    # MOI.initialize(evaluator, [:Grad, :Jac, :Hess])
-    # @test MOI.hessian_lagrangian_structure(evaluator) == [(1, 1), (2, 2), (2, 1)]
-    # H = fill(NaN, 3)
-    # MOI.eval_hessian_lagrangian(evaluator, H, [1.0, 1.0], 1.5, Float64[])
-    # @test H == 1.5 .* [802, 200, -400]
+    @test :Hess in MOI.features_available(evaluator)
+    MOI.initialize(evaluator, [:Grad, :Jac, :Hess])
+    @test MOI.hessian_lagrangian_structure(evaluator) == [(1, 1), (2, 2), (2, 1)]
+    H = fill(NaN, 3)
+    MOI.eval_hessian_lagrangian(evaluator, H, [1.0, 1.0], 1.5, Float64[])
+    @test H == 1.5 .* [802, 200, -400]
     return
 end