Add support for multi-indexed constraints (#69)

* MOI fix

* multi_con

* Double index constraints

* first attempt -- doesn't look promising

* new idx introduced

* Working. Need to add a test

* added test

* dead end?

* collect approach

* test fix

* test fix

* typo fix

* fixes

* works except for roc and oneapi

* fixed oneapi and amdgpu issues

* rollback changes in ext packages

---------

Co-authored-by: Sungho Shin <[email protected]>

Author: sshin23

ext/ExaModelsCUDA.jl

@@ -3,7 +3,6 @@ module ExaModelsCUDA
 import ExaModels: ExaModels, NLPModels
 import CUDA: CUDA, CUDABackend, CuArray
 
-ExaModels.ExaCore(backend::CUDABackend) = ExaModels.ExaCore(Float64, backend)
 ExaModels.convert_array(v, backend::CUDABackend) = CuArray(v)
 
 end

ext/ExaModelsMOI.jl

@@ -332,11 +332,16 @@ function _exafy(e::MOI.ScalarNonlinearFunction, p = ())
 end
 
 function _exafy(e::MOI.ScalarAffineFunction{T}, p = ()) where {T}
-    return sum(begin
-        c1, p = _exafy(term, p)
-        c1
-    end for term in e.terms) + ExaModels.ParIndexed(ExaModels.ParSource(), length(p) + 1),
-    (p..., e.constant)
+    ec = if !isempty(e.terms)
+        sum(begin
+                c1, p = _exafy(term, p)
+                c1
+            end for term in e.terms) + ExaModels.ParIndexed(ExaModels.ParSource(), length(p) + 1)
+    else
+        ExaModels.ParIndexed(ExaModels.ParSource(), length(p) + 1)
+    end
+    
+    return ec, (p..., e.constant)
 end
 
 function _exafy(e::MOI.ScalarAffineTerm{T}, p = ()) where {T}

ext/ExaModelsOneAPI.jl

@@ -7,15 +7,15 @@ function ExaModels.append!(
     a::A,
     b::Base.Generator{UnitRange{I}},
     lb,
-) where {I,A<:oneAPI.oneVector}
+) where {I,A<:oneAPI.oneArray}
     la = length(a)
     aa = similar(a, la + lb)
     copyto!(view(aa, 1:la), a)
     map!(b.f, view(aa, (la+1):(la+lb)), b.iter)
     return aa
 end
 
-function ExaModels.append!(backend, a::A, b::Base.Generator, lb) where {A<:oneAPI.oneVector}
+function ExaModels.append!(backend, a::A, b::Base.Generator, lb) where {A<:oneAPI.oneArray}
     la = length(a)
     aa = similar(a, la + lb)
     copyto!(view(aa, 1:la), a)
@@ -29,7 +29,7 @@ function ExaModels.append!(
     a::A,
     b::V,
     lb,
-) where {A<:oneAPI.oneVector,V<:AbstractVector}
+) where {A<:oneAPI.oneArray,V<:AbstractArray}
     la = length(a)
     aa = similar(a, la + lb)
     copyto!(view(aa, 1:la), a)
@@ -38,7 +38,7 @@ function ExaModels.append!(
 end
 
 
-function ExaModels.append!(backend, a::A, b::Number, lb) where {A<:oneAPI.oneVector}
+function ExaModels.append!(backend, a::A, b::Number, lb) where {A<:oneAPI.oneArray}
     la = length(a)
     aa = similar(a, la + lb)
     copyto!(view(aa, 1:la), a)
@@ -48,16 +48,16 @@ end
 
 ExaModels.convert_array(v, backend::oneAPI.oneAPIBackend) = oneAPI.oneArray(v)
 
-ExaModels.sort!(array::A; lt = isless) where {A<:oneAPI.oneVector} =
+ExaModels.sort!(array::A; lt = isless) where {A<:oneAPI.oneArray} =
     copyto!(array, sort!(Array(array); lt = lt))
 
 # below is type piracy
-function Base.findall(f::F, bitarray::A) where {F<:Function,A<:oneAPI.oneVector}
+function Base.findall(f::F, bitarray::A) where {F<:Function,A<:oneAPI.oneArray}
     a = Array(bitarray)
     b = findall(f, a)
     c = similar(bitarray, eltype(b), length(b))
     return copyto!(c, b)
 end
-Base.findall(bitarray::A) where {A<:oneAPI.oneVector} = Base.findall(identity, bitarray)
+Base.findall(bitarray::A) where {A<:oneAPI.oneArray} = Base.findall(identity, bitarray)
 
 end # module

src/nlp.jl

@@ -223,7 +223,7 @@ julia> result = ipopt(m; print_level=0)    # solve the problem
 
 ```
 """
-function ExaModel(c::C; kwargs...) where {C<:ExaCore}
+function ExaModel(c::C; prod = nothing) where {C<:ExaCore}
     return ExaModel(
         c.obj,
         c.con,
@@ -254,15 +254,16 @@ end
 )
 
 function append!(backend, a, b::Base.Generator, lb)
-
+    b = _adapt_gen(b)
+    
     la = length(a)
     resize!(a, la + lb)
     map!(b.f, view(a, (la+1):(la+lb)), convert_array(b.iter, backend))
     return a
 end
 
 function append!(backend, a, b::Base.Generator{UnitRange{I}}, lb) where {I}
-
+    
     la = length(a)
     resize!(a, la + lb)
     map!(b.f, view(a, (la+1):(la+lb)), b.iter)
@@ -367,6 +368,7 @@ Objective
 ```
 """
 function objective(c::C, gen) where {C<:ExaCore}
+    gen = _adapt_gen(gen)
     f = SIMDFunction(gen, c.nobj, c.nnzg, c.nnzh)
     pars = gen.iter
 
@@ -426,8 +428,9 @@ function constraint(
     start = zero(T),
     lcon = zero(T),
     ucon = zero(T),
-) where {T,C<:ExaCore{T}}
-
+    ) where {T,C<:ExaCore{T}}
+    
+    gen = _adapt_gen(gen)
     f = SIMDFunction(gen, c.ncon, c.nnzj, c.nnzh)
     pars = gen.iter
 
@@ -487,6 +490,8 @@ function _constraint(c, f, pars, start, lcon, ucon)
 end
 
 function constraint!(c::C, c1, gen::Base.Generator) where {C<:ExaCore}
+    
+    gen = _adapt_gen(gen)
     f = SIMDFunction(gen, offset0(c1, 0), c.nnzj, c.nnzh)
     pars = gen.iter
 
@@ -621,7 +626,7 @@ function hess_coord!(
     x::AbstractVector,
     y::AbstractVector,
     hess::AbstractVector;
-    obj_weight = one(eltype(x)),
+   obj_weight = one(eltype(x)),
 )
 
     fill!(hess, zero(eltype(hess)))
@@ -680,6 +685,22 @@ end
 @inbounds @inline offset0(a::C, i) where {C<:ConstraintAug} = offset0(a.f, a.itr, i)
 @inbounds @inline offset0(f::F, itr, i) where {P<:Pair,F<:SIMDFunction{P}} =
     f.o0 + f.f.first(itr[i], nothing)
+@inbounds @inline offset0(f::F, itr, i) where {T<:Tuple,P<:Pair{T},F<:SIMDFunction{P}} = f.o0 + idxx(coord(itr, i, f.f.first), Base.size(itr))
+
+@inline idx(itr, I) = @inbounds itr[I]
+@inline idx(itr::Base.Iterators.ProductIterator{V}, I) where V =  _idx(I-1, itr.iterators, Base.size(itr))
+@inline function _idx(n, (vec1, vec...), (si1, si...))
+    d, r = divrem(n, si1)
+    return (vec1[r + 1], _idx(d, vec, si)...)
+end
+@inline _idx(n, (vec,), ::Tuple{Int}) = @inbounds vec[n + 1]
+
+@inline idxx(coord, si) = _idxx(coord, si, 1) + 1
+@inline _idxx(coord, si, a) = a * (coord[1] - 1) + _idxx(coord[2:end], si[2:end], a*si[1])
+@inline _idxx(::Tuple{}, ::Tuple{}, a) = 0
+
+@inline coord(itr, i, (f,fs...)) = (f(idx(itr,i), nothing), coord(itr, i, fs)...)
+@inline coord(itr, i, ::Tuple{}) = ()
 
 for (thing, val) in [(:solution, 1), (:multipliers_L, 0), (:multipliers_U, 2)]
     @eval begin
@@ -755,3 +776,7 @@ function multipliers(result::SolverCore.AbstractExecutionStats, y::Constraint)
     len = length(y.itr)
     return view(result.multipliers, o+1:o+len)
 end
+
+
+_adapt_gen(gen) = Base.Generator(gen.f, collect(gen.iter))
+_adapt_gen(gen::Base.Generator{P}) where P <: Union{AbstractArray,AbstractRange} = gen

src/simdfunction.jl

@@ -1,4 +1,4 @@
-@inline (a::Pair{P,S} where {P<:AbstractNode,S<:AbstractNode})(i, x) = a.second(i, x)
+@inline (a::Pair{P,S} where {P,S<:AbstractNode})(i, x) = a.second(i, x)
 
 """
     Compressor{I}

test/NLPTest/NLPTest.jl

@@ -123,24 +123,24 @@ end
 
 function runtests()
     @testset "NLP test" begin
-        for (name, args) in NLP_TEST_ARGUMENTS
-            @testset "$name $args" begin
+        for backend in BACKENDS
+            @testset "$backend" begin
+                for (name, args) in NLP_TEST_ARGUMENTS
+                    @testset "$name $args" begin
 
-                exa_model = getfield(@__MODULE__, Symbol("_exa_$(name)_model"))
-                jump_model = getfield(@__MODULE__, Symbol("_jump_$(name)_model"))
+                        exa_model = getfield(@__MODULE__, Symbol("_exa_$(name)_model"))
+                        jump_model = getfield(@__MODULE__, Symbol("_jump_$(name)_model"))
 
-                m, vars0, cons0 = exa_model(nothing, args)
-                m0 = WrapperNLPModel(m)
+                        m, vars0, cons0 = exa_model(nothing, args)
+                        m0 = WrapperNLPModel(m)
 
-                m, vars2, cons2 = jump_model(nothing, args)
-                m2 = MathOptNLPModel(m)
+                        m, vars2, cons2 = jump_model(nothing, args)
+                        m2 = MathOptNLPModel(m)
+                        
+                        set_optimizer(m, MadNLP.Optimizer)
+                        set_optimizer_attribute(m, "print_level", MadNLP.ERROR)
+                        optimize!(m)
 
-                set_optimizer(m, MadNLP.Optimizer)
-                set_optimizer_attribute(m, "print_level", MadNLP.ERROR)
-                optimize!(m)
-
-                for backend in BACKENDS
-                    @testset "$backend" begin
 
                         m, vars1, cons1 = exa_model(backend, args)
                         m1 = WrapperNLPModel(m)
@@ -164,11 +164,18 @@ function runtests()
                                 end
                             end
                         end
-
                         result1 = madnlp(m1; print_level = MadNLP.ERROR)
                         test_api(result1, vars1, cons1, vars2, cons2)
                     end
                 end
+
+                m3 = WrapperNLPModel(exa_luksan_vlcek_model(nothing, 3; M = 2))
+                m4 = jump_luksan_vlcek_model(nothing, 3; M = 2)
+
+                @testset "Multi-column constraints" begin
+                    test_nlp(m3, m4; full = false)
+                end
+
             end
         end
     end

test/NLPTest/luksan.jl

@@ -1,47 +1,51 @@
-function luksan_vlcek_obj(x, i)
-    return 100 * (x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2
+function luksan_vlcek_obj(x, i, j)
+    return 100 * (x[i-1, j]^2 - x[i, j])^2 + (x[i-1, j] - 1)^2
 end
 
-function luksan_vlcek_con(x, i)
-    return 3x[i+1]^3 + 2 * x[i+2] - 5 + sin(x[i+1] - x[i+2])sin(x[i+1] + x[i+2]) + 4x[i+1] -
-           x[i]exp(x[i] - x[i+1]) - 3
+function luksan_vlcek_con1(x, i, j)
+    return 3x[i+1, j]^3 + 2 * x[i+2, j] - 5
+end
+function luksan_vlcek_con2(x, i, j)
+    return sin(x[i+1, j] - x[i+2, j])sin(x[i+1, j] + x[i+2, j]) + 4x[i+1, j] -
+        x[i, j]exp(x[i, j] - x[i+1, j]) - 3
 end
 
 function luksan_vlcek_x0(i)
     return mod(i, 2) == 1 ? -1.2 : 1.0
 end
 
-function _exa_luksan_vlcek_model(backend, N)
+function _exa_luksan_vlcek_model(backend, N; M = 1)
 
     c = ExaCore(backend = backend)
-    x = variable(c, N; start = (luksan_vlcek_x0(i) for i = 1:N))
-    s = constraint(c, luksan_vlcek_con(x, i) for i = 1:N-2)
-    objective(c, luksan_vlcek_obj(x, i) for i = 2:N)
+    x = variable(c, N, M; start = [luksan_vlcek_x0(i) for i = 1:N, j=1:M])
+    s = constraint(c, luksan_vlcek_con1(x, i, j) for i = 1:N-2, j=1:M)
+    constraint!(c, s, (i,j) => luksan_vlcek_con2(x, i, j) for i = 1:N-2, j=1:M)
+    objective(c, luksan_vlcek_obj(x, i, j) for i = 2:N, j=1:M)
 
     return ExaModel(c; prod = true), (x,), (s,)
 end
 
-function exa_luksan_vlcek_model(backend, N)
-    m, vars, cons = _exa_luksan_vlcek_model(backend, N)
+function exa_luksan_vlcek_model(backend, N; M = 1)
+    m, vars, cons = _exa_luksan_vlcek_model(backend, N;M =  M)
     return m
 end
 
-function _jump_luksan_vlcek_model(backend, N)
+function _jump_luksan_vlcek_model(backend, N; M = 1)
     jm = JuMP.Model()
 
-    JuMP.@variable(jm, x[i = 1:N], start = mod(i, 2) == 1 ? -1.2 : 1.0)
+    JuMP.@variable(jm, x[i = 1:N, j=1:M], start = mod(i, 2) == 1 ? -1.2 : 1.0)
     JuMP.@NLconstraint(
         jm,
-        s[i = 1:N-2],
-        3x[i+1]^3 + 2x[i+2] - 5 + sin(x[i+1] - x[i+2])sin(x[i+1] + x[i+2]) + 4x[i+1] -
-        x[i]exp(x[i] - x[i+1]) - 3 == 0.0
+        s[i = 1:N-2, j=1:M],
+        3x[i+1,j]^3 + 2x[i+2,j] - 5 + sin(x[i+1,j] - x[i+2,j])sin(x[i+1,j] + x[i+2,j]) + 4x[i+1,j] -
+        x[i,j]exp(x[i,j] - x[i+1,j]) - 3 == 0.0
     )
-    JuMP.@NLobjective(jm, Min, sum(100(x[i-1]^2 - x[i])^2 + (x[i-1] - 1)^2 for i = 2:N))
+    JuMP.@NLobjective(jm, Min, sum(100(x[i-1,j]^2 - x[i,j])^2 + (x[i-1,j] - 1)^2 for i = 2:N, j=1:M))
 
     return jm, (x,), (s,)
 end
 
-function jump_luksan_vlcek_model(backend, N)
-    jm, vars, cons = _jump_luksan_vlcek_model(backend, N)
+function jump_luksan_vlcek_model(backend, N; M = 1)
+    jm, vars, cons = _jump_luksan_vlcek_model(backend, N; M = M)
     return MathOptNLPModel(jm)
 end