automatically generated lagrangian allocates

[baggepinnen]

The automatically generated lagrangian allocates a lot of memory, killing performance
https://github.com/SciML/OptimizationBase.jl/blob/main/src/OptimizationDISparseExt.jl#L121C2-L137C12

It would be nice if it could use a cache instead of recreating the _res = zeros(eltype(x), num_cons) on each invokation

[gdalle]

I think this could work with contexts of type DifferentiationInterface.Cache, but it is a less well-tested part of the package so beware of the dragons. Wanna try it out @Vaibhavdixit02?

[Vaibhavdixit02]

The issue there was knowing the types upfront, I haven't looked at DifferentiationInterface.Cache maybe you have handled it already @gdalle

[gdalle]

What do you mean by knowing the types upfront?

[baggepinnen]

To preallocate the cache it's beneficial to know the type of the dual number that the function will be called with. Generally, the function will be called with more than one type, so another strategy is to make use of a more flexible cache object that allocates upon first call with a new type.

[ChrisRackauckas]

PreallocationTools handles this by just having the cache of the right or larger size, since the type doesn't matter if you can reinterpret, but what is important is knowing the total memory size.

[gdalle]

So the idea of DI.Cache with ForwardDiff is precisely to allocate something with the correct dual type for the differentiation operator in question. It isn't optimized yet but it works, and preallocation will be easy to add. However it doesn't handle multiple types by default.

[gdalle]

In this case it might look like:

function lagrangian(θ, σ, λ, p, res)
    return σ * f.f(θ, p) + dot(λ, f.cons(res, θ, p))
end

And then you would call

DI.hessian(lagrangian, backend, θ, DI.Constant(σ), DI.Constant(λ), DI.Constant(p), DI.Cache(res))

At the moment it will not bring benefits because the Cache is re-allocated at every call. However it is a rather straightforward PR to allocate it only once during preparation. I just never got around to it because there were more pressing matters and nobody was using Cache. If there is interest, I can try to get it done soon

[gdalle]

In the upcoming DI release (v0.6.44) you can do this without any allocation when ForwardDiff is used:

using Chairmarks, LinearAlgebra
using DifferentiationInterface, SparseConnectivityTracer, SparseMatrixColorings
using ForwardDiff: ForwardDiff

function constraints!(c, x)
    c .= x .- 1
    return nothing
end

objective(x) = sum(abs2, x)

function lagrangian!(x, λ, c)
    constraints!(c, x)
    f = objective(x)
    return f + dot(λ, c)
end

backend = AutoSparse(
    AutoForwardDiff();
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(),
)

julia> x, c, λ = rand(10), zeros(10), rand(10);

julia> prep = prepare_hessian(lagrangian!, backend, x, Constant(λ), Cache(c));

julia> H = hessian(lagrangian!, prep, backend, x, Constant(λ), Cache(c));

julia> @b hessian!($lagrangian!, $H, $prep, $backend, $x, Constant($λ), Cache($c))  # alloc-free
354.671 ns