Work towards “in-place first” internally

[kellertuer]

For the high-level interfaces it is nice that

gradient_descent(M, f, grad_f)

works for the allocating gradient grad_f(M,p) by default.
This aims to be easier to use.

Internally this is a bit chaotic since

• one always has to distinguish allocating and in-place everywhere.
• an objective can only be one of a kind. For example for the Hessian, both gradient and Hessian have to be the same kind of function (alloc vs nonalloc)

Since the first step is nice for first-time users, keeping that default would be nice; but internally always preferring in-place would be nice.

One possibiliy is maybe to “wrap allocating early” in a function wrapper

struct AllocationFunctionWrapper{T,F}
function::F
end
  # Gradient example:
  function (f:: AllocationFunctionWrapper)(M, X, p)
  # evaluate allocation function and copy over
  return copyto!(M, X, p, f.function(M,p)
end

where T could be a type to determine different things to do (as opposed to copying over the gradient), though most often it is probably just a tangent vector or a point.
Then internally all Objectivbes could always assume having in-place functions; their constructors could maybe be tolerant still.

This is an issue more for the long-time, since it would require a humungous rework.

[mateuszbaran]

Nearly all internals should only ever use the allocating in-place variant IMO so I think it wouldn't be all that chaotic. The chaos would be limited to figuring out user input so internals can work with that easily. Manifolds.jl handles the in-place-first approach quite well for example.

EDIT: corrected "allocating" to "in-place".

[kellertuer]

I would have thought the internal functions should always work in-place – not allocating?

And I meant it is chaotic now, in such a new form it is probably less chaotic; sorry if that was not stated clearly.

[mateuszbaran]

Yes, internally everything should work in place; right now it's quite inconsistent, even in the new RLM code.

[kellertuer]

Ah I might have misread your first sentence in the previous comment, maybe?

Nearly all internals should only ever use the allocating variant

read to me like you want to switch to alloc; but you want to switch to inlace consistently.

Yes, currently it is a bit chaotic, indeed, since (a) all methods are always implemented twice (a alloc, b in-place fields in the objective) and for both alloc and in-place calls; so overall often 4 times.
This should be improved to

1. 1 implementation of in-place fields of an objective (like it is always objective.grad!(M, X, p))
2. if a user provides an allocating variant (like grad(M, p)) it should be wrapped as described above (on the high-level interface only)
3. maybe in a few places one could do a second generic implementation (similar to exp(M, p, X)) that allocates and calls the function from 1) – probably even on a very abstract objective level.

[mateuszbaran]

Ah, sorry, I mistyped in the first comment. I meant in-place of course. Your proposed improvement sounds good.

[kellertuer]

chrchr yeah that was a confusing typo then, but good that we have clarified that and that we do agree!

[SteffenPL]

Going purely to in-place would make it impossible though to use SMatrix/SVector representation though?

I don't know if there was prior discussion, but especially when matrix multiplications and such are chained, writing a in-place function with mul!(...) etc requires cache or lazy computing, whereas a SMatrix/SVector implementation is close to the math...

(The code below unfortunately falls onto MMatrix anyway, but it would be neat if SMatrix would work.)

Example

using LinearAlgebra
using Manifolds
using Manopt
using StaticArrays

function half_squared_distance(S, A)
    return 0.5 * sum(abs2, S .- A)
end

function closest_symplectic(A; p0 = typeof(A)(I), maxiter = 1000, tol = 1e-10)
    m, k = size(A)
    @assert m == k && iseven(m) "A must be a square 2n x 2n matrix"

    M = SymplecticMatrices(m)
    J = SymplecticElement(A)
    f(M, S) = half_squared_distance(S, A)
    function grad_f(M, S)
        Y = S - A
        JS = J * S
        return Y * S' * S + JS * Y' * JS
    end
    objective = ManifoldGradientObjective(f, grad_f; evaluation = AllocatingEvaluation())
    result = gradient_descent(
        M, objective, p0;
        stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol),
        retraction_method = CayleyRetraction(),
        debug = [],
    )
    return result
end

function symplectic_residual(S)
    J = SymplecticElement(S)
    return norm(S' * (J * S) - J)
end

A = SMatrix{4,4}(randn(4, 4))
p0 = SMatrix{4,4,Float64}(I)
@time S = closest_symplectic(A; p0 = p0, maxiter = 100, tol = 1e-4);

@show norm(S - A)
@show symplectic_residual(S)

[kellertuer]

No we would not go “in-place only”.

We would

• get rid of the current parameter at the objective
• only store in-place functions
• wrap allocating functions accordingly

Probably more like a “localize” these two kinds a bit more and hence being closer to what we already do in Manifolds.jl
For best of cases I do not even expect the high-level interfaces to change – but internally it would both simplify a lot of code and make several places more efficient.

[SteffenPL]

Mh, correct me if I'm wrong, but with if a function f_grad(M,p) gets internally translated to

function f_grad!(M, X, p) 
   X .= f_grad(M,p)  
   return nothing
end

then we do get another copy compared to X = f_grad(M,p) in case of SMatrix? (But probable, a copy is not that relevant in the big picture 😃)

[mateuszbaran]

Going purely to in-place would make it impossible though to use SMatrix/SVector representation though?

In vast majority of cases caches is the right way for optimization due to its iterative nature. For small arrays MArray caches offer similar performance to SArray, and the only case I know where SArray have a clear advantage is CUDA kernel code generation but that's very niche and IMO justifies having a separate implementation. AFAIK the SciML ecosystem settled on a similar model.

User code would still be able to use SArray but there would be a wrapper that converts to MArray-based optimization.

then we do get another copy compared to X = f_grad(M,p) in case of SMatrix? (But probable, a copy is not that relevant in the big picture 😃)

There are scenarios where this additional copy would reduce performance but you need significantly different implementation of optimization algorithms to exploit that. That other implementation would make things much slower in the large array case, which is the primary target here.

[SteffenPL]

Ah, I see! Thanks :) (I am on the GPU side indeed but not really performance critical at this point.) Thanks for the details by the way!

[kellertuer]

...the internal function would of course return X and not nothing, but the main different is that instead of that e.g. the ManifoldGradientObjective being the one that “knows” (by type parameter) that its gradient is in-place or not, the “gradient itself would know”, so it localises a bit more and that could also bring an advantage in other cases.

...but to be honest I am definetly not the performance expert here. I mainly annoy everyone with interface design, documentation and testing, I feel that is my main job ;)