Schur design constraints and sparse Schur complement solver

[michel2323]

This is another pass over the Schur complement-based solver while adding design constraint support and replacing the dense Schur complement solve with sparse CUDSS on the GPU.

I'm sorry that it is again a large diff. But the tag-based detection of cross and design only constraints and the change in the structure implied a new variable ordering.

I've added our classic SCOPF formulation that we used in ECP and elsewhere to ExaModelsPower.jl (no PR yet). This was very useful to debug. Currently, it passes all IEEE case 9 and 118 with the line contingencies we always used.

One major change was also that I did not relax the equality constraints in the GPU solver. This is fixed now.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic master) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl b/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl
index 3375891..ecd2ea8 100644
--- a/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl
+++ b/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl
@@ -9,7 +9,7 @@ function MadNLP.MadNLPOptions{T}(
     kkt_system = dense_callback ? MadNLP.DenseCondensedKKTSystem : MadNLP.SparseCondensedKKTSystem,
     linear_solver = MadNLPGPU.LapackROCmSolver,
     tol = MadNLP.get_tolerance(T,kkt_system),
-    bound_relax_factor = (kkt_system <: MadNLP.SparseCondensedKKTSystem) ? tol : T(1.0e-8),
+        bound_relax_factor = (kkt_system <: MadNLP.SparseCondensedKKTSystem) ? tol : T(1.0e-8),
 ) where {T, VT <: ROCVector{T}}
     return MadNLP.MadNLPOptions{T}(
         tol = tol,
@@ -28,9 +28,9 @@ function MadNLP.create_kkt_system(
     ) where {T, VT <: ROCVector{T}}
     return error(
         "SchurComplementCondensedKKTSystem is not supported on ROCm/AMDGPU: there is no ROCm " *
-        "Schur implementation (the batched per-scenario factorization and Schur reduction are " *
-        "cuDSS / CUBLAS-specific). Use a CUDA device for the GPU Schur path, or solve the Schur " *
-        "system on CPU."
+            "Schur implementation (the batched per-scenario factorization and Schur reduction are " *
+            "cuDSS / CUBLAS-specific). Use a CUDA device for the GPU Schur path, or solve the Schur " *
+            "system on CPU."
     )
 end
 
diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl
index 75a06b4..9c1c364 100644
--- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl
+++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl
@@ -250,14 +250,18 @@ function _compute_det_scatter(kkt::GPUSchurComplementCondensedKKTSystem)
     n_iA = kkt.n_per_s_ineq_Akk
     nz = Array(kkt.gpu_ineq_Akk_nzpos)
     slotA = n_iA > 0 ? Int[((i - 1) ÷ n_iA) * nnzb + nz[i] for i in eachindex(nz)] : Int[]
-    Akk = _segment_scatter(slotA,
-        Array(kkt.gpu_ineq_Akk_jcoo1), Array(kkt.gpu_ineq_Akk_jcoo2), Array(kkt.gpu_ineq_Akk_bufidx))
+    Akk = _segment_scatter(
+        slotA,
+        Array(kkt.gpu_ineq_Akk_jcoo1), Array(kkt.gpu_ineq_Akk_jcoo2), Array(kkt.gpu_ineq_Akk_bufidx)
+    )
     # ineq → C_dk: linear slot into (blk × m × ns) = v + (d-1)*blk + (k-1)*blk*m
     n_iC = kkt.n_per_s_ineq_Cdk
     vv = Array(kkt.gpu_ineq_Cdk_col); dd = Array(kkt.gpu_ineq_Cdk_row)
     slotC = n_iC > 0 ? Int[vv[i] + (dd[i] - 1) * blk + ((i - 1) ÷ n_iC) * blk * m for i in eachindex(vv)] : Int[]
-    Cdk = _segment_scatter(slotC,
-        Array(kkt.gpu_ineq_Cdk_jcoo_d), Array(kkt.gpu_ineq_Cdk_jcoo_v), Array(kkt.gpu_ineq_Cdk_bufidx))
+    Cdk = _segment_scatter(
+        slotC,
+        Array(kkt.gpu_ineq_Cdk_jcoo_d), Array(kkt.gpu_ineq_Cdk_jcoo_v), Array(kkt.gpu_ineq_Cdk_bufidx)
+    )
     # ineq → S: merge scenario + design-only contributions, slot = nzpos
     slotS = vcat(Array(kkt.schur_ineq_S_nzpos), Array(kkt.schur_design_ineq_S_nzpos))
     j1S = vcat(Array(kkt.schur_ineq_S_jcoo1), Array(kkt.schur_design_ineq_S_jcoo1))
@@ -305,9 +309,11 @@ function _compute_det_scatter(kkt::GPUSchurComplementCondensedKKTSystem)
     njac = length(kkt.jac)
     jt_lin = _segment_lin_scatter(Array(kkt.jt_csc_map), collect(1:njac))
 
-    return (Akk = Akk, Cdk = Cdk, S = S, hessAkk = hessAkk, hessCdk = hessCdk, hessS = hessS,
-            hess_lin = hess_lin, hess_full = hess_full, hess_full_perm = hess_full_perm,
-            jt_lin = jt_lin)
+    return (
+        Akk = Akk, Cdk = Cdk, S = S, hessAkk = hessAkk, hessCdk = hessCdk, hessS = hessS,
+        hess_lin = hess_lin, hess_full = hess_full, hess_full_perm = hess_full_perm,
+        jt_lin = jt_lin,
+    )
 end
 
 function _get_det_scatter(kkt::GPUSchurComplementCondensedKKTSystem)
@@ -317,7 +323,7 @@ end
 
 # --- Dispatch: GPU path when callback uses CuVector ---
 function MadNLP.create_kkt_system(
-    ::Type{MadNLP.SchurComplementCondensedKKTSystem},
+        ::Type{MadNLP.SchurComplementCondensedKKTSystem},
     cb::MadNLP.SparseCallback{T, VT},
     linear_solver::Type;
     opt_linear_solver=MadNLP.default_options(linear_solver),
@@ -479,7 +485,7 @@ function MadNLP.create_kkt_system(
     gpu_hess_Cdk_row = CuVector{Int}(coupled_inv[flat.all_hess_Cdk_row])
     gpu_hess_Cdk_col    = CuVector{Int}(flat.all_hess_Cdk_col)
     gpu_pr_diag_global  = CuVector{Int}(flat.all_pr_diag_global)
-    gpu_pr_diag_nzpos   = CuVector{Int}(flat.all_pr_diag_nzpos)
+    gpu_pr_diag_nzpos = CuVector{Int}(flat.all_pr_diag_nzpos)
     gpu_ineq_Akk_nzpos  = CuVector{Int}(flat.all_ineq_Akk_nzpos)
     gpu_ineq_Akk_jcoo1  = CuVector{Int}(flat.all_ineq_Akk_jcoo1)
     gpu_ineq_Akk_jcoo2  = CuVector{Int}(flat.all_ineq_Akk_jcoo2)
@@ -620,7 +626,7 @@ end
 # iteration. Passing `det` to a specialized inner method makes every `det.X` access (and the
 # kernel launches taking them) type-stable after a single dynamic dispatch at the call.
 MadNLP.build_kkt!(kkt::GPUSchurComplementCondensedKKTSystem) = _schur_build_kkt!(kkt, _get_det_scatter(kkt))
-function _schur_build_kkt!(kkt::GPUSchurComplementCondensedKKTSystem{T}, det) where T
+function _schur_build_kkt!(kkt::GPUSchurComplementCondensedKKTSystem{T}, det) where {T}
     ns = kkt.ns
     nv = kkt.nv
     nd = kkt.nd
@@ -767,7 +773,7 @@ end
 
 # --- solve_kkt! ---
 function MadNLP.solve_kkt!(
-    kkt::GPUSchurComplementCondensedKKTSystem{T},
+        kkt::GPUSchurComplementCondensedKKTSystem{T},
     w::MadNLP.AbstractKKTVector{T},
 ) where T
 
@@ -866,20 +872,20 @@ end
 # hit many times per Richardson step) out of the hot body so the SpMV dispatches statically.
 function MadNLP.mul!(
     w::MadNLP.AbstractKKTVector{T, VT},
-    kkt::GPUSchurComplementCondensedKKTSystem{T},
+        kkt::GPUSchurComplementCondensedKKTSystem{T},
     x::MadNLP.AbstractKKTVector,
     alpha = one(T),
     beta = zero(T),
-) where {T, VT <: CuVector{T}}
+    ) where {T, VT <: CuVector{T}}
     return _schur_mul!(w, kkt, x, alpha, beta, _get_det_scatter(kkt).hess_full)
 end
 function _schur_mul!(
-    w::MadNLP.AbstractKKTVector{T, VT},
-    kkt::GPUSchurComplementCondensedKKTSystem{T},
-    x::MadNLP.AbstractKKTVector,
-    alpha,
-    beta,
-    hess_full,
+        w::MadNLP.AbstractKKTVector{T, VT},
+        kkt::GPUSchurComplementCondensedKKTSystem{T},
+        x::MadNLP.AbstractKKTVector,
+        alpha,
+        beta,
+        hess_full,
 ) where {T, VT <: CuVector{T}}
     n = MadNLP.num_variables(kkt)
 
diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl
index b2810dc..247f972 100644
--- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl
+++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl
@@ -4,7 +4,7 @@ MadNLP.@kwdef mutable struct CudssSolverOptions <: MadNLP.AbstractOptions
     cudss_ordering::ORDERING = DEFAULT_ORDERING
     cudss_perm::Vector{Cint} = Cint[]
     cudss_ir::Int = 0
-    cudss_ir_tol::Float64 = 1e-8
+    cudss_ir_tol::Float64 = 1.0e-8
     cudss_pivot_threshold::Float64 = 0.0
     cudss_pivot_epsilon::Float64 = 0.0
     cudss_matching_alg::String = "default"
diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl
index 85d5651..a5578b0 100644
--- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl
+++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl
@@ -30,11 +30,11 @@ end
 # `scen_var_global` (flat, length ns*nv, scenario-major), since a scenario's
 # variables need not be contiguous.
 @kernel function _extract_rhs_kernel!(
-    rhs_k,
-    @Const(wx),
-    @Const(scen_var_global),
-    @Const(nv),
-    @Const(ns),
+        rhs_k,
+        @Const(wx),
+        @Const(scen_var_global),
+        @Const(nv),
+        @Const(ns),
 )
     i = @index(Global)
     k = (i - 1) ÷ nv + 1
@@ -47,11 +47,11 @@ end
 
 # Write back per-scenario solution to global wx
 @kernel function _writeback_rhs_kernel!(
-    wx,
-    @Const(rhs_k),
-    @Const(scen_var_global),
-    @Const(nv),
-    @Const(ns),
+        wx,
+        @Const(rhs_k),
+        @Const(scen_var_global),
+        @Const(nv),
+        @Const(ns),
 )
     i = @index(Global)
     k = (i - 1) ÷ nv + 1
diff --git a/src/IPM/options.jl b/src/IPM/options.jl
index fddd006..8494544 100644
--- a/src/IPM/options.jl
+++ b/src/IPM/options.jl
@@ -214,14 +214,14 @@ function MadNLPOptions{T}(
     callback = dense_callback ? DenseCallback : SparseCallback,
     kkt_system = dense_callback ? DenseCondensedKKTSystem : SparseKKTSystem,
     linear_solver = dense_callback ? LapackCPUSolver : default_sparse_solver(nlp),
-    tol = get_tolerance(T,kkt_system),
-    # Condensed KKT systems (Sparse + Schur) relax equalities into the barrier — the SAME
-    # systems selected for RelaxEquality below. Flooring those slacks at a tiny 1e-8 box blows up
-    # the condensation weights z/s and ruins the conditioning of the (SPD) factorization, so relax
-    # by `tol` there (feasible to ~tol, but well-conditioned). Non-condensed systems keep equalities
-    # exact and use the tight 1e-8. Must match the CUDA / ROCm extension constructors so a problem
-    # behaves identically on CPU and GPU.
-    bound_relax_factor = (kkt_system <: CondensedKKTSystems) ? tol : T(1.0e-8),
+        tol = get_tolerance(T, kkt_system),
+        # Condensed KKT systems (Sparse + Schur) relax equalities into the barrier — the SAME
+        # systems selected for RelaxEquality below. Flooring those slacks at a tiny 1e-8 box blows up
+        # the condensation weights z/s and ruins the conditioning of the (SPD) factorization, so relax
+        # by `tol` there (feasible to ~tol, but well-conditioned). Non-condensed systems keep equalities
+        # exact and use the tight 1e-8. Must match the CUDA / ROCm extension constructors so a problem
+        # behaves identically on CPU and GPU.
+        bound_relax_factor = (kkt_system <: CondensedKKTSystems) ? tol : T(1.0e-8),
 ) where {T}
     return MadNLPOptions{T}(
         tol = tol,
@@ -234,7 +234,7 @@ end
 
 get_tolerance(::Type{T},::Type{KKT}) where {T, KKT} = 10^round(log10(eps(T))/2)
 get_tolerance(::Type{T},::Type{SparseCondensedKKTSystem}) where T = 10^(round(log10(eps(T))/4))
-get_tolerance(::Type{T},::Type{SchurComplementCondensedKKTSystem}) where T = 10^(round(log10(eps(T))/4))
+get_tolerance(::Type{T}, ::Type{SchurComplementCondensedKKTSystem}) where {T} = 10^(round(log10(eps(T)) / 4))
 
 default_sparse_solver(nlp::AbstractNLPModel) = MumpsSolver
 
diff --git a/src/KKT/Schur/schur.jl b/src/KKT/Schur/schur.jl
index b93c2e0..32d5f97 100644
--- a/src/KKT/Schur/schur.jl
+++ b/src/KKT/Schur/schur.jl
@@ -107,7 +107,7 @@ struct SchurComplementCondensedKKTSystem{
     # Per-scenario sparse augmented blocks (lower triangle only)
     A_kk::Vector{SparseMatrixCSC{T, Int32}}
     C_dk::Vector{MT}                # ns × (m × blk_size) — dense, reduced to the m COUPLED
-                                    # design rows (design vars that never couple are exactly zero)
+    # design rows (design vars that never couple are exactly zero)
 
     # Sparse first-stage Schur complement (lower-triangular CSC, size nd × nd): the
     # design block S_dd plus the condensed design inequalities. SPD. Only `m` design
@@ -395,11 +395,11 @@ function _build_schur_symbolic(
     # and SPD — there is no equality saddle / bordered block.
     isempty(ind_eq) || error(
         "SchurComplementCondensedKKTSystem is RelaxEquality-only, but got $(length(ind_eq)) " *
-        "constraint(s) kept as equalities. The bordered EnforceEquality saddle that the old " *
-        "`SchurComplementKKTSystem` used was removed: the first-stage Schur complement is now SPD " *
-        "and requires every constraint to be relaxed into the barrier. Pass " *
-        "`equality_treatment=MadNLP.RelaxEquality` (the default when " *
-        "`kkt_system=SchurComplementCondensedKKTSystem`) — do not override it with `EnforceEquality`."
+            "constraint(s) kept as equalities. The bordered EnforceEquality saddle that the old " *
+            "`SchurComplementKKTSystem` used was removed: the first-stage Schur complement is now SPD " *
+            "and requires every constraint to be relaxed into the barrier. Pass " *
+            "`equality_treatment=MadNLP.RelaxEquality` (the default when " *
+            "`kkt_system=SchurComplementCondensedKKTSystem`) — do not override it with `EnforceEquality`."
     )
     ind_ineq_set = Set(Int.(ind_ineq))
 
@@ -425,7 +425,7 @@ function _build_schur_symbolic(
         if n_in_k != nc_ineq_per_s
             error(
                 "SchurComplementCondensedKKTSystem requires uniform per-scenario constraint counts. " *
-                "Scenario 1 has ineq=$nc_ineq_per_s; scenario $k has ineq=$n_in_k."
+                    "Scenario 1 has ineq=$nc_ineq_per_s; scenario $k has ineq=$n_in_k."
             )
         end
     end
@@ -921,7 +921,7 @@ function _flatten_block_maps(block_maps::Vector{ScenarioBlockMap})
     bm1 = block_maps[1]
     n_per_s_hess_Akk = length(bm1.hess_Akk_coo)
     n_per_s_hess_Cdk = length(bm1.hess_Cdk_coo)
-    n_per_s_pr_diag  = length(bm1.pr_diag_global)
+    n_per_s_pr_diag = length(bm1.pr_diag_global)
     n_per_s_ineq_Akk = length(bm1.ineq_Akk_nzpos)
     n_per_s_ineq_Cdk = length(bm1.ineq_Cdk_row)
     n_per_s_ineq_S   = length(bm1.ineq_S_row)
@@ -965,7 +965,7 @@ function _flatten_block_maps(block_maps::Vector{ScenarioBlockMap})
 end
 
 function create_kkt_system(
-    ::Type{SchurComplementCondensedKKTSystem},
+        ::Type{SchurComplementCondensedKKTSystem},
     cb::SparseCallback{T,VT},
     linear_solver::Type;
     opt_linear_solver=default_options(linear_solver),
@@ -978,9 +978,9 @@ function create_kkt_system(
         schur_var_scen = nothing,
         schur_con_scen = nothing,
     schur_scenario_linear_solver::Type=MumpsSolver,
-    schur_scenario_opt_linear_solver=default_options(schur_scenario_linear_solver),
-    schur_opt_linear_solver=nothing,
-    kwargs...,
+        schur_scenario_opt_linear_solver = default_options(schur_scenario_linear_solver),
+        schur_opt_linear_solver = nothing,
+        kwargs...,
 ) where {T, VT}
 
     isempty(kwargs) || Base.@warn(
@@ -1047,14 +1047,14 @@ function create_kkt_system(
     # (nd ≫ m) tractable on CPU.
     block_maps = [
         let bm = sym.block_maps[k]
-            ScenarioBlockMap(
-                bm.hess_Akk_coo, bm.hess_Akk_nzpos,
-                bm.hess_Cdk_coo, coupled_inv[bm.hess_Cdk_row], bm.hess_Cdk_col,
-                bm.ineq_Akk_nzpos, bm.ineq_Akk_jcoo1, bm.ineq_Akk_jcoo2, bm.ineq_Akk_bufidx,
-                coupled_inv[bm.ineq_Cdk_row], bm.ineq_Cdk_col, bm.ineq_Cdk_jcoo_d, bm.ineq_Cdk_jcoo_v, bm.ineq_Cdk_bufidx,
-                bm.ineq_S_row, bm.ineq_S_col, bm.ineq_S_jcoo1, bm.ineq_S_jcoo2, bm.ineq_S_bufidx,
-                bm.pr_diag_global, bm.pr_diag_nzpos,
-            )
+                ScenarioBlockMap(
+                    bm.hess_Akk_coo, bm.hess_Akk_nzpos,
+                    bm.hess_Cdk_coo, coupled_inv[bm.hess_Cdk_row], bm.hess_Cdk_col,
+                    bm.ineq_Akk_nzpos, bm.ineq_Akk_jcoo1, bm.ineq_Akk_jcoo2, bm.ineq_Akk_bufidx,
+                    coupled_inv[bm.ineq_Cdk_row], bm.ineq_Cdk_col, bm.ineq_Cdk_jcoo_d, bm.ineq_Cdk_jcoo_v, bm.ineq_Cdk_bufidx,
+                    bm.ineq_S_row, bm.ineq_S_col, bm.ineq_S_jcoo1, bm.ineq_S_jcoo2, bm.ineq_S_bufidx,
+                    bm.pr_diag_global, bm.pr_diag_nzpos,
+                )
         end for k in 1:ns
     ]
 
@@ -1260,7 +1260,7 @@ function factorize_kkt!(kkt::SchurComplementCondensedKKTSystem)
 end
 
 function solve_kkt!(
-    kkt::SchurComplementCondensedKKTSystem,
+        kkt::SchurComplementCondensedKKTSystem,
     w::AbstractKKTVector{T},
 ) where T
 
@@ -1359,7 +1359,7 @@ function solve_kkt!(
 end
 
 # KKT matrix-vector product for iterative refinement
-function mul!(w::AbstractKKTVector{T}, kkt::SchurComplementCondensedKKTSystem{T}, x::AbstractKKTVector, alpha = one(T), beta = zero(T)) where T
+function mul!(w::AbstractKKTVector{T}, kkt::SchurComplementCondensedKKTSystem{T}, x::AbstractKKTVector, alpha = one(T), beta = zero(T)) where {T}
     n = num_variables(kkt)
     ns_ineq = kkt.n_ineq
     wx = @view(primal(w)[1:n])
diff --git a/test/schur_test.jl b/test/schur_test.jl
index 83181ce..67ea477 100644
--- a/test/schur_test.jl
+++ b/test/schur_test.jl
@@ -184,8 +184,8 @@ using MadNLPTests
         # scenario var and the design var.
         hess_I_ok = Int32[1, 2, 3]; hess_J_ok = Int32[1, 2, 3]
         jac_I_ok  = Int32[1, 1, 2, 2]; jac_J_ok = Int32[1, 3, 2, 3]
-        ind_eq    = Int32[]
-        ind_ineq  = Int32[1, 2]
+        ind_eq = Int32[]
+        ind_ineq = Int32[1, 2]
         @test MadNLP._build_schur_symbolic(
             Float64, n, m, ns, nv, nd, nc,
             hess_I_ok, hess_J_ok, jac_I_ok, jac_J_ok, ind_eq, ind_ineq,

[codecov]

Codecov Report

❌ Patch coverage is 94.67085% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.13%. Comparing base (4f4fb9a) to head (dd0d099).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
src/KKT/Schur/schur.jl	94.58%	17 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #623      +/-   ##
==========================================
+ Coverage   88.91%   89.13%   +0.21%     
==========================================
  Files          55       55              
  Lines        5015     5125     +110     
==========================================
+ Hits         4459     4568     +109     
- Misses        556      557       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.