Optimize the conversion between JuMP and ExaModels

[joaquimg]

Given the model:

using ExaModels, JuMP, CUDA

model = Model()

@variable(model, ...)

@constraint(model, c1[i = ...], expr1... == 0.0)

@constraint(model, c2[i = ...], expr1... == 0.0)

@objective(model, Min, expr3...)

that would be converted with:

em = ExaModel(model; backend = CUDABackend())

It would be possible to use JuMP attributes to communicate part of the structure:

JuMP.set_attribute.(c1, ExaModels.ConstraintTag(), "c1")
JuMP.set_attribute.(c2, ExaModels.ConstraintTag(), "c2")

This way all the constraints in the c1 list would be tagged with "c1" (analogous for c2).
Hence, the step of grouping constraints could be done only within the tag groups, which would minimize the work.

[sshin23]

Hmm, that's an interesting idea.

We still don't have a guarantee that each JuMP constraint will lead to a single type of ExaModel constraint. E.g.,

m = JuMP.Model()
@variable(model, x[1:10])
@constraint(model, [i=1:10], mod(i,2) == 0 ? sin(x[i]) : cos(x[i]) == 0)

will lead to two types of ExaModels constraint.

Not sure how much difference tagging the constraints will make. We first convert the JuMP expression into an ExaModels expression
https://github.com/exanauts/ExaModels.jl/blob/main/ext/ExaModelsMOI.jl#L253-L410
and then classify them. Once they are converted into an ExaModel expression, classification is pretty quick. This tagging wouldn't help make conversion faster

[joaquimg]

I admit that I do not fully understand the JuMP->Exa converter. My comments come from watching some presentations on ExaModels. I might have gotten some bits wrong.

About your example, tagging would be a user responsibility. In your example, the user would have to:

@constraint(model, c[i=1:10], mod(i,2) == 0 ? sin(x[i]) : cos(x[i]) == 0)
JuMP.set_attribute.(c[collect(i for i in 1:10 if mod(i,2) == 0)], ExaModels.ConstraintTag(), "c1")
JuMP.set_attribute.(c[collect(i for i in 1:10 if mod(i,2) != 0)], ExaModels.ConstraintTag(), "c2")

About converting and classifying, if conversion is the bottleneck, this would not help much.

On the other hand, there is hope if the two operations are interwined, that is, some precomputations of classifying occur at conversion.

Again, I might have gotten it wrong. I see there are _exafy and update_bin. If exafy is the bottleneck, there is no hope for tagging. But if the bottleneck is update_bin, which has to compare a new expression with previously created ones in tree, there is hope. Because smaller trees with similar candidates might be much faster to walk through than a single big tree.

[sshin23]

I haven't done profiling, but I'm quite sure exafy (the conversion) is the bottleneck; the JuMP expression trees are not concrete, so conversion tends to be slow. But once it is converted to ExaModel, classification can be done pretty much by just looking at the type of expression tree.

In the long term, to have the best of both worlds (flexibility of JuMP and speed of ExaModels), I think a proper approach would be to implement a special constraint syntax where the user can specify templatized constraint (and objective) and set ExaModels as a backend for doing AD for that part of the model.

[odow]

@blegat and I are going to take a look at this. I think we can take a lot of inspiration from the code in SymbolicAD to speed this up. The existing code can identify shared structure. We probably need some sort of extension to detect shared structure in summations. x-ref lanl-ansi/MathOptSymbolicAD.jl#22

See here:

https://github.com/jump-dev/MathOptInterface.jl/blob/03d9ada0465ecac462f78bd5fcd39ec471003fe9/src/Nonlinear/SymbolicAD/SymbolicAD.jl#L798-L1202

[blegat]

I think it's always nice to keep an explicit path. If the users uses GenOpt, he tells explicitly where the repeated structure are so we could differentiate that part with ExaModels and the other constraints with ArrayDiff or ReverseAD. We could also have an utilities inspired from MathOptSymbolicAD that tries to reconstruct the repeated structure as an MOI layer that can takes classical MOI input and produce output containing GenOpt constraints.

[odow]

I need this for LANL. My only goal would be to improve the speed of https://github.com/exanauts/ExaModels.jl/blob/main/ext/ExaModelsMOI.jl to make it competitive for AC-OPF problems.

I don't want to add any dependencies on GenOpt or ArrayDiff etc.

[blegat]

I think it's cleaner from a long-term perspective to have an explicit interface and then view things as preprocessing steps that then go through these explicit interfaces. For now, it might be indeed be easier to first experiment with ideas that are specific to ExaModelsMOI.jl. We can see how generalizable it is in a later stage.