Add AbstractKKTVector structure (#167)

* add AbstractKKTVector
* rewrite RichardsonIterator
* add documentation for LinearSolver's interface
* rework interface for solve_refine!

.ci/ci.jl

@@ -11,6 +11,8 @@ if ARGS[1] == "full"
     pkgs = ["MadNLPHSL","MadNLPPardiso","MadNLPMumps","MadNLPGPU","MadNLPKrylov"]
 elseif ARGS[1] == "basic"
     pkgs = ["MadNLPMumps","MadNLPKrylov"]
+elseif ARGS[1] == "gpu"
+    pkgs = ["MadNLPGPU"]
 else
     error("proper argument should be given - full or basic")
 end

lib/MadNLPGPU/src/kernels.jl

@@ -123,7 +123,7 @@ end
     DenseKKTSystem kernels
 =#
 
-function MadNLP.mul!(y::AbstractVector, kkt::MadNLP.DenseKKTSystem{T, VT, MT}, x::AbstractVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
+function LinearAlgebra.mul!(y::AbstractVector, kkt::MadNLP.DenseKKTSystem{T, VT, MT}, x::AbstractVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
     # Load buffers
     haskey(kkt.etc, :hess_w1) || (kkt.etc[:hess_w1] = CuVector{T}(undef, size(kkt.aug_com, 1)))
     haskey(kkt.etc, :hess_w2) || (kkt.etc[:hess_w2] = CuVector{T}(undef, size(kkt.aug_com, 1)))
@@ -136,6 +136,9 @@ function MadNLP.mul!(y::AbstractVector, kkt::MadNLP.DenseKKTSystem{T, VT, MT}, x
     LinearAlgebra.mul!(d_y, kkt.aug_com, d_x)
     copyto!(y, d_y)
 end
+function LinearAlgebra.mul!(y::MadNLP.ReducedKKTVector, kkt::MadNLP.DenseKKTSystem{T, VT, MT}, x::MadNLP.ReducedKKTVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
+    LinearAlgebra.mul!(MadNLP.full(y), kkt, MadNLP.full(x))
+end
 
 @kernel function _build_dense_kkt_system_kernel!(
     dest, hess, jac, pr_diag, du_diag, diag_hess, ind_ineq, con_scale, n, m, ns
@@ -252,7 +255,7 @@ function MadNLP._build_condensed_kkt_system!(
     wait(ev)
 end
 
-function MadNLP.mul!(y::AbstractVector, kkt::MadNLP.DenseCondensedKKTSystem{T, VT, MT}, x::AbstractVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
+function LinearAlgebra.mul!(y::AbstractVector, kkt::MadNLP.DenseCondensedKKTSystem{T, VT, MT}, x::AbstractVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
     if length(y) == length(x) == size(kkt.aug_com, 1)
         # Load buffers
         haskey(kkt.etc, :hess_w1) || (kkt.etc[:hess_w1] = CuVector{T}(undef, size(kkt.aug_com, 1)))
@@ -279,6 +282,9 @@ function MadNLP.mul!(y::AbstractVector, kkt::MadNLP.DenseCondensedKKTSystem{T, V
         copyto!(y, d_y)
     end
 end
+function LinearAlgebra.mul!(y::MadNLP.ReducedKKTVector, kkt::MadNLP.DenseCondensedKKTSystem{T, VT, MT}, x::MadNLP.ReducedKKTVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
+    LinearAlgebra.mul!(MadNLP.full(y), kkt, MadNLP.full(x))
+end
 
 function MadNLP.jprod_ineq!(y::AbstractVector, kkt::MadNLP.DenseCondensedKKTSystem{T, VT, MT}, x::AbstractVector) where {T, VT<:CuVector{T}, MT<:CuMatrix{T}}
     # Create buffers

src/IPM/IPM.jl

@@ -31,38 +31,14 @@ mutable struct InteriorPointSolver{KKTSystem} <: AbstractInteriorPointSolver
 
     jacl::Vector{Float64}
 
-    d::Vector{Float64}
-    dx::StrideOneVector{Float64}
-    dl::StrideOneVector{Float64}
-    dzl::StrideOneVector{Float64}
-    dzu::StrideOneVector{Float64}
-
-    p::Vector{Float64}
-    px::StrideOneVector{Float64}
-    pl::StrideOneVector{Float64}
-
-    pzl::Union{Nothing,StrideOneVector{Float64}}
-    pzu::Union{Nothing,StrideOneVector{Float64}}
-
-    _w1::Vector{Float64}
-    _w1x::StrideOneVector{Float64}
-    _w1l::StrideOneVector{Float64}
-    _w1zl::Union{Nothing,StrideOneVector{Float64}}
-    _w1zu::Union{Nothing,StrideOneVector{Float64}}
-
-    _w2::Vector{Float64}
-    _w2x::StrideOneVector{Float64}
-    _w2l::StrideOneVector{Float64}
-    _w2zl::Union{Nothing,StrideOneVector{Float64}}
-    _w2zu::Union{Nothing,StrideOneVector{Float64}}
-
-    _w3::Vector{Float64}
-    _w3x::StrideOneVector{Float64}
-    _w3l::StrideOneVector{Float64}
-
-    _w4::Vector{Float64}
-    _w4x::StrideOneVector{Float64}
-    _w4l::StrideOneVector{Float64}
+    d::AbstractKKTVector{Float64, Vector{Float64}}
+    p::AbstractKKTVector{Float64, Vector{Float64}}
+
+    _w1::AbstractKKTVector{Float64, Vector{Float64}}
+    _w2::AbstractKKTVector{Float64, Vector{Float64}}
+
+    _w3::AbstractKKTVector{Float64, Vector{Float64}}
+    _w4::AbstractKKTVector{Float64, Vector{Float64}}
 
     x_trial::Vector{Float64}
     c_trial::Vector{Float64}
@@ -201,42 +177,18 @@ function InteriorPointSolver{KKTSystem}(nlp::AbstractNLPModel, opt::Options;
 
     aug_vec_length = is_reduced(kkt) ? n+m : n+m+nlb+nub
 
-    _w1 = zeros(aug_vec_length) # fixes the random failure for inertia-free + Unreduced
-    _w1x= view(_w1,1:n)
-    _w1l= view(_w1,n+1:n+m)
-    _w1zl = is_reduced(kkt) ? nothing : view(_w1,n+m+1:n+m+nlb)
-    _w1zu = is_reduced(kkt) ? nothing : view(_w1,n+m+nlb+1:n+m+nlb+nub)
-
-
-    _w2 = Vector{Float64}(undef,aug_vec_length)
-    _w2x= view(_w2,1:n)
-    _w2l= view(_w2,n+1:n+m)
-    _w2zl = is_reduced(kkt) ? nothing : view(_w2,n+m+1:n+m+nlb)
-    _w2zu = is_reduced(kkt) ? nothing : view(_w2,n+m+nlb+1:n+m+nlb+nub)
-
-    _w3 = zeros(aug_vec_length) # fixes the random failure for inertia-free + Unreduced
-    _w3x= view(_w3,1:n)
-    _w3l= view(_w3,n+1:n+m)
-    _w4 = zeros(aug_vec_length) # need to initialize to zero due to mul!
-    _w4x= view(_w4,1:n)
-    _w4l= view(_w4,n+1:n+m)
-
+    _w1 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
+    _w2 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
+    _w3 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
+    _w4 = is_reduced(kkt) ? ReducedKKTVector(n, m) : UnreducedKKTVector(n, m, nlb, nub)
 
     jacl = zeros(n) # spblas may throw an error if not initialized to zero
 
-    d = Vector{Float64}(undef,aug_vec_length)
-    dx= view(d,1:n)
-    dl= view(d,n+1:n+m)
-    dzl= is_reduced(kkt) ? Vector{Float64}(undef,nlb) : view(d,n+m+1:n+m+nlb)
-    dzu= is_reduced(kkt) ? Vector{Float64}(undef,nub) : view(d,n+m+nlb+1:n+m+nlb+nub)
-    dx_lr = view(dx,ind_cons.ind_lb)
-    dx_ur = view(dx,ind_cons.ind_ub)
+    d = UnreducedKKTVector(n, m, nlb, nub)
+    dx_lr = view(d.xp, ind_cons.ind_lb) # TODO
+    dx_ur = view(d.xp, ind_cons.ind_ub) # TODO
 
-    p = Vector{Float64}(undef,aug_vec_length)
-    px= view(p,1:n)
-    pl= view(p,n+1:n+m)
-    pzl= is_reduced(kkt) ? Vector{Float64}(undef,nlb) : view(p,n+m+1:n+m+nlb)
-    pzu= is_reduced(kkt) ? Vector{Float64}(undef,nub) : view(p,n+m+nlb+1:n+m+nlb+nub)
+    p = UnreducedKKTVector(n, m, nlb, nub)
 
     obj_scale = [1.0]
     con_scale = ones(m)
@@ -246,11 +198,10 @@ function InteriorPointSolver{KKTSystem}(nlp::AbstractNLPModel, opt::Options;
     cnt.linear_solver_time =
         @elapsed linear_solver = opt.linear_solver.Solver(get_kkt(kkt) ; option_dict=option_linear_solver,logger=logger)
 
-    n_kkt = size(get_kkt(kkt), 1)
+    n_kkt = size(kkt, 1)
+    buffer_vec = similar(full(d), n_kkt)
     @trace(logger,"Initializing iterative solver.")
-    iterator = opt.iterator.Solver(
-        similar(d, n_kkt),
-        (b, x)->mul!(b, kkt, x), (x)->solve!(linear_solver, x) ; option_dict=option_linear_solver)
+    iterator = opt.iterator(linear_solver, kkt, buffer_vec)
 
     @trace(logger,"Initializing fixed variable treatment scheme.")
 
@@ -263,8 +214,8 @@ function InteriorPointSolver{KKTSystem}(nlp::AbstractNLPModel, opt::Options;
     return InteriorPointSolver{KKTSystem}(nlp,kkt,opt,cnt,logger,
         n,m,nlb,nub,x,l,zl,zu,xl,xu,0.,f,c,
         jacl,
-        d,dx,dl,dzl,dzu,p,px,pl,pzl,pzu,
-        _w1,_w1x,_w1l,_w1zl,_w1zu,_w2,_w2x,_w2l,_w2zl,_w2zu,_w3,_w3x,_w3l,_w4,_w4x,_w4l,
+        d, p,
+        _w1, _w2, _w3, _w4,
         x_trial,c_trial,0.,x_slk,c_slk,rhs,
         ind_cons.ind_ineq,ind_cons.ind_fixed,ind_cons.ind_llb,ind_cons.ind_uub,
         x_lr,x_ur,xl_r,xu_r,zl_r,zu_r,dx_lr,dx_ur,x_trial_lr,x_trial_ur,

src/IPM/callbacks.jl

@@ -51,10 +51,10 @@ function eval_lag_hess_wrapper!(ipp::InteriorPointSolver, kkt::AbstractKKTSystem
     nlp = ipp.nlp
     cnt = ipp.cnt
     @trace(ipp.logger,"Evaluating Lagrangian Hessian.")
-    ipp._w1l .= l.*ipp.con_scale
+    dual(ipp._w1) .= l.*ipp.con_scale
     hess = get_hessian(kkt)
     cnt.eval_function_time += @elapsed hess_coord!(
-        nlp, view(x,1:get_nvar(nlp)), ipp._w1l, view(hess,1:get_nnzh(nlp));
+        nlp, view(x,1:get_nvar(nlp)), dual(ipp._w1), view(hess, 1:get_nnzh(nlp));
         obj_weight = (get_minimize(nlp) ? 1. : -1.) * (is_resto ? 0.0 : ipp.obj_scale[]))
     compress_hessian!(kkt)
     cnt.lag_hess_cnt+=1
@@ -80,10 +80,10 @@ function eval_lag_hess_wrapper!(ipp::InteriorPointSolver, kkt::AbstractDenseKKTS
     nlp = ipp.nlp
     cnt = ipp.cnt
     @trace(ipp.logger,"Evaluating Lagrangian Hessian.")
-    ipp._w1l .= l.*ipp.con_scale
+    dual(ipp._w1) .= l.*ipp.con_scale
     hess = get_hessian(kkt)
     cnt.eval_function_time += @elapsed hess_dense!(
-        nlp, view(x,1:get_nvar(nlp)), ipp._w1l, hess;
+        nlp, view(x,1:get_nvar(nlp)), dual(ipp._w1), hess;
         obj_weight = (get_minimize(nlp) ? 1. : -1.) * (is_resto ? 0.0 : ipp.obj_scale[]))
     compress_hessian!(kkt)
     cnt.lag_hess_cnt+=1

src/IPM/factorization.jl

@@ -5,32 +5,44 @@ function factorize_wrapper!(ips::InteriorPointSolver)
     ips.cnt.linear_solver_time += @elapsed factorize!(ips.linear_solver)
 end
 
-function solve_refine_wrapper!(ips::InteriorPointSolver, x,b)
+function solve_refine_wrapper!(
+    ips::InteriorPointSolver,
+    x::AbstractKKTVector,
+    b::AbstractKKTVector,
+)
     cnt = ips.cnt
     @trace(ips.logger,"Iterative solution started.")
-    fixed_variable_treatment_vec!(b, ips.ind_fixed)
+    fixed_variable_treatment_vec!(full(b), ips.ind_fixed)
+
+    cnt.linear_solver_time += @elapsed begin
+        result = solve_refine!(x, ips.iterator, b)
+    end
 
-    cnt.linear_solver_time += @elapsed (result = solve_refine!(x, ips.iterator, b))
     if result == :Solved
         solve_status =  true
     else
         if improve!(ips.linear_solver)
             cnt.linear_solver_time += @elapsed begin
                 factorize!(ips.linear_solver)
-                solve_status = (solve_refine!(x, ips.iterator, b) == :Solved ? true : false)
+                ret = solve_refine!(x, ips.iterator, b)
+                solve_status = (ret == :Solved)
             end
         else
             solve_status = false
         end
     end
-    fixed_variable_treatment_vec!(x, ips.ind_fixed)
+    fixed_variable_treatment_vec!(full(x), ips.ind_fixed)
     return solve_status
 end
 
-function solve_refine_wrapper!(ips::InteriorPointSolver{<:DenseCondensedKKTSystem}, x, b)
+function solve_refine_wrapper!(
+    ips::InteriorPointSolver{<:DenseCondensedKKTSystem},
+    x::AbstractKKTVector,
+    b::AbstractKKTVector,
+)
     cnt = ips.cnt
     @trace(ips.logger,"Iterative solution started.")
-    fixed_variable_treatment_vec!(b, ips.ind_fixed)
+    fixed_variable_treatment_vec!(full(b), ips.ind_fixed)
 
     kkt = ips.kkt
 
@@ -39,26 +51,26 @@ function solve_refine_wrapper!(ips::InteriorPointSolver{<:DenseCondensedKKTSyste
     n_condensed = n + n_eq
 
     # load buffers
-    b_c = view(ips._w1, 1:n_condensed)
-    x_c = view(ips._w2, 1:n_condensed)
-    jv_x = view(ips._w3, 1:ns) # for jprod
-    jv_t = ips._w4x            # for jtprod
-    v_c = ips._w4l
+    b_c = view(full(ips._w1), 1:n_condensed)
+    x_c = view(full(ips._w2), 1:n_condensed)
+    jv_x = view(full(ips._w3), 1:ns) # for jprod
+    jv_t = primal(ips._w4)             # for jtprod
+    v_c = dual(ips._w4)
 
     Σs = get_slack_regularization(kkt)
     α = get_scaling_inequalities(kkt)
 
     # Decompose right hand side
-    bx = view(b, 1:n)
-    bs = view(b, n+1:n+ns)
-    by = view(b, kkt.ind_eq_shifted)
-    bz = view(b, kkt.ind_ineq_shifted)
+    bx = view(full(b), 1:n)
+    bs = view(full(b), n+1:n+ns)
+    by = view(full(b), kkt.ind_eq_shifted)
+    bz = view(full(b), kkt.ind_ineq_shifted)
 
     # Decompose results
-    xx = view(x, 1:n)
-    xs = view(x, n+1:n+ns)
-    xy = view(x, kkt.ind_eq_shifted)
-    xz = view(x, kkt.ind_ineq_shifted)
+    xx = view(full(x), 1:n)
+    xs = view(full(x), n+1:n+ns)
+    xy = view(full(x), kkt.ind_eq_shifted)
+    xz = view(full(x), kkt.ind_ineq_shifted)
 
     v_c .= 0.0
     v_c[kkt.ind_ineq] .= (Σs .* bz .+ α .* bs) ./ α.^2
@@ -77,7 +89,7 @@ function solve_refine_wrapper!(ips::InteriorPointSolver{<:DenseCondensedKKTSyste
     xz .= sqrt.(Σs) ./ α .* jv_x .- Σs .* bz ./ α.^2 .- bs ./ α
     xs .= (bs .+ α .* xz) ./ Σs
 
-    fixed_variable_treatment_vec!(x, ips.ind_fixed)
+    fixed_variable_treatment_vec!(full(x), ips.ind_fixed)
     return solve_status
 end
 

src/IPM/kernels.jl

@@ -30,61 +30,61 @@ end
 
 # Set RHS
 function set_aug_rhs!(ips::InteriorPointSolver, kkt::AbstractKKTSystem, c)
-    ips.px.=.-ips.f.+ips.mu./(ips.x.-ips.xl).-ips.mu./(ips.xu.-ips.x).-ips.jacl
-    ips.pl.=.-c
+    primal(ips.p) .= .-ips.f.+ips.mu./(ips.x.-ips.xl).-ips.mu./(ips.xu.-ips.x).-ips.jacl
+    dual(ips.p)   .= .-c
 end
 
 function set_aug_rhs!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem, c)
-    ips.px.=.-ips.f.+ips.zl.-ips.zu.-ips.jacl
-    ips.pl.=.-c
-    ips.pzl.=(ips.xl_r-ips.x_lr).*kkt.l_lower .+ ips.mu./kkt.l_lower
-    ips.pzu.=(ips.xu_r-ips.x_ur).*kkt.u_lower .- ips.mu./kkt.u_lower
+    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu.-ips.jacl
+    dual(ips.p) .= .-c
+    dual_lb(ips.p) .= (ips.xl_r-ips.x_lr).*kkt.l_lower .+ ips.mu./kkt.l_lower
+    dual_ub(ips.p) .= (ips.xu_r-ips.x_ur).*kkt.u_lower .- ips.mu./kkt.u_lower
 end
 
 function set_aug_rhs_ifr!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem,c)
-    ips._w1x .= 0.
-    ips._w1l .= .-c
-    ips._w1zl.= 0.
-    ips._w1zu.= 0.
+    primal(ips._w1) .= 0.0
+    dual(ips._w1) .= .-c
+    dual_lb(ips._w1) .= 0.0
+    dual_ub(ips._w1) .= 0.0
 end
 
 # Set RHS RR
 function set_aug_rhs_RR!(
     ips::InteriorPointSolver, kkt::AbstractKKTSystem, RR::RobustRestorer, rho,
 )
-    ips.px.=.-RR.f_R.-ips.jacl.+RR.mu_R./(ips.x.-ips.xl).-RR.mu_R./(ips.xu.-ips.x)
-    ips.pl.=.-ips.c.+RR.pp.-RR.nn.+(RR.mu_R.-(rho.-ips.l).*RR.pp)./RR.zp.-(RR.mu_R.-(rho.+ips.l).*RR.nn)./RR.zn
+    primal(ips.p) .= .-RR.f_R.-ips.jacl.+RR.mu_R./(ips.x.-ips.xl).-RR.mu_R./(ips.xu.-ips.x)
+    dual(ips.p) .= .-ips.c.+RR.pp.-RR.nn.+(RR.mu_R.-(rho.-ips.l).*RR.pp)./RR.zp.-(RR.mu_R.-(rho.+ips.l).*RR.nn)./RR.zn
 end
 
 # Finish
 function finish_aug_solve!(ips::InteriorPointSolver, kkt::AbstractKKTSystem, mu)
-    ips.dzl.= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
-    ips.dzu.= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
+    dual_lb(ips.d) .= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
+    dual_ub(ips.d) .= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
 end
 
 function finish_aug_solve!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem, mu)
-    ips.dzl.*=.-kkt.l_lower
-    ips.dzu.*=kkt.u_lower
-    ips.dzl.= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
-    ips.dzu.= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
+    dual_lb(ips.d) .*= .-kkt.l_lower
+    dual_ub(ips.d) .*= kkt.u_lower
+    dual_lb(ips.d) .= (mu.-ips.zl_r.*ips.dx_lr)./(ips.x_lr.-ips.xl_r).-ips.zl_r
+    dual_ub(ips.d) .= (mu.+ips.zu_r.*ips.dx_ur)./(ips.xu_r.-ips.x_ur).-ips.zu_r
 end
 
 # Initial
 function set_initial_rhs!(ips::InteriorPointSolver, kkt::AbstractKKTSystem)
-    ips.px .= .-ips.f.+ips.zl.-ips.zu
-    ips.pl .= 0.0
+    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu
+    dual(ips.p) .= 0.0
 end
 function set_initial_rhs!(ips::InteriorPointSolver, kkt::SparseUnreducedKKTSystem)
-    ips.px .= .-ips.f.+ips.zl.-ips.zu
-    ips.pl .= 0.0
-    ips.pzl.= 0.0
-    ips.pzu.= 0.0
+    primal(ips.p) .= .-ips.f.+ips.zl.-ips.zu
+    dual(ips.p) .= 0.0
+    dual_lb(ips.p) .= 0.0
+    dual_ub(ips.p) .= 0.0
 end
 
 # Set ifr
 function set_aug_rhs_ifr!(ips::InteriorPointSolver, kkt::AbstractKKTSystem)
-    ips._w1x .= 0.0
-    ips._w1l .= .-ips.c
+    primal(ips._w1) .= 0.0
+    dual(ips._w1) .= .-ips.c
 end
 
 # Finish RR