Add ExaModels.SIMDMode: ExaModels as an MOI.Nonlinear AD backend - #307
Add ExaModels.SIMDMode: ExaModels as an MOI.Nonlinear AD backend#307blegat wants to merge 4 commits into
Conversation
The purpose of this rewrite is to: * simplify the implementation by removing the linked list * thoroughly document and comment the file to explain what is happening * add various improvements: split additive terms into separate constraint augmentations, support :(x - y), skip constants if they are 0.0, re-use variables in symbolic expressions * fix various bugs in the MOI wrapper * add many more tests
|
This would still require a CPU <-> GPU switch at each iteration. Part of the reason for ExaModels.Optimizer is that the entire problem gets copied to the GPU, instead of just evaluating the derivatives. |
| objective::Any | ||
| constraints::Vector{Tuple{Any,Any}} | ||
| SIMDNonlinearModel() = new(nothing, Tuple{Any,Any}[]) | ||
| end |
There was a problem hiding this comment.
Why does this PR need a new model object? Can't we just build the ExaModel from the existing MOI.Nonlinear.Model?
There was a problem hiding this comment.
Good question, the issue is that it is easier to do the conversion at once but MOI.Nonlinear.add_constraint gives things one by one. For example, we need ordered_variables which is given only at MOI.Nonlinear.Evaluator. This is kind of the supports_incremental_interface we have in MOI. So it means ExaModels's AD backend should receive the model as a whole. I'm starting to think the interface should simply be
ad_model = MOI.Nonlinear.evaluator(ad_backend, model, ordered_variables, features)
# Now we can call the callbacks. `ad_backend` can do
# whathever he wants as long as the callbacks workThen so here ExaModels can implement copy_to so it does not need to create any new custom objects.
If often happens that we only want to give part of the model to ad_model but that's very easy to do with MOI.Utilities.Filter, CCOpt is doing this for complementarity constraints.
Incremental interface is nice for solver supporting modifications but it's annoying in other cases. We've already simplified our life greatly by allowing solvers to just implement MOI.copy_to, I think we just need to do the same here.
There was a problem hiding this comment.
But if the AD model does not support incremental mode, it also means that any solvers needing an AD model cannot support it so I'll first check if there is no way to just work around it
There was a problem hiding this comment.
Actually, your PR #290 makes it easy to support incremental interface, I'll just assume ordered_variables is trivial for now
b700d45 to
ed08937
Compare
Instead of ExaModels pretending to be a solver via ExaModels.Optimizer, any MOI solver that supports MOI.AutomaticDifferentiationBackend can now evaluate its nonlinear model through ExaModels by setting the backend to ExaModels.SIMDMode(; device). The extension defines: - SIMDNonlinearModel: the model built by MOI.Nonlinear.model(::SIMDMode), wrapped only in ModelWithOracles because ExaModels consumes affine and quadratic functions natively. Constraints and the objective are translated to SIMD-grouped bins as they are added, reusing the incremental update_bin! machinery of ExaModels.Optimizer (including its term-splitting of affine, quadratic, and +/- expressions into separate deduplicated bins). - SIMDEvaluator: assembles the ExaCore from the bins during MOI.initialize, once the number of variables is known, and delegates the MOI evaluator callbacks (including the Jacobian and Hessian products, with ExaModel(prod = true) when requested) to the NLPModels API of the resulting ExaModel. - constraint_bounds/num_constraints/constraint_linearity/ objective_linearity queries, and exploits_structure(::SIMDMode) = true so that consumers pass functions unparsed. The bins reference the variables by their raw index, so the ordered_variables of the evaluator must be MOI.VariableIndex.(1:n); MOI.initialize errors otherwise. User-defined operators, MOI.Nonlinear parameters, and expressions are not supported and error explicitly.
2 participants
Instead of ExaModels pretending to be a solver via ExaModels.Optimizer, any MOI solver that supports
MOI.AutomaticDifferentiationBackendcan now evaluate its nonlinear model through ExaModels by setting the backend to ExaModels.SIMDMode(; device). Written by Claude Fable 5, still need to check. Needs