dead end?

Author: sshin23

ext/ExaModelsKernelAbstractions.jl

@@ -477,7 +477,7 @@ function ExaModels.sgradient!(
     f,
     x,
     adj,
-) where {B<:KernelAbstractions.Backend}
+    ) where {B<:KernelAbstractions.Backend}
     return kerg(backend)(y, f.f, f.itr, x, adj; ndrange = size(f.itr))
 end
 

src/nlp.jl

@@ -490,14 +490,14 @@ function _constraint(c, f, pars, start, lcon, ucon)
 end
 
 function constraint!(c::C, c1, gen::Base.Generator) where {C<:ExaCore}
-    f = SIMDFunction(gen, offset0(c1, 0), c.nnzj, c.nnzh)
+    f = SIMDFunction(gen, offset0(c1, 0), c.nnzj, c.nnzh; tsize = Base.size(c1.itr))
     pars = gen.iter
 
     _constraint!(c, f, pars)
 end
 
 function constraint!(c::C, c1, expr, pars) where {C<:ExaCore}
-    f = _simdfunction(expr, offset0(c1, 0), c.nnzj, c.nnzh)
+    f = _simdfunction(expr, offset0(c1, 0), c.nnzj, c.nnzh; tsize = Base.size(c1.itr))
 
     _constraint!(c, f, pars)
 end
@@ -697,7 +697,7 @@ coord(itr, i, ::Tuple{}) = ()
 @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))
+@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), f.tsize)
 
 for (thing, val) in [(:solution, 1), (:multipliers_L, 0), (:multipliers_U, 2)]
     @eval begin

src/simdfunction.jl

@@ -13,7 +13,7 @@ struct Compressor{I}
 end
 @inline (i::Compressor{I})(n) where {I} = @inbounds i.inner[n]
 
-struct SIMDFunction{F,C1,C2}
+struct SIMDFunction{F,C1,C2,T}
     f::F
     comp1::C1
     comp2::C2
@@ -22,6 +22,7 @@ struct SIMDFunction{F,C1,C2}
     o2::Int
     o1step::Int
     o2step::Int
+    tsize::T
 end
 
 """
@@ -35,14 +36,14 @@ Returns a `SIMDFunction` using the `gen`.
 - `o1`: offset for the derivative evalution
 - `o2`: offset for the second-order derivative evalution
 """
-function SIMDFunction(gen::Base.Generator, o0 = 0, o1 = 0, o2 = 0)
+function SIMDFunction(gen::Base.Generator, o0 = 0, o1 = 0, o2 = 0; tsize = ())
 
     f = gen.f(Par(eltype(gen.iter)))
 
-    _simdfunction(f, o0, o1, o2)
+    _simdfunction(f, o0, o1, o2; tsize = tsize)
 end
 
-function _simdfunction(f, o0, o1, o2)
+function _simdfunction(f, o0, o1, o2; tsize = ())
     d = f(Identity(), AdjointNodeSource(nothing))
     y1 = []
     ExaModels.grpass(d, nothing, y1, nothing, 0, NaN)
@@ -59,5 +60,5 @@ function _simdfunction(f, o0, o1, o2)
     o2step = length(a2)
     c2 = Compressor(Tuple(findfirst(isequal(i), a2) for i in y2))
 
-    SIMDFunction(f, c1, c2, o0, o1, o2, o1step, o2step)
+    SIMDFunction(f, c1, c2, o0, o1, o2, o1step, o2step, tsize)
 end