diff --git a/docs/src/lib/kkt.md b/docs/src/lib/kkt.md index fed4aa2ce..7db3dd5f4 100644 --- a/docs/src/lib/kkt.md +++ b/docs/src/lib/kkt.md @@ -79,10 +79,13 @@ DenseCondensedKKTSystem For two-stage stochastic programs with a block-arrowhead structure, MadNLP provides a specialized KKT system that eliminates the per-scenario -blocks and reduces the linear system to a dense `nd × nd` Schur complement -on the design variables. +blocks and reduces the linear system to a **sparse** (lower-triangular CSC) +`nd × nd` Schur complement on the design variables. All equality constraints +are relaxed into the barrier (RelaxEquality), so the per-scenario blocks and +the first-stage Schur complement are symmetric positive definite; the reduction +`Σ_k C_dk A_kk⁻¹ C_dk'` fills only the coupled-design × coupled-design block. ```@docs -SchurComplementKKTSystem +SchurComplementCondensedKKTSystem ``` diff --git a/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl b/lib/MadNLPGPU/ext/MadNLPGPUAMDGPUExt/rocm.jl index 1c997e87e..3375891ab 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 = tol, + bound_relax_factor = (kkt_system <: MadNLP.SparseCondensedKKTSystem) ? tol : T(1.0e-8), ) where {T, VT <: ROCVector{T}} return MadNLP.MadNLPOptions{T}( tol = tol, @@ -20,6 +20,20 @@ function MadNLP.MadNLPOptions{T}( ) end +function MadNLP.create_kkt_system( + ::Type{MadNLP.SchurComplementCondensedKKTSystem}, + cb::MadNLP.SparseCallback{T, VT}, + linear_solver::Type; + kwargs..., + ) 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." + ) +end + #= SparseMatrixCSC to ROCSparseMatrixCSC =# diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda.jl index bcbbba545..db03cbf52 100644 --- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda.jl +++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda.jl @@ -9,7 +9,9 @@ function MadNLP.MadNLPOptions{T}( kkt_system = dense_callback ? MadNLP.DenseCondensedKKTSystem : MadNLP.SparseCondensedKKTSystem, linear_solver = dense_callback ? LapackCUDASolver : CUDSSSolver, tol = MadNLP.get_tolerance(T, kkt_system), - bound_relax_factor = (kkt_system == MadNLP.SparseCondensedKKTSystem) ? tol : T(1.0e-8), + # See MadNLP.MadNLPOptions in src/IPM/options.jl for why condensed systems (Sparse + Schur) + # relax by `tol` while the rest use 1e-8. Kept identical on CPU / CUDA / ROCm. + bound_relax_factor = (kkt_system <: MadNLP.CondensedKKTSystems) ? tol : T(1.0e-8), ) where {T, VT <: CuVector{T}} return MadNLP.MadNLPOptions{T}( tol = tol, diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl index 44c32ca45..75a06b419 100644 --- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl +++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cuda_schur.jl @@ -1,15 +1,22 @@ ########################################################### -##### CUDA wrappers for SchurComplementKKTSystem ########## +##### CUDA wrappers for SchurComplementCondensedKKTSystem ########## ########################################################### """ - GPUSchurComplementKKTSystem + GPUSchurComplementCondensedKKTSystem GPU-native Schur complement KKT system, the CUDA counterpart of CPU -[`MadNLP.SchurComplementKKTSystem`](@ref). Uses a single batched +[`MadNLP.SchurComplementCondensedKKTSystem`](@ref). Uses a single batched `CuSparseMatrixCSC` holding all `ns` scenario blocks (factored by `CUDSSSolver` -with uniform batching) and CUBLAS strided batched GEMM for the Schur complement -accumulation. +with uniform batching) and CUBLAS strided batched GEMM for the Schur reduction. + +The first-stage Schur complement `S` is itself assembled as a **sparse** +lower-triangular `CuSparseMatrixCSC` (the reduction `Σ_k C_dk A_kk⁻¹ C_dk'` fills +only the coupled-design × coupled-design block) and factored by a second +`CUDSSSolver` (`nbatch=1`, which also reports inertia). The coupling block `C_dk` +is stored reduced to its `m` coupled design columns. The sparsity pattern is +static across the IPM loop: cuDSS analysis runs once at construction; each +iteration only refreshes `S.nzVal` and refactorizes. Variable layout: `[v_1, ..., v_ns, d]`, `v_k ∈ R^nv`, `d ∈ R^nd`. Constraint layout: `[c_1, ..., c_ns]`, `c_k ∈ R^nc`. @@ -20,13 +27,13 @@ scenarios concatenated in scenario order. Each `n_per_s_*` field stores the per-scenario length; the matching `gpu_*` vector then has length `ns * n_per_s_*` and is indexed directly by the kernel global index. """ -struct GPUSchurComplementKKTSystem{ +struct GPUSchurComplementCondensedKKTSystem{ T, VT <: AbstractVector{T}, MT <: AbstractMatrix{T}, QN, - LS, # linear solver for Schur complement S (LapackCUDASolver) - LS2, # batched scenario solver (CUDSSSolver) + LS, # linear solver for the sparse Schur complement S (CUDSSSolver, nbatch=1) + LS2, # batched scenario solver (CUDSSSolver, nbatch=ns) COO_T, # SparseMatrixCOO type for hess_raw COO_JT, # SparseMatrixCOO type for jt_coo CSC_T, # CuSparseMatrixCSC type @@ -63,28 +70,32 @@ struct GPUSchurComplementKKTSystem{ nv::Int nd::Int nc::Int - nc_eq_per_s::Int nc_ineq_per_s::Int - blk_size::Int + blk_size::Int # == nv (per-scenario condensed block size) + m::Int # number of coupled design vars (Schur fill width) # Batched scenario block A_kk_batched::CSC_T nnz_per_scenario::Int - # Dense blocks - C_dk_batched::CuArray{T, 3} # (blk_size, nd, ns) — per-scenario cross-block stored - # with scenario-var dim first, so each slice - # C_dk[:, :, k] is directly usable as the cuDSS - # multi-RHS input for `A_kk \ C_dk'`. - aug_com::MT # (nd, nd) — Schur complement S + # Reduced coupling blocks: only the m design columns that couple to a scenario. + # `(blk_size, m, ns)`, scenario-var dim first, so each slice `C_dk[:, :, k]` is + # the cuDSS multi-RHS input for `A_kk \ C_dk_red'`. + C_dk_batched::CuArray{T, 3} + schur_csc::CSC_T # lower-triangular sparse Schur complement S (nd × nd, SPD) + schur_block_batched::CuArray{T, 3} # (m, m, ns) — per-scenario C_dk_red' A_kk⁻¹ C_dk_red' + coupled_design_local::VI # length m — design-local indices that couple to scenarios + schur_fill_nzpos::CuMatrix{Int} # (m, m) — nzval position of the Schur-fill block entries # Buffers diag_buffer::VT buffer::VT - wy_eq_buf::VT # ns*nc_eq_per_s — preserves eq duals across J*Δx round-trip - rhs_d::VT + rhs_d::VT # nd + rhs_d_red::VT # m — reduced design coupling RHS + rhs_d_red_batched::CuArray{T, 3} # (m, 1, ns) — strided-batched reduction GEMM output + rhs_d_red_mat::MT # (m, ns) — rhs_d_red broadcast across scenarios (back-sub) rhs_k_batched::MT # (blk_size, ns) - tmp_blk_nd_batched::CuArray{T, 3} # (blk_size, nd, ns) + tmp_blk_nd_batched::CuArray{T, 3} # (blk_size, m, ns) solve_buffer::VT # blk_size * ns # Flattened GPU index maps (all scenarios concatenated in scenario order). @@ -102,19 +113,6 @@ struct GPUSchurComplementKKTSystem{ gpu_pr_diag_global::VI gpu_pr_diag_nzpos::VI - n_per_s_du_diag::Int - gpu_du_diag_global::VI - gpu_du_diag_nzpos::VI - - n_per_s_jeq_Akk::Int - gpu_jeq_Akk_coo::VI - gpu_jeq_Akk_nzpos::VI - - n_per_s_jeq_Cdk::Int - gpu_jeq_Cdk_coo::VI - gpu_jeq_Cdk_row::VI - gpu_jeq_Cdk_col::VI - n_per_s_ineq_Akk::Int gpu_ineq_Akk_nzpos::VI gpu_ineq_Akk_jcoo1::VI @@ -128,20 +126,27 @@ struct GPUSchurComplementKKTSystem{ gpu_ineq_Cdk_jcoo_v::VI gpu_ineq_Cdk_bufidx::VI - n_per_s_ineq_S::Int - gpu_ineq_S_row::VI - gpu_ineq_S_col::VI - gpu_ineq_S_jcoo1::VI - gpu_ineq_S_jcoo2::VI - gpu_ineq_S_bufidx::VI - - hess_S_coo::VI - hess_S_row::VI - hess_S_col::VI - - eq_global_indices::VI - - # Inequality/equality/bound indices + # Sparse-Schur scatter maps: scatter into `schur_csc.nzVal` at precomputed nzval + # positions (lower-triangle CSC; cuDSS symmetrizes). Lists are lower-only so each + # slot is hit once. + schur_hess_coo::VI # design Hessian: COO index into `hess` + schur_hess_nzpos::VI # → nzval slot + schur_diag_nzpos::VI # pr_diag design diagonal (paired with design_var_global) + schur_ineq_S_nzpos::VI # scenario ineq → S (flat over all scenarios) + schur_ineq_S_jcoo1::VI + schur_ineq_S_jcoo2::VI + schur_ineq_S_bufidx::VI + schur_design_ineq_S_nzpos::VI # design ineq → S + schur_design_ineq_S_jcoo1::VI + schur_design_ineq_S_jcoo2::VI + schur_design_ineq_S_bufidx::VI + + # Tag-driven global index lists (design/scenario vars need not be contiguous). + design_var_global::VI # length nd — global index of each design var + scen_var_global::VI # length ns*nv (scenario-major) — scenario var globals + + # Inequality/equality/bound indices (n_eq == 0 / ind_eq empty under RelaxEquality, + # kept for the generic KKT interface). n_eq::Int ind_eq::VI n_ineq::Int @@ -150,20 +155,169 @@ struct GPUSchurComplementKKTSystem{ ind_ub::VI # Solvers - scenario_solver::LS2 + scenario_solver::LS2 # batched per-scenario blocks (CUDSSSolver, nbatch=ns) + # Schur complement (CUDSSSolver, nbatch=1); the IPM queries this field for inertia/factorize. linear_solver::LS - # Multi-RHS descriptors over `C_dk_batched` / `tmp_blk_nd_batched`, used to - # solve A_kk * tmp[:, :, k] = C_dk[:, :, k] for all k in a single cuDSS - # call (nd right-hand sides per scenario, ns batch). Zero-copy: both are - # `cudss_update`'d onto the existing `(blk, nd, ns)` buffers each iteration. - scenario_x_multi::CUDSS.CudssMatrix{T} - scenario_b_multi::CUDSS.CudssMatrix{T} + # Lazily-built deterministic condensation-scatter structure (NamedTuple of _DetScatter / + # _DetLinScatter), memoized here so it is owned by — and freed with — the kkt. + det_cache::Base.RefValue{Any} +end + +# cuDSS sparse-LDL is less accurate than a direct factorization. The relaxed Schur +# path solves each (SPD) per-scenario block with it and accumulates A_kk⁻¹ into the +# first-stage complement, so per-block error compounds — iterative refinement is the +# remedy, applied to both the batched per-scenario blocks and the nbatch=1 complement. +# Empirically on case9 two-stage SCOPF: ir=1 diverges to NaN, ir=3 reaches the optimum +# neighbourhood but stalls (inf_du ~1e-1), ir=5 converges cleanly to tol=1e-4 (obj +# matches the CPU / `:single` solve). Default to 5; stiffer/larger problems may need +# more — override via `--cudss-ir` / the `schur_*_opt_linear_solver` kwargs. +# +# This inner cuDSS IR is COMPLEMENTARY to — not a replacement for — the outer full-KKT +# `RichardsonIterator` (src/LinearSolvers/backsolve.jl), which already refines every Newton +# step against the true KKT residual. The outer IR's approximate inverse is one Schur +# reduction, whose accuracy depends on these per-scenario solves, so accurate inner solves +# are a prerequisite for the outer IR to converge (cf. Petra–Schenk–Lubin 2014). When the +# inner solve genuinely fails (e.g. IR on a numerically singular block), the cuDSS wrappers +# raise a Solve/FactorizationException → ERROR_IN_STEP_COMPUTATION rather than crashing. +const SCHUR_DEFAULT_CUDSS_IR = 5 +function _default_schur_cudss_options() + opt = MadNLP.default_options(CUDSSSolver) + opt.cudss_ir = SCHUR_DEFAULT_CUDSS_IR + opt.cudss_ir_tol = 0.0 # disarm cuDSS 0.8's IR_FAILED gate; keep the refinement steps + return opt +end + +# --- Deterministic condensation scatter: sorted-by-slot segment structure ------------- +# The Σ-amplified inequality-condensation scatters were atomic-add (nondeterministic +# summation order). For a given target array, group the contributions by their linear target +# slot so each slot can be summed by a single thread in a fixed order (see _det_quad_scatter! +# in kernels_schur.jl). Built once per kkt from the static index maps and cached. +struct _DetScatter{VI} # quadratic condensation: value = diag_buffer*jac*jac + jc1::VI + jc2::VI + buf::VI + segstart::VI + segslot::VI +end +struct _DetLinScatter{VI} # linear scatter: value = src[src_idx] + src_idx::VI + segstart::VI + segslot::VI +end + +# Build the contiguous per-slot segment boundaries for a stable sort of `slots`. +function _segments(slots::Vector{Int}) + perm = sortperm(slots) + s = slots[perm] + segslot = Int[]; segstart = Int[] + @inbounds for t in eachindex(s) + if t == 1 || s[t] != s[t - 1] + push!(segslot, s[t]); push!(segstart, t) + end + end + push!(segstart, length(s) + 1) + return perm, segstart, segslot +end + +function _segment_scatter(slots::Vector{Int}, jc1::Vector{Int}, jc2::Vector{Int}, buf::Vector{Int}) + if isempty(slots) + z = CuVector{Int}(undef, 0) + return _DetScatter(z, z, z, CuVector{Int}([1]), z) + end + perm, segstart, segslot = _segments(slots) + return _DetScatter( + CuVector{Int}(jc1[perm]), CuVector{Int}(jc2[perm]), CuVector{Int}(buf[perm]), + CuVector{Int}(segstart), CuVector{Int}(segslot), + ) +end + +function _segment_lin_scatter(slots::Vector{Int}, src_idx::Vector{Int}) + if isempty(slots) + z = CuVector{Int}(undef, 0) + return _DetLinScatter(z, CuVector{Int}([1]), z) + end + perm, segstart, segslot = _segments(slots) + return _DetLinScatter(CuVector{Int}(src_idx[perm]), CuVector{Int}(segstart), CuVector{Int}(segslot)) +end + +# Build the (A_kk, C_dk, S) deterministic scatter structures for `kkt` from its static index +# maps. Computed once and memoized in `kkt.det_cache` (a Ref field owned by the kkt, so it is +# freed with the kkt — no global cache, no leak). +function _compute_det_scatter(kkt::GPUSchurComplementCondensedKKTSystem) + ns = kkt.ns; nnzb = kkt.nnz_per_scenario; blk = kkt.blk_size; m = kkt.m + # ineq → A_kk: linear slot into the batched nzVal = (k-1)*nnzb + nzpos + 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)) + # 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)) + # 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)) + j2S = vcat(Array(kkt.schur_ineq_S_jcoo2), Array(kkt.schur_design_ineq_S_jcoo2)) + bfS = vcat(Array(kkt.schur_ineq_S_bufidx), Array(kkt.schur_design_ineq_S_bufidx)) + S = _segment_scatter(slotS, j1S, j2S, bfS) + + # Hessian scatters are ALSO large-valued (constraint Hessian × the relaxed-equality duals + # λ) and the Hessian COO has many duplicates → the non-atomic `+=` races. Make them + # deterministic linear scatters too. (pr_diag scatters target distinct diagonal slots with + # no collision, so they are already deterministic and left as-is.) + n_hA = kkt.n_per_s_hess_Akk + hnz = Array(kkt.gpu_hess_Akk_nzpos) + slotHA = n_hA > 0 ? Int[((i - 1) ÷ n_hA) * nnzb + hnz[i] for i in eachindex(hnz)] : Int[] + hessAkk = _segment_lin_scatter(slotHA, Array(kkt.gpu_hess_Akk_coo)) + n_hC = kkt.n_per_s_hess_Cdk + hv = Array(kkt.gpu_hess_Cdk_col); hd = Array(kkt.gpu_hess_Cdk_row) + slotHC = n_hC > 0 ? Int[hv[i] + (hd[i] - 1) * blk + ((i - 1) ÷ n_hC) * blk * m for i in eachindex(hv)] : Int[] + hessCdk = _segment_lin_scatter(slotHC, Array(kkt.gpu_hess_Cdk_coo)) + hessS = _segment_lin_scatter(Array(kkt.schur_hess_nzpos), Array(kkt.schur_hess_coo)) + + # --- Correct Hessian for the iterative-refinement mul! ----------------------------- + # The refinement matvec needs the TRUE full symmetric Hessian H. Two GPU pitfalls make + # the naive `mul!(.., Symmetric(hess_csc,:L), ..)` wrong (it only poisons mul!, not the + # Schur step — but mul! is the Richardson reference operator, so a wrong mul! stalls the + # outer refinement and inflates the IPM iteration count ~10x): + # (1) the generic GPU `transfer!` is a non-summing scatter `view(nzVal,map) .= V`, so + # duplicate Hessian COO entries (very common from AD) are dropped, not summed; + # (2) CUSPARSE ignores the `Symmetric(.,:L)` wrapper and multiplies only the stored + # lower triangle, dropping the strict-upper contribution. + # Fix both here: (1) a deterministic segmented SUM that fills hess_csc.nzVal correctly, + # and (2) a full-symmetric CSC `hess_full` so a plain general SpMV computes the true H*x. + nnzc = length(kkt.hess) + hess_lin = _segment_lin_scatter(Array(kkt.hess_csc_map), collect(1:nnzc)) + hess_full, hess_full_perm = MadNLP.get_tril_to_full(kkt.hess_csc) + fill!(hess_full.nzVal, zero(eltype(hess_full.nzVal))) + + # --- Correct Jacobian assembly for compress_jacobian! ------------------------------- + # The generic GPU `transfer!` (`view(jt_csc.nzVal, map) .= jac`) is a NON-summing scatter, so + # duplicate Jacobian COO entries (same (constraint, variable) coordinate — common with + # ExaModels AD) are dropped last-write-wins. The condensation kernels and the Hessian scatter + # SUM duplicates, so jt_csc must too, else jtprod!/mul!/dual recovery use a *different* + # Jacobian than the assembled A_kk/C_dk/S blocks. Build a deterministic segmented SUM keyed by + # jt_csc nzval slot (source = the jac COO values). + 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) +end + +function _get_det_scatter(kkt::GPUSchurComplementCondensedKKTSystem) + kkt.det_cache[] === nothing && (kkt.det_cache[] = _compute_det_scatter(kkt)) + return kkt.det_cache[] end # --- Dispatch: GPU path when callback uses CuVector --- function MadNLP.create_kkt_system( - ::Type{MadNLP.SchurComplementKKTSystem}, + ::Type{MadNLP.SchurComplementCondensedKKTSystem}, cb::MadNLP.SparseCallback{T, VT}, linear_solver::Type; opt_linear_solver=MadNLP.default_options(linear_solver), @@ -173,9 +327,35 @@ function MadNLP.create_kkt_system( schur_nv::Int=0, schur_nd::Int=0, schur_nc::Int=0, - schur_scenario_opt_linear_solver=MadNLP.default_options(CUDSSSolver), + schur_var_scen = nothing, + schur_con_scen = nothing, + schur_scenario_linear_solver::Type = CUDSSSolver, + schur_scenario_opt_linear_solver = _default_schur_cudss_options(), + schur_opt_linear_solver = _default_schur_cudss_options(), + kwargs..., ) where {T, VT <: CuVector{T}} + isempty(kwargs) || Base.@warn( + "GPUSchurComplementCondensedKKTSystem ignores unsupported kkt_options: " * + join(string.(keys(kwargs)), ", ") + ) + schur_scenario_linear_solver === CUDSSSolver || Base.@warn( + "GPUSchurComplementCondensedKKTSystem always factorizes the per-scenario blocks with a " * + "batched cuDSS solver; the requested `schur_scenario_linear_solver=" * + "$(schur_scenario_linear_solver)` is ignored. Pass `schur_scenario_opt_linear_solver` " * + "to configure the cuDSS scenario solver." + ) + # The first-stage Schur complement is also cuDSS (nbatch=1, so it reports inertia); the + # positional `linear_solver`/`opt_linear_solver` chosen by the IPM are NOT used here. Warn + # loudly when the user overrode `linear_solver` (e.g. LapackCUDASolver or a custom + # AbstractLinearSolver) instead of silently substituting cuDSS. The first-stage cuDSS options + # are configured via `schur_opt_linear_solver`. + linear_solver === CUDSSSolver || Base.@warn( + "GPUSchurComplementCondensedKKTSystem factorizes the first-stage Schur complement with " * + "cuDSS; the requested `linear_solver=$(linear_solver)` (and its `opt_linear_solver`) " * + "is ignored. Configure the first-stage cuDSS solver via `schur_opt_linear_solver`." + ) + n = cb.nvar m = cb.ncon ns_ineq = length(cb.ind_ineq) @@ -183,7 +363,8 @@ function MadNLP.create_kkt_system( nlb = length(cb.ind_lb) nub = length(cb.ind_ub) - ns, nv, nd, nc = MadNLP._resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc) + dims = MadNLP._resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc, schur_var_scen, schur_con_scen) + ns, nv, nd, nc = dims.ns, dims.nv, dims.nd, dims.nc # --- Get sparsity patterns on CPU --- jac_sparsity_I = Vector{Int32}(undef, cb.nnzj) @@ -228,12 +409,18 @@ function MadNLP.create_kkt_system( hess_sparsity_I, hess_sparsity_J, jac_sparsity_I, jac_sparsity_J, cpu_ind_eq, cpu_ind_ineq, + dims.var_scen, dims.con_scen, ) - nc_eq_per_s = sym.nc_eq_per_s nc_ineq_per_s = sym.nc_ineq_per_s blk_size = sym.blk_size nnz_per_scenario = sym.nnz_per_scenario akk_csc_cpu = sym.akk_csc_template + m_coupled = sym.m_coupled + coupled_inv = sym.coupled_inv + + # Flatten the per-scenario variable global indices (scenario-major) for the + # extract/writeback kernels, which can no longer assume contiguous stripes. + scen_var_global_flat = reduce(vcat, sym.scen_var_global; init = Int[]) # Flatten per-scenario block_maps into the per-field concatenated CPU vectors # the GPU struct stores (one upload per field below). @@ -248,12 +435,19 @@ function MadNLP.create_kkt_system( batched_colPtr, batched_rowVal, batched_nzVal, (blk_size, blk_size), ) - # --- Dense arrays on GPU --- - aug_com = CuMatrix{T}(undef, nd, nd) - fill!(aug_com, zero(T)) - C_dk_batched = CuArray{T, 3}(undef, blk_size, nd, ns) - fill!(C_dk_batched, zero(T)) - tmp_blk_nd_batched = CuArray{T, 3}(undef, blk_size, nd, ns) + # --- Sparse Schur complement (lower-triangular CSC) + reduced coupling buffers --- + schur_colPtr = CuVector{Cint}(Vector{Cint}(sym.schur_csc_colptr)) + schur_rowVal = CuVector{Cint}(Vector{Cint}(sym.schur_csc_rowval)) + schur_nzVal = CUDACore.fill(zero(T), sym.schur_nnz) + schur_csc = cuSPARSE.CuSparseMatrixCSC{T, Cint}( + schur_colPtr, schur_rowVal, schur_nzVal, (nd, nd), + ) + # Reduced coupling: only the m design columns that couple to a scenario. + C_dk_batched = CUDACore.fill(zero(T), blk_size, m_coupled, ns) + tmp_blk_nd_batched = CuArray{T, 3}(undef, blk_size, m_coupled, ns) + schur_block_batched = CuArray{T, 3}(undef, m_coupled, m_coupled, ns) + gpu_schur_fill_nzpos = CuMatrix{Int}(sym.schur_fill_nzpos) + gpu_coupled_design_local = CuVector{Int}(sym.coupled_design_local) # --- Diagonal vectors on GPU --- reg = CuVector{T}(undef, n + ns_ineq) @@ -270,11 +464,10 @@ function MadNLP.create_kkt_system( # --- Buffers --- diag_buffer = CuVector{T}(undef, max(ns_ineq, 1)) buffer = CuVector{T}(undef, m) - # Size matches the eq_global_indices gather/scatter shape used in - # solve_kkt! step 7. ns * nc_eq_per_s == n_eq_total under the uniform- - # scenario invariant validated in _build_schur_symbolic. - wy_eq_buf = CuVector{T}(undef, ns * nc_eq_per_s) - rhs_d = CuVector{T}(undef, nd) + rhs_d = CuVector{T}(undef, nd) + rhs_d_red = CuVector{T}(undef, m_coupled) + rhs_d_red_batched = CuArray{T, 3}(undef, m_coupled, 1, ns) + rhs_d_red_mat = CuMatrix{T}(undef, m_coupled, ns) rhs_k_batched = CuMatrix{T}(undef, blk_size, ns) solve_buffer = CuVector{T}(undef, blk_size * ns) @@ -282,108 +475,97 @@ function MadNLP.create_kkt_system( gpu_hess_Akk_coo = CuVector{Int}(flat.all_hess_Akk_coo) gpu_hess_Akk_nzpos = CuVector{Int}(flat.all_hess_Akk_nzpos) gpu_hess_Cdk_coo = CuVector{Int}(flat.all_hess_Cdk_coo) - gpu_hess_Cdk_row = CuVector{Int}(flat.all_hess_Cdk_row) + # Reduced C_dk: remap the design-column targets (1:nd) to compact columns (1:m). + 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_du_diag_global = CuVector{Int}(flat.all_du_diag_global) - gpu_du_diag_nzpos = CuVector{Int}(flat.all_du_diag_nzpos) - gpu_jeq_Akk_coo = CuVector{Int}(flat.all_jeq_Akk_coo) - gpu_jeq_Akk_nzpos = CuVector{Int}(flat.all_jeq_Akk_nzpos) - gpu_jeq_Cdk_coo = CuVector{Int}(flat.all_jeq_Cdk_coo) - gpu_jeq_Cdk_row = CuVector{Int}(flat.all_jeq_Cdk_row) - gpu_jeq_Cdk_col = CuVector{Int}(flat.all_jeq_Cdk_col) 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) gpu_ineq_Akk_bufidx = CuVector{Int}(flat.all_ineq_Akk_bufidx) - gpu_ineq_Cdk_row = CuVector{Int}(flat.all_ineq_Cdk_row) + gpu_ineq_Cdk_row = CuVector{Int}(coupled_inv[flat.all_ineq_Cdk_row]) gpu_ineq_Cdk_col = CuVector{Int}(flat.all_ineq_Cdk_col) gpu_ineq_Cdk_jcoo_d = CuVector{Int}(flat.all_ineq_Cdk_jcoo_d) gpu_ineq_Cdk_jcoo_v = CuVector{Int}(flat.all_ineq_Cdk_jcoo_v) gpu_ineq_Cdk_bufidx = CuVector{Int}(flat.all_ineq_Cdk_bufidx) - gpu_ineq_S_row = CuVector{Int}(flat.all_ineq_S_row) - gpu_ineq_S_col = CuVector{Int}(flat.all_ineq_S_col) - gpu_ineq_S_jcoo1 = CuVector{Int}(flat.all_ineq_S_jcoo1) - gpu_ineq_S_jcoo2 = CuVector{Int}(flat.all_ineq_S_jcoo2) - gpu_ineq_S_bufidx = CuVector{Int}(flat.all_ineq_S_bufidx) - gpu_hess_S_coo = CuVector{Int}(sym.hess_S_coo) - gpu_hess_S_row = CuVector{Int}(sym.hess_S_row) - gpu_hess_S_col = CuVector{Int}(sym.hess_S_col) - gpu_eq_global_indices = CuVector{Int}(sym.eq_global_flat) + + # Sparse-Schur nzpos scatter maps (lower-only) + their value sources. + gpu_schur_hess_coo = CuVector{Int}(sym.schur_hess_coo) + gpu_schur_hess_nzpos = CuVector{Int}(sym.schur_hess_nzpos) + gpu_schur_diag_nzpos = CuVector{Int}(sym.schur_diag_nzpos) + gpu_schur_ineq_S_nzpos = CuVector{Int}(sym.schur_ineq_S_nzpos) + gpu_schur_ineq_S_jcoo1 = CuVector{Int}(sym.schur_ineq_S_jcoo1) + gpu_schur_ineq_S_jcoo2 = CuVector{Int}(sym.schur_ineq_S_jcoo2) + gpu_schur_ineq_S_bufidx = CuVector{Int}(sym.schur_ineq_S_bufidx) + gpu_schur_design_ineq_S_nzpos = CuVector{Int}(sym.schur_design_ineq_S_nzpos) + gpu_schur_design_ineq_S_jcoo1 = CuVector{Int}(sym.schur_design_ineq_S_jcoo1) + gpu_schur_design_ineq_S_jcoo2 = CuVector{Int}(sym.schur_design_ineq_S_jcoo2) + gpu_schur_design_ineq_S_bufidx = CuVector{Int}(sym.schur_design_ineq_S_bufidx) + + # Tag-driven global index lists. + gpu_design_var_global = CuVector{Int}(sym.design_var_global) + gpu_scen_var_global = CuVector{Int}(scen_var_global_flat) # --- Create solvers --- quasi_newton = MadNLP.create_quasi_newton(hessian_approximation, cb, n; options=qn_options) - # cuDSS is the only sparse batched solver wired into this path; users tune it - # via `schur_scenario_opt_linear_solver` rather than swapping the type. + # Both the per-scenario blocks and the Schur complement are factorized by cuDSS. + # The per-scenario blocks are batched (nbatch=ns); the Schur complement is a single + # lower-triangular CSC (nbatch=1), so cuDSS reports its inertia. The dense cuSOLVER + # `sytrf` path is replaced. `linear_solver`/`opt_linear_solver` are unused on GPU. + # The batched scenario solver is analyzed once for its single-RHS shape (nbatch=ns) inside + # the CUDSSSolver constructor, and every scenario solve in build_kkt!/solve_kkt! is + # single-RHS (the multi-RHS ubatch solve is broken above nbatch≈14, see build_kkt!), so no + # extra multi-RHS descriptors/analysis are needed. scenario_solver = CUDSSSolver(A_kk_batched; opt=schur_scenario_opt_linear_solver) - _linear_solver = linear_solver(aug_com; opt=opt_linear_solver) - - # --- Multi-RHS cuDSS descriptors for batched A_kk \ C_dk' --- - # The descriptors hold the (blk × nd) shape per scenario and point at the - # existing buffers each iteration. - scenario_b_multi = CUDSS.CudssMatrix(T, blk_size, nd; nbatch=ns) - scenario_x_multi = CUDSS.CudssMatrix(T, blk_size, nd; nbatch=ns) - # Re-analyze for the multi-RHS shape so cuDSS plans enough workspace. - # This is the LARGEST RHS shape we will ever solve with on this handle; - # the later single-RHS solve at `solve_kkt!` step 3 reuses the same handle - # with a (blk × 1) × ns descriptor, which relies on the invariant that - # "analysis planned for a larger RHS accepts a smaller RHS at solve time". - # If that breaks on a future cuDSS version (e.g. a strict shape check or - # per-column IR state), the single-RHS path would need its own handle or - # padding. The `schur_cudss_ir` test exercises `cudss_ir > 0` — the config - # most likely to surface such a regression — as a tripwire. - CUDSS.cudss( - "analysis", scenario_solver.inner, - scenario_x_multi, scenario_b_multi; - asynchronous=scenario_solver.opt.cudss_asynchronous, - ) + schur_solver = CUDSSSolver(schur_csc; opt = schur_opt_linear_solver) # analysis runs once - return GPUSchurComplementKKTSystem( + return GPUSchurComplementCondensedKKTSystem( hess, jac, hess_raw, jt_coo, hess_csc, hess_csc_map, jt_csc, jt_csc_map, quasi_newton, reg, pr_diag, du_diag, l_diag, u_diag, l_lower, u_lower, - ns, nv, nd, nc, nc_eq_per_s, nc_ineq_per_s, blk_size, + ns, nv, nd, nc, nc_ineq_per_s, blk_size, m_coupled, A_kk_batched, nnz_per_scenario, - C_dk_batched, aug_com, - diag_buffer, buffer, wy_eq_buf, rhs_d, rhs_k_batched, tmp_blk_nd_batched, solve_buffer, + C_dk_batched, schur_csc, schur_block_batched, gpu_coupled_design_local, gpu_schur_fill_nzpos, + diag_buffer, buffer, rhs_d, rhs_d_red, rhs_d_red_batched, rhs_d_red_mat, rhs_k_batched, tmp_blk_nd_batched, solve_buffer, flat.n_per_s_hess_Akk, gpu_hess_Akk_coo, gpu_hess_Akk_nzpos, flat.n_per_s_hess_Cdk, gpu_hess_Cdk_coo, gpu_hess_Cdk_row, gpu_hess_Cdk_col, flat.n_per_s_pr_diag, gpu_pr_diag_global, gpu_pr_diag_nzpos, - flat.n_per_s_du_diag, gpu_du_diag_global, gpu_du_diag_nzpos, - flat.n_per_s_jeq_Akk, gpu_jeq_Akk_coo, gpu_jeq_Akk_nzpos, - flat.n_per_s_jeq_Cdk, gpu_jeq_Cdk_coo, gpu_jeq_Cdk_row, gpu_jeq_Cdk_col, flat.n_per_s_ineq_Akk, gpu_ineq_Akk_nzpos, gpu_ineq_Akk_jcoo1, gpu_ineq_Akk_jcoo2, gpu_ineq_Akk_bufidx, flat.n_per_s_ineq_Cdk, gpu_ineq_Cdk_row, gpu_ineq_Cdk_col, gpu_ineq_Cdk_jcoo_d, gpu_ineq_Cdk_jcoo_v, gpu_ineq_Cdk_bufidx, - flat.n_per_s_ineq_S, gpu_ineq_S_row, gpu_ineq_S_col, gpu_ineq_S_jcoo1, gpu_ineq_S_jcoo2, gpu_ineq_S_bufidx, - gpu_hess_S_coo, gpu_hess_S_row, gpu_hess_S_col, - gpu_eq_global_indices, + gpu_schur_hess_coo, gpu_schur_hess_nzpos, + gpu_schur_diag_nzpos, + gpu_schur_ineq_S_nzpos, gpu_schur_ineq_S_jcoo1, gpu_schur_ineq_S_jcoo2, gpu_schur_ineq_S_bufidx, + gpu_schur_design_ineq_S_nzpos, gpu_schur_design_ineq_S_jcoo1, gpu_schur_design_ineq_S_jcoo2, gpu_schur_design_ineq_S_bufidx, + gpu_design_var_global, gpu_scen_var_global, n_eq_total, CuVector{Int}(cpu_ind_eq), ns_ineq, CuVector{Int}(cpu_ind_ineq), CuVector{Int}(cpu_ind_lb), CuVector{Int}(cpu_ind_ub), - scenario_solver, _linear_solver, - scenario_x_multi, scenario_b_multi, + scenario_solver, schur_solver, + Base.RefValue{Any}(nothing), ) end # --- Trivial accessors --- -MadNLP.num_variables(kkt::GPUSchurComplementKKTSystem) = size(kkt.hess_csc, 1) +MadNLP.num_variables(kkt::GPUSchurComplementCondensedKKTSystem) = size(kkt.hess_csc, 1) -function MadNLP.get_slack_regularization(kkt::GPUSchurComplementKKTSystem) +function MadNLP.get_slack_regularization(kkt::GPUSchurComplementCondensedKKTSystem) n = MadNLP.num_variables(kkt) return view(kkt.pr_diag, n+1:n+kkt.n_ineq) end -function MadNLP.is_inertia_correct(kkt::GPUSchurComplementKKTSystem, num_pos, num_zero, num_neg) - return (num_zero == 0) && (num_pos == size(kkt.aug_com, 1)) +function MadNLP.is_inertia_correct(kkt::GPUSchurComplementCondensedKKTSystem, num_pos, num_zero, num_neg) + # RelaxEquality-only: the first-stage Schur complement is SPD (nd positive + # eigenvalues, no negative or zero ones). + return (num_zero == 0) && (num_pos == kkt.nd) && (num_neg == 0) end -MadNLP.should_regularize_dual(kkt::GPUSchurComplementKKTSystem, num_pos, num_zero, num_neg) = true +MadNLP.should_regularize_dual(kkt::GPUSchurComplementCondensedKKTSystem, num_pos, num_zero, num_neg) = true -MadNLP.nnz_jacobian(kkt::GPUSchurComplementKKTSystem) = MadNLP.nnz(kkt.jt_coo) +MadNLP.nnz_jacobian(kkt::GPUSchurComplementCondensedKKTSystem) = MadNLP.nnz(kkt.jt_coo) -function MadNLP.jtprod!(y::VT, kkt::GPUSchurComplementKKTSystem, x::VT) where {VT <: CuVector} +function MadNLP.jtprod!(y::VT, kkt::GPUSchurComplementCondensedKKTSystem, x::VT) where {VT <: CuVector} nx = MadNLP.num_variables(kkt) ns_ineq = kkt.n_ineq yx = view(y, 1:nx) @@ -393,24 +575,62 @@ function MadNLP.jtprod!(y::VT, kkt::GPUSchurComplementKKTSystem, x::VT) where {V return end -function MadNLP.compress_jacobian!(kkt::GPUSchurComplementKKTSystem) - MadNLP.transfer!(kkt.jt_csc, kkt.jt_coo, kkt.jt_csc_map) +MadNLP.compress_jacobian!(kkt::GPUSchurComplementCondensedKKTSystem) = + _schur_compress_jacobian!(kkt, _get_det_scatter(kkt)) # function barrier (see build_kkt!) +function _schur_compress_jacobian!(kkt::GPUSchurComplementCondensedKKTSystem, det) + # NOT the generic non-summing `transfer!` (`view(jt_csc.nzVal, map) .= jac`), which drops + # duplicate Jacobian COO entries (last-write-wins). Deterministically SUM duplicates into + # jt_csc.nzVal so jtprod!/mul!/dual recovery use the SAME Jacobian as the assembled + # A_kk/C_dk/S blocks (which also sum duplicates). See the `jt_lin` comment in + # `_compute_det_scatter`. Matches the CPU `_transfer!`, which sums with `+=`. + backend = CUDABackend() + nz = kkt.jt_csc.nzVal + fill!(nz, zero(eltype(nz))) + if length(det.jt_lin.segslot) > 0 + _det_lin_scatter!(backend)( + nz, kkt.jac, det.jt_lin.src_idx, det.jt_lin.segstart, det.jt_lin.segslot; + ndrange = length(det.jt_lin.segslot), + ) + end + return end -function MadNLP.compress_hessian!(kkt::GPUSchurComplementKKTSystem) - MadNLP.transfer!(kkt.hess_csc, kkt.hess_raw, kkt.hess_csc_map) +MadNLP.compress_hessian!(kkt::GPUSchurComplementCondensedKKTSystem) = + _schur_compress_hessian!(kkt, _get_det_scatter(kkt)) # function barrier (see build_kkt!) +function _schur_compress_hessian!(kkt::GPUSchurComplementCondensedKKTSystem, det) + # NOT the generic non-summing `transfer!`: deterministically SUM duplicate Hessian COO + # entries into hess_csc.nzVal, then materialize the full symmetric `hess_full` used by the + # refinement mul!. (See the hess_full comment in `_compute_det_scatter`.) + backend = CUDABackend() + nz = kkt.hess_csc.nzVal + fill!(nz, zero(eltype(nz))) + if length(det.hess_lin.segslot) > 0 + _det_lin_scatter!(backend)( + nz, kkt.hess, det.hess_lin.src_idx, det.hess_lin.segstart, det.hess_lin.segslot; + ndrange = length(det.hess_lin.segslot), + ) + end + copyto!(det.hess_full.nzVal, det.hess_full_perm) + return end # --- build_kkt! --- -function MadNLP.build_kkt!(kkt::GPUSchurComplementKKTSystem{T}) where T +# Function barrier: `_get_det_scatter` returns the memoized scatter as `::Any` (its concrete +# NamedTuple type is huge), so accessing `det.X` in-body would dynamic-dispatch every IPM +# 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 ns = kkt.ns nv = kkt.nv nd = kkt.nd + m = kkt.m n = MadNLP.num_variables(kkt) blk = kkt.blk_size backend = CUDABackend() nzval = kkt.A_kk_batched.nzVal nnz_s = kkt.nnz_per_scenario + Snz = kkt.schur_csc.nzVal # Compute condensing diagonal for inequalities if kkt.n_ineq > 0 @@ -419,40 +639,39 @@ function MadNLP.build_kkt!(kkt::GPUSchurComplementKKTSystem{T}) where T kkt.diag_buffer .= Sigma_s ./ (one(T) .- Sigma_d .* Sigma_s) end - # Zero out assembly targets - fill!(kkt.aug_com, zero(T)) + # Zero out assembly targets (sparse Schur nzval + scenario block + reduced coupling) + fill!(Snz, zero(T)) fill!(nzval, zero(T)) fill!(kkt.C_dk_batched, zero(T)) - # Initialize S from design-design Hessian - n_hess_S = length(kkt.hess_S_coo) - if n_hess_S > 0 - _init_S_hess_kernel!(backend)( - kkt.aug_com, kkt.hess, kkt.hess_S_coo, kkt.hess_S_row, kkt.hess_S_col; - ndrange=n_hess_S, + # Scatter design-design Hessian (deterministic) and pr_diag diagonal into the sparse + # Schur nzval. Design Hessian is large-valued + has duplicates → deterministic linear + # scatter; pr_diag targets distinct diagonal slots → already race-free, left atomic. + if length(det.hessS.segslot) > 0 + _det_lin_scatter!(backend)( + Snz, kkt.hess, det.hessS.src_idx, det.hessS.segstart, det.hessS.segslot; + ndrange = length(det.hessS.segslot), ) end if nd > 0 - _init_S_diag_kernel!(backend)(kkt.aug_com, kkt.pr_diag, ns * nv, nd; ndrange=nd) + _scatter_to_csc_atomic!(backend)( + Snz, kkt.pr_diag, kkt.design_var_global, kkt.schur_diag_nzpos; ndrange = nd, + ) end - # Scatter Hessian diagonal → A_kk (flattened maps) - n_hA = kkt.n_per_s_hess_Akk - if n_hA > 0 - _scatter_to_Akk_batched!(backend)( - nzval, kkt.hess, kkt.gpu_hess_Akk_coo, kkt.gpu_hess_Akk_nzpos, - nnz_s, n_hA; - ndrange=ns * n_hA, + # Scatter Hessian diagonal → A_kk (deterministic linear scatter) + if length(det.hessAkk.segslot) > 0 + _det_lin_scatter!(backend)( + nzval, kkt.hess, det.hessAkk.src_idx, det.hessAkk.segstart, det.hessAkk.segslot; + ndrange = length(det.hessAkk.segslot), ) end - # Scatter Hessian coupling → C_dk - n_hC = kkt.n_per_s_hess_Cdk - if n_hC > 0 - _scatter_to_Cdk_batched!(backend)( - kkt.C_dk_batched, kkt.hess, kkt.gpu_hess_Cdk_coo, - kkt.gpu_hess_Cdk_col, kkt.gpu_hess_Cdk_row, n_hC; - ndrange=ns * n_hC, + # Scatter Hessian coupling → C_dk (deterministic linear scatter) + if length(det.hessCdk.segslot) > 0 + _det_lin_scatter!(backend)( + reshape(kkt.C_dk_batched, :), kkt.hess, det.hessCdk.src_idx, det.hessCdk.segstart, det.hessCdk.segslot; + ndrange = length(det.hessCdk.segslot), ) end @@ -466,118 +685,98 @@ function MadNLP.build_kkt!(kkt::GPUSchurComplementKKTSystem{T}) where T ) end - # Scatter du_diag → A_kk diagonal - n_du = kkt.n_per_s_du_diag - if n_du > 0 - _scatter_to_Akk_batched!(backend)( - nzval, kkt.du_diag, kkt.gpu_du_diag_global, kkt.gpu_du_diag_nzpos, - nnz_s, n_du; - ndrange=ns * n_du, - ) - end - - # Scatter equality Jacobian → A_kk - n_jA = kkt.n_per_s_jeq_Akk - if n_jA > 0 - _scatter_to_Akk_batched!(backend)( - nzval, kkt.jac, kkt.gpu_jeq_Akk_coo, kkt.gpu_jeq_Akk_nzpos, - nnz_s, n_jA; - ndrange=ns * n_jA, - ) - end - - # Scatter equality Jacobian coupling → C_dk - n_jC = kkt.n_per_s_jeq_Cdk - if n_jC > 0 - _scatter_to_Cdk_batched!(backend)( - kkt.C_dk_batched, kkt.jac, kkt.gpu_jeq_Cdk_coo, - kkt.gpu_jeq_Cdk_col, kkt.gpu_jeq_Cdk_row, n_jC; - ndrange=ns * n_jC, - ) - end - - # Inequality condensation → A_kk - n_iA = kkt.n_per_s_ineq_Akk - if n_iA > 0 - _ineq_condense_Akk_kernel!(backend)( + # Inequality condensation (Σ-amplified) → A_kk / C_dk / S, DETERMINISTIC (one thread per + # output slot, fixed-order sum). Replaces the atomic-add scatters whose nondeterministic + # summation order perturbed the blocks ~1e-7 run-to-run and made convergence a roulette. + if length(det.Akk.segslot) > 0 + _det_quad_scatter!(backend)( nzval, kkt.jac, kkt.diag_buffer, - kkt.gpu_ineq_Akk_nzpos, kkt.gpu_ineq_Akk_jcoo1, - kkt.gpu_ineq_Akk_jcoo2, kkt.gpu_ineq_Akk_bufidx, - nnz_s, n_iA; - ndrange=ns * n_iA, + det.Akk.jc1, det.Akk.jc2, det.Akk.buf, det.Akk.segstart, det.Akk.segslot; + ndrange = length(det.Akk.segslot), ) end - - # Inequality condensation → C_dk - n_iC = kkt.n_per_s_ineq_Cdk - if n_iC > 0 - _ineq_condense_Cdk_kernel!(backend)( - kkt.C_dk_batched, kkt.jac, kkt.diag_buffer, - kkt.gpu_ineq_Cdk_col, kkt.gpu_ineq_Cdk_row, - kkt.gpu_ineq_Cdk_jcoo_d, kkt.gpu_ineq_Cdk_jcoo_v, - kkt.gpu_ineq_Cdk_bufidx, n_iC; - ndrange=ns * n_iC, + if length(det.Cdk.segslot) > 0 + _det_quad_scatter!(backend)( + reshape(kkt.C_dk_batched, :), kkt.jac, kkt.diag_buffer, + det.Cdk.jc1, det.Cdk.jc2, det.Cdk.buf, det.Cdk.segstart, det.Cdk.segslot; + ndrange = length(det.Cdk.segslot), ) end - - # Inequality condensation → S (atomic adds) - n_iS = kkt.n_per_s_ineq_S - if n_iS > 0 - _ineq_condense_S_kernel!(backend)( - kkt.aug_com, kkt.jac, kkt.diag_buffer, - kkt.gpu_ineq_S_row, kkt.gpu_ineq_S_col, - kkt.gpu_ineq_S_jcoo1, kkt.gpu_ineq_S_jcoo2, - kkt.gpu_ineq_S_bufidx; - ndrange=ns * n_iS, + # Scenario + design-only inequalities → S, merged into one deterministic scatter. + if length(det.S.segslot) > 0 + _det_quad_scatter!(backend)( + Snz, kkt.jac, kkt.diag_buffer, + det.S.jc1, det.S.jc2, det.S.buf, det.S.segstart, det.S.segslot; + ndrange = length(det.S.segslot), ) end # Factorize all scenario blocks in one batched cuDSS call MadNLP.factorize!(kkt.scenario_solver) - # Compute tmp = A_kk^{-1} * C_dk' in one batched multi-RHS cuDSS solve: - # for each scenario k, solve the (blk × blk) system with nd right-hand sides - # packed as the columns of C_dk_batched[:, :, k]. The (blk, nd, ns) layouts - # of both buffers line up directly with a cuDSS batched dense descriptor, so - # this is zero-copy — we only retarget the matrix descriptors each iteration. - CUDSS.cudss_update(kkt.scenario_b_multi, kkt.C_dk_batched) - CUDSS.cudss_update(kkt.scenario_x_multi, kkt.tmp_blk_nd_batched) - CUDSS.cudss( - "solve", kkt.scenario_solver.inner, - kkt.scenario_x_multi, kkt.scenario_b_multi; - asynchronous=kkt.scenario_solver.opt.cudss_asynchronous, - ) - - # S -= Σ_k C_dk[:,:,k]' * tmp[:,:,k] - # Reshape the (blk, nd, ns) buffers as (blk*ns, nd): column-major flattening - # collapses the per-scenario contributions into a single GEMM, since - # (C_2d' * tmp_2d)[a, b] = Σ_{i,k} C[i,a,k] * tmp[i,b,k] = Σ_k (C_dk[:,:,k]' * tmp[:,:,k])[a, b]. - # Avoids both the (nd × nd × ns) S_contrib buffer and the per-iteration - # `dropdims(sum(...))` allocation that the batched-GEMM path required. - C_2d = reshape(kkt.C_dk_batched, blk * ns, nd) - tmp_2d = reshape(kkt.tmp_blk_nd_batched, blk * ns, nd) - mul!(kkt.aug_com, C_2d', tmp_2d, -one(T), one(T)) + # Schur reduction restricted to the coupled-design block. Only the m design + # columns that couple to a scenario contribute (the rest of C_dk is exactly + # zero), so the reduction fills only the coupled × coupled sub-block. + if m > 0 + # tmp_red = A_kk⁻¹ * C_dk_red' (m right-hand sides per scenario). + # + # cuDSS's uniform-batch ("ubatch") solve is BROKEN for MULTI-RHS (nrhs > 1) once the + # batch count nbatch ≳ 14: it returns garbage (off by ~1e13). Verified on case118 + # (ns=176) and reproduced in PURE CUDSS with a synthetic uniform batch — single-RHS + # (nrhs=1) is always correct, multi-RHS fails above the threshold regardless of the + # matrix (a synthetic pure-CUDSS batch reproduces it). cuDSS's own batched tests only + # cover single-RHS, so this path was never exercised upstream. + # So solve the m coupled columns one at a time with the (correct) single-RHS batched + # solve — column j across all ns scenarios — mirroring the CPU path's column-by-column + # `A_kk \ C_dk'`. Slower than one multi-RHS call, but correct. (Remove this loop and + # restore the batched `scenario_*_multi` solve once cuDSS fixes multi-RHS ubatch.) + for j in 1:m + copyto!(reshape(kkt.solve_buffer, blk, ns), view(kkt.C_dk_batched, :, j, :)) + # `solve_linear_system!` already converts a cuDSS failure into a SolveException (after + # logging the status), which the IPM maps to ERROR_IN_STEP_COMPUTATION. The old outer + # try/catch that re-wrapped it as a FactorizationException was redundant double-handling + # (and its `CUDSS.CUDSSError` branch was dead — never reached the outer scope). + MadNLP.solve_linear_system!(kkt.scenario_solver, kkt.solve_buffer) + copyto!(view(kkt.tmp_blk_nd_batched, :, j, :), reshape(kkt.solve_buffer, blk, ns)) + end + + # D[:,:,k] = C_dk_red[:,:,k]' * tmp_red[:,:,k] (m×m) for all k in ONE batched + # GEMM (each batch slice is independent, so this is the correct batched form — + # NOT the fused-reshape trick that scrambles d/k for ns>1). Then scatter the + # lower-or-diagonal half as `S -= D` into the Schur nzval at the precomputed + # fill positions. + cuBLAS.gemm_strided_batched!( + 'T', 'N', one(T), kkt.C_dk_batched, kkt.tmp_blk_nd_batched, + zero(T), kkt.schur_block_batched, + ) + # Deterministic: one thread per lower-triangle (a,b) slot sums -D[a,b,k] over k. + _det_schur_block!(backend)( + Snz, kkt.schur_block_batched, kkt.schur_fill_nzpos, m, ns; ndrange = m * m, + ) + end return end # --- factorize_kkt! --- -function MadNLP.factorize_kkt!(kkt::GPUSchurComplementKKTSystem) +function MadNLP.factorize_kkt!(kkt::GPUSchurComplementCondensedKKTSystem) + # cuDSS reads `nonzeros(schur_solver.tril)` (=== schur_csc.nzVal, populated by + # build_kkt!) and runs factorization (first iter) / refactorization (later). return MadNLP.factorize!(kkt.linear_solver) end # --- solve_kkt! --- function MadNLP.solve_kkt!( - kkt::GPUSchurComplementKKTSystem{T}, + kkt::GPUSchurComplementCondensedKKTSystem{T}, w::MadNLP.AbstractKKTVector{T}, ) where T ns = kkt.ns nv = kkt.nv nd = kkt.nd + m = kkt.m n = MadNLP.num_variables(kkt) blk = kkt.blk_size - nc_eq = kkt.nc_eq_per_s backend = CUDABackend() wx = view(MadNLP.full(w), 1:n) @@ -595,58 +794,64 @@ function MadNLP.solve_kkt!( mul!(wx, kkt.jt_csc, kkt.buffer, one(T), one(T)) end - # Step 2: Extract per-scenario RHS blocks via kernel + # Step 2: Extract per-scenario RHS blocks via kernel. if blk * ns > 0 _extract_rhs_kernel!(backend)( - kkt.rhs_k_batched, wx, wy, kkt.eq_global_indices, - nv, nc_eq, ns, blk; + kkt.rhs_k_batched, wx, kkt.scen_var_global, nv, ns; ndrange=blk * ns, ) end - copyto!(kkt.rhs_d, view(wx, ns*nv+1:ns*nv+nd)) + @views kkt.rhs_d[1:nd] .= wx[kkt.design_var_global] # Step 3: Forward elimination — batched solve copyto!(kkt.solve_buffer, vec(kkt.rhs_k_batched)) MadNLP.solve_linear_system!(kkt.scenario_solver, kkt.solve_buffer) copyto!(vec(kkt.rhs_k_batched), kkt.solve_buffer) - # rhs_d -= Σ_k C_dk[:,:,k]' * rhs_k[:,k]: same column-major reshape trick as - # build_kkt!. C_2d' * vec(rhs_k_batched) == Σ_k C_dk[:,:,k]' * rhs_k[:,k]. - # One launch instead of ns. - C_2d = reshape(kkt.C_dk_batched, blk * ns, nd) - mul!(kkt.rhs_d, C_2d', vec(kkt.rhs_k_batched), -one(T), one(T)) + # rhs_d[coupled] -= Σ_k C_dk_red[:,:,k]' * rhs_k[:,k]. Only the coupled design rows receive a + # contribution (non-coupled C_dk columns are exactly zero). One strided-batched GEMM computes + # the per-scenario (m×1) products C_dk_red[:,:,k]' rhs_k[:,k]; sum over the batch dim gives the + # accumulated reduced RHS, which is then scatter-subtracted. (Replaces the per-scenario `mul!` + # loop = ns kernel launches per solve.) + if m > 0 + cuBLAS.gemm_strided_batched!( + 'T', 'N', one(T), kkt.C_dk_batched, reshape(kkt.rhs_k_batched, blk, 1, ns), + zero(T), kkt.rhs_d_red_batched, + ) + sum!(reshape(kkt.rhs_d_red, m, 1, 1), kkt.rhs_d_red_batched) # reduce over scenarios + @views kkt.rhs_d[kkt.coupled_design_local] .-= kkt.rhs_d_red + end - # Step 4: Solve Schur complement + # Step 4: Solve the first-stage Schur complement system (size nd, SPD) with cuDSS. MadNLP.solve_linear_system!(kkt.linear_solver, kkt.rhs_d) - # Step 5: Back-substitution. rhs_k[:,k] -= tmp[:,:,k] * rhs_d for all k; - # reshape tmp as (blk*ns, nd) so a single GEMV updates the flattened - # rhs_k_batched in place. - tmp_2d = reshape(kkt.tmp_blk_nd_batched, blk * ns, nd) - mul!(vec(kkt.rhs_k_batched), tmp_2d, kkt.rhs_d, -one(T), one(T)) + # Step 5: Back-substitution. rhs_k[:,k] -= tmp_red[:,:,k] * rhs_d[coupled] per scenario. The + # reduced design update is shared across scenarios, so broadcast it into an (m×ns) buffer and + # do the whole back-substitution as one strided-batched GEMM. + if m > 0 + @views kkt.rhs_d_red .= kkt.rhs_d[kkt.coupled_design_local] + kkt.rhs_d_red_mat .= reshape(kkt.rhs_d_red, m, 1) + cuBLAS.gemm_strided_batched!( + 'N', 'N', -one(T), kkt.tmp_blk_nd_batched, reshape(kkt.rhs_d_red_mat, m, 1, ns), + one(T), reshape(kkt.rhs_k_batched, blk, 1, ns), + ) + end # Step 6: Write back to w via kernel if blk * ns > 0 _writeback_rhs_kernel!(backend)( - wx, wy, kkt.rhs_k_batched, kkt.eq_global_indices, - nv, nc_eq, ns, blk; + wx, kkt.rhs_k_batched, kkt.scen_var_global, nv, ns; ndrange=blk * ns, ) end - copyto!(view(wx, ns*nv+1:ns*nv+nd), kkt.rhs_d) + @views wx[kkt.design_var_global] .= kkt.rhs_d[1:nd] - # Step 7: Recover inequality duals and slacks + # Step 7: Recover inequality duals and slacks (all constraints are inequalities + # under RelaxEquality, so there are no equality duals to preserve). if kkt.n_ineq > 0 - # Stash eq duals; mul! below overwrites all of wy. - # Vectorized GPU gather/scatter — no @allowscalar. - copyto!(kkt.wy_eq_buf, view(wy, kkt.eq_global_indices)) - - # J * Δx + # J * Δx (overwrites all of wy) mul!(wy, kkt.jt_csc', wx) - # Restore equality duals - view(wy, kkt.eq_global_indices) .= kkt.wy_eq_buf - # Inequality dual recovery wy[kkt.ind_ineq] .= kkt.diag_buffer .* wy[kkt.ind_ineq] .- kkt.buffer[kkt.ind_ineq] ws .= (ws .+ view(wy, kkt.ind_ineq)) ./ Sigma_s @@ -657,12 +862,24 @@ function MadNLP.solve_kkt!( end # --- mul! for iterative refinement --- +# Function barrier: hoist the `::Any` det access (its `hess_full` is the SpMV reference operator, +# 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::GPUSchurComplementKKTSystem{T}, + kkt::GPUSchurComplementCondensedKKTSystem{T}, x::MadNLP.AbstractKKTVector, alpha = one(T), beta = zero(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, ) where {T, VT <: CuVector{T}} n = MadNLP.num_variables(kkt) @@ -675,7 +892,9 @@ function MadNLP.mul!( xz = @view(MadNLP.dual(x)[kkt.ind_ineq]) wx .= beta .* wx - mul!(wx, Symmetric(kkt.hess_csc, :L), xx, alpha, one(T)) + # Use the correctly-summed full-symmetric Hessian (general SpMV) — NOT + # Symmetric(hess_csc,:L), which CUSPARSE multiplies as lower-triangle-only. + mul!(wx, hess_full, xx, alpha, one(T)) m = size(kkt.jt_csc, 2) if m > 0 @@ -690,9 +909,11 @@ function MadNLP.mul!( return w end -function MadNLP.mul_hess_blk!(wx::VT, kkt::GPUSchurComplementKKTSystem{T}, t) where {T, VT <: CuVector{T}} +MadNLP.mul_hess_blk!(wx::VT, kkt::GPUSchurComplementCondensedKKTSystem{T}, t) where {T, VT <: CuVector{T}} = + _schur_mul_hess_blk!(wx, kkt, t, _get_det_scatter(kkt).hess_full) # function barrier (see mul!) +function _schur_mul_hess_blk!(wx::VT, kkt::GPUSchurComplementCondensedKKTSystem{T}, t, hess_full) where {T, VT <: CuVector{T}} n = MadNLP.num_variables(kkt) - mul!(@view(wx[1:n]), Symmetric(kkt.hess_csc, :L), @view(t[1:n])) + mul!(@view(wx[1:n]), hess_full, @view(t[1:n])) fill!(@view(wx[n+1:end]), 0) wx .+= t .* kkt.pr_diag end diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/cudss.jl index 28508a54c..b2810dc5c 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 = 1.0e-8 # currently ignored by cuDSS + cudss_ir_tol::Float64 = 1e-8 cudss_pivot_threshold::Float64 = 0.0 cudss_pivot_epsilon::Float64 = 0.0 cudss_matching_alg::String = "default" @@ -47,7 +47,12 @@ function set_cudss_options!(solver::CUDSS.CudssSolver, opt::CudssSolverOptions) CUDSS.cudss_set(solver, "pivot_threshold", opt.cudss_pivot_threshold) end if opt.cudss_matching - if pkgversion(CUDSS) < v"0.8" + if solver.matrix.nbatch > 1 + # cuDSS matching (`matching_alg`) fails the *analysis* phase with + # CUDSS_STATUS_NOT_SUPPORTED on a uniform-batch solver (through cuDSS 0.8), + # so never enable it there (e.g. the two-stage per-scenario batch solver). + Base.@warn "cuDSS matching is not supported on uniform-batch (ubatch) solvers; ignoring `cudss_matching = true` for this batched solver." maxlog = 1 + elseif pkgversion(CUDSS) < v"0.8" CUDSS.cudss_set(solver, "use_matching", 1) if opt.cudss_matching_alg != "default" CUDSS.cudss_set(solver, "matching_alg", opt.cudss_matching_alg) @@ -90,6 +95,7 @@ mutable struct CUDSSSolver{T, V} <: MadNLP.AbstractLinearSolver{T} x_gpu::CUDSS.CudssMatrix{T} b_gpu::CUDSS.CudssMatrix{T} buffer::V + diag::V opt::CudssSolverOptions logger::MadNLP.MadNLPLogger @@ -151,19 +157,41 @@ function CUDSSSolver( # Always allocate it to support dynamic updates to opt.cudss_ir buffer = CuVector{T}(undef, n * nbatch) + # Scratch for the factor diagonal, used to recover the inertia when matching is on + # (cuDSS misreports it as (0, 0)); only the nbatch == 1 path ever queries inertia. + diag = CuVector{T}(undef, n) + return CUDSSSolver( solver, csc, - x_gpu, b_gpu, buffer, + x_gpu, b_gpu, buffer, diag, opt, logger, ) end function MadNLP.factorize!(M::CUDSSSolver) CUDSS.cudss_update(M.inner.matrix, nonzeros(M.tril)) - if M.inner.fresh_factorization - CUDSS.cudss("factorization", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) - else - CUDSS.cudss("refactorization", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) + # A cuDSS error during (re)factorization is translated into a FactorizationException + # (→ ERROR_IN_STEP_COMPUTATION) rather than crashing with a raw CUDSSError. This does + # not interfere with the inertia→regularization recovery, which only runs when + # factorize! returns normally and reports indefiniteness via `inertia` (info != 0 + # yields a dummy inertia). + try + if M.inner.fresh_factorization + CUDSS.cudss("factorization", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) + else + CUDSS.cudss("refactorization", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) + end + catch e + # Log the cuDSS status before erasing it: FactorizationException is field-less, so an OOM + # (CUDSS_STATUS_ALLOC_FAILED), a not-initialized handle, or an internal error would + # otherwise be indistinguishable from numerical breakdown once the IPM maps this to + # ERROR_IN_STEP_COMPUTATION. + if e isa CUDSS.CUDSSError + @warn(M.logger, "cuDSS (re)factorization failed with status $(e.code); reporting a FactorizationException (-> ERROR_IN_STEP_COMPUTATION).") + throw(FactorizationException()) + else + rethrow(e) + end end return M end @@ -176,7 +204,22 @@ function MadNLP.solve_linear_system!(M::CUDSSSolver{T, V}, xb::V) where {T, V} CUDSS.cudss_update(M.b_gpu, xb) end CUDSS.cudss_update(M.x_gpu, xb) - CUDSS.cudss("solve", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) + # A cuDSS failure here — e.g. iterative refinement on a numerically singular system + # reporting CUDSS_STATUS_EXECUTION_FAILED — is a step-computation failure, not a bug. + # Translate it into a MadNLP SolveException so the IPM maps it to + # ERROR_IN_STEP_COMPUTATION (-3) instead of crashing with a raw CUDSSError. + try + CUDSS.cudss("solve", M.inner, M.x_gpu, M.b_gpu, asynchronous = M.opt.cudss_asynchronous) + catch e + # As in factorize!: attach the cuDSS status before converting to the field-less + # SolveException, so OOM vs numerical breakdown is diagnosable. + if e isa CUDSS.CUDSSError + @warn(M.logger, "cuDSS solve failed with status $(e.code); reporting a SolveException (-> ERROR_IN_STEP_COMPUTATION).") + throw(SolveException()) + else + rethrow(e) + end + end return xb end @@ -202,6 +245,21 @@ function MadNLP.inertia(M::CUDSSSolver) elseif M.opt.cudss_algorithm == MadNLP.LDL # N.B.: cuDSS does not always return the correct inertia. if info == 0 + if M.opt.cudss_matching + # cuDSS (through 0.8) reports inertia (0, 0) whenever matching is enabled, + # even though the factorization is correct — trusting it sends + # InertiaBased/InertiaAuto into an endless regularization bump and then + # restoration. Recover the inertia from the sign counts of the factor + # diagonal instead. Caveat: exact only for 1×1 pivots, which holds for the + # (quasi-definite) condensed KKT family this solver targets. + CUDSS.cudss_set(M.inner, "diag", M.diag) + CUDSS.cudss_get(M.inner, "diag") + d = Array(M.diag) + z = zero(eltype(d)) + npos = count(>(z), d) + nneg = count(<(z), d) + return (npos, n - npos - nneg, nneg) + end (k, l) = CUDSS.cudss_get(M.inner, "inertia") @assert 0 ≤ k + l ≤ n return (k, n - k - l, l) diff --git a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl index 6f56c0372..85d565181 100644 --- a/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl +++ b/lib/MadNLPGPU/ext/MadNLPGPUCUDAExt/kernels_schur.jl @@ -24,151 +24,131 @@ @inbounds nzval[offset + nzpos_flat[i]] += src[src_idx[i]] end -# Scatter src[src_idx[i]] into C_dk[v_idx, d_idx, k] (blk_size × nd × ns). -# Layout has scenario-var dim first so each C_dk[:, :, k] is a contiguous -# (blk × nd) matrix — directly usable as a cuDSS multi-RHS dense matrix. -# Used for: Hessian coupling entries, equality Jacobian coupling entries. -@kernel function _scatter_to_Cdk_batched!( - C_dk, - @Const(src), - @Const(src_idx), - @Const(v_idx_flat), # scenario-var index (1..blk) - @Const(d_idx_flat), # design-var index (1..nd) - @Const(n_per_s), -) - i = @index(Global) - k = (i - 1) ÷ n_per_s + 1 - @inbounds C_dk[v_idx_flat[i], d_idx_flat[i], k] += src[src_idx[i]] -end - -# Inequality condensation → A_kk (lower triangle). -@kernel function _ineq_condense_Akk_kernel!( - nzval, - @Const(jac), - @Const(diag_buffer), - @Const(nzpos_flat), - @Const(jcoo1_flat), - @Const(jcoo2_flat), - @Const(bufidx_flat), - @Const(nnz_per_block), - @Const(n_per_s), +# Extract per-scenario RHS from global wx → rhs_k_batched (nv × ns). Under +# RelaxEquality the per-scenario block has no equality rows (blk_size == nv), so +# only scenario variables are gathered. Their global indices come from +# `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), ) i = @index(Global) - offset = ((i - 1) ÷ n_per_s) * nnz_per_block - @inbounds nzval[offset + nzpos_flat[i]] += diag_buffer[bufidx_flat[i]] * - jac[jcoo1_flat[i]] * jac[jcoo2_flat[i]] + k = (i - 1) ÷ nv + 1 + local_idx = (i - 1) % nv + 1 + if k <= ns + @inbounds gi = scen_var_global[(k - 1) * nv + local_idx] + @inbounds rhs_k[local_idx, k] = wx[gi] + end end -# Inequality condensation → C_dk (blk_size × nd × ns; see _scatter_to_Cdk_batched!). -@kernel function _ineq_condense_Cdk_kernel!( - C_dk, - @Const(jac), - @Const(diag_buffer), - @Const(v_idx_flat), # scenario-var index (1..blk) - @Const(d_idx_flat), # design-var index (1..nd) - @Const(jcoo_d_flat), - @Const(jcoo_v_flat), - @Const(bufidx_flat), - @Const(n_per_s), +# 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), ) i = @index(Global) - k = (i - 1) ÷ n_per_s + 1 - @inbounds C_dk[v_idx_flat[i], d_idx_flat[i], k] += diag_buffer[bufidx_flat[i]] * - jac[jcoo_d_flat[i]] * jac[jcoo_v_flat[i]] + k = (i - 1) ÷ nv + 1 + local_idx = (i - 1) % nv + 1 + if k <= ns + @inbounds gi = scen_var_global[(k - 1) * nv + local_idx] + @inbounds wx[gi] = rhs_k[local_idx, k] + end end -# Inequality condensation → S (atomic adds since multiple scenarios target the -# same nd × nd dense matrix). -@kernel function _ineq_condense_S_kernel!( - S, - @Const(jac), - @Const(diag_buffer), - @Const(row_flat), - @Const(col_flat), - @Const(jcoo1_flat), - @Const(jcoo2_flat), - @Const(bufidx_flat), -) +# ===== Sparse-Schur (cuDSS) assembly kernel ====================================== +# Scatters into the lower-triangular CSC `nzval` of the sparse Schur complement at a +# precomputed nzval position. Used for the design pr_diag diagonal, which targets distinct +# diagonal slots with no collision (so a plain atomic add is race-free and deterministic; the +# Σ-amplified / duplicate-prone contributions go through the deterministic scatters below). +@kernel function _scatter_to_csc_atomic!( + nzval, + @Const(src), + @Const(src_idx), + @Const(nzpos), + ) i = @index(Global) - @inbounds begin - val = diag_buffer[bufidx_flat[i]] * jac[jcoo1_flat[i]] * jac[jcoo2_flat[i]] - CUDACore.@atomic S[row_flat[i], col_flat[i]] += val - end + @inbounds CUDACore.@atomic nzval[nzpos[i]] += src[src_idx[i]] end -# Initialize S from design-design Hessian -@kernel function _init_S_hess_kernel!( - S, - @Const(hess), - @Const(coo_indices), - @Const(row_indices), - @Const(col_indices), -) - idx = @index(Global) - @inbounds S[row_indices[idx], col_indices[idx]] += hess[coo_indices[idx]] -end +# ===== Deterministic (atomic-free) condensation scatters ======================== +# The atomic-add condensation scatters above are correct but their summation ORDER is +# nondeterministic. With the RelaxEquality condensation weights Σ ≈ 1e8, that reorders the +# accumulation enough to perturb the assembled blocks by ~1e-7 run-to-run; amplified by the +# (legitimate) KKT conditioning this tips the IPM step-acceptance residual across its +# threshold → nondeterministic convergence. These deterministic variants assign ONE thread +# per output slot, which sums that slot's contributions in a fixed (sorted) order with a plain +# `+=` (no two threads touch the same slot within a launch, and launches are sequential), so +# the assembled blocks are reproducible run-to-run — matching the deterministic CPU assembly. +# +# `segstart`/`segslot` describe, for one target array, the contiguous contribution ranges per +# slot in a stable-by-slot ordering: segment `s` owns sorted contributions +# `segstart[s] : segstart[s+1]-1`, all targeting linear index `segslot[s]`. -# Add pr_diag_dd to S diagonal -@kernel function _init_S_diag_kernel!( - S, - @Const(pr_diag), - @Const(diag_offset), - @Const(nd), -) - idx = @index(Global) - if idx <= nd - @inbounds S[idx, idx] += pr_diag[diag_offset + idx] +# Linear scatter value = src[src_idx] → any target nzval, deterministic (one thread per slot). +@kernel function _det_lin_scatter!( + nzval, + @Const(src), + @Const(src_idx), + @Const(segstart), + @Const(segslot), + ) + s = @index(Global) + @inbounds begin + acc = zero(eltype(nzval)) + for j in segstart[s]:(segstart[s + 1] - 1) + acc += src[src_idx[j]] + end + nzval[segslot[s]] += acc end end -# Extract per-scenario RHS from global wx/wy → rhs_k_batched (blk_size × ns) -@kernel function _extract_rhs_kernel!( - rhs_k, - @Const(wx), - @Const(wy), - @Const(eq_global_indices), - @Const(nv), - @Const(nc_eq), - @Const(ns), - @Const(blk_size), -) - i = @index(Global) - k = (i - 1) ÷ blk_size + 1 - local_idx = (i - 1) % blk_size + 1 - if k <= ns - if local_idx <= nv - gi = (k - 1) * nv + local_idx - @inbounds rhs_k[local_idx, k] = wx[gi] - else - eq_idx = local_idx - nv - flat_idx = (k - 1) * nc_eq + eq_idx - @inbounds rhs_k[local_idx, k] = wy[eq_global_indices[flat_idx]] +# Inequality condensation (Σ-amplified) → any target nzval (A_kk / C_dk / S), deterministic. +@kernel function _det_quad_scatter!( + nzval, + @Const(jac), + @Const(diag_buffer), + @Const(jc1), + @Const(jc2), + @Const(buf), + @Const(segstart), + @Const(segslot), + ) + s = @index(Global) + @inbounds begin + acc = zero(eltype(nzval)) + for j in segstart[s]:(segstart[s + 1] - 1) + acc += diag_buffer[buf[j]] * jac[jc1[j]] * jac[jc2[j]] end + nzval[segslot[s]] += acc end end -# Write back per-scenario solution to global wx/wy -@kernel function _writeback_rhs_kernel!( - wx, - wy, - @Const(rhs_k), - @Const(eq_global_indices), - @Const(nv), - @Const(nc_eq), - @Const(ns), - @Const(blk_size), -) +# Schur reduction block → CSC nzval, deterministic: one thread per lower-triangle (a,b) slot +# sums -D[a,b,k] over the ns scenarios in fixed order. +@kernel function _det_schur_block!( + nzval, + @Const(D), + @Const(fill_nzpos), + @Const(m), + @Const(ns), + ) i = @index(Global) - k = (i - 1) ÷ blk_size + 1 - local_idx = (i - 1) % blk_size + 1 - if k <= ns - if local_idx <= nv - gi = (k - 1) * nv + local_idx - @inbounds wx[gi] = rhs_k[local_idx, k] - else - eq_idx = local_idx - nv - flat_idx = (k - 1) * nc_eq + eq_idx - @inbounds wy[eq_global_indices[flat_idx]] = rhs_k[local_idx, k] + a = (i - 1) % m + 1 + b = (i - 1) ÷ m + 1 + if a >= b + @inbounds begin + acc = zero(eltype(nzval)) + for k in 1:ns + acc += -D[a, b, k] + end + nzval[fill_nzpos[a, b]] += acc end end end diff --git a/lib/MadNLPGPU/test/schur_cuda_test.jl b/lib/MadNLPGPU/test/schur_cuda_test.jl index 79e4317ec..36e6aac51 100644 --- a/lib/MadNLPGPU/test/schur_cuda_test.jl +++ b/lib/MadNLPGPU/test/schur_cuda_test.jl @@ -1,12 +1,13 @@ using Test using LinearAlgebra +using SparseArrays using CUDA using CUDSS using MadNLP using MadNLPGPU using MadNLPTests -@testset "GPUSchurComplementKKTSystem" begin +@testset "GPUSchurComplementCondensedKKTSystem" begin @testset "Basic convergence — quadratic with coupling" begin ns, nv, nd, nc = 3, 1, 1, 1 @@ -26,8 +27,8 @@ using MadNLPTests result = madnlp( nlp; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCUDASolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -54,17 +55,27 @@ using MadNLPTests ref = madnlp( nlp_cpu; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCPUSolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) + # The equality constraints (lcon == ucon == 0) are relaxed into slacks, whose + # barrier weight σ_s ~ 1/μ² blows up near convergence and enters each condensed + # per-scenario block as σ_s·JᵀJ — a rank-1 term that drives cond → 1e16. cuDSS + # iterative refinement cannot refine a numerically singular system and raises + # CUDSS_STATUS_IR_FAILED, so pin the batched scenario solver to ir=0 (best-effort + # direct solve); MadNLP's outer Richardson refinement handles the residual. + gpu_opts = schur_opts(; ns, nv, nd, nc) + scenario_opt = MadNLP.default_options(CUDSSSolver) + scenario_opt.cudss_ir = 0 + gpu_opts[:schur_scenario_opt_linear_solver] = scenario_opt gpu_result = madnlp( nlp_gpu; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCUDASolver, - kkt_options = schur_opts(; ns, nv, nd, nc), + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, + kkt_options = gpu_opts, print_level = MadNLP.ERROR, ) @@ -96,8 +107,8 @@ using MadNLPTests result = madnlp( nlp; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCUDASolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -142,8 +153,8 @@ using MadNLPTests result = madnlp( nlp; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCUDASolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, kkt_options = opts, print_level = MadNLP.ERROR, ) @@ -168,8 +179,8 @@ using MadNLPTests result = madnlp( nlp; callback = MadNLP.SparseCallback, - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCUDASolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -179,4 +190,111 @@ using MadNLPTests @test isapprox(sol[2], 7.0; atol = 1.0e-3) @test isapprox(sol[3], 5.0; atol = 1.0e-3) end + + @testset "Design-only constraints — match CPU reference" begin + # Non-contiguous layout + design-only equality and inequality constraints. + for (ns, nv, nd) in ((2, 2, 2), (3, 2, 3)) + n = ns * nv + nd + qp_cpu, _, _, kkt_opts = + build_twostage_qp_general(zeros(Float64, n); ns, nv, nd, permute = true) + qp_gpu, _, _, _ = + build_twostage_qp_general(CUDA.zeros(Float64, n); ns, nv, nd, permute = true) + + # Pin both cuDSS solvers to ir=0: relaxed-equality slacks drive the condensed + # systems numerically singular near convergence (σ_s·JᵀJ, σ_s ~ 1/μ²), where + # cuDSS iterative refinement raises CUDSS_STATUS_IR_FAILED. Here the design-only + # equality lands in the first-stage Schur complement, so the complement solver + # is the one that goes singular (the scenario blocks can too). The best-effort + # direct solve plus MadNLP's outer Richardson refinement suffice. + scenario_opt = MadNLP.default_options(CUDSSSolver) + scenario_opt.cudss_ir = 0 + complement_opt = MadNLP.default_options(CUDSSSolver) + complement_opt.cudss_ir = 0 + kkt_opts[:schur_scenario_opt_linear_solver] = scenario_opt + kkt_opts[:schur_opt_linear_solver] = complement_opt + + ref = madnlp(qp_cpu; linear_solver = LapackCPUSolver, print_level = MadNLP.ERROR) + # The condensed default `bound_relax_factor = tol` (= 1e-4 here) relaxes the + # design-only EQUALITIES into ±2e-4 boxes; the optimizer exploits that slack and + # lands ~2·tol off the exact-equality reference — this testset compares against + # ground truth, so pin the tight relaxation explicitly. + gpu_result = madnlp( + qp_gpu; + callback = MadNLP.SparseCallback, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = CUDSSSolver, + kkt_options = kkt_opts, + bound_relax_factor = 1.0e-8, + print_level = MadNLP.ERROR, + ) + + @test ref.status == MadNLP.SOLVE_SUCCEEDED + @test gpu_result.status == MadNLP.SOLVE_SUCCEEDED + @test isapprox(gpu_result.objective, ref.objective; atol = 1.0e-6) + @test isapprox(Array(gpu_result.solution), Array(ref.solution); atol = 1.0e-4) + end + end + + # Assemble one KKT matrix (jac+hess at x0, pr_diag=1, du_diag=-1e-8) and return it. + # RelaxEquality (the Schur default): all constraints condense, so a coupling + # equality becomes a coupling inequality whose J'ΣJ splits across A_kk / C_dk / S_dd. + function _assembled_kkt(qp, solver, kkt_opts) + cb = MadNLP.create_callback( + MadNLP.SparseCallback, qp; equality_treatment = MadNLP.RelaxEquality, + ) + kkt = MadNLP.create_kkt_system(MadNLP.SchurComplementCondensedKKTSystem, cb, solver; kkt_opts...) + x0 = cb.nlp.meta.x0 + y0 = cb.nlp.meta.y0 + MadNLP._eval_jac_wrapper!(cb, x0, MadNLP.get_jacobian(kkt)) + MadNLP.compress_jacobian!(kkt) + MadNLP._eval_lag_hess_wrapper!(cb, x0, y0, MadNLP.get_hessian(kkt)) + MadNLP.compress_hessian!(kkt) + fill!(kkt.pr_diag, 1.0) + fill!(kkt.du_diag, -1.0e-8) + MadNLP.build_kkt!(kkt) + return kkt + end + + @testset "Sparse Schur (cuDSS) — GPU schur_csc matches CPU schur_csc" begin + # Both the CPU and GPU first-stage Schur complements are sparse lower-triangular + # CSCs (same symbolic pattern, size nd × nd, SPD). Densified + symmetrized they + # must agree. Directly catches triangle / sign / nzpos / double-count bugs in the + # sparse assembly, including the relaxed coupling-equality → C_dk cross-term. + for (ns, nv, nd) in ((2, 2, 2), (3, 2, 3), (4, 1, 2)) + n = ns * nv + nd + qp_cpu, _, _, kkt_opts = + build_twostage_qp_general(zeros(Float64, n); ns, nv, nd, permute = true) + qp_gpu, _, _, _ = + build_twostage_qp_general(CUDA.zeros(Float64, n); ns, nv, nd, permute = true) + + kkt_cpu = _assembled_kkt(qp_cpu, MadNLP.MumpsSolver, kkt_opts) + kkt_gpu = _assembled_kkt(qp_gpu, CUDSSSolver, kkt_opts) + + A_dense = Array(Symmetric(Matrix(kkt_cpu.schur_csc), :L)) + S = kkt_gpu.schur_csc + nS = size(S, 1) + S_lower = SparseMatrixCSC( + nS, nS, Array(S.colPtr), Array(S.rowVal), Array(S.nzVal), + ) + S_full = Array(Symmetric(Matrix(S_lower), :L)) + + @test kkt_gpu.m <= nd + @test isapprox(S_full, A_dense; atol = 1.0e-7, rtol = 1.0e-7) + end + end + + @testset "Sparse Schur (cuDSS) — reduced coupling buffers sized to m ≤ nd" begin + # The reduced C_dk / tmp / Schur-block buffers are width m (the coupled-design + # count), not nd — the memory/compute win. m ≤ nd always; for SCOPF m ≪ nd. + ns, nv, nd = 3, 2, 3 + n = ns * nv + nd + qp_gpu, _, _, kkt_opts = + build_twostage_qp_general(CUDA.zeros(Float64, n); ns, nv, nd, permute = true) + kkt = _assembled_kkt(qp_gpu, CUDSSSolver, kkt_opts) + @test 1 <= kkt.m <= nd + @test size(kkt.C_dk_batched, 2) == kkt.m + @test size(kkt.tmp_blk_nd_batched, 2) == kkt.m + @test size(kkt.schur_block_batched, 1) == kkt.m + @test size(kkt.schur_block_batched, 2) == kkt.m + end end diff --git a/lib/MadNLPTests/src/Instances/twostage_qp.jl b/lib/MadNLPTests/src/Instances/twostage_qp.jl index 8275ca857..7192493d7 100644 --- a/lib/MadNLPTests/src/Instances/twostage_qp.jl +++ b/lib/MadNLPTests/src/Instances/twostage_qp.jl @@ -2,7 +2,7 @@ TwoStageQP{T, VT, MT, VI} <: NLPModels.AbstractNLPModel{T, VT} Diagonal-Hessian two-stage stochastic QP, used to exercise -`SchurComplementKKTSystem` (CPU and GPU) without depending on ExaModels. +`SchurComplementCondensedKKTSystem` (CPU and GPU) without depending on ExaModels. Variable layout: `x = [v_{1,1},...,v_{1,nv}, v_{2,1},...,v_{ns,nv}, d_1,...,d_nd]` Constraint layout: `c = [c_{1,1},...,c_{1,nc}, c_{2,1},...,c_{ns,nc}]` @@ -181,11 +181,180 @@ function build_twostage_qp( end """ - schur_opts(; ns, nv, nd, nc) -> Dict + schur_opts(; ns, nv, nd, nc, var_scen=nothing, con_scen=nothing) -> Dict -Convenience: build the `kkt_options` dict for `SchurComplementKKTSystem` -with the four scenario dimensions. +Convenience: build the `kkt_options` dict for `SchurComplementCondensedKKTSystem` with the +four scenario dimensions. When `var_scen`/`con_scen` are supplied (per-variable / +per-constraint scenario tags, 0 = design), they are passed through so the Schur +build partitions by tag instead of assuming the contiguous `[v_1..v_ns, d]` / +`[c_1..c_ns]` layout — required for models with non-contiguous orderings and/or +design-only constraints. """ -schur_opts(; ns, nv, nd, nc) = Dict{Symbol, Any}( - :schur_ns => ns, :schur_nv => nv, :schur_nd => nd, :schur_nc => nc, -) +function schur_opts(; ns, nv, nd, nc, var_scen = nothing, con_scen = nothing) + opts = Dict{Symbol, Any}( + :schur_ns => ns, :schur_nv => nv, :schur_nd => nd, :schur_nc => nc, + ) + var_scen === nothing || (opts[:schur_var_scen] = var_scen) + con_scen === nothing || (opts[:schur_con_scen] = con_scen) + return opts +end + +""" + build_twostage_qp_general(x0_template = Vector{Float64}(); + ns, nv, nd, permute = true) -> (qp, var_scen, con_scen, kkt_opts) + +Build a strictly-convex two-stage QP that exercises the *general* Schur path: +design-only **equality** and **inequality** constraints, plus a **non-contiguous** +global ordering (design variables are NOT last and each scenario's variables are +scattered). Strict convexity (positive-diagonal Hessian) gives a unique optimum, so +the Schur solve must match the reference `SparseKKTSystem` solve exactly. + +Per scenario: `nv` variables, one equality and one inequality constraint coupling +that scenario's variables to the first `n_coupled` design variables (default all +`nd`). Globally: one design-only equality (`Σ d = 1`) and one design-only +inequality (`-5 ≤ d_1 - d_2 ≤ 5`, inactive). Passing `n_coupled < nd` leaves +`nd - n_coupled` design variables **uncoupled** (they enter only the objective and +the design-only constraints), so the Schur fill width `m = n_coupled < nd` — this +exercises the reduced-width coupling path. With `permute = true` the canonical +contiguous layout is reordered to +`[design vars, then component-major scenario vars]` / `[design cons, then scenario +cons]`; the returned `var_scen` / `con_scen` describe that order. + +Returns the `TwoStageQP`, the tag vectors, and a ready `kkt_options` dict. +""" +function build_twostage_qp_general( + x0_template::AbstractVector{T} = Vector{Float64}(undef, 0); + ns::Int, nv::Int, nd::Int, permute::Bool = true, n_coupled::Int = nd, + ) where {T} + nd >= 2 || error("build_twostage_qp_general needs nd >= 2 for the design constraints") + 1 <= n_coupled <= nd || error("build_twostage_qp_general needs 1 <= n_coupled <= nd") + + # Per-scenario: 1 equality + 1 inequality. Design: 1 equality + 1 inequality. + nc = 2 + n = ns * nv + nd + m = ns * nc + 2 + + # --- Canonical (contiguous) assembly: vars [v_1..v_ns, d], cons [scen.., design..] --- + cvar(k, j) = (k - 1) * nv + j # canonical scenario var index + cdes(l) = ns * nv + l # canonical design var index + cscon(k, t) = (k - 1) * nc + t # canonical scenario constraint (t=1 eq, 2 ineq) + cdcon(t) = ns * nc + t # canonical design constraint (t=1 eq, 2 ineq) + + H = zeros(T, n) + g = zeros(T, n) + lvar = fill(T(-50), n) + uvar = fill(T(50), n) + A = zeros(T, m, n) + lcon = zeros(T, m) + ucon = zeros(T, m) + + @inbounds for k in 1:ns, j in 1:nv + i = cvar(k, j) + H[i] = 2 + T(0.1) * (j + k) + g[i] = -2 * (T(j) + T(k)) # pulls v[k,j] toward a positive target + end + @inbounds for l in 1:nd + i = cdes(l) + H[i] = 2 + T(0.1) * l + g[i] = T(0.5) * l + end + + @inbounds for k in 1:ns + # scenario equality: Σ_j v[k,j] + Σ_{l≤n_coupled} d[l] = 0. Only the first `n_coupled` + # design vars appear in a scenario constraint, so exactly those couple into C_dk; design + # vars l > n_coupled are "design-only" (they enter only the objective and the design-only + # constraints below). With n_coupled < nd this exercises the m < nd reduced Schur path. + re = cscon(k, 1) + for j in 1:nv + A[re, cvar(k, j)] = 1 + end + for l in 1:n_coupled + A[re, cdes(l)] = 1 + end + lcon[re] = 0; ucon[re] = 0 + # scenario inequality (wide → inactive): -10 ≤ Σ_j v[k,j] - d[1] ≤ 10 + ri = cscon(k, 2) + for j in 1:nv + A[ri, cvar(k, j)] = 1 + end + A[ri, cdes(1)] = -1 + lcon[ri] = -10; ucon[ri] = 10 + end + # design equality: Σ_l d[l] = 1 + let re = cdcon(1) + for l in 1:nd + A[re, cdes(l)] = 1 + end + lcon[re] = 1; ucon[re] = 1 + end + # design inequality (wide → inactive): -5 ≤ d[1] - d[2] ≤ 5 + let ri = cdcon(2) + A[ri, cdes(1)] = 1; A[ri, cdes(2)] = -1 + lcon[ri] = -5; ucon[ri] = 5 + end + + canon_var_scen = vcat((fill(k, nv) for k in 1:ns)..., fill(0, nd)) + canon_con_scen = vcat((fill(k, nc) for k in 1:ns)..., 0, 0) + + # --- Permutation: design first, then component-major scenario vars/cons --- + if permute + var_order = Int[] + for l in 1:nd + push!(var_order, cdes(l)) + end + for j in 1:nv, k in 1:ns + push!(var_order, cvar(k, j)) + end + con_order = Int[] + push!(con_order, cdcon(1)); push!(con_order, cdcon(2)) + for t in 1:nc, k in 1:ns + push!(con_order, cscon(k, t)) + end + else + var_order = collect(1:n) + con_order = collect(1:m) + end + + H_p = H[var_order]; g_p = g[var_order] + lvar_p = lvar[var_order]; uvar_p = uvar[var_order] + A_p = A[con_order, var_order] + lcon_p = lcon[con_order]; ucon_p = ucon[con_order] + var_scen = canon_var_scen[var_order] + con_scen = canon_con_scen[con_order] + + # --- COO structure (dense Jacobian nonzeros; diagonal Hessian) --- + jrows = Int[]; jcols = Int[]; jvals = T[] + @inbounds for i in 1:m, j in 1:n + if A_p[i, j] != 0 + push!(jrows, i); push!(jcols, j); push!(jvals, A_p[i, j]) + end + end + hrows = collect(1:n); hcols = collect(1:n); hvals = copy(H_p) + + to_device(v_h::AbstractVector) = (v = similar(x0_template, length(v_h)); copyto!(v, v_h); v) + to_device(M_h::AbstractMatrix) = (M = similar(x0_template, size(M_h)...); copyto!(M, M_h); M) + + meta = NLPModels.NLPModelMeta( + n; + ncon = m, + nnzj = length(jrows), + nnzh = n, + x0 = to_device(zeros(T, n)), + y0 = to_device(zeros(T, m)), + lvar = to_device(lvar_p), + uvar = to_device(uvar_p), + lcon = to_device(lcon_p), + ucon = to_device(ucon_p), + minimize = true, + ) + + qp = TwoStageQP( + meta, NLPModels.Counters(), + to_device(H_p), to_device(A_p), to_device(g_p), + to_device(hvals), to_device(jvals), + hrows, hcols, jrows, jcols, + ) + + kkt_opts = schur_opts(; ns, nv, nd, nc, var_scen, con_scen) + return qp, var_scen, con_scen, kkt_opts +end diff --git a/lib/MadNLPTests/src/MadNLPTests.jl b/lib/MadNLPTests/src/MadNLPTests.jl index 390c397c9..c4289cf3f 100644 --- a/lib/MadNLPTests/src/MadNLPTests.jl +++ b/lib/MadNLPTests/src/MadNLPTests.jl @@ -13,7 +13,7 @@ import JuMP: Model, @variable, @constraint, @objective, optimize!, set_attribute MOI, termination_status, LowerBoundRef, UpperBoundRef, value, dual, fix import NLPModelsJuMP -export test_madnlp, solcmp, TwoStageQP, build_twostage_qp, schur_opts +export test_madnlp, solcmp, TwoStageQP, build_twostage_qp, build_twostage_qp_general, schur_opts function solcmp(x,sol;atol=1e-4,rtol=1e-4) aerr = norm(x-sol,Inf) diff --git a/src/IPM/IPM.jl b/src/IPM/IPM.jl index 7cd07a4c4..e7e413758 100644 --- a/src/IPM/IPM.jl +++ b/src/IPM/IPM.jl @@ -201,7 +201,7 @@ function MadNLPSolver(nlp::AbstractNLPModel{T,VT}; kwargs...) where {T, VT} dx_ur = view(d.xp, cb.ind_ub) # TODO inertia_correction_method = if ipm_opt.inertia_correction_method == InertiaAuto - is_inertia(kkt.linear_solver)::Bool ? InertiaBased : InertiaFree + is_inertia(kkt)::Bool ? InertiaBased : InertiaFree else ipm_opt.inertia_correction_method end diff --git a/src/IPM/options.jl b/src/IPM/options.jl index 5553ed6fa..fddd00654 100644 --- a/src/IPM/options.jl +++ b/src/IPM/options.jl @@ -4,6 +4,8 @@ parse_option(::Type{Module},str::String) = eval(Symbol(str)) parse_option(::Type{<:AbstractUserCallback}, f::Any) = f parse_option(type::Type{T},i::Int64) where {T<:Enum} = type(i) +const CondensedKKTSystems = Union{SparseCondensedKKTSystem, SchurComplementCondensedKKTSystem} + function set_options!(opt::AbstractOptions, options) other_options = Dict{Symbol, Any}() for (key, val) in options @@ -143,8 +145,8 @@ tau\\_min | 0.99 | lower bound on fraction- # NLP options kappa_d::T = 1e-5 - fixed_variable_treatment::Type = kkt_system <: MadNLP.SparseCondensedKKTSystem ? MadNLP.RelaxBound : MadNLP.MakeParameter - equality_treatment::Type = kkt_system <: MadNLP.SparseCondensedKKTSystem ? MadNLP.RelaxEquality : MadNLP.EnforceEquality + fixed_variable_treatment::Type = kkt_system <: CondensedKKTSystems ? MadNLP.RelaxBound : MadNLP.MakeParameter + equality_treatment::Type = kkt_system <: CondensedKKTSystems ? MadNLP.RelaxEquality : MadNLP.EnforceEquality bound_relax_factor::T = 1e-8 jacobian_constant::Bool = false hessian_constant::Bool = false @@ -157,7 +159,7 @@ tau\\_min | 0.99 | lower bound on fraction- # initialization options dual_initialized::Bool = false - dual_initialization_method::Type = kkt_system <: MadNLP.SparseCondensedKKTSystem ? DualInitializeSetZero : DualInitializeLeastSquares + dual_initialization_method::Type = kkt_system <: CondensedKKTSystems ? DualInitializeSetZero : DualInitializeLeastSquares constr_mult_init_max::T = 1e3 bound_push::T = 1e-2 bound_fac::T = 1e-2 @@ -212,25 +214,41 @@ 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) + 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, callback = callback, kkt_system = kkt_system, linear_solver = linear_solver, + bound_relax_factor = bound_relax_factor, ) 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)) default_sparse_solver(nlp::AbstractNLPModel) = MumpsSolver function check_option_sanity(options) - is_kkt_dense = options.kkt_system <: AbstractDenseKKTSystem || options.kkt_system <: SchurComplementKKTSystem + is_kkt_dense = options.kkt_system <: AbstractDenseKKTSystem is_hess_approx_dense = options.hessian_approximation <: Union{BFGS, DampedBFGS} - if input_type(options.linear_solver) == :csc && is_kkt_dense + # `SchurComplementCondensedKKTSystem` is a *sparse-callback* KKT system: on both CPU and + # GPU it factorizes a sparse lower-triangular first-stage Schur complement (MumpsSolver / + # any `:csc` solver on CPU, cuDSS on GPU), so a `:csc` linear solver is valid for it — only + # the genuinely dense KKT systems require a dense linear solver here. It also has no dense + # quasi-Newton wiring, so it is NOT classified as dense above: a dense BFGS/DampedBFGS + # approximation is rejected below with a clear options error instead of failing later with + # an obscure constructor/dispatch error. + if input_type(options.linear_solver) == :csc && options.kkt_system <: AbstractDenseKKTSystem error("[options] Sparse Linear solver is not supported in dense mode.\n"* "Please use a dense linear solver or change `kkt_system` ") end diff --git a/src/KKT/KKTsystem.jl b/src/KKT/KKTsystem.jl index 7e6277d71..bd92cc0ae 100644 --- a/src/KKT/KKTsystem.jl +++ b/src/KKT/KKTsystem.jl @@ -233,6 +233,16 @@ get_kkt(kkt::AbstractKKTSystem) = kkt.aug_com get_jacobian(kkt::AbstractKKTSystem) = kkt.jac get_hessian(kkt::AbstractKKTSystem) = kkt.hess +""" + is_inertia(kkt::AbstractKKTSystem) + +Whether `kkt` can report the inertia of the full KKT matrix, and hence whether +`InertiaAuto` may select the inertia-based regularization. Defaults to the capability of +the KKT system's linear solver; KKT systems whose matrix is factorized by more than one +solver (e.g. the Schur path) must override this to require all of them. +""" +is_inertia(kkt::AbstractKKTSystem) = is_inertia(kkt.linear_solver) + """ is inertia_correct(kkt::AbstractKKTSystem, num_pos, num_zero, num_neg) diff --git a/src/KKT/Schur/schur.jl b/src/KKT/Schur/schur.jl index a3eb36768..6790e5932 100644 --- a/src/KKT/Schur/schur.jl +++ b/src/KKT/Schur/schur.jl @@ -15,15 +15,6 @@ struct ScenarioBlockMap hess_Cdk_row::Vector{Int} hess_Cdk_col::Vector{Int} - # Equality Jacobian, scenario vars → A_kk lower triangle only - jeq_Akk_coo::Vector{Int} - jeq_Akk_nzpos::Vector{Int} - - # Equality Jacobian, design vars → C_dk - jeq_Cdk_coo::Vector{Int} - jeq_Cdk_row::Vector{Int} # design var index (1:nd) - jeq_Cdk_col::Vector{Int} # nv + eq_local_idx - # Inequality condensation → A_kk (lower triangle) ineq_Akk_nzpos::Vector{Int} ineq_Akk_jcoo1::Vector{Int} @@ -47,14 +38,10 @@ struct ScenarioBlockMap # Diagonal positions in A_kk nzval for pr_diag (nv entries) pr_diag_global::Vector{Int} pr_diag_nzpos::Vector{Int} - - # Diagonal positions in A_kk nzval for du_diag (nc_eq entries) - du_diag_global::Vector{Int} - du_diag_nzpos::Vector{Int} end """ - SchurComplementKKTSystem{T, VT, MT, QN, LS, LS2, VI} <: AbstractCondensedKKTSystem{T, VT, MT, QN} + SchurComplementCondensedKKTSystem{T, VT, MT, QN, LS, LS2, VI} <: AbstractCondensedKKTSystem{T, VT, MT, QN} KKT system exploiting block-arrowhead structure from two-stage stochastic programs via Schur complement decomposition, using sparse COO/CSC storage for the global @@ -63,13 +50,18 @@ Hessian and Jacobian and sparse per-scenario block solvers. Variable layout: `[v_1, ..., v_ns, d]` where `v_k ∈ R^nv`, `d ∈ R^nd`. Constraint layout: `[c_1, ..., c_ns]` where `c_k ∈ R^nc`. -The augmented per-scenario block `A_k` (size `blk_size × blk_size`) is stored -as a sparse lower-triangular `SparseMatrixCSC` and factored by a configurable -sparse solver (default `MumpsSolver` — each `A_k` is symmetric indefinite, not -SQD in general). The coupling blocks `C_dk` remain dense. The Schur complement -`S = aug_com` (size `nd × nd`) is dense. +All equality constraints are relaxed to inequalities (`RelaxEquality`) and +condensed, so the per-scenario block `A_k` (size `nv × nv`) and the first-stage +Schur complement `schur_csc` (size `nd × nd`) are both symmetric *positive +definite* — there is no bordered equality saddle. `A_k` is a sparse +lower-triangular `SparseMatrixCSC` (`= H_k + pr_diag + Σ_i J_i' D_i J_i`) factored +by a configurable sparse solver (default `MumpsSolver`). The coupling blocks +`C_dk` remain dense. The reduction `Σ_k C_dk A_kk⁻¹ C_dk'` fills only the `m×m` +coupled-design block of `schur_csc`, factored by a sparse symmetric solver +(`MumpsSolver` by default; any `:csc` symmetric solver that reports inertia, +e.g. `Ma57Solver`/`LDLSolver`, works). """ -struct SchurComplementKKTSystem{ +struct SchurComplementCondensedKKTSystem{ T, VT <: AbstractVector{T}, MT <: AbstractMatrix{T}, @@ -109,37 +101,56 @@ struct SchurComplementKKTSystem{ nv::Int nd::Int nc::Int - nc_eq_per_s::Int nc_ineq_per_s::Int - blk_size::Int # = nv + nc_eq_per_s + blk_size::Int # = nv (per-scenario condensed block size) # Per-scenario sparse augmented blocks (lower triangle only) A_kk::Vector{SparseMatrixCSC{T, Int32}} - C_dk::Vector{MT} # ns × (nd × blk_size) — dense - - # Schur complement (what the dense linear solver sees) - aug_com::MT # nd × nd + C_dk::Vector{MT} # ns × (m × blk_size) — dense, reduced to the m COUPLED + # 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 + # vars couple to a scenario, so the reduction `Σ_k C_dk A_kk⁻¹ C_dk'` fills only + # the `m×m` coupled-design block. + schur_csc::SparseMatrixCSC{T, Int32} + m::Int # coupled design vars (Schur fill width) + coupled_design_local::Vector{Int} # length m — design-local indices that couple + schur_fill_nzpos::Matrix{Int} # (m, m) — nzval positions of the Schur-fill block + schur_block::MT # (m, m) buffer — per-scenario C_dk_red' A_kk⁻¹ C_dk_red' # Buffers diag_buffer::VT # n_ineq — condensing diagonal buffer::VT # m_total — general - wy_eq_buf::VT # ns*nc_eq_per_s — preserves eq duals across J*Δx round-trip - rhs_d::VT # nd — design RHS + rhs_d::VT # nd — design vars + rhs_d_red::VT # m — reduced (coupled) design RHS accumulator rhs_k::Vector{VT} # ns × blk_size — scenario RHS buffers - tmp_blk_nd::Vector{MT} # ns × (blk_size × nd) + tmp_blk_nd::Vector{MT} # ns × (blk_size × m) — A_kk⁻¹ C_dk_red' (reduced coupled cols) solve_buffers::Vector{VT} # ns × blk_size — per-scenario column-by-column solve buffers - # Precomputed index maps block_maps::Vector{ScenarioBlockMap} - hess_S_coo::Vector{Int} # COO indices for design-design Hessian - hess_S_row::Vector{Int} # S row (1:nd) - hess_S_col::Vector{Int} # S col (1:nd) - - # Flat per-scenario equality indices (length ns * nc_eq_per_s). - # Scenario k's indices live at (k-1)*nc_eq_per_s+1 : k*nc_eq_per_s. - eq_global_indices::Vector{Int} - # Inequality/equality/bound index info + # Sparse-Schur scatter maps: scatter into `schur_csc.nzval` at precomputed nzval + # positions (lower-triangle; the symmetric solver reads the lower triangle). Lists + # are lower-only so each slot is hit once. + schur_hess_coo::Vector{Int} # design Hessian: COO index into `hess` + schur_hess_nzpos::Vector{Int} # → nzval slot + schur_diag_nzpos::Vector{Int} # pr_diag design diagonal (paired with design_var_global) + schur_ineq_S_nzpos::Vector{Int} # scenario ineq → S (flat over all scenarios) + schur_ineq_S_jcoo1::Vector{Int} + schur_ineq_S_jcoo2::Vector{Int} + schur_ineq_S_bufidx::Vector{Int} + schur_design_ineq_S_nzpos::Vector{Int} # design ineq → S + schur_design_ineq_S_jcoo1::Vector{Int} + schur_design_ineq_S_jcoo2::Vector{Int} + schur_design_ineq_S_bufidx::Vector{Int} + + # Tag-driven global index lists (design/scenario vars need not be contiguous). + design_var_global::Vector{Int} # length nd — global index of each design var + scen_var_global::Vector{Vector{Int}} # ns × nv — global index of each scenario var + + # Inequality/equality/bound index info (n_eq == 0 / ind_eq empty under RelaxEquality, + # kept for the generic KKT interface). n_eq::Int ind_eq::VI n_ineq::Int @@ -147,91 +158,143 @@ struct SchurComplementKKTSystem{ ind_lb::VI ind_ub::VI + # Aggregated inertia of the per-scenario blocks A_k, refreshed by `build_kkt!`. + # Haynsworth additivity gives In(K) = Σ_k In(A_k) + In(S), but the IPM only ever + # queries the linear solver for In(S) (`solver.jl`: `inertia(kkt.linear_solver)`). + # `is_inertia_correct` reads these so the test covers all of In(K). + block_num_neg::Base.RefValue{Int} # Σ_k #negative eigenvalues of A_k + block_num_zero::Base.RefValue{Int} # Σ_k #zero eigenvalues of A_k + block_num_indef::Base.RefValue{Int} # #{k : A_k not positive definite} — diagnostics + # Solvers scenario_solvers::Vector{LS2} linear_solver::LS # for Schur complement S (dense) end """ - _resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc) - -> (ns, nv, nd, nc) - -Resolve two-stage stochastic dimensions for a SchurComplementKKTSystem. + _schur_tags_from_callback(cb) -> (ns, var_scen, con_scen) | nothing -If `schur_ns == 0`, attempt to auto-detect from `cb.nlp.tags::TwoStageTags` -(an ExaModel convention): `tags.var_scenario[i] == 0` flags design variables, -`== 1` flags scenario variables; same encoding for `con_scenario`. +Read per-variable/per-constraint scenario-assignment vectors from the NLP model +behind `cb`, if it exposes a two-stage tag. Supports the ExaModels convention +(`cb.nlp.tag::TwoStageExaModelTag` with fields `nscen`, `var_scen`, `con_scen`) +and a legacy/test interface (`cb.nlp.tags` with `ns`, `var_scenario`, +`con_scenario`). Returns `nothing` when no recognizable tag is present. -Asserts that the resolved dimensions are consistent with `n` and `m`. +In both conventions `var_scen[i] == 0` flags a design variable and `== k` flags +scenario `k`; same encoding for constraints. """ -function _resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc) - if schur_ns == 0 && hasproperty(cb.nlp, :tags) - tags = cb.nlp.tags +function _schur_tags_from_callback(cb) + nlp = cb.nlp + if hasproperty(nlp, :tag) + tag = nlp.tag + if hasproperty(tag, :nscen) && hasproperty(tag, :var_scen) && hasproperty(tag, :con_scen) + return Int(tag.nscen), Vector{Int}(Array(tag.var_scen)), Vector{Int}(Array(tag.con_scen)) + end + end + if hasproperty(nlp, :tags) + tags = nlp.tags if hasproperty(tags, :ns) && hasproperty(tags, :var_scenario) && hasproperty(tags, :con_scenario) - schur_ns = tags.ns - var_scen = Array(tags.var_scenario) - con_scen = Array(tags.con_scenario) - - # Single-pass histograms over the scenario tags. Index 1 is the design - # bucket (tag 0); 1+k is scenario k. Cheaper than the previous three - # `count(==(·))` passes and lets us validate per-scenario uniformity - # in the same pass — important because a malformed model whose global - # aggregates happen to satisfy n == ns*nv + nd would otherwise drive - # downstream symbolic build into garbage territory. - var_hist = zeros(Int, schur_ns + 1) - for tag in var_scen - t = Int(tag) - (0 <= t <= schur_ns) || error( - "var_scenario tag $t out of range [0, $schur_ns]; " * - "0 = design, 1..$schur_ns = scenario index." - ) - @inbounds var_hist[t + 1] += 1 - end - con_hist = zeros(Int, schur_ns + 1) - for tag in con_scen - t = Int(tag) - (0 <= t <= schur_ns) || error( - "con_scenario tag $t out of range [0, $schur_ns]; " * - "1..$schur_ns = scenario index." - ) - @inbounds con_hist[t + 1] += 1 - end + return Int(tags.ns), Vector{Int}(Array(tags.var_scenario)), Vector{Int}(Array(tags.con_scenario)) + end + end + return nothing +end - schur_nd = var_hist[1] - schur_nv = var_hist[2] - schur_nc = con_hist[2] +""" + _resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc, + schur_var_scen=nothing, schur_con_scen=nothing) + -> (; ns, nv, nd, nc, nc_design, var_scen, con_scen) + +Resolve two-stage stochastic dimensions AND the per-variable/per-constraint +scenario-assignment vectors for a `SchurComplementCondensedKKTSystem`. The vectors encode +`var_scen[i] ∈ {0..ns}` (0 = design) and `con_scen[j] ∈ {0..ns}`; the symbolic +build partitions the Hessian/Jacobian by these tags rather than by index +arithmetic, so design variables and a scenario's variables need NOT be contiguous +in the global ordering. + +Resolution priority: +1. explicit `schur_var_scen` / `schur_con_scen` (passed via `kkt_options`); +2. else a two-stage tag on the model (see [`_schur_tags_from_callback`](@ref)); +3. else synthesize the contiguous layout `[v_1..v_ns, d]` / `[c_1..c_ns]` from + `schur_ns/nv/nd/nc` (backward-compatible with hand-built two-stage models). + +Design-only constraints (`con_scen == 0`) are supported; `nc_design` counts them. +Asserts consistency with `n` and `m`, and per-scenario uniformity. +""" +function _resolve_schur_dims( + cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc, + schur_var_scen = nothing, schur_con_scen = nothing, + ) + var_scen = nothing + con_scen = nothing + ns = schur_ns + + if schur_var_scen !== nothing && schur_con_scen !== nothing + var_scen = Vector{Int}(Array(schur_var_scen)) + con_scen = Vector{Int}(Array(schur_con_scen)) + ns = schur_ns > 0 ? schur_ns : maximum(var_scen) + elseif schur_ns == 0 + tag_info = _schur_tags_from_callback(cb) + if tag_info !== nothing + ns, var_scen, con_scen = tag_info + end + end - # Reject design-only constraints — the Schur reduction has no slot - # for them. - con_hist[1] == 0 || error( - "$(con_hist[1]) constraints have con_scenario tag 0; " * - "design-only constraints are not supported by SchurComplementKKTSystem." - ) + if var_scen === nothing + # No tags/vectors available: synthesize the contiguous layout from the + # explicit per-stage dimensions. Reproduces the legacy index-arithmetic + # behaviour exactly (design variables last, scenario stripes contiguous). + @assert schur_ns > 0 "schur_ns must be specified and positive (or pass schur_var_scen/schur_con_scen, or use a two-stage tag for auto-detection)" + @assert schur_nv > 0 "schur_nv must be specified and positive" + @assert schur_nd > 0 "schur_nd must be specified and positive" + ns = schur_ns + var_scen = vcat((fill(k, schur_nv) for k in 1:ns)..., fill(0, schur_nd)) + con_scen = vcat((fill(k, schur_nc) for k in 1:ns)...) + end - # Per-scenario uniformity. Cheap loop, fires before any symbolic - # work and points at the offending scenario. - for k in 2:schur_ns - @inbounds nv_k = var_hist[k + 1] - @inbounds nc_k = con_hist[k + 1] - nv_k == schur_nv || error( - "Scenario $k has $nv_k variables; scenario 1 has $schur_nv. " * - "SchurComplementKKTSystem requires uniform per-scenario sizes." - ) - nc_k == schur_nc || error( - "Scenario $k has $nc_k constraints; scenario 1 has $schur_nc. " * - "SchurComplementKKTSystem requires uniform per-scenario sizes." - ) - end - end + @assert ns > 0 "resolved scenario count must be positive" + length(var_scen) == n || error("var_scen has length $(length(var_scen)); expected n=$n") + length(con_scen) == m || error("con_scen has length $(length(con_scen)); expected m=$m") + + # Single-pass histograms over the scenario tags. Index 1 is the design bucket + # (tag 0); 1+k is scenario k. Validates per-scenario uniformity in the same + # pass — a malformed model whose global aggregates happen to satisfy + # n == ns*nv + nd would otherwise drive the symbolic build into garbage. + var_hist = zeros(Int, ns + 1) + for t in var_scen + (0 <= t <= ns) || error("var_scen tag $t out of range [0, $ns]; 0 = design, 1..$ns = scenario index.") + @inbounds var_hist[t + 1] += 1 + end + con_hist = zeros(Int, ns + 1) + for t in con_scen + (0 <= t <= ns) || error("con_scen tag $t out of range [0, $ns]; 0 = design, 1..$ns = scenario index.") + @inbounds con_hist[t + 1] += 1 end - @assert schur_ns > 0 "schur_ns must be specified and positive (or use TwoStageTags for auto-detection)" - @assert schur_nv > 0 "schur_nv must be specified and positive" - @assert schur_nd > 0 "schur_nd must be specified and positive" - @assert n == schur_ns * schur_nv + schur_nd "Variable count mismatch: n=$n != ns*nv+nd=$(schur_ns*schur_nv+schur_nd)" - @assert m == schur_ns * schur_nc "Constraint count mismatch: m=$m != ns*nc=$(schur_ns*schur_nc)" + nd = var_hist[1] + nv = var_hist[2] + nc = con_hist[2] + nc_design = con_hist[1] - return schur_ns, schur_nv, schur_nd, schur_nc + for k in 2:ns + @inbounds nv_k = var_hist[k + 1] + @inbounds nc_k = con_hist[k + 1] + nv_k == nv || error( + "Scenario $k has $nv_k variables; scenario 1 has $nv. " * + "SchurComplementCondensedKKTSystem requires uniform per-scenario sizes." + ) + nc_k == nc || error( + "Scenario $k has $nc_k constraints; scenario 1 has $nc. " * + "SchurComplementCondensedKKTSystem requires uniform per-scenario sizes." + ) + end + + @assert nv > 0 "resolved per-scenario variable count nv must be positive" + @assert nd > 0 "resolved design variable count nd must be positive" + @assert n == ns * nv + nd "Variable count mismatch: n=$n != ns*nv+nd=$(ns * nv + nd)" + @assert m == ns * nc + nc_design "Constraint count mismatch: m=$m != ns*nc+nc_design=$(ns * nc + nc_design)" + + return (; ns, nv, nd, nc, nc_design, var_scen, con_scen) end # --- Index-driven scatter helpers used by build_kkt! --- @@ -266,22 +329,25 @@ end hess_I, hess_J, jac_I, jac_J, ind_eq, ind_ineq) -> NamedTuple -Pure-CPU symbolic construction shared by the CPU `SchurComplementKKTSystem` -and GPU `GPUSchurComplementKKTSystem` constructors. Inputs are CPU-resident +Pure-CPU symbolic construction shared by the CPU `SchurComplementCondensedKKTSystem` +and GPU `GPUSchurComplementCondensedKKTSystem` constructors. Inputs are CPU-resident sparsity arrays (the GPU side downloads its sparsity before calling). **Assumes uniform per-scenario structure** (same A_kk pattern for every k). This matches typical two-stage stochastic models and is what the GPU batched cuDSS path requires; CPU follows the same assumption. +RelaxEquality-only: `ind_eq` must be empty (all constraints are inequalities), +so the per-scenario blocks and the first-stage Schur complement are condensed/SPD. + Returns a NamedTuple with: -- `eq_per_scenario`, `ineq_per_scenario` — `Vector{Vector{Int}}` of global indices -- `nc_eq_per_s`, `nc_ineq_per_s`, `blk_size` — per-scenario sizes +- `ineq_per_scenario` — `Vector{Vector{Int}}` of global inequality indices +- `nc_ineq_per_s`, `blk_size` — per-scenario sizes (`blk_size == nv`) - `block_maps::Vector{ScenarioBlockMap}` — per-scenario index maps - `hess_S_coo`, `hess_S_row`, `hess_S_col` — design-design Hessian COO maps - `akk_csc_template::SparseMatrixCSC{T, Int32}` — shared A_kk sparsity (zero values) - `nnz_per_scenario::Int` -- `eq_global_flat::Vector{Int}` — flattened eq indices in scenario order (for GPU) +- the sparse-Schur pattern (`schur_csc_*`, `schur_*_nzpos`, `coupled_design_local`, …) """ function _build_schur_symbolic( ::Type{T}, @@ -290,50 +356,89 @@ function _build_schur_symbolic( jac_I::AbstractVector{<:Integer}, jac_J::AbstractVector{<:Integer}, ind_eq::AbstractVector{<:Integer}, ind_ineq::AbstractVector{<:Integer}, + var_scen = nothing, + con_scen = nothing, ) where {T} n_hess = length(hess_I) n_jac = length(jac_I) - # --- Classify constraints per scenario --- - ind_eq_set = Set(ind_eq) - ind_ineq_set = Set(ind_ineq) - - eq_per_scenario = Vector{Vector{Int}}(undef, ns) - ineq_per_scenario = Vector{Vector{Int}}(undef, ns) + # Synthesize the contiguous tag layout (`[v_1..v_ns, d]`, `[c_1..c_ns]`) when + # no tags are supplied. This reproduces the legacy index-arithmetic behaviour + # exactly, so hand-built two-stage models and the symbolic unit tests keep + # working unchanged. + if var_scen === nothing + var_scen = vcat((fill(k, nv) for k in 1:ns)..., fill(0, nd)) + end + if con_scen === nothing + con_scen = vcat((fill(k, nc) for k in 1:ns)...) + end + var_scen = Vector{Int}(Array(var_scen)) + con_scen = Vector{Int}(Array(con_scen)) + + # --- Tag-driven global index lists & local maps --- + # Partition variables/constraints by their scenario tag (0 = design) instead + # of by contiguous index arithmetic, so design variables and a scenario's + # variables need NOT be contiguous in the global ordering. `var_local[i]` is + # the position of variable `i` within its own block (1:nd for design, + # 1:nv for its scenario), assigned in increasing global-index order so it is + # monotone within a block — which keeps the lower-triangular A_kk layout valid. + design_var_global = Int[] # length nd + scen_var_global = [Int[] for _ in 1:ns] # each length nv + var_local = zeros(Int, n) + for i in 1:n + o = var_scen[i] + if o == 0 + push!(design_var_global, i) + var_local[i] = length(design_var_global) + else + push!(scen_var_global[o], i) + var_local[i] = length(scen_var_global[o]) + end + end - for k in 1:ns - cr = (k-1)*nc+1 : k*nc - eq_per_scenario[k] = Int[] - ineq_per_scenario[k] = Int[] - for gi in cr - if gi in ind_eq_set - push!(eq_per_scenario[k], gi) - end - if gi in ind_ineq_set - push!(ineq_per_scenario[k], gi) - end + # --- Classify constraints per scenario / design (tag-driven) --- + # RelaxEquality-only: every constraint is an inequality (ind_eq must be empty), + # so the per-scenario blocks and the first-stage Schur complement are condensed + # 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`." + ) + ind_ineq_set = Set(Int.(ind_ineq)) + + ineq_per_scenario = [Int[] for _ in 1:ns] + design_ineq = Int[] # design-only inequality constraint global rows + + for gi in 1:m + (gi in ind_ineq_set) || continue + o = con_scen[gi] + if o == 0 + push!(design_ineq, gi) + else + push!(ineq_per_scenario[o], gi) end end # Scenario 1 sets the canonical per-scenario constraint count; reject any # scenario that disagrees, since downstream code (and the GPU batched layout) # assumes uniform per-scenario shape. - nc_eq_per_s = length(eq_per_scenario[1]) nc_ineq_per_s = length(ineq_per_scenario[1]) for k in 2:ns - n_eq_k = length(eq_per_scenario[k]) n_in_k = length(ineq_per_scenario[k]) - if n_eq_k != nc_eq_per_s || n_in_k != nc_ineq_per_s + if n_in_k != nc_ineq_per_s error( - "SchurComplementKKTSystem requires uniform per-scenario constraint counts. " * - "Scenario 1 has (eq=$nc_eq_per_s, ineq=$nc_ineq_per_s); " * - "scenario $k has (eq=$n_eq_k, ineq=$n_in_k)." + "SchurComplementCondensedKKTSystem requires uniform per-scenario constraint counts. " * + "Scenario 1 has ineq=$nc_ineq_per_s; scenario $k has ineq=$n_in_k." ) end end - blk_size = nv + nc_eq_per_s + blk_size = nv # Lookup: global ineq index → diag_buffer index ineq_to_bufidx = Dict{Int,Int}() @@ -350,10 +455,7 @@ function _build_schur_symbolic( push!(entries, (ci, col)) end - d_start = ns * nv + 1 - d_end = ns * nv + nd - - # --- Classify Hessian COO entries --- + # --- Classify Hessian COO entries (tag-driven) --- hess_S_coo = Int[] hess_S_row = Int[] hess_S_col = Int[] @@ -364,43 +466,44 @@ function _build_schur_symbolic( for ci in 1:n_hess ri = Int(hess_I[ci]) - rj = Int(hess_J[ci]) # lower triangle: ri >= rj - - # Both design vars → S entry (write both triangles for dense S) - if ri >= d_start && ri <= d_end && rj >= d_start && rj <= d_end - di = ri - d_start + 1 - dj = rj - d_start + 1 + rj = Int(hess_J[ci]) # global lower triangle: ri >= rj + oi = var_scen[ri] + oj = var_scen[rj] + + if oi == 0 && oj == 0 + # Both design vars → S entry (write both triangles for dense S). + di = var_local[ri] + dj = var_local[rj] push!(hess_S_coo, ci); push!(hess_S_row, di); push!(hess_S_col, dj) if di != dj push!(hess_S_coo, ci); push!(hess_S_row, dj); push!(hess_S_col, di) end hess_classified[ci] = true - continue - end - - # One design + one scenario var → coupling block - if ri >= d_start && ri <= d_end && rj < d_start - di = ri - d_start + 1 - k = div(rj - 1, nv) + 1 - if k >= 1 && k <= ns - vj = rj - (k-1)*nv - push!(hess_per_scenario_coupling[k], (ci, di, vj)) - hess_classified[ci] = true + elseif (oi == 0) != (oj == 0) + # One design + one scenario var → coupling block. Either global index + # may be the design one now (design vars are not necessarily last). + if oi == 0 + di = var_local[ri]; k = oj; vj = var_local[rj] + else + di = var_local[rj]; k = oi; vj = var_local[ri] end - continue - end - - # Both within the same scenario → A_kk diagonal block - if ri < d_start && rj < d_start - ki = div(ri - 1, nv) + 1 - kj = div(rj - 1, nv) + 1 - if ki == kj && ki >= 1 && ki <= ns - li = ri - (ki-1)*nv - lj = rj - (ki-1)*nv - push!(hess_per_scenario_diag[ki], (ci, li, lj)) - hess_classified[ci] = true + push!(hess_per_scenario_coupling[k], (ci, di, vj)) + hess_classified[ci] = true + elseif oi == oj + # Both within the same scenario → A_kk diagonal block. `var_local` is + # monotone in the global index within a scenario, so li >= lj here; + # normalize defensively to keep the lower-triangular layout. + k = oi + li = var_local[ri] + lj = var_local[rj] + if li < lj + li, lj = lj, li end + push!(hess_per_scenario_diag[k], (ci, li, lj)) + hess_classified[ci] = true end + # else: both scenario but oi != oj → cross-scenario coupling, left + # unclassified and reported below. end # Anything left unclassified is a Hessian entry that doesn't fit the @@ -412,19 +515,14 @@ function _build_schur_symbolic( sample = first(bad, min(5, length(bad))) details = join(("(row=$(Int(hess_I[ci])), col=$(Int(hess_J[ci])))" for ci in sample), ", ") error( - "$n_bad_hess Hessian COO entries do not fit the SchurComplementKKTSystem " * + "$n_bad_hess Hessian COO entries do not fit the SchurComplementCondensedKKTSystem " * "block-arrowhead pattern (likely cross-scenario coupling). First few: " * details ) end # --- Build shared A_kk template from scenario 1 --- - eq_cons_1 = eq_per_scenario[1] ineq_cons_1 = ineq_per_scenario[1] - eq_local_1 = Dict{Int,Int}() - for (ci, gi) in enumerate(eq_cons_1) - eq_local_1[gi] = ci - end akk_entries = Dict{Tuple{Int,Int}, Nothing}() @@ -436,24 +534,12 @@ function _build_schur_symbolic( for i in 1:nv akk_entries[(i, i)] = nothing end - # Equality Jacobian: row = nv + eq_local, col = local_var - for gi in eq_cons_1 - for (_, col) in get(jac_by_constraint, gi, Tuple{Int,Int}[]) - if col >= 1 && col <= nv - akk_entries[(nv + eq_local_1[gi], col)] = nothing - end - end - end - # du_diag for eq constraints - for ci in 1:length(eq_cons_1) - akk_entries[(nv + ci, nv + ci)] = nothing - end # Inequality condensation fill-in (lower triangle pairs of scenario vars) for gi in ineq_cons_1 local_vars = Int[] for (_, col) in get(jac_by_constraint, gi, Tuple{Int,Int}[]) - if col >= 1 && col <= nv - push!(local_vars, col) + if var_scen[col] == 1 + push!(local_vars, var_local[col]) end end for a in local_vars, b in local_vars @@ -487,7 +573,7 @@ function _build_schur_symbolic( # scenario 1's template) surfaces as a meaningful error instead of KeyError. @inline akk_pos(key, k) = let p = get(akk_lookup, key, 0) p == 0 && error( - "SchurComplementKKTSystem: scenario $k has an A_kk entry at local " * + "SchurComplementCondensedKKTSystem: scenario $k has an A_kk entry at local " * "(row=$(key[1]), col=$(key[2])) absent from scenario 1's template. " * "Per-scenario Hessian/Jacobian sparsity must be uniform." ) @@ -498,19 +584,11 @@ function _build_schur_symbolic( # --- Build per-scenario ScenarioBlockMaps --- block_maps = Vector{ScenarioBlockMap}(undef, ns) - eq_global_flat = Int[] jac_classified = falses(n_jac) for k in 1:ns - vr_start = (k-1)*nv + 1 - eq_cons = eq_per_scenario[k] ineq_cons = ineq_per_scenario[k] - eq_local = Dict{Int,Int}() - for (ci, gi) in enumerate(eq_cons) - eq_local[gi] = ci - end - # Hessian diagonal → A_kk hess_Akk_coo_vec = Int[] hess_Akk_nzpos_vec = Int[] @@ -529,47 +607,14 @@ function _build_schur_symbolic( push!(hess_Cdk_col_vec, vj) end - # Equality Jacobian → A_kk and C_dk - jeq_Akk_coo_vec = Int[] - jeq_Akk_nzpos_vec = Int[] - jeq_Cdk_coo_vec = Int[] - jeq_Cdk_row_vec = Int[] - jeq_Cdk_col_vec = Int[] - - for gi in eq_cons - local_eq = eq_local[gi] - for (coo_idx, col) in get(jac_by_constraint, gi, Tuple{Int,Int}[]) - if col >= vr_start && col < vr_start + nv - local_var = col - vr_start + 1 - push!(jeq_Akk_coo_vec, coo_idx) - push!(jeq_Akk_nzpos_vec, akk_pos((nv + local_eq, local_var), k)) - jac_classified[coo_idx] = true - elseif col >= d_start && col <= d_end - di = col - d_start + 1 - push!(jeq_Cdk_coo_vec, coo_idx) - push!(jeq_Cdk_row_vec, di) - push!(jeq_Cdk_col_vec, nv + local_eq) - jac_classified[coo_idx] = true - end - end - end - # pr_diag → A_kk diagonal pr_diag_global_vec = Int[] pr_diag_nzpos_vec = Int[] for i in 1:nv - push!(pr_diag_global_vec, vr_start + i - 1) + push!(pr_diag_global_vec, scen_var_global[k][i]) push!(pr_diag_nzpos_vec, akk_pos((i, i), k)) end - # du_diag → A_kk diagonal - du_diag_global_vec = Int[] - du_diag_nzpos_vec = Int[] - for (ci, gi) in enumerate(eq_cons) - push!(du_diag_global_vec, gi) - push!(du_diag_nzpos_vec, akk_pos((nv + ci, nv + ci), k)) - end - # Inequality condensation ineq_Akk_nzpos_vec = Int[] ineq_Akk_jcoo1_vec = Int[] @@ -594,13 +639,15 @@ function _build_schur_symbolic( d_entries = Tuple{Int,Int}[] for (coo_idx, col) in get(jac_by_constraint, gi, Tuple{Int,Int}[]) - if col >= vr_start && col < vr_start + nv - push!(v_entries, (coo_idx, col - vr_start + 1)) + oc = var_scen[col] + if oc == k + push!(v_entries, (coo_idx, var_local[col])) jac_classified[coo_idx] = true - elseif col >= d_start && col <= d_end - push!(d_entries, (coo_idx, col - d_start + 1)) + elseif oc == 0 + push!(d_entries, (coo_idx, var_local[col])) jac_classified[coo_idx] = true end + # else: cross-scenario column, left unclassified and reported below. end # A_kk: lower-triangle pairs of scenario vars @@ -635,16 +682,42 @@ function _build_schur_symbolic( block_maps[k] = ScenarioBlockMap( hess_Akk_coo_vec, hess_Akk_nzpos_vec, hess_Cdk_coo_vec, hess_Cdk_row_vec, hess_Cdk_col_vec, - jeq_Akk_coo_vec, jeq_Akk_nzpos_vec, - jeq_Cdk_coo_vec, jeq_Cdk_row_vec, jeq_Cdk_col_vec, ineq_Akk_nzpos_vec, ineq_Akk_jcoo1_vec, ineq_Akk_jcoo2_vec, ineq_Akk_bufidx_vec, ineq_Cdk_row_vec, ineq_Cdk_col_vec, ineq_Cdk_jcoo_d_vec, ineq_Cdk_jcoo_v_vec, ineq_Cdk_bufidx_vec, ineq_S_row_vec, ineq_S_col_vec, ineq_S_jcoo1_vec, ineq_S_jcoo2_vec, ineq_S_bufidx_vec, pr_diag_global_vec, pr_diag_nzpos_vec, - du_diag_global_vec, du_diag_nzpos_vec, ) + end - append!(eq_global_flat, eq_cons) + # --- Design-only constraint maps --- + # Design-only constraints touch ONLY design variables; they are condensed into + # S_dd just like scenario inequalities. Referencing a scenario variable is + # unrepresentable. + nc_design_ineq = length(design_ineq) + + # Design inequality condensation → S_dd (design × design, full — both triangles). + design_ineq_S_row = Int[] + design_ineq_S_col = Int[] + design_ineq_S_jcoo1 = Int[] + design_ineq_S_jcoo2 = Int[] + design_ineq_S_bufidx = Int[] + for gi in design_ineq + bidx = ineq_to_bufidx[gi] + d_entries = Tuple{Int, Int}[] + for (coo_idx, col) in get(jac_by_constraint, gi, Tuple{Int, Int}[]) + var_scen[col] == 0 || error( + "Design-only inequality (row $gi) references scenario variable (col $col)." + ) + push!(d_entries, (coo_idx, var_local[col])) + jac_classified[coo_idx] = true + end + for (coo_a, da) in d_entries, (coo_b, db) in d_entries + push!(design_ineq_S_row, da) + push!(design_ineq_S_col, db) + push!(design_ineq_S_jcoo1, coo_a) + push!(design_ineq_S_jcoo2, coo_b) + push!(design_ineq_S_bufidx, bidx) + end end # Catch Jacobian entries whose column doesn't match the constraint's own @@ -655,7 +728,7 @@ function _build_schur_symbolic( sample = first(bad, min(5, length(bad))) details = join(("(row=$(Int(jac_I[ci])), col=$(Int(jac_J[ci])))" for ci in sample), ", ") error( - "$n_bad_jac Jacobian COO entries do not fit the SchurComplementKKTSystem " * + "$n_bad_jac Jacobian COO entries do not fit the SchurComplementCondensedKKTSystem " * "block-arrowhead pattern (column belongs to a different scenario). First few: " * details ) @@ -668,25 +741,140 @@ function _build_schur_symbolic( bm1 = block_maps[1] fields_to_check = ( :hess_Akk_coo, :hess_Cdk_coo, - :jeq_Akk_coo, :jeq_Cdk_coo, :ineq_Akk_nzpos, :ineq_Cdk_row, :ineq_S_row, - :pr_diag_global, :du_diag_global, + :pr_diag_global, ) for k in 2:ns, f in fields_to_check n1 = length(getfield(bm1, f)) nk = length(getfield(block_maps[k], f)) if nk != n1 error( - "SchurComplementKKTSystem requires uniform per-scenario sparsity. " * + "SchurComplementCondensedKKTSystem requires uniform per-scenario sparsity. " * "Scenario 1 has $n1 entries in $f; scenario $k has $nk." ) end end + # ===== Static sparse-Schur pattern (shared CPU/GPU) ============================= + # The first-stage Schur complement `schur_csc` (size nd × nd) is sparse: the + # reduction Σ_k C_dk A_kk⁻¹ C_dk' only fills the coupled-design × coupled-design + # block (design vars that actually couple to a scenario). Compute, once: the set of + # coupled design vars (uniform across scenarios), the lower-triangular sparsity + # pattern, and per-contribution nzval-position maps used to assemble and factorize + # `schur_csc` (the GPU additionally uploads them to build a CuSparseMatrixCSC). + + # (1) Coupled design vars per scenario (design-local indices appearing in C_dk). + coupled_per_s = [Set{Int}() for _ in 1:ns] + for k in 1:ns + for (_, di, _) in hess_per_scenario_coupling[k] + push!(coupled_per_s[k], di) + end + for di in block_maps[k].ineq_Cdk_row + push!(coupled_per_s[k], di) + end + end + for k in 2:ns + coupled_per_s[k] == coupled_per_s[1] || error( + "SchurComplementCondensedKKTSystem (sparse): scenario $k couples a different set of design " * + "variables than scenario 1; the sparse Schur path requires uniform coupling." + ) + end + coupled_design_local = sort!(collect(coupled_per_s[1])) # length m + m_coupled = length(coupled_design_local) + coupled_inv = zeros(Int, nd) # design-local → compact col (1:m) or 0 + for (c, di) in enumerate(coupled_design_local) + coupled_inv[di] = c + end + + # (2) Lower-triangular sparsity pattern: union of every contribution, folded to (max,min). + _lo(r, c) = r >= c ? (r, c) : (c, r) + schur_set = Dict{Tuple{Int, Int}, Nothing}() + for t in eachindex(hess_S_row) # design Hessian + schur_set[_lo(hess_S_row[t], hess_S_col[t])] = nothing + end + for i in 1:nd # design pr_diag diagonal + schur_set[(i, i)] = nothing + end + for k in 1:ns # scenario ineq → S (positions uniform) + bm = block_maps[k] + for t in eachindex(bm.ineq_S_row) + schur_set[_lo(bm.ineq_S_row[t], bm.ineq_S_col[t])] = nothing + end + end + for t in eachindex(design_ineq_S_row) # design ineq → S + schur_set[_lo(design_ineq_S_row[t], design_ineq_S_col[t])] = nothing + end + for a in coupled_design_local, b in coupled_design_local # Schur fill block + schur_set[_lo(a, b)] = nothing + end + + schur_nnz = length(schur_set) + schur_I = Vector{Int32}(undef, schur_nnz) + schur_J = Vector{Int32}(undef, schur_nnz) + let t = 0 + for ((r, c), _) in schur_set + @assert r >= c "sparse Schur pattern entry ($r,$c) is not lower-triangular" + t += 1 + schur_I[t] = Int32(r) + schur_J[t] = Int32(c) + end + end + schur_csc_template, schur_coo_map = coo_to_csc( + SparseMatrixCOO(nd, nd, schur_I, schur_J, zeros(T, schur_nnz)) + ) + # (row,col) → nzval position, from the canonical (column-major) CSC ordering. + schur_pos = Dict{Tuple{Int, Int}, Int}() + for t in 1:schur_nnz + schur_pos[(Int(schur_I[t]), Int(schur_J[t]))] = Int(schur_coo_map[t]) + end + _nzpos(r, c) = schur_pos[_lo(r, c)] + + # (3) Per-contribution nzpos maps, lower-triangle only (each lower slot hit once; + # cuDSS reads the lower triangle and symmetrizes). + schur_hess_coo = Int[] # design Hessian + schur_hess_nzpos = Int[] + for t in eachindex(hess_S_row) + hess_S_row[t] >= hess_S_col[t] || continue # keep lower only + push!(schur_hess_coo, hess_S_coo[t]) + push!(schur_hess_nzpos, _nzpos(hess_S_row[t], hess_S_col[t])) + end + schur_diag_nzpos = Int[_nzpos(i, i) for i in 1:nd] # pr_diag (with design_var_global) + + schur_ineq_S_nzpos = Int[] # scenario ineq → S (flat, lower only) + schur_ineq_S_jcoo1 = Int[] + schur_ineq_S_jcoo2 = Int[] + schur_ineq_S_bufidx = Int[] + for k in 1:ns + bm = block_maps[k] + for t in eachindex(bm.ineq_S_row) + bm.ineq_S_row[t] >= bm.ineq_S_col[t] || continue + push!(schur_ineq_S_nzpos, _nzpos(bm.ineq_S_row[t], bm.ineq_S_col[t])) + push!(schur_ineq_S_jcoo1, bm.ineq_S_jcoo1[t]) + push!(schur_ineq_S_jcoo2, bm.ineq_S_jcoo2[t]) + push!(schur_ineq_S_bufidx, bm.ineq_S_bufidx[t]) + end + end + + schur_design_ineq_S_nzpos = Int[] # design ineq → S (lower only) + schur_design_ineq_S_jcoo1 = Int[] + schur_design_ineq_S_jcoo2 = Int[] + schur_design_ineq_S_bufidx = Int[] + for t in eachindex(design_ineq_S_row) + design_ineq_S_row[t] >= design_ineq_S_col[t] || continue + push!(schur_design_ineq_S_nzpos, _nzpos(design_ineq_S_row[t], design_ineq_S_col[t])) + push!(schur_design_ineq_S_jcoo1, design_ineq_S_jcoo1[t]) + push!(schur_design_ineq_S_jcoo2, design_ineq_S_jcoo2[t]) + push!(schur_design_ineq_S_bufidx, design_ineq_S_bufidx[t]) + end + + # (4) Schur-fill nzpos (m×m); column-major flat, lower-or-diagonal entries valid. + schur_fill_nzpos = zeros(Int, m_coupled, m_coupled) + for a in 1:m_coupled, b in 1:m_coupled + schur_fill_nzpos[a, b] = _nzpos(coupled_design_local[a], coupled_design_local[b]) + end + return ( - eq_per_scenario = eq_per_scenario, ineq_per_scenario = ineq_per_scenario, - nc_eq_per_s = nc_eq_per_s, nc_ineq_per_s = nc_ineq_per_s, blk_size = blk_size, block_maps = block_maps, @@ -695,7 +883,37 @@ function _build_schur_symbolic( hess_S_col = hess_S_col, akk_csc_template = akk_csc_template, nnz_per_scenario = nnz_per_scenario, - eq_global_flat = eq_global_flat, + # Tag-driven global index lists (replace the old contiguous arithmetic). + var_scen = var_scen, + con_scen = con_scen, + design_var_global = design_var_global, + scen_var_global = scen_var_global, + # Design-only constraint block. + nc_design_ineq = nc_design_ineq, + design_ineq_S_row = design_ineq_S_row, + design_ineq_S_col = design_ineq_S_col, + design_ineq_S_jcoo1 = design_ineq_S_jcoo1, + design_ineq_S_jcoo2 = design_ineq_S_jcoo2, + design_ineq_S_bufidx = design_ineq_S_bufidx, + # Sparse-Schur (GPU cuDSS) pattern + nzpos maps. + m_coupled = m_coupled, + coupled_design_local = coupled_design_local, + coupled_inv = coupled_inv, + schur_csc_colptr = schur_csc_template.colptr, + schur_csc_rowval = schur_csc_template.rowval, + schur_nnz = length(schur_csc_template.nzval), + schur_hess_coo = schur_hess_coo, + schur_hess_nzpos = schur_hess_nzpos, + schur_diag_nzpos = schur_diag_nzpos, + schur_ineq_S_nzpos = schur_ineq_S_nzpos, + schur_ineq_S_jcoo1 = schur_ineq_S_jcoo1, + schur_ineq_S_jcoo2 = schur_ineq_S_jcoo2, + schur_ineq_S_bufidx = schur_ineq_S_bufidx, + schur_design_ineq_S_nzpos = schur_design_ineq_S_nzpos, + schur_design_ineq_S_jcoo1 = schur_design_ineq_S_jcoo1, + schur_design_ineq_S_jcoo2 = schur_design_ineq_S_jcoo2, + schur_design_ineq_S_bufidx = schur_design_ineq_S_bufidx, + schur_fill_nzpos = schur_fill_nzpos, ) end @@ -712,9 +930,6 @@ function _flatten_block_maps(block_maps::Vector{ScenarioBlockMap}) 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_du_diag = length(bm1.du_diag_global) - n_per_s_jeq_Akk = length(bm1.jeq_Akk_coo) - n_per_s_jeq_Cdk = length(bm1.jeq_Cdk_coo) 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) @@ -735,19 +950,6 @@ function _flatten_block_maps(block_maps::Vector{ScenarioBlockMap}) all_pr_diag_global = cat_int(bm -> bm.pr_diag_global), all_pr_diag_nzpos = cat_int(bm -> bm.pr_diag_nzpos), - n_per_s_du_diag = n_per_s_du_diag, - all_du_diag_global = cat_int(bm -> bm.du_diag_global), - all_du_diag_nzpos = cat_int(bm -> bm.du_diag_nzpos), - - n_per_s_jeq_Akk = n_per_s_jeq_Akk, - all_jeq_Akk_coo = cat_int(bm -> bm.jeq_Akk_coo), - all_jeq_Akk_nzpos = cat_int(bm -> bm.jeq_Akk_nzpos), - - n_per_s_jeq_Cdk = n_per_s_jeq_Cdk, - all_jeq_Cdk_coo = cat_int(bm -> bm.jeq_Cdk_coo), - all_jeq_Cdk_row = cat_int(bm -> bm.jeq_Cdk_row), - all_jeq_Cdk_col = cat_int(bm -> bm.jeq_Cdk_col), - n_per_s_ineq_Akk = n_per_s_ineq_Akk, all_ineq_Akk_nzpos = cat_int(bm -> bm.ineq_Akk_nzpos), all_ineq_Akk_jcoo1 = cat_int(bm -> bm.ineq_Akk_jcoo1), @@ -771,7 +973,7 @@ function _flatten_block_maps(block_maps::Vector{ScenarioBlockMap}) end function create_kkt_system( - ::Type{SchurComplementKKTSystem}, + ::Type{SchurComplementCondensedKKTSystem}, cb::SparseCallback{T,VT}, linear_solver::Type; opt_linear_solver=default_options(linear_solver), @@ -781,9 +983,24 @@ function create_kkt_system( schur_nv::Int=0, schur_nd::Int=0, schur_nc::Int=0, + 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..., ) where {T, VT} + isempty(kwargs) || Base.@warn( + "SchurComplementCondensedKKTSystem (CPU) ignores unsupported kkt_options: " * + join(string.(keys(kwargs)), ", ") + ) + schur_opt_linear_solver === nothing || Base.@warn( + "SchurComplementCondensedKKTSystem (CPU): `schur_opt_linear_solver` is a GPU-only option " * + "(first-stage cuDSS options) and is ignored on CPU; the first-stage Schur complement " * + "is solved by `linear_solver` with `opt_linear_solver`." + ) + n = cb.nvar m = cb.ncon ns_ineq = length(cb.ind_ineq) @@ -791,7 +1008,8 @@ function create_kkt_system( nlb = length(cb.ind_lb) nub = length(cb.ind_ub) - ns, nv, nd, nc = _resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc) + dims = _resolve_schur_dims(cb, n, m, schur_ns, schur_nv, schur_nd, schur_nc, schur_var_scen, schur_con_scen) + ns, nv, nd, nc = dims.ns, dims.nv, dims.nd, dims.nc # --- Get sparsity patterns --- jac_sparsity_I = Vector{Int32}(undef, cb.nnzj) @@ -822,20 +1040,45 @@ function create_kkt_system( T, n, m, ns, nv, nd, nc, hess_sparsity_I, hess_sparsity_J, jac_sparsity_I, jac_sparsity_J, - cb.ind_eq, cb.ind_ineq, + Array(cb.ind_eq), Array(cb.ind_ineq), + dims.var_scen, dims.con_scen, ) - nc_eq_per_s = sym.nc_eq_per_s nc_ineq_per_s = sym.nc_ineq_per_s blk_size = sym.blk_size - block_maps = sym.block_maps + m_coupled = sym.m_coupled + coupled_inv = sym.coupled_inv # design-local → compact coupled col (1:m) or 0 + + # Reduce the per-scenario coupling maps to the m COUPLED design rows: remap each C_dk row + # (a design-local index in 1:nd) to its compact column in 1:m. Every C_dk row is coupled by + # construction, so `coupled_inv` is nonzero there. Mirrors the GPU reduction — the coupling + # blocks and their solves are width m, not nd, which is what makes SCOPF-shaped problems + # (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, + ) + end for k in 1:ns + ] # Per-scenario A_kk: independent copies of the shared template (same sparsity, fresh nzval). A_kk_vec = [copy(sym.akk_csc_template) for _ in 1:ns] - # --- Dense matrices --- - aug_com = Matrix{T}(undef, nd, nd) - C_dk = [Matrix{T}(undef, nd, blk_size) for _ in 1:ns] - tmp_blk_nd = [Matrix{T}(undef, blk_size, nd) for _ in 1:ns] + # --- Sparse Schur complement + reduced dense coupling blocks --- + # schur_csc is the first-stage block (lower-triangular, size nd). The reduction + # `Σ_k C_dk A_kk⁻¹ C_dk'` fills only the m×m coupled block, so C_dk is stored reduced to its + # m coupled design rows (m × blk_size) and tmp_blk_nd = A_kk⁻¹ C_dk_red' is (blk_size × m). + schur_csc = SparseMatrixCSC{T, Int32}( + nd, nd, sym.schur_csc_colptr, sym.schur_csc_rowval, zeros(T, sym.schur_nnz), + ) + schur_block = Matrix{T}(undef, m_coupled, m_coupled) + C_dk = [Matrix{T}(undef, m_coupled, blk_size) for _ in 1:ns] + tmp_blk_nd = [Matrix{T}(undef, blk_size, m_coupled) for _ in 1:ns] # --- Diagonal vectors --- reg = VT(undef, n + ns_ineq) @@ -849,62 +1092,115 @@ function create_kkt_system( # --- Buffers --- diag_buffer = VT(undef, ns_ineq) buffer = VT(undef, m) - # Size from the per-scenario invariant rather than the global eq count: the - # uniform-scenario validation in `_build_schur_symbolic` guarantees these - # are equal, but tying the buffer to `ns * nc_eq_per_s` keeps the - # round-trip in `solve_kkt!` (gather/scatter via the eq index list) shape- - # consistent by construction. - wy_eq_buf = VT(undef, ns * nc_eq_per_s) - rhs_d = VT(undef, nd) + rhs_d = VT(undef, nd) + rhs_d_red = VT(undef, m_coupled) rhs_k = [VT(undef, blk_size) for _ in 1:ns] solve_buffers = [VT(undef, blk_size) for _ in 1:ns] # --- Init --- - fill!(aug_com, zero(T)) fill!(pr_diag, zero(T)) fill!(du_diag, zero(T)) # --- Create solvers --- quasi_newton = create_quasi_newton(hessian_approximation, cb, n; options=qn_options) - scenario_solvers = [schur_scenario_linear_solver(A_kk_vec[k]) for k in 1:ns] - _linear_solver = linear_solver(aug_com; opt = opt_linear_solver) + scenario_solvers = [schur_scenario_linear_solver(A_kk_vec[k]; opt = schur_scenario_opt_linear_solver) for k in 1:ns] + _linear_solver = linear_solver(schur_csc; opt = opt_linear_solver) - return SchurComplementKKTSystem( + return SchurComplementCondensedKKTSystem( hess, jac, hess_raw, jt_coo, hess_csc, hess_csc_map, jt_csc, jt_csc_map, quasi_newton, reg, pr_diag, du_diag, l_diag, u_diag, l_lower, u_lower, ns, nv, nd, nc, - nc_eq_per_s, nc_ineq_per_s, blk_size, + nc_ineq_per_s, blk_size, A_kk_vec, C_dk, - aug_com, - diag_buffer, buffer, wy_eq_buf, rhs_d, rhs_k, tmp_blk_nd, solve_buffers, + schur_csc, m_coupled, sym.coupled_design_local, sym.schur_fill_nzpos, schur_block, + diag_buffer, buffer, rhs_d, rhs_d_red, rhs_k, tmp_blk_nd, solve_buffers, block_maps, - sym.hess_S_coo, sym.hess_S_row, sym.hess_S_col, - sym.eq_global_flat, + sym.schur_hess_coo, sym.schur_hess_nzpos, + sym.schur_diag_nzpos, + sym.schur_ineq_S_nzpos, sym.schur_ineq_S_jcoo1, sym.schur_ineq_S_jcoo2, sym.schur_ineq_S_bufidx, + sym.schur_design_ineq_S_nzpos, sym.schur_design_ineq_S_jcoo1, sym.schur_design_ineq_S_jcoo2, sym.schur_design_ineq_S_bufidx, + sym.design_var_global, sym.scen_var_global, n_eq, cb.ind_eq, ns_ineq, cb.ind_ineq, cb.ind_lb, cb.ind_ub, + Ref(0), Ref(0), Ref(0), scenario_solvers, _linear_solver, ) end -num_variables(kkt::SchurComplementKKTSystem) = size(kkt.hess_csc, 1) +num_variables(kkt::SchurComplementCondensedKKTSystem) = size(kkt.hess_csc, 1) -function get_slack_regularization(kkt::SchurComplementKKTSystem) +function get_slack_regularization(kkt::SchurComplementCondensedKKTSystem) n = num_variables(kkt) ns_ineq = kkt.n_ineq return view(kkt.pr_diag, n+1:n+ns_ineq) end -function is_inertia_correct(kkt::SchurComplementKKTSystem, num_pos, num_zero, num_neg) - return (num_zero == 0) && (num_pos == size(kkt.aug_com, 1)) +""" + is_inertia(kkt::SchurComplementCondensedKKTSystem) + +The Schur path can only certify In(K) if *every* factorization contributing to the +Haynsworth sum reports inertia — the complement solver and all `ns` block solvers. +""" +is_inertia(kkt::SchurComplementCondensedKKTSystem) = + is_inertia(kkt.linear_solver) && all(is_inertia, kkt.scenario_solvers) + +function is_inertia_correct(kkt::SchurComplementCondensedKKTSystem, num_pos, num_zero, num_neg) + # `(num_pos, num_zero, num_neg)` is In(S) alone — the IPM reads it off + # `kkt.linear_solver`, which factorizes only the nd×nd first-stage complement. By + # Haynsworth additivity In(K) = Σ_k In(A_k) + In(S), so In(S) == (nd, 0, 0) certifies + # the whole condensed KKT matrix ONLY if every block A_k is positive definite too. + # + # This is not a conservative omission: an indefinite A_k makes its contribution + # −C_k A_k⁻¹ C_kᵀ to S a *positive* shift, which can leave S positive definite. S then + # passes, δ_w is never raised, and the IPM accepts a direction with wrong curvature. + # The counts below are refreshed by `build_kkt!`, which factorizes the blocks; the IPM + # calls `factorize_wrapper!` (→ `build_kkt!`) before every inertia query, so they are + # always in sync with the In(S) passed in. + is_inertia(kkt) || throw(InertiaException()) + return (num_zero == 0) && (num_pos == kkt.nd) && (num_neg == 0) && + (kkt.block_num_neg[] == 0) && (kkt.block_num_zero[] == 0) +end + +""" + _accumulate_block_inertia!(kkt::SchurComplementCondensedKKTSystem) + +Sum the per-scenario block inertias into `kkt.block_num_neg` / `block_num_zero` / +`block_num_indef`. Called by `build_kkt!` right after the blocks are factorized; the +counts are a by-product of that factorization (MUMPS stores them in INFOG(12)), so this +only reads them back. Solvers that do not report inertia are skipped — `is_inertia(kkt)` +is then false and `is_inertia_correct` refuses to certify anything. + +Set `MADNLP_SCHUR_BLOCK_INERTIA` in the environment to log the counts per factorization. +""" +function _accumulate_block_inertia!(kkt::SchurComplementCondensedKKTSystem) + num_neg = 0 + num_zero = 0 + num_indef = 0 + worst = 0 + for s in kkt.scenario_solvers + is_inertia(s) || continue + (_, z, neg) = inertia(s) + num_neg += neg + num_zero += z + (neg > 0 || z > 0) && (num_indef += 1) + worst = max(worst, neg) + end + kkt.block_num_neg[] = num_neg + kkt.block_num_zero[] = num_zero + kkt.block_num_indef[] = num_indef + if haskey(ENV, "MADNLP_SCHUR_BLOCK_INERTIA") + Base.@info "schur block inertia" blocks_indefinite = num_indef total_neg_pivots = num_neg total_zero_pivots = num_zero max_neg_pivots = worst ns = kkt.ns + end + return end -should_regularize_dual(kkt::SchurComplementKKTSystem, num_pos, num_zero, num_neg) = true +should_regularize_dual(kkt::SchurComplementCondensedKKTSystem, num_pos, num_zero, num_neg) = true -function jtprod!(y::AbstractVector, kkt::SchurComplementKKTSystem, x::AbstractVector) +function jtprod!(y::AbstractVector, kkt::SchurComplementCondensedKKTSystem, x::AbstractVector) nx = num_variables(kkt) ns_ineq = kkt.n_ineq yx = view(y, 1:nx) @@ -914,20 +1210,21 @@ function jtprod!(y::AbstractVector, kkt::SchurComplementKKTSystem, x::AbstractVe return end -function compress_jacobian!(kkt::SchurComplementKKTSystem) +function compress_jacobian!(kkt::SchurComplementCondensedKKTSystem) transfer!(kkt.jt_csc, kkt.jt_coo, kkt.jt_csc_map) end -function compress_hessian!(kkt::SchurComplementKKTSystem) +function compress_hessian!(kkt::SchurComplementCondensedKKTSystem) transfer!(kkt.hess_csc, kkt.hess_raw, kkt.hess_csc_map) end -nnz_jacobian(kkt::SchurComplementKKTSystem) = nnz(kkt.jt_coo) +nnz_jacobian(kkt::SchurComplementCondensedKKTSystem) = nnz(kkt.jt_coo) -function build_kkt!(kkt::SchurComplementKKTSystem{T, VT, MT}) where {T, VT, MT} +function build_kkt!(kkt::SchurComplementCondensedKKTSystem{T, VT, MT}) where {T, VT, MT} ns = kkt.ns nv = kkt.nv nd = kkt.nd + m = kkt.m n = num_variables(kkt) blk = kkt.blk_size @@ -938,34 +1235,33 @@ function build_kkt!(kkt::SchurComplementKKTSystem{T, VT, MT}) where {T, VT, MT} kkt.diag_buffer .= Sigma_s ./ (one(T) .- Sigma_d .* Sigma_s) end - # Initialize Schur complement S = H_dd + diag(pr_diag_dd) - S = kkt.aug_com - fill!(S, zero(T)) - _scatter_add!(S, kkt.hess, kkt.hess_S_coo, kkt.hess_S_row, kkt.hess_S_col) - @inbounds for i in 1:nd - S[i, i] += kkt.pr_diag[ns*nv+i] - end + # Initialize the sparse first-stage bordered block: scatter the design Hessian and + # the design pr_diag diagonal into the lower-triangular CSC `nz` by precomputed + # position (design variables need not be the last nd global indices). + nz = kkt.schur_csc.nzval + fill!(nz, zero(T)) + _scatter_add!(nz, kkt.hess, kkt.schur_hess_coo, kkt.schur_hess_nzpos) + _scatter_add!(nz, kkt.pr_diag, kkt.design_var_global, kkt.schur_diag_nzpos) # Phase 1 (parallel): assemble per-scenario blocks, factorize, compute A_kk^{-1} * C_dk'. # `@blas_safe_threads` runs `Threads.@threads` over scenarios while pinning # BLAS to a single thread per task to avoid oversubscription with the - # per-scenario `mul!` / `factorize!` calls inside the loop. + # per-scenario `mul!` / `factorize!` calls inside the loop. Only A_kk[k]/C_dk[k]/ + # tmp[k] are written here — the shared sparse `nz` is touched sequentially below. @blas_safe_threads for k in 1:ns bm = kkt.block_maps[k] A_kk = kkt.A_kk[k] C_dk = kkt.C_dk[k] - nz = A_kk.nzval + aknz = A_kk.nzval - fill!(nz, zero(T)) + fill!(aknz, zero(T)) fill!(C_dk, zero(T)) - _scatter_add!(nz, kkt.hess, bm.hess_Akk_coo, bm.hess_Akk_nzpos) + _scatter_add!(aknz, kkt.hess, bm.hess_Akk_coo, bm.hess_Akk_nzpos) _scatter_add!(C_dk, kkt.hess, bm.hess_Cdk_coo, bm.hess_Cdk_row, bm.hess_Cdk_col) - _scatter_add!(nz, kkt.pr_diag, bm.pr_diag_global, bm.pr_diag_nzpos) - _scatter_add!(nz, kkt.du_diag, bm.du_diag_global, bm.du_diag_nzpos) - _scatter_add!(nz, kkt.jac, bm.jeq_Akk_coo, bm.jeq_Akk_nzpos) - _scatter_add!(C_dk, kkt.jac, bm.jeq_Cdk_coo, bm.jeq_Cdk_row, bm.jeq_Cdk_col) - _scatter_quad_add!(nz, kkt.jac, kkt.diag_buffer, + _scatter_add!(aknz, kkt.pr_diag, bm.pr_diag_global, bm.pr_diag_nzpos) + _scatter_quad_add!( + aknz, kkt.jac, kkt.diag_buffer, bm.ineq_Akk_nzpos, bm.ineq_Akk_jcoo1, bm.ineq_Akk_jcoo2, bm.ineq_Akk_bufidx) _scatter_quad_add!(C_dk, kkt.jac, kkt.diag_buffer, bm.ineq_Cdk_row, bm.ineq_Cdk_col, @@ -974,38 +1270,63 @@ function build_kkt!(kkt::SchurComplementKKTSystem{T, VT, MT}) where {T, VT, MT} # Factor A_kk factorize!(kkt.scenario_solvers[k]) - # Compute tmp = A_kk^{-1} * C_dk' (blk × nd) + # Compute tmp_red = A_kk^{-1} * C_dk_red' (blk × m): only the m coupled columns — + # the other nd-m design columns of C_dk are exactly zero, so we never solve them. buf = kkt.solve_buffers[k] - for j in 1:nd + tmp = kkt.tmp_blk_nd[k] + for j in 1:m @inbounds for i in 1:blk - buf[i] = C_dk[j, i] # C_dk' column j + buf[i] = C_dk[j, i] # C_dk_red' column j = row j of the reduced C_dk end solve_linear_system!(kkt.scenario_solvers[k], buf) @inbounds for i in 1:blk - kkt.tmp_blk_nd[k][i, j] = buf[i] + tmp[i, j] = buf[i] end end end - # Phase 2 (sequential): accumulate into shared Schur complement S - for k in 1:ns - bm = kkt.block_maps[k] - _scatter_quad_add!(S, kkt.jac, kkt.diag_buffer, - bm.ineq_S_row, bm.ineq_S_col, - bm.ineq_S_jcoo1, bm.ineq_S_jcoo2, bm.ineq_S_bufidx) - # S -= C_dk * A_kk^{-1} * C_dk' - mul!(S, kkt.C_dk[k], kkt.tmp_blk_nd[k], -one(T), one(T)) + # Aggregate the block inertias for `is_inertia_correct`. Each solver computed this as a + # by-product of the `factorize!` above (MUMPS: INFOG(12)), so this is a read, not work. + _accumulate_block_inertia!(kkt) + + # Phase 2 (sequential): scatter the scenario inequality condensation into `nz` + # (flattened over scenarios), then the Schur reduction. The reduction + # `Σ_k C_dk A_kk⁻¹ C_dk'` is nonzero only on the m coupled design vars, so it is + # formed from the coupled rows/cols of the dense C_dk/tmp as an `m×m` symmetric + # block and the lower half scattered as `-D` into `nz`. + _scatter_quad_add!( + nz, kkt.jac, kkt.diag_buffer, + kkt.schur_ineq_S_nzpos, kkt.schur_ineq_S_jcoo1, + kkt.schur_ineq_S_jcoo2, kkt.schur_ineq_S_bufidx + ) + if m > 0 + for k in 1:ns + # C_dk_red (m×blk) and tmp_red (blk×m) are already reduced to the coupled block, so + # this is a plain contiguous GEMM (BLAS) — no Vector{Int}-indexed views, unlike the + # old full-width nd path which fell back to the generic (non-BLAS) mul!. + mul!(kkt.schur_block, kkt.C_dk[k], kkt.tmp_blk_nd[k]) + @inbounds for b in 1:m, a in b:m + nz[kkt.schur_fill_nzpos[a, b]] -= kkt.schur_block[a, b] + end + end end + # Design-only inequalities: condense into S_dd at their precomputed + # (lower-triangle) positions. + _scatter_quad_add!( + nz, kkt.jac, kkt.diag_buffer, + kkt.schur_design_ineq_S_nzpos, kkt.schur_design_ineq_S_jcoo1, + kkt.schur_design_ineq_S_jcoo2, kkt.schur_design_ineq_S_bufidx + ) return end -function factorize_kkt!(kkt::SchurComplementKKTSystem) +function factorize_kkt!(kkt::SchurComplementCondensedKKTSystem) return factorize!(kkt.linear_solver) end function solve_kkt!( - kkt::SchurComplementKKTSystem, + kkt::SchurComplementCondensedKKTSystem, w::AbstractKKTVector{T}, ) where T @@ -1013,6 +1334,7 @@ function solve_kkt!( nv = kkt.nv nd = kkt.nd nc = kkt.nc + m = kkt.m n = num_variables(kkt) blk = kkt.blk_size @@ -1032,23 +1354,16 @@ function solve_kkt!( mul!(wx, kkt.jt_csc, kkt.buffer, one(T), one(T)) end - # Step 2: Extract per-scenario RHS blocks - # NOTE: writes to disjoint slices of wx/wy → safe to @blas_safe_threads, - # but per-iteration work is just a few scalar copies; profile before threading. - nc_eq = kkt.nc_eq_per_s + # Step 2: Extract per-scenario RHS blocks via the tag-driven global index lists. @inbounds for k in 1:ns - vr_start = (k-1)*nv + sv = kkt.scen_var_global[k] rhs = kkt.rhs_k[k] for i in 1:nv - rhs[i] = wx[vr_start + i] - end - eq_base = (k-1)*nc_eq - for ci in 1:nc_eq - rhs[nv+ci] = wy[kkt.eq_global_indices[eq_base + ci]] + rhs[i] = wx[sv[i]] end end @inbounds for i in 1:nd - kkt.rhs_d[i] = wx[ns*nv+i] + kkt.rhs_d[i] = wx[kkt.design_var_global[i]] end # Step 3: Forward elimination @@ -1056,46 +1371,47 @@ function solve_kkt!( @blas_safe_threads for k in 1:ns solve_linear_system!(kkt.scenario_solvers[k], kkt.rhs_k[k]) end - # Phase 2 (sequential): accumulate into shared rhs_d - for k in 1:ns - mul!(kkt.rhs_d, kkt.C_dk[k], kkt.rhs_k[k], -one(T), one(T)) + # Phase 2 (sequential): rhs_d[coupled] -= Σ_k C_dk_red[k] * rhs_k[k]. Only the m coupled + # design rows receive a contribution (the reduced C_dk holds exactly those rows). Accumulate + # into the contiguous reduced buffer (BLAS GEMV), then scatter-subtract into rhs_d. + if m > 0 + fill!(kkt.rhs_d_red, zero(T)) + for k in 1:ns + mul!(kkt.rhs_d_red, kkt.C_dk[k], kkt.rhs_k[k], one(T), one(T)) + end + @views kkt.rhs_d[kkt.coupled_design_local] .-= kkt.rhs_d_red end - # Step 4: Solve Schur complement + # Step 4: Solve the first-stage Schur complement system (size nd, SPD) solve_linear_system!(kkt.linear_solver, kkt.rhs_d) - # Step 5: Back-substitution (parallel — reads shared rhs_d, writes per-scenario rhs_k) - @blas_safe_threads for k in 1:ns - mul!(kkt.rhs_k[k], kkt.tmp_blk_nd[k], kkt.rhs_d, -one(T), one(T)) + # Step 5: Back-substitution. rhs_k[k] -= tmp_red[k] * rhs_d[coupled] per scenario (only the + # coupled design entries feed back; tmp_red is the reduced blk×m block). + if m > 0 + @views kkt.rhs_d_red .= kkt.rhs_d[kkt.coupled_design_local] + @blas_safe_threads for k in 1:ns + mul!(kkt.rhs_k[k], kkt.tmp_blk_nd[k], kkt.rhs_d_red, -one(T), one(T)) + end end - # Step 6: Write back to w (same threading note as Step 2 above) + # Step 6: Write back to w @inbounds for k in 1:ns - vr_start = (k-1)*nv + sv = kkt.scen_var_global[k] rhs = kkt.rhs_k[k] for i in 1:nv - wx[vr_start + i] = rhs[i] - end - eq_base = (k-1)*nc_eq - for ci in 1:nc_eq - wy[kkt.eq_global_indices[eq_base + ci]] = rhs[nv+ci] + wx[sv[i]] = rhs[i] end end @inbounds for i in 1:nd - wx[ns*nv+i] = kkt.rhs_d[i] + wx[kkt.design_var_global[i]] = kkt.rhs_d[i] end - # Step 7: Recover inequality duals and slacks + # Step 7: Recover inequality duals and slacks (all constraints are inequalities + # under RelaxEquality, so there are no equality duals to preserve). if kkt.n_ineq > 0 - # Stash eq duals; mul! below overwrites all of wy - copyto!(kkt.wy_eq_buf, view(wy, kkt.ind_eq)) - - # J * Δx via sparse: (jt_csc)' * wx + # J * Δx via sparse: (jt_csc)' * wx (overwrites all of wy) mul!(wy, kkt.jt_csc', wx) - # Restore equality duals - view(wy, kkt.ind_eq) .= kkt.wy_eq_buf - # Inequality dual recovery @inbounds for idx in 1:length(kkt.ind_ineq) gi = kkt.ind_ineq[idx] @@ -1109,7 +1425,7 @@ function solve_kkt!( end # KKT matrix-vector product for iterative refinement -function mul!(w::AbstractKKTVector{T}, kkt::SchurComplementKKTSystem{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]) @@ -1138,7 +1454,7 @@ function mul!(w::AbstractKKTVector{T}, kkt::SchurComplementKKTSystem{T}, x::Abst return w end -function mul_hess_blk!(wx, kkt::SchurComplementKKTSystem, t) +function mul_hess_blk!(wx, kkt::SchurComplementCondensedKKTSystem, t) n = num_variables(kkt) mul!(@view(wx[1:n]), Symmetric(kkt.hess_csc, :L), @view(t[1:n])) fill!(@view(wx[n+1:end]), 0) diff --git a/src/MadNLP.jl b/src/MadNLP.jl index 4d9299d63..cac6e8bd7 100644 --- a/src/MadNLP.jl +++ b/src/MadNLP.jl @@ -14,7 +14,7 @@ import SolverCore: getStatus, AbstractOptimizationSolver, AbstractExecutionStats import LDLFactorizations import MUMPS_seq_jll, OpenBLAS32_jll -export MadNLPSolver, MadNLPOptions, LDLSolver, LapackCPUSolver, MumpsSolver, MadNLPExecutionStats, madnlp, solve!, madsuite, SchurComplementKKTSystem +export MadNLPSolver, MadNLPOptions, LDLSolver, LapackCPUSolver, MumpsSolver, MadNLPExecutionStats, madnlp, solve!, madsuite, SchurComplementCondensedKKTSystem Base.USE_GPL_LIBS && export UmfpackSolver, CHOLMODSolver function __init__() @@ -44,4 +44,20 @@ madsuite(::Val{:madnlp}, args...; kwargs...) = madnlp(args...; kwargs...) global Optimizer +# Backwards-compatible (DEPRECATED) alias. `SchurComplementKKTSystem` was renamed to +# `SchurComplementCondensedKKTSystem` when the bordered EnforceEquality saddle was replaced by a +# condensed (RelaxEquality) SPD Schur complement. The rename tracks a *silent behaviour change* — +# the option defaults flipped (equality_treatment EnforceEquality→RelaxEquality, +# fixed_variable_treatment MakeParameter→RelaxBound, a looser `tol`/`bound_relax_factor`) — so +# warn on use instead of aliasing silently, and point users at the new name. +Base.@deprecate_binding( + SchurComplementKKTSystem, + SchurComplementCondensedKKTSystem, + true, + ": it was renamed to `SchurComplementCondensedKKTSystem` and its defaults changed (now " * + "RelaxEquality + RelaxBound with `tol`-relaxed bounds, instead of the old EnforceEquality " * + "bordered saddle). Switch to `SchurComplementCondensedKKTSystem` and re-check your " * + "convergence tolerances.", +) + end # end module diff --git a/test/schur_test.jl b/test/schur_test.jl index f15941927..83181ced1 100644 --- a/test/schur_test.jl +++ b/test/schur_test.jl @@ -3,7 +3,7 @@ using LinearAlgebra using MadNLP using MadNLPTests -@testset "SchurComplementKKTSystem" begin +@testset "SchurComplementCondensedKKTSystem" begin @testset "Basic convergence — quadratic with coupling" begin # min sum_k (v_k - θ_k)^2 + (d - 1)^2 @@ -26,8 +26,8 @@ using MadNLPTests result = madnlp( qp; - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCPUSolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -60,11 +60,16 @@ using MadNLPTests end ref = madnlp(mk(); linear_solver = LapackCPUSolver, print_level = MadNLP.ERROR) + # The condensed default `bound_relax_factor = tol` (= 1e-4 here) relaxes the + # EQUALITY constraint into a ±2e-4 box, landing ~2·tol off the exact-equality + # reference — this testset compares against ground truth, so pin the tight + # relaxation explicitly. schur = madnlp( mk(); - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCPUSolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, kkt_options = schur_opts(; ns, nv, nd, nc), + bound_relax_factor = 1.0e-8, print_level = MadNLP.ERROR, ) @@ -98,8 +103,8 @@ using MadNLPTests result = madnlp( qp; - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCPUSolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -125,8 +130,8 @@ using MadNLPTests result = madnlp( qp; - kkt_system = SchurComplementKKTSystem, - linear_solver = LapackCPUSolver, + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, kkt_options = schur_opts(; ns, nv, nd, nc), print_level = MadNLP.ERROR, ) @@ -137,12 +142,14 @@ using MadNLPTests end @testset "Autodetect dims via tags" begin - # Fake `cb` exposing the tag interface that _resolve_schur_dims expects. + # Fake `cb` exposing the legacy tag interface that _resolve_schur_dims reads. mkcb(tags) = (; nlp = (; tags)) # Happy case: ns=2, nv=1, nd=1, nc=1 → var_scen [1, 2, 0], con_scen [1, 2] ok = (; ns = 2, var_scenario = [1, 2, 0], con_scenario = [1, 2]) - @test MadNLP._resolve_schur_dims(mkcb(ok), 3, 2, 0, 0, 0, 0) == (2, 1, 1, 1) + r = MadNLP._resolve_schur_dims(mkcb(ok), 3, 2, 0, 0, 0, 0) + @test (r.ns, r.nv, r.nd, r.nc, r.nc_design) == (2, 1, 1, 1, 0) + @test r.var_scen == [1, 2, 0] && r.con_scen == [1, 2] # Out-of-range variable tag bad_tag = (; ns = 2, var_scenario = [1, 5, 0], con_scenario = [1, 2]) @@ -156,22 +163,29 @@ using MadNLPTests nu_con = (; ns = 2, var_scenario = [1, 2, 0], con_scenario = [1, 2, 2]) @test_throws ErrorException MadNLP._resolve_schur_dims(mkcb(nu_con), 3, 3, 0, 0, 0, 0) - # Design-only constraint (con tag 0) is rejected. + # Design-only constraint (con tag 0) is now SUPPORTED and counted in nc_design. d_only = (; ns = 2, var_scenario = [1, 2, 0], con_scenario = [1, 2, 0]) - @test_throws ErrorException MadNLP._resolve_schur_dims(mkcb(d_only), 3, 3, 0, 0, 0, 0) + r2 = MadNLP._resolve_schur_dims(mkcb(d_only), 3, 3, 0, 0, 0, 0) + @test (r2.ns, r2.nv, r2.nd, r2.nc, r2.nc_design) == (2, 1, 1, 1, 1) + + # Explicit kkt_options vectors take priority over dims and tags. + r3 = MadNLP._resolve_schur_dims(mkcb(ok), 3, 3, 0, 0, 0, 0, [0, 1, 2], [0, 1, 2]) + @test (r3.ns, r3.nv, r3.nd, r3.nc, r3.nc_design) == (2, 1, 1, 1, 1) + @test r3.var_scen == [0, 1, 2] end @testset "Layout validation" begin - # ns=2, nv=1, nd=1, nc=1: vars [v1, v2, d], cons [c1, c2] + # ns=2, nv=1, nd=1, nc=1: vars [v1, v2, d], cons [c1, c2]. RelaxEquality-only: + # all constraints are inequalities (ind_eq must be empty). ns, nv, nd, nc = 2, 1, 1, 1 n, m = ns*nv + nd, ns*nc # Happy case as a baseline: diagonal Hessian, each constraint touches its own - # scenario var and the design var, all equality. + # 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[1, 2] - ind_ineq = Int32[] + 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, @@ -193,7 +207,7 @@ using MadNLPTests hess_I_ok, hess_J_ok, jac_I_bad, jac_J_bad, ind_eq, ind_ineq, ) - # Non-uniform eq/ineq counts: c_1 is equality, c_2 is inequality. + # RelaxEquality-only: a non-empty `ind_eq` is rejected. @test_throws ErrorException MadNLP._build_schur_symbolic( Float64, n, m, ns, nv, nd, nc, hess_I_ok, hess_J_ok, jac_I_ok, jac_J_ok, Int32[1], Int32[2], @@ -213,7 +227,103 @@ using MadNLPTests jac_J_nu = Int32[1, 2, 5, 3, 4, 5] @test_throws ErrorException MadNLP._build_schur_symbolic( Float64, n2, m2, ns2, nv2, nd2, nc2, - hess_I_nu, hess_J_nu, jac_I_nu, jac_J_nu, Int32[1, 2], Int32[], + hess_I_nu, hess_J_nu, jac_I_nu, jac_J_nu, Int32[], Int32[1, 2], + ) + end + + @testset "Design-only constraints — match SparseKKT reference" begin + # Non-contiguous layout (design vars NOT last, scenario vars scattered) with + # design-only equality AND inequality constraints. Strict convexity ⇒ unique + # optimum, so Schur must match the default sparse KKT solve. + for (ns, nv, nd) in ((2, 2, 2), (3, 2, 3), (4, 1, 2)) + qp, var_scen, con_scen, kkt_opts = + build_twostage_qp_general(; ns, nv, nd, permute = true) + qp_ref = build_twostage_qp_general(; ns, nv, nd, permute = true)[1] + + ref = madnlp(qp_ref; linear_solver = LapackCPUSolver, print_level = MadNLP.ERROR) + # Pin the tight relaxation for the ground-truth compare (see "Match SparseKKT + # reference" above): the condensed default relaxes the design-only equalities + # by ±2·tol. + schur = madnlp( + qp; + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, + kkt_options = kkt_opts, + bound_relax_factor = 1.0e-8, + print_level = MadNLP.ERROR, + ) + + @test ref.status == MadNLP.SOLVE_SUCCEEDED + @test schur.status == MadNLP.SOLVE_SUCCEEDED + @test isapprox(schur.objective, ref.objective; atol = 1.0e-6) + @test isapprox(schur.solution, ref.solution; atol = 1.0e-4) + end + end + + @testset "Uncoupled design variables (m < nd) — match SparseKKT reference" begin + # Only the first `n_coupled` design vars couple to scenarios; the remaining + # `nd - n_coupled` are design-only, so the Schur fill width m = n_coupled < nd. + # Exercises the reduced-width (coupled-subset) elimination path. The weakly- + # constrained uncoupled design directions need the tight tol to match the exact- + # equality reference (the default condensed tol of 1e-4 leaves them ~3e-4 off). + for (ns, nv, nd, n_coupled) in ((3, 2, 3, 1), (4, 2, 4, 2), (2, 3, 5, 1)) + qp, var_scen, con_scen, kkt_opts = + build_twostage_qp_general(; ns, nv, nd, permute = true, n_coupled) + qp_ref = build_twostage_qp_general(; ns, nv, nd, permute = true, n_coupled)[1] + + ref = madnlp(qp_ref; linear_solver = LapackCPUSolver, print_level = MadNLP.ERROR) + schur = madnlp( + qp; + kkt_system = SchurComplementCondensedKKTSystem, + linear_solver = MadNLP.MumpsSolver, + kkt_options = kkt_opts, + bound_relax_factor = 1.0e-8, + tol = 1.0e-8, + print_level = MadNLP.ERROR, + ) + + @test ref.status == MadNLP.SOLVE_SUCCEEDED + @test schur.status == MadNLP.SOLVE_SUCCEEDED + @test isapprox(schur.objective, ref.objective; atol = 1.0e-6) + @test isapprox(schur.solution, ref.solution; atol = 1.0e-4) + end + end + + @testset "Design-only constraints — symbolic build" begin + # ns=2, nv=1, nd=2: vars [v1, v2, d1, d2]. RelaxEquality-only, so all cons + # are inequalities that condense into A_kk / C_dk / S_dd: + # c1 (scen 1): v1 + d1 c2 (scen 2): v2 + d1 + # c3 (design): d1 + d2 c4 (design): d1 + ns, nv, nd, nc = 2, 1, 2, 1 + n, m = ns * nv + nd, 4 + var_scen = [1, 2, 0, 0] + con_scen = [1, 2, 0, 0] + hess_I = Int32[1, 2, 3, 4]; hess_J = Int32[1, 2, 3, 4] + # jac rows: c1:(v1=1,d1=3), c2:(v2=2,d1=3), c3:(d1=3,d2=4), c4:(d1=3) + jac_I = Int32[1, 1, 2, 2, 3, 3, 4] + jac_J = Int32[1, 3, 2, 3, 3, 4, 3] + ind_eq = Int32[] + ind_ineq = Int32[1, 2, 3, 4] + + sym = MadNLP._build_schur_symbolic( + Float64, n, m, ns, nv, nd, nc, + hess_I, hess_J, jac_I, jac_J, ind_eq, ind_ineq, + var_scen, con_scen, + ) + @test sym.nc_design_ineq == 2 # c3, c4 are design-only + @test sym.design_var_global == [3, 4] + # Design inequalities condense into S_dd (design-local cols: d1→1, d2→2): + # c3 over {1,2} ⇒ (1,1),(1,2),(2,1),(2,2); c4 over {1} ⇒ (1,1). + S_entries = Set(zip(sym.design_ineq_S_row, sym.design_ineq_S_col)) + @test (1, 1) in S_entries && (2, 2) in S_entries + @test (1, 2) in S_entries && (2, 1) in S_entries + + # A design-only constraint that reaches a scenario variable must error. + jac_J_bad = Int32[1, 3, 2, 3, 3, 1, 3] # c3 now touches v1 (col 1, scenario 1) + @test_throws ErrorException MadNLP._build_schur_symbolic( + Float64, n, m, ns, nv, nd, nc, + hess_I, hess_J, jac_I, jac_J_bad, ind_eq, ind_ineq, + var_scen, con_scen, ) end end