From 7173e43cc323a07ca2cad361224b33b664504cdb Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 11:09:22 -0400 Subject: [PATCH 01/51] initial --- Project.toml | 14 +- ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 7 + ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 178 +++++++ ext/MadIPMCUDAExt/cuda_wrapper.jl | 70 ++- ext/MadIPMCUDAExt/operators.jl | 40 +- src/MadIPM.jl | 14 +- src/batch/madipm/kernels.jl | 417 +++++++++++++++++ src/batch/madipm/solver.jl | 591 ++++++++++++++++++++++++ src/batch/madnlp/callback.jl | 228 +++++++++ src/batch/madnlp/initialization.jl | 104 +++++ src/batch/madnlp/kernels.jl | 17 + src/batch/madnlp/kkt.jl | 377 +++++++++++++++ src/batch/madnlp/linear_solver.jl | 54 +++ src/batch/madnlp/nlpmodels.jl | 202 ++++++++ src/batch/madnlp/rhs.jl | 71 +++ src/batch/structure.jl | 207 +++++++++ src/batch/utils.jl | 182 ++++++++ 17 files changed, 2766 insertions(+), 7 deletions(-) create mode 100644 ext/MadIPMCUDAExt/cuda_batch_kernels.jl create mode 100644 src/batch/madipm/kernels.jl create mode 100644 src/batch/madipm/solver.jl create mode 100644 src/batch/madnlp/callback.jl create mode 100644 src/batch/madnlp/initialization.jl create mode 100644 src/batch/madnlp/kernels.jl create mode 100644 src/batch/madnlp/kkt.jl create mode 100644 src/batch/madnlp/linear_solver.jl create mode 100644 src/batch/madnlp/nlpmodels.jl create mode 100644 src/batch/madnlp/rhs.jl create mode 100644 src/batch/structure.jl create mode 100644 src/batch/utils.jl diff --git a/Project.toml b/Project.toml index 1ab66531..afae98fb 100644 --- a/Project.toml +++ b/Project.toml @@ -14,16 +14,21 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [weakdeps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +[sources] +QuadraticModels = {rev = "mk/rhsbatch", url = "https://github.com/klamike/QuadraticModels.jl.git"} + [extensions] -MadIPMCUDAExt = ["CUDA", "KernelAbstractions", "MadNLPGPU"] +MadIPMCUDAExt = ["CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] MadIPMMathOptInterfaceExt = "MathOptInterface" [compat] CUDA = "5.4.0" +CUDSS = "0.6" KernelAbstractions = "0.9" LDLFactorizations = "0.10.1" LinearAlgebra = "1.10" @@ -39,8 +44,13 @@ Test = "1.10" julia = "1.10" [extras] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["CUDA", "KernelAbstractions", "MadNLPTests", "MadNLPGPU", "MathOptInterface", "Test"] +test = ["CUDA", "CUDSS", "KernelAbstractions", "MadNLPTests", "MadNLPGPU", "MathOptInterface", "Test"] diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index aec17c5a..0c8ca828 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -6,13 +6,20 @@ using NLPModels using QuadraticModels using CUDA using CUDA.CUSPARSE +using CUDSS using KernelAbstractions import QuadraticModels: SparseMatrixCOO import MadIPM +import MadNLP include("cuda_wrapper.jl") +include("cuda_batch_kernels.jl") include("operators.jl") +function MadIPM._csc_with_nzval(A::CUSPARSE.CuSparseMatrixCSC, nzval, n) + return CUSPARSE.CuSparseMatrixCSC(A.colPtr, A.rowVal, nzval, (n, n)) +end + @kernel function _fill_sparse_structure!(rows, cols, Ap, Aj, Ax) i = @index(Global, Linear) for c in Ap[i]:Ap[i+1]-1 diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl new file mode 100644 index 00000000..6018974f --- /dev/null +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -0,0 +1,178 @@ +@kernel function _batch_set_con_scale_sparse_kernel!( + con_scale, + @Const(ptr), + @Const(inds), + @Const(jac_buffer), +) + (index, j) = @index(Global, NTuple) + @inbounds begin + rng = ptr[index]:ptr[index+1]-1 + for k in rng + (row, i) = inds[k] + con_scale[row, j] = max(con_scale[row, j], abs(jac_buffer[i, j])) + end + end +end + +function MadNLP._set_con_scale_sparse!( + con_scale::CuMatrix{T}, + jac_I::CuVector{<:Integer}, + jac_buffer::CuMatrix{T}, +) where T + ind_jac = CuVector{Int}(1:length(jac_I)) + inds = map((i, j) -> (i, j), jac_I, ind_jac) + !isempty(inds) && sort!(inds) + ptr = getptr(inds; by = ((x1, x2), (y1, y2)) -> x1 != y1) + if length(ptr) > 1 + backend = CUDABackend() + _batch_set_con_scale_sparse_kernel!(backend)( + con_scale, + ptr, + inds, + jac_buffer; + ndrange = (length(ptr) - 1, size(con_scale, 2)), + ) + KernelAbstractions.synchronize(backend) + end + return con_scale +end + +@kernel function _block_argmin_kernel!(out_val, out_idx, @Const(parent_data), offset, nrows) + tid = @index(Local, Linear) + j = @index(Group, Linear) + gs = @groupsize()[1] + T = eltype(out_val) + + sval = @localmem T (64,) + sidx = @localmem Int32 (64,) + + # Strided scan + local_min = T(Inf) + local_idx = Int32(0) + @inbounds begin + i = Int32(tid) + while i <= nrows + v = parent_data[offset + i, j] + if v < local_min + local_min = v + local_idx = i + end + i += Int32(gs) + end + sval[tid] = local_min + sidx[tid] = local_idx + end + @synchronize() + + # Tree reduction + @inbounds begin + stride = Int32(gs) >> Int32(1) + while stride > Int32(0) + if Int32(tid) <= stride + if sval[tid + stride] < sval[tid] + sval[tid] = sval[tid + stride] + sidx[tid] = sidx[tid + stride] + end + end + @synchronize() + stride >>= Int32(1) + end + + if tid == 1 + out_val[1, j] = sval[1] + out_idx[1, j] = sidx[1] + end + end +end + +@kernel function _mehrotra_correction_kernel!( + alpha_p, alpha_d, + @Const(mu), + @Const(val_xl), @Const(idx_xl), @Const(val_xu), @Const(idx_xu), + @Const(val_zl), @Const(idx_zl), @Const(val_zu), @Const(idx_zu), + @Const(d_vals), @Const(x_vals), @Const(xl_vals), @Const(xu_vals), + @Const(zl_vals), @Const(zu_vals), + @Const(ind_lb), @Const(ind_ub), + dlb_off, dub_off, gamma_f, +) + j = @index(Global, Linear) + T = eltype(alpha_p) + + mu_j = mu[1, j] + max_ap = alpha_p[1, j] + max_ad = alpha_d[1, j] + + # primal step + corrected_p = one(T) + @inbounds if max_ap < one(T) + i_xl = idx_xl[1, j] + i_xu = idx_xu[1, j] + if val_xl[1, j] <= val_xu[1, j] && i_xl > Int32(0) + idx = ind_lb[i_xl] + zl_stepped = zl_vals[idx, j] + max_ad * d_vals[dlb_off + i_xl, j] + corrected_p = (x_vals[idx, j] - xl_vals[idx, j] - mu_j / zl_stepped) / (-d_vals[idx, j]) + elseif i_xu > Int32(0) + idx = ind_ub[i_xu] + zu_stepped = zu_vals[idx, j] + max_ad * d_vals[dub_off + i_xu, j] + corrected_p = (xu_vals[idx, j] - x_vals[idx, j] - mu_j / zu_stepped) / d_vals[idx, j] + end + end + @inbounds alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) + + # dual step + corrected_d = one(T) + @inbounds if max_ad < one(T) + i_zl = idx_zl[1, j] + i_zu = idx_zu[1, j] + if val_zl[1, j] <= val_zu[1, j] && i_zl > Int32(0) + idx = ind_lb[i_zl] + x_gap = x_vals[idx, j] + max_ap * d_vals[idx, j] - xl_vals[idx, j] + corrected_d = -(zl_vals[idx, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] + elseif i_zu > Int32(0) + idx = ind_ub[i_zu] + x_gap = xu_vals[idx, j] - x_vals[idx, j] - max_ap * d_vals[idx, j] + corrected_d = -(zu_vals[idx, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] + end + end + @inbounds alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) +end + +const _MEHROTRA_BLOCK = 64 + +function MadIPM._argmin_columns!( + out_val::CuMatrix{T}, out_idx::CuMatrix{Int32}, + parent_data::CuMatrix{T}, offset::Int, nrows::Int; threads_per_column = _MEHROTRA_BLOCK +) where T + ncols = size(out_val, 2) + if ncols > 0 && nrows > 0 + backend = CUDABackend() + _block_argmin_kernel!(backend, threads_per_column)( + out_val, out_idx, parent_data, Int32(offset), Int32(nrows); + ndrange = threads_per_column * ncols, + ) + KernelAbstractions.synchronize(backend) + end +end + +function MadIPM._mehrotra_correct_steps!( + alpha_p::CuMatrix{T}, alpha_d::CuMatrix{T}, mu, + val_xl, idx_xl, val_xu, idx_xu, + val_zl, idx_zl, val_zu, idx_zu, + d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, + ind_lb, ind_ub, dlb_off::Int, dub_off::Int, gamma_f, +) where T + bs = size(alpha_p, 2) + if bs > 0 + backend = CUDABackend() + _mehrotra_correction_kernel!(backend)( + alpha_p, alpha_d, mu, + val_xl, idx_xl, val_xu, idx_xu, + val_zl, idx_zl, val_zu, idx_zu, + d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, + ind_lb, ind_ub, + Int32(dlb_off), Int32(dub_off), gamma_f; + ndrange = bs, + ) + KernelAbstractions.synchronize(backend) + end +end diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 56e173e3..b4ba5e08 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -1,5 +1,6 @@ using MadNLPGPU import LinearAlgebra: BlasFloat +import MadNLP: getptr @kernel function _transfer_to_map!(dest, to_map, src) k = @index(Global, Linear) @@ -14,10 +15,14 @@ function MadNLP.transfer!( src::MadNLP.SparseMatrixCOO{Tv}, map::CuVector{Int}, ) where {Tv} - fill!(nonzeros(dest), zero(Tv)) + return MadNLP._transfer!(dest.nzVal, src.V, map) +end + +function MadNLP._transfer!(dest::CuVector{T}, src::CuVector{T}, map::CuVector{Int}) where T + fill!(dest, zero(T)) if length(map) > 0 backend = CUDABackend() - _transfer_to_map!(backend)(nonzeros(dest), map, src.V; ndrange=length(map)) + _transfer_to_map!(backend)(dest, map, src; ndrange=length(map)) KernelAbstractions.synchronize(backend) end return @@ -184,3 +189,64 @@ MadIPM.sparse_csc_format(::Type{<:CuArray}) = CuSparseMatrixCSC MadIPM._colptr(A::CuSparseMatrixCSC) = A.colPtr MadIPM._rowval(A::CuSparseMatrixCSC) = A.rowVal MadIPM._nzval(A::CuSparseMatrixCSC) = A.nzVal + +function MadIPM._coo_to_scatter( + coo_I, nrows::Int, n_entries::Int, + proto_I, nzVals::CuMatrix{T}, batch_size::Int, +) where T + if n_entries == 0 + scatter = CUSPARSE.CuSparseMatrixCSC( + CuVector{Int32}([1]), + CuVector{Int32}(undef, 0), + CuVector{T}(undef, 0), + (nrows, 0), + ) + op = MadIPMOperator(scatter; spmm_ncols=batch_size) + buffer = similar(nzVals, 0, batch_size) + return op, buffer + end + coo_J = similar(proto_I, n_entries) + coo_J .= Int32(1):Int32(n_entries) + coo_V = similar(nzVals, n_entries) + fill!(coo_V, one(T)) + scatter, _ = MadNLP.coo_to_csc( + MadNLP.SparseMatrixCOO(nrows, n_entries, coo_I, coo_J, coo_V), + ) + fill!(MadIPM._nzval(scatter), one(T)) + op = MadIPMOperator(scatter; spmm_ncols=batch_size) + buffer = similar(nzVals, n_entries, batch_size) + fill!(buffer, zero(T)) + return op, buffer +end + +# we introduce a new constructor that takes the nzvals as a matrix explicitly +function MadNLPGPU.CUDSSSolver( + aug_com::CUSPARSE.CuSparseMatrixCSC{T,Cint}, + nzvals_mat::CuMatrix{T}, + n::Int; + opt::MadNLPGPU.CudssSolverOptions = MadNLPGPU.CudssSolverOptions(), +) where T + batch_nzVal = vec(nzvals_mat) + batch_aug_com = CUSPARSE.CuSparseMatrixCSC( + aug_com.colPtr, aug_com.rowVal, batch_nzVal, size(aug_com), + ) + solver = MadNLPGPU.CUDSSSolver(batch_aug_com; opt=opt) + solver.tril.nzVal = batch_nzVal + return solver +end + +MadIPM.is_factorized(::MadNLPGPU.CUDSSSolver) = true + +function MadIPM._active_factorize!(s::MadNLPGPU.CUDSSSolver, na::Int) + CUDSS.cudss_set(s.inner, "ubatch_size", na) + MadNLP.factorize!(s) + return +end + +function MadIPM._active_solve!(s::MadNLPGPU.CUDSSSolver{T}, rhs::CuVector{T}, na::Int, n::Int) where T + rhs_active = unsafe_wrap(CuArray{T, 2}, pointer(rhs), (n, na)) + CUDSS.cudss_update(s.b_gpu, rhs_active) + CUDSS.cudss_update(s.x_gpu, rhs_active) + CUDSS.cudss("solve", s.inner, s.x_gpu, s.b_gpu, asynchronous=s.opt.cudss_asynchronous) + return +end diff --git a/ext/MadIPMCUDAExt/operators.jl b/ext/MadIPMCUDAExt/operators.jl index a1b7687b..7ba3d79c 100644 --- a/ext/MadIPMCUDAExt/operators.jl +++ b/ext/MadIPMCUDAExt/operators.jl @@ -7,6 +7,7 @@ mutable struct MadIPMOperator{T,M,M2} <: AbstractMatrix{T} transa::Char descA::CUSPARSE.CuSparseMatrixDescriptor buffer::CuVector{UInt8} + spmm_buffer::CuVector{UInt8} alpha::Base.RefValue{T} beta::Base.RefValue{T} end @@ -19,7 +20,7 @@ for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T}), :BlasFloat), (:(CuSparseMatrixCSC{T}), :BlasFloat), (:(CuSparseMatrixCOO{T}), :BlasFloat)) @eval begin - function MadIPMOperator(A::$SparseMatrixType; transa::Char='N', symmetric::Bool=false) where T <: $BlasType + function MadIPMOperator(A::$SparseMatrixType; transa::Char='N', symmetric::Bool=false, spmm_ncols::Int=0) where T <: $BlasType m, n = size(A) alpha = Ref{T}(one(T)) beta = Ref{T}(zero(T)) @@ -39,11 +40,37 @@ for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T}), :BlasFloat), M2 = typeof(mat) alpha = Ref{T}(one(T)) beta = Ref{T}(zero(T)) - return MadIPMOperator{T,M,M2}(T, m, n, A, mat, transa, descA, buffer, alpha, beta) + spmm_buffer = if spmm_ncols > 0 + descB = CUSPARSE.CuDenseMatrixDescriptor(T, n, spmm_ncols) + descC = CUSPARSE.CuDenseMatrixDescriptor(T, m, spmm_ncols) + spmm_buf_size = Ref{Csize_t}() + spmm_algo = CUSPARSE.CUSPARSE_SPMM_ALG_DEFAULT + CUSPARSE.cusparseSpMM_bufferSize(CUSPARSE.handle(), transa, 'N', alpha, descA, descB, beta, descC, T, spmm_algo, spmm_buf_size) + buf = CuVector{UInt8}(undef, spmm_buf_size[]) + if CUSPARSE.version() ≥ v"12.3" + CUSPARSE.cusparseSpMM_preprocess(CUSPARSE.handle(), transa, 'N', alpha, descA, descB, beta, descC, T, spmm_algo, buf) + end + buf + else + CuVector{UInt8}(undef, 0) + end + return MadIPMOperator{T,M,M2}(T, m, n, A, mat, transa, descA, buffer, spmm_buffer, alpha, beta) end end end +function LinearAlgebra.mul!(Y::CuMatrix{T}, A::MadIPMOperator{T}, X::CuMatrix{T}) where T <: BlasFloat + (size(Y, 1) != A.m) && throw(DimensionMismatch("size(Y,1) != A.m")) + (size(X, 1) != A.n) && throw(DimensionMismatch("size(X,1) != A.n")) + descX = CUSPARSE.CuDenseMatrixDescriptor(X) + descY = CUSPARSE.CuDenseMatrixDescriptor(Y) + CUSPARSE.cusparseSpMM( + CUSPARSE.handle(), A.transa, 'N', + A.alpha, A.descA, descX, A.beta, descY, + T, CUSPARSE.CUSPARSE_SPMM_ALG_DEFAULT, A.spmm_buffer, + ) +end + function LinearAlgebra.mul!(y::CuVector{T}, A::MadIPMOperator{T}, x::CuVector{T}) where T <: BlasFloat (length(y) != A.m) && throw(DimensionMismatch("length(y) != A.m")) (length(x) != A.n) && throw(DimensionMismatch("length(x) != A.n")) @@ -52,3 +79,12 @@ function LinearAlgebra.mul!(y::CuVector{T}, A::MadIPMOperator{T}, x::CuVector{T} algo = CUSPARSE.CUSPARSE_SPMV_ALG_DEFAULT CUSPARSE.cusparseSpMV(CUSPARSE.handle(), A.transa, A.alpha, A.descA, descX, A.beta, descY, T, algo, A.buffer) end + +function LinearAlgebra.mul!(Y::CuMatrix{T}, A::MadIPMOperator{T}, X::CuMatrix{T}, α::Number, β::Number) where T <: BlasFloat + A.alpha[] = T(α) + A.beta[] = T(β) + mul!(Y, A, X) + A.alpha[] = one(T) + A.beta[] = zero(T) + return Y +end diff --git a/src/MadIPM.jl b/src/MadIPM.jl index 85fb907e..ec839a1e 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -16,7 +16,19 @@ include("KKT/normalkkt.jl") include("linear_solver.jl") include("solver.jl") -export MPCSolver, madipm +include("batch/utils.jl") +include("batch/madnlp/rhs.jl") +include("batch/madnlp/callback.jl") +include("batch/madnlp/kkt.jl") +include("batch/structure.jl") +include("batch/madnlp/linear_solver.jl") +include("batch/madnlp/kernels.jl") +include("batch/madnlp/initialization.jl") +include("batch/madnlp/nlpmodels.jl") +include("batch/madipm/kernels.jl") +include("batch/madipm/solver.jl") + +export MPCSolver, madipm, madipm_batch, madipm_foreach MadNLP.madsuite(::Val{:madipm}, args...; kwargs...) = madipm(args...; kwargs...) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl new file mode 100644 index 00000000..69ac1653 --- /dev/null +++ b/src/batch/madipm/kernels.jl @@ -0,0 +1,417 @@ +function set_initial_primal_rhs!(solver::AbstractBatchMPCSolver) + p = solver.p + fill!(MadNLP.full(p), 0.0) + py = MadNLP.dual(p) + b = MadNLP.full(solver.c) + + py .= .- b + return +end + +function set_initial_dual_rhs!(solver::AbstractBatchMPCSolver) + p = solver.p + fill!(MadNLP.full(p), 0.0) + px = MadNLP.primal(p) + c = MadNLP.primal(solver.f) + + px .= .- c + return +end + +function set_predictive_rhs!(solver::AbstractBatchMPCSolver, kkt::AbstractBatchKKTSystem) + px = MadNLP.primal(solver.p) + py = MadNLP.dual(solver.p) + pzl = MadNLP.dual_lb(solver.p) + pzu = MadNLP.dual_ub(solver.p) + f = MadNLP.primal(solver.f) + c = MadNLP.full(solver.c) + zl = MadNLP.full(solver.zl) + zu = MadNLP.full(solver.zu) + jacl = MadNLP.full(solver.jacl) + xl_r = lower(solver.xl) + x_lr = lower(solver.x) + zl_r = lower(solver.zl) + xu_r = upper(solver.xu) + x_ur = upper(solver.x) + zu_r = upper(solver.zu) + + fill!(MadNLP.full(solver.p), 0.0) + + px .= .-f .+ zl .- zu .- jacl + py .= .-c + pzl .= (xl_r .- x_lr) .* zl_r + pzu .= (xu_r .- x_ur) .* zu_r + return +end + +function set_correction_rhs!(bs::AbstractBatchMPCSolver, kkt::AbstractBatchKKTSystem, mu, correction_lb, correction_ub, ind_lb, ind_ub) + px = MadNLP.primal(bs.p) + py = MadNLP.dual(bs.p) + pzl = MadNLP.dual_lb(bs.p) + pzu = MadNLP.dual_ub(bs.p) + f = MadNLP.primal(bs.f) + c = MadNLP.full(bs.c) + zl = MadNLP.full(bs.zl) + zu = MadNLP.full(bs.zu) + jacl = MadNLP.full(bs.jacl) + xl_r = lower(bs.xl) + x_lr = lower(bs.x) + zl_r = lower(bs.zl) + xu_r = upper(bs.xu) + x_ur = upper(bs.x) + zu_r = upper(bs.zu) + + px .= .-f .+ zl .- zu .- jacl + py .= .-c + pzl .= (xl_r .- x_lr) .* zl_r .+ mu .- correction_lb + pzu .= (xu_r .- x_ur) .* zu_r .- mu .- correction_ub + return +end + +function get_correction!( + batch_solver::AbstractBatchMPCSolver, + correction_lb, + correction_ub +) + dlb = MadNLP.dual_lb(batch_solver.d) + dub = MadNLP.dual_ub(batch_solver.d) + + dx_lr = xp_lr(batch_solver.d) + dx_ur = xp_ur(batch_solver.d) + + correction_lb .= dx_lr .* dlb + correction_ub .= dx_ur .* dub + return +end + +function set_aug_diagonal_reg!(kkt, solver::AbstractBatchMPCSolver) + xl_r = lower(solver.xl) + x_lr = lower(solver.x) + zl_r = lower(solver.zl) + xu_r = upper(solver.xu) + x_ur = upper(solver.x) + zu_r = upper(solver.zu) + + kkt.reg .= solver.del_w + du_diag(kkt) .= solver.del_c + + kkt.l_diag .= xl_r .- x_lr + kkt.u_diag .= x_ur .- xu_r + + kkt.l_lower .= zl_r + kkt.u_lower .= zu_r + + pr_diag(kkt) .= kkt.reg + pr_diag_lb = view(kkt.nzVals, _get_ind_lb(solver), :) + pr_diag_ub = view(kkt.nzVals, _get_ind_ub(solver), :) + pr_diag_lb .-= kkt.l_lower ./ kkt.l_diag + pr_diag_ub .-= kkt.u_lower ./ kkt.u_diag + return +end + +function get_complementarity_measure!(solver::AbstractBatchMPCSolver) + ws = solver.workspace + nlb, nub = solver.d.nlb, solver.d.nub + T = eltype(ws.mu_curr) + + if nlb + nub == 0 + fill!(ws.mu_curr, zero(T)) + return ws.mu_curr + end + + xl_r = lower(solver.xl) + x_lr = lower(solver.x) + zl_r = lower(solver.zl) + xu_r = upper(solver.xu) + x_ur = upper(solver.x) + zu_r = upper(solver.zu) + + _scratch_lb = MadNLP.dual_lb(solver._w2) + @. _scratch_lb = (x_lr - xl_r) * zl_r + sum!(ws.sum_lb, _scratch_lb) + + _scratch_ub = MadNLP.dual_ub(solver._w2) + @. _scratch_ub = (x_ur - xu_r) * zu_r + sum!(ws.sum_ub, _scratch_ub) + + @. ws.mu_curr = (ws.sum_lb + ws.sum_ub) / (nlb + nub) + return ws.mu_curr +end + +function get_affine_complementarity_measure!(solver::AbstractBatchMPCSolver, alpha_p, alpha_d) + ws = solver.workspace + nlb, nub = solver.d.nlb, solver.d.nub + T = eltype(ws.mu_affine) + + if nlb + nub == 0 + fill!(ws.mu_affine, zero(T)) + return ws.mu_affine + end + + xl_r = lower(solver.xl) + x_lr = lower(solver.x) + zl_r = lower(solver.zl) + xu_r = upper(solver.xu) + x_ur = upper(solver.x) + zu_r = upper(solver.zu) + dx_lr = xp_lr(solver.d) + dx_ur = xp_ur(solver.d) + dzlb = MadNLP.dual_lb(solver.d) + dzub = MadNLP.dual_ub(solver.d) + + _scratch_lb = MadNLP.dual_lb(solver._w2) + @. _scratch_lb = (x_lr + alpha_p * dx_lr - xl_r) * (zl_r + alpha_d * dzlb) + sum!(ws.sum_lb, _scratch_lb) + + _scratch_ub = MadNLP.dual_ub(solver._w2) + @. _scratch_ub = (xu_r - x_ur - alpha_p * dx_ur) * (zu_r + alpha_d * dzub) + sum!(ws.sum_ub, _scratch_ub) + + @. ws.mu_affine = (ws.sum_lb + ws.sum_ub) / (nlb + nub) + return ws.mu_affine +end + +function update_barrier!(::Mehrotra, solver::AbstractBatchMPCSolver, mu_affine) + ws = solver.workspace + T = eltype(ws.mu_curr) + + has_inequalities = (length(_get_ind_llb(solver)) + length(_get_ind_uub(solver))) > 0 + + get_complementarity_measure!(solver) + + if has_inequalities + @. ws.mu_batch = clamp((ws.mu_affine / ws.mu_curr) ^ 3, T(1e-6), T(10.0)) + @. ws.mu_batch = max(solver.opt.mu_min, ws.mu_batch * ws.mu_curr) + else + @. ws.mu_batch = max(solver.opt.mu_min, ws.mu_curr) + end + return +end + +function get_fraction_to_boundary_step!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + x, xl, xu = batch_solver.x, batch_solver.xl, batch_solver.xu + zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d + nlb, nub = d.nlb, d.nub + T = eltype(ws.alpha_p) + inf_val = T(Inf) + + if nlb > 0 + _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) + _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + + map!( + (dx, xl, x) -> dx < 0 ? (xl - x) / dx : inf_val, + _scratch_lb, _dx_lr, _xl_r, _x_lr, + ) + minimum!(ws.alpha_xl, _scratch_lb) + @. ws.alpha_xl *= ws.tau + + map!( + (dz, z) -> dz < 0 ? -z / dz : inf_val, + _scratch_lb, _dzlb, _zl_r, + ) + minimum!(ws.alpha_zl, _scratch_lb) + @. ws.alpha_zl *= ws.tau + else + fill!(ws.alpha_xl, one(T)) + fill!(ws.alpha_zl, one(T)) + end + + if nub > 0 + _dx_ur = xp_ur(d); _xu_r = upper(xu); _x_ur = upper(x) + _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + + map!( + (dx, xu, x) -> dx > 0 ? (xu - x) / dx : inf_val, + _scratch_ub, _dx_ur, _xu_r, _x_ur, + ) + minimum!(ws.alpha_xu, _scratch_ub) + @. ws.alpha_xu *= ws.tau + + map!( + (dz, z) -> (dz < 0) & (z + dz < 0) ? -z / dz : inf_val, + _scratch_ub, _dzub, _zu_r, + ) + minimum!(ws.alpha_zu, _scratch_ub) + @. ws.alpha_zu *= ws.tau + else + fill!(ws.alpha_xu, one(T)) + fill!(ws.alpha_zu, one(T)) + end + + ws.alpha_p .= min.(ws.alpha_xl, ws.alpha_xu, one(T)) + ws.alpha_d .= min.(ws.alpha_zl, ws.alpha_zu, one(T)) + return +end + +function set_tau!(rule::ConservativeStep, batch_solver::AbstractBatchMPCSolver) + fill!(batch_solver.workspace.tau, rule.tau) +end +function set_tau!(rule::AdaptiveStep, batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + ws.tau .= max.(1 .- ws.mu_batch, rule.tau_min) +end +function update_step!(rule::Union{ConservativeStep, AdaptiveStep}, batch_solver::AbstractBatchMPCSolver) + set_tau!(rule, batch_solver) + get_fraction_to_boundary_step!(batch_solver) + return +end + +function _argmin_columns!( + out_val::AbstractMatrix{T}, out_idx::AbstractMatrix{Int32}, + parent_data::AbstractMatrix{T}, offset::Int, nrows::Int, +) where T + @inbounds for j in axes(out_val, 2) + min_val = T(Inf) + min_idx = Int32(0) + for i in 1:nrows + v = parent_data[offset + i, j] + if v < min_val + min_val = v + min_idx = Int32(i) + end + end + out_val[1, j] = min_val + out_idx[1, j] = min_idx + end +end + +function _mehrotra_correct_steps!( + alpha_p, alpha_d, mu, + val_xl, idx_xl, val_xu, idx_xu, + val_zl, idx_zl, val_zu, idx_zu, + d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, + ind_lb, ind_ub, dlb_off::Int, dub_off::Int, gamma_f, +) + T = eltype(alpha_p) + @inbounds for j in axes(alpha_p, 2) + mu_j = mu[1, j] + max_ap = alpha_p[1, j] + max_ad = alpha_d[1, j] + + # primal step + corrected_p = one(T) + if max_ap < one(T) + i_xl = idx_xl[1, j] + i_xu = idx_xu[1, j] + if val_xl[1, j] <= val_xu[1, j] && i_xl > Int32(0) + idx = ind_lb[i_xl] + zl_stepped = zl_vals[idx, j] + max_ad * d_vals[dlb_off + i_xl, j] + corrected_p = (x_vals[idx, j] - xl_vals[idx, j] - mu_j / zl_stepped) / (-d_vals[idx, j]) + elseif i_xu > Int32(0) + idx = ind_ub[i_xu] + zu_stepped = zu_vals[idx, j] + max_ad * d_vals[dub_off + i_xu, j] + corrected_p = (xu_vals[idx, j] - x_vals[idx, j] - mu_j / zu_stepped) / d_vals[idx, j] + end + end + alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) + + # dual step + corrected_d = one(T) + if max_ad < one(T) + i_zl = idx_zl[1, j] + i_zu = idx_zu[1, j] + if val_zl[1, j] <= val_zu[1, j] && i_zl > Int32(0) + idx = ind_lb[i_zl] + x_gap = x_vals[idx, j] + max_ap * d_vals[idx, j] - xl_vals[idx, j] + corrected_d = -(zl_vals[idx, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] + elseif i_zu > Int32(0) + idx = ind_ub[i_zu] + x_gap = xu_vals[idx, j] - x_vals[idx, j] - max_ap * d_vals[idx, j] + corrected_d = -(zu_vals[idx, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] + end + end + alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) + end +end + +function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + x, xl, xu = batch_solver.x, batch_solver.xl, batch_solver.xu + zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d + nlb, nub = d.nlb, d.nub + T = eltype(ws.alpha_p) + gamma_f = T(rule.gamma_f) + gamma_a = one(T) / (one(T) - gamma_f) + + fill!(ws.tau, one(T)) + get_fraction_to_boundary_step!(batch_solver) + + get_affine_complementarity_measure!(batch_solver, ws.alpha_p, ws.alpha_d) + mu_full = ws.mu_curr + @. mu_full = ws.mu_affine / gamma_a + + w2_vals = batch_solver._w2.values + dlb_off = d.n + d.m + dub_off = d.n + d.m + d.nlb + + if nlb > 0 + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) + _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) + + map!((dx, xl, x) -> dx < 0 ? (xl - x) / dx : T(Inf), _scratch_lb, _dx_lr, _xl_r, _x_lr) + _argmin_columns!(ws.alpha_xl, ws.idx_xl, w2_vals, dlb_off, nlb) + + map!((dz, z) -> dz < 0 ? -z / dz : T(Inf), _scratch_lb, _dzlb, _zl_r) + _argmin_columns!(ws.alpha_zl, ws.idx_zl, w2_vals, dlb_off, nlb) + else + fill!(ws.alpha_xl, one(T)); fill!(ws.idx_xl, Int32(0)) + fill!(ws.alpha_zl, one(T)); fill!(ws.idx_zl, Int32(0)) + end + + if nub > 0 + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + _dx_ur = xp_ur(d); _xu_r = upper(xu); _x_ur = upper(x) + _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) + + map!((dx, xu, x) -> dx > 0 ? (xu - x) / dx : T(Inf), _scratch_ub, _dx_ur, _xu_r, _x_ur) + _argmin_columns!(ws.alpha_xu, ws.idx_xu, w2_vals, dub_off, nub) + + map!((dz, z) -> (dz < 0) & (z + dz < 0) ? -z / dz : T(Inf), _scratch_ub, _dzub, _zu_r) + _argmin_columns!(ws.alpha_zu, ws.idx_zu, w2_vals, dub_off, nub) + else + fill!(ws.alpha_xu, one(T)); fill!(ws.idx_xu, Int32(0)) + fill!(ws.alpha_zu, one(T)); fill!(ws.idx_zu, Int32(0)) + end + + _mehrotra_correct_steps!( + ws.alpha_p, ws.alpha_d, mu_full, + ws.alpha_xl, ws.idx_xl, ws.alpha_xu, ws.idx_xu, + ws.alpha_zl, ws.idx_zl, ws.alpha_zu, ws.idx_zu, + d.values, x.values, xl.values, xu.values, zl.values, zu.values, + d.ind_lb, d.ind_ub, dlb_off, dub_off, gamma_f, + ) + + return +end + +function init_regularization!(solver::AbstractBatchMPCSolver, ::NoRegularization) + fill!(solver.del_w, 1.0) + fill!(solver.del_c, 0.0) +end +function update_regularization!(solver::AbstractBatchMPCSolver, ::NoRegularization) + fill!(solver.del_w, 0.0) + fill!(solver.del_c, 0.0) +end +function init_regularization!(solver::AbstractBatchMPCSolver, reg::FixedRegularization) + fill!(solver.del_w, 1.0) + fill!(solver.del_c, reg.delta_d) +end +function update_regularization!(solver::AbstractBatchMPCSolver, reg::FixedRegularization) + fill!(solver.del_w, reg.delta_p) + fill!(solver.del_c, reg.delta_d) +end +function init_regularization!(solver::AbstractBatchMPCSolver, reg::AdaptiveRegularization) + fill!(solver.del_w, 1.0) + fill!(solver.del_c, reg.delta_d) +end +function update_regularization!(solver::AbstractBatchMPCSolver, reg::AdaptiveRegularization) + reg.delta_p = max(reg.delta_p / 10.0, reg.delta_min) + reg.delta_d = min(reg.delta_d / 10.0, -reg.delta_min) + fill!(solver.del_w, reg.delta_p) + fill!(solver.del_c, reg.delta_d) +end diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl new file mode 100644 index 00000000..18664fdd --- /dev/null +++ b/src/batch/madipm/solver.jl @@ -0,0 +1,591 @@ +function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T + bkkt = batch_solver.kkt + bs = batch_solver.batch_size + n = batch_solver.d.n + m = batch_solver.d.m + + bx, bxl, bxu = batch_solver.x, batch_solver.xl, batch_solver.xu + bzl, bzu = batch_solver.zl, batch_solver.zu + x = MadNLP.primal(bx) + l, u = MadNLP.full(bxl), MadNLP.full(bxu) + lb, ub = lower(bxl), upper(bxu) + zl, zu = lower(bzl), upper(bzu) + xl, xu = lower(bx), upper(bx) + # use jacl as a buffer + res = MadNLP.full(batch_solver.jacl) + + # Add initial primal-dual regularization + bkkt.reg .= batch_solver.del_w + pr_diag(bkkt) .= batch_solver.del_w + du_diag(bkkt) .= batch_solver.del_c + + # Step 0: factorize initial KKT system + MadNLP.factorize_wrapper!(batch_solver) + + # Step 1: Compute initial primal variable as x0 = x + dx, with dx the + # least square solution of the system A * dx = (b - A*x) + set_initial_primal_rhs!(batch_solver) + solve_system!(batch_solver.d, batch_solver, batch_solver.p) + # x0 = x + dx + x .+= MadNLP.primal(batch_solver.d) + + # Step 2: Compute initial dual variable as the least square solution of A' * y = -f + set_initial_dual_rhs!(batch_solver) + solve_system!(batch_solver.d, batch_solver, batch_solver.p) + MadNLP.full(batch_solver.y) .= MadNLP.dual(batch_solver.d) + + # Step 3: init bounds multipliers using c + A' * y - zl + zu = 0 + # A' * y + MadNLP.jtprod!(res, bkkt, batch_solver.y) + # A'*y + c + res .+= MadNLP.primal(batch_solver.f) + # Initialize bounds multipliers + map!( + (r_, l_, u_, zl_) -> begin + val = if isfinite(l_) && isfinite(u_) + 0.5 * r_ + elseif isfinite(l_) + r_ + else + zl_ + end + val + end, + MadNLP.full(batch_solver.zl), res, l, u, MadNLP.full(batch_solver.zl), + ) + map!( + (r_, l_, u_, zu_) -> begin + val = if isfinite(l_) && isfinite(u_) + -0.5 * r_ + elseif isfinite(u_) + -r_ + else + zu_ + end + val + end, + MadNLP.full(batch_solver.zu), res, l, u, MadNLP.full(batch_solver.zu), + ) + + ws = batch_solver.workspace + nlb_init, nub_init = batch_solver.d.nlb, batch_solver.d.nub + _s1 = ws.alpha_xl # (1,bs) scratch + _s2 = ws.alpha_xu # (1,bs) scratch + + # delta_x = max(0, -1.5 * min(xl-lb, 0), -1.5 * min(ub-xu, 0)) + if nlb_init > 0 + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + @. _scratch_lb = xl - lb + minimum!(_s1, _scratch_lb) + @. _s1 = min(_s1, zero(T)) # clamp: init=0.0 behavior + else + fill!(_s1, zero(T)) + end + if nub_init > 0 + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + @. _scratch_ub = ub - xu + minimum!(_s2, _scratch_ub) + @. _s2 = min(_s2, zero(T)) + else + fill!(_s2, zero(T)) + end + delta_x = ws.mu_batch # (1,bs) scratch + @. delta_x = max(zero(T), T(-1.5) * _s1, T(-1.5) * _s2) + + # delta_s = max(0, -1.5 * min(zl, 0), -1.5 * min(zu, 0)) + if nlb_init > 0 + minimum!(_s1, zl) + @. _s1 = min(_s1, zero(T)) + else + fill!(_s1, zero(T)) + end + if nub_init > 0 + minimum!(_s2, zu) + @. _s2 = min(_s2, zero(T)) + else + fill!(_s2, zero(T)) + end + delta_s = ws.mu_curr # (1,bs) scratch + @. delta_s = max(zero(T), T(-1.5) * _s1, T(-1.5) * _s2) + + xl .+= delta_x + xu .-= delta_x + zl .+= 1.0 .+ delta_s + zu .+= 1.0 .+ delta_s + + # μ = sum((xl-lb)*zl) + sum((ub-xu)*zu) + μ = ws.mu_affine # (1,bs) scratch + fill!(μ, zero(T)) + if nlb_init > 0 + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + @. _scratch_lb = (xl - lb) * zl + sum!(ws.sum_lb, _scratch_lb) + μ .+= ws.sum_lb + end + if nub_init > 0 + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + @. _scratch_ub = (ub - xu) * zu + sum!(ws.sum_ub, _scratch_ub) + μ .+= ws.sum_ub + end + + # delta_x2 = μ / (2 * (sum(zl) + sum(zu))) + if nlb_init > 0 + sum!(ws.sum_lb, zl) + else + fill!(ws.sum_lb, zero(T)) + end + if nub_init > 0 + sum!(ws.sum_ub, zu) + else + fill!(ws.sum_ub, zero(T)) + end + delta_x2 = _s1 # reuse (1,bs) scratch + @. delta_x2 = μ / (2 * (ws.sum_lb + ws.sum_ub)) + + # delta_s2 = μ / (2 * (sum(xl-lb) + sum(ub-xu))) + if nlb_init > 0 + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + @. _scratch_lb = xl - lb + sum!(ws.sum_lb, _scratch_lb) + else + fill!(ws.sum_lb, zero(T)) + end + if nub_init > 0 + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + @. _scratch_ub = ub - xu + sum!(ws.sum_ub, _scratch_ub) + else + fill!(ws.sum_ub, zero(T)) + end + delta_s2 = _s2 # reuse (1,bs) scratch + @. delta_s2 = μ / (2 * (ws.sum_lb + ws.sum_ub)) + + xl .+= delta_x2 + xu .-= delta_x2 + zl .+= delta_s2 + zu .+= delta_s2 + + # Use Ipopt's heuristic to project x back on the interval [l, u] + kappa = batch_solver.opt.bound_fac + map!( + (l_, u_, x_) -> begin + out = if x_ < l_ + pl = min(kappa * max(1.0, l_), kappa * (u_ - l_)) + l_ + pl + elseif u_ < x_ + pu = min(kappa * max(1.0, u_), kappa * (u_ - l_)) + u_ - pu + else + x_ + end + out + end, + x, l, u, x, + ) + return +end + +function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T + opt = batch_solver.opt + bcb = batch_solver.bcb + ws = batch_solver.workspace + + MadNLP.initialize!( + bcb, + batch_solver.x, + batch_solver.xl, + batch_solver.xu, + MadNLP.full(batch_solver.y), + MadNLP.full(batch_solver.rhs), + bcb.ind_ineq; + tol=opt.bound_relax_factor, + bound_push=opt.bound_push, + bound_fac=opt.bound_fac, + ) + fill!(MadNLP.full(batch_solver.jacl), zero(T)) + + if opt.scaling + MadNLP.set_scaling!( + bcb, + batch_solver.x, + batch_solver.xl, + batch_solver.xu, + MadNLP.full(batch_solver.y), + MadNLP.full(batch_solver.rhs), + bcb.ind_ineq, + T(opt.nlp_scaling_max_gradient), + ) + end + + MadNLP.initialize!(batch_solver.kkt) + init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + init_starting_point!(batch_solver) + + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(T)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, zero(T)) + fill!(ws.inf_du, zero(T)) + fill!(ws.inf_compl, zero(T)) + fill!(ws.dual_obj, zero(T)) + fill!(ws.alpha_p, zero(T)) + fill!(ws.alpha_d, zero(T)) + t_now = time() + batch_solver.batch_cnt.start_time[] = t_now + fill!(batch_solver.batch_cnt.k, 0) + batch_solver.batch_cnt.linear_solver_time[] = 0.0 + batch_solver.batch_cnt.eval_function_time[] = 0.0 + batch_solver.batch_cnt.obj_cnt[] = 0 + batch_solver.batch_cnt.obj_grad_cnt[] = 0 + batch_solver.batch_cnt.con_cnt[] = 0 + + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + + return +end + +function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) where T + ws = batch_solver.workspace + opt = batch_solver.opt + bcnt = batch_solver.batch_cnt + x, xl, xu = batch_solver.x, batch_solver.xl, batch_solver.xu + zl, zu = batch_solver.zl, batch_solver.zu + bs = batch_solver.batch_size + nlb, nub = batch_solver.d.nlb, batch_solver.d.nub + + f_vals = MadNLP.full(batch_solver.f) + zl_vals = MadNLP.full(zl) + zu_vals = MadNLP.full(zu) + jacl_vals = MadNLP.full(batch_solver.jacl) + y_vals = MadNLP.full(batch_solver.y) + rhs_vals = MadNLP.full(batch_solver.rhs) + + # inf_pr[i] = norm(c[:, i], Inf) / max(1, norm_b[i]) + ws.inf_pr .= maximum(abs, MadNLP.full(batch_solver.c); dims=1) + @. ws.inf_pr /= max(one(T), ws.norm_b) + + # inf_du[i] = max|f-zl+zu+jacl| / max(1, norm_c[i]) + _scratch_n = MadNLP.primal(batch_solver._w2) + @. _scratch_n = abs(f_vals - zl_vals + zu_vals + jacl_vals) + maximum!(ws.inf_du, _scratch_n) + @. ws.inf_du /= max(one(T), ws.norm_c) + + # inf_compl[i] = get_optimality_gap / max(1, norm_c[i]) + if nlb > 0 + x_lr = lower(x); xl_r = lower(xl); zl_r = lower(zl) + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + @. _scratch_lb = abs(x_lr - xl_r) * zl_r + maximum!(ws.sum_lb, _scratch_lb) + else + fill!(ws.sum_lb, zero(T)) + end + if nub > 0 + xu_r = upper(xu); x_ur = upper(x); zu_r = upper(zu) + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + @. _scratch_ub = abs(xu_r - x_ur) * zu_r + maximum!(ws.sum_ub, _scratch_ub) + else + fill!(ws.sum_ub, zero(T)) + end + @. ws.inf_compl = max(ws.sum_lb, ws.sum_ub) / max(one(T), ws.norm_c) + @. ws.best_complementarity = min(ws.best_complementarity, ws.inf_compl) + + _scratch_m = MadNLP.dual(batch_solver._w2) + @. _scratch_m = y_vals * rhs_vals + sum!(ws.dual_obj, _scratch_m) + ws.dual_obj .*= -one(T) + if nlb > 0 + zl_r = lower(zl); xl_r = lower(xl) + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + @. _scratch_lb = zl_r * xl_r + sum!(ws.sum_lb, _scratch_lb) + ws.dual_obj .+= ws.sum_lb + end + if nub > 0 + zu_r = upper(zu); xu_r = upper(xu) + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + @. _scratch_ub = zu_r * xu_r + sum!(ws.sum_ub, _scratch_ub) + ws.dual_obj .-= ws.sum_ub + end + + ds = T(opt.divergence_scale) + copyto!(ws.term_converged, + vec(@. max(ws.inf_pr, ws.inf_du, ws.inf_compl) <= opt.tol)) + copyto!(ws.term_infeasible, + vec(@. (ws.inf_compl > opt.divergence_tol * ws.best_complementarity + ) & (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))))) + copyto!(ws.term_diverging, + vec(@. ws.obj_val < -(opt.divergence_tol * max(ds, abs(ws.dual_obj), one(T))))) + + walltime_hit = time() - bcnt.start_time[] >= opt.max_wall_time + @inbounds for i in 1:bs + ws.status[i] != MadNLP.REGULAR && continue + if ws.term_converged[i] + ws.status[i] = MadNLP.SOLVE_SUCCEEDED + elseif ws.term_infeasible[i] + ws.status[i] = MadNLP.INFEASIBLE_PROBLEM_DETECTED + elseif ws.term_diverging[i] + ws.status[i] = MadNLP.DIVERGING_ITERATES + elseif bcnt.k[i] >= opt.max_iter + ws.status[i] = MadNLP.MAXIMUM_ITERATIONS_EXCEEDED + elseif walltime_hit + ws.status[i] = MadNLP.MAXIMUM_WALLTIME_EXCEEDED + end + end + return +end + +function solve_system!( + d::BatchUnreducedKKTVector{T}, + batch_solver::AbstractBatchMPCSolver{T}, + p::BatchUnreducedKKTVector{T}, +) where T + opt = batch_solver.opt + copyto!(MadNLP.full(d), MadNLP.full(p)) + MadNLP.solve_kkt!(batch_solver.kkt, batch_solver) + + w = batch_solver._w1 + copyto!(MadNLP.full(w), MadNLP.full(p)) + mul!(w, batch_solver.kkt, d, -one(T), one(T)) + + bkkt = batch_solver.kkt + bs = bkkt.batch_size + @inbounds for i in 1:bs + if bkkt.batch_map[i] == 0 + view(MadNLP.full(w), :, i) .= zero(T) + view(MadNLP.full(p), :, i) .= zero(T) + end + end + norm_w = norm(MadNLP.full(w), Inf) + norm_p = norm(MadNLP.full(p), Inf) + + residual_ratio = norm_w / max(one(T), norm_p) + MadNLP.@debug( + batch_solver.logger, + @sprintf("Residual after linear solve: %6.2e", residual_ratio), + ) + if isnan(residual_ratio) || (opt.check_residual && (residual_ratio > opt.tol_linear_solve)) + throw(MadNLP.SolveException()) + end + return d +end + +function increment_k!(batch_solver::AbstractBatchMPCSolver) + bcnt = batch_solver.batch_cnt + ws = batch_solver.workspace + for i in 1:batch_solver.batch_size + if ws.status[i] == MadNLP.REGULAR + bcnt.k[i] += 1 + end + end +end + +function update_solution!(stats::BatchExecutionStats, batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + bcb = batch_solver.bcb + x, zl, zu = batch_solver.x, batch_solver.zl, batch_solver.zu + + stats.status .= ws.status + stats.iter .= batch_solver.batch_cnt.k + + MadNLP.unpack_x!(stats.solution, bcb, x) + MadNLP.unpack_y!(stats.multipliers, bcb, MadNLP.full(batch_solver.y)) + MadNLP.unpack_z!(stats.multipliers_L, bcb, MadNLP.variable(zl)) + MadNLP.unpack_z!(stats.multipliers_U, bcb, MadNLP.variable(zu)) + stats.objective .= MadNLP.unpack_obj(bcb, ws.obj_val) + MadNLP.unpack_cons!(stats.constraints, bcb, MadNLP.full(batch_solver.c), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, MadNLP.slack(x)) + + stats.dual_feas .= vec(ws.inf_du) + stats.primal_feas .= vec(ws.inf_pr) + stats.total_time .= batch_solver.batch_cnt.total_time + return stats +end + +function affine_direction!(solver::AbstractBatchMPCSolver) + set_predictive_rhs!(solver, solver.kkt) + solve_system!(solver.d, solver, solver.p) +end + +function prediction_step!(solver::AbstractBatchMPCSolver) + ws = solver.workspace + affine_direction!(solver) + + fill!(ws.tau, one(eltype(ws.tau))) + get_fraction_to_boundary_step!(solver) + zero_inactive_step!(solver) + get_affine_complementarity_measure!(solver, ws.alpha_p, ws.alpha_d) + get_correction!(solver, MadNLP.full(solver.correction_lb), MadNLP.full(solver.correction_ub)) + update_barrier!(solver.opt.barrier_update, solver, ws.mu_affine) + return +end + +function mehrotra_correction_direction!(solver::AbstractBatchMPCSolver) + set_correction_rhs!(solver, solver.kkt, solver.workspace.mu_batch, MadNLP.full(solver.correction_lb), MadNLP.full(solver.correction_ub), nothing, nothing) + solve_system!(solver.d, solver, solver.p) + return +end + +function factorize_system!(batch_solver::AbstractBatchMPCSolver) + update_regularization!(batch_solver, batch_solver.opt.regularization) + max_trials = 3 + for _ in 1:max_trials + set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) + MadNLP.factorize_wrapper!(batch_solver) + is_factorized(batch_solver.kkt.batch_solver) && break # exit once all are factorized + batch_solver.del_w .*= 100.0 + batch_solver.del_c .*= 100.0 + end + return +end + + +function apply_step!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + x, xl, xu = batch_solver.x, batch_solver.xl, batch_solver.xu + zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d + batch_size = batch_solver.batch_size + nlb, nub = d.nlb, d.nub + n, m = d.n, d.m + + # x += alpha_p * dx + MadNLP.full(x) .+= ws.alpha_p .* MadNLP.primal(d) + + # y += alpha_d * d_dual + MadNLP.full(batch_solver.y) .+= ws.alpha_d .* MadNLP.dual(d) + + # zl_r += alpha_d * dzl, zu_r += alpha_d * dzu + if nlb > 0 + lower(zl) .+= ws.alpha_d .* MadNLP.dual_lb(d) + end + if nub > 0 + upper(zu) .+= ws.alpha_d .* MadNLP.dual_ub(d) + end + + MadNLP.adjust_boundary!(lower(x), lower(xl), upper(x), upper(xu), ws.mu_batch) + increment_k!(batch_solver) + return +end + +function evaluate_model!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + bcb = batch_solver.bcb + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return +end + +function mpc_step!(batch_solver::AbstractBatchMPCSolver) + factorize_system!(batch_solver) + prediction_step!(batch_solver) + mehrotra_correction_direction!(batch_solver) + update_step!(batch_solver.opt.step_rule, batch_solver) + zero_inactive_step!(batch_solver) + apply_step!(batch_solver) + evaluate_model!(batch_solver) +end + +function _update_active_mask!(batch_solver::AbstractBatchMPCSolver{T}) where T + ws = batch_solver.workspace + bmap = batch_solver.kkt.batch_map + copyto!(ws.active_mask, reshape(T.(bmap .!= 0), 1, :)) +end + +function mpc!(batch_solver::AbstractBatchMPCSolver) + while true + MadNLP.print_iter(batch_solver) + update_termination_criteria!(batch_solver) + update_active_set!(batch_solver.kkt, batch_solver.workspace.status) + batch_solver.kkt.active_batch_size[] == 0 && return + _update_active_mask!(batch_solver) + mpc_step!(batch_solver) + end +end + +function solve!(batch_solver::AbstractBatchMPCSolver{T}) where T + ws = batch_solver.workspace + bcb = batch_solver.bcb + bs = batch_solver.batch_size + + nvar_nlp = bcb.nlp.meta.nvar + ncon = bcb.ncon + VT = typeof(ws.bf) + MT = typeof(MadNLP.full(batch_solver.x)) + stats = BatchExecutionStats(MT, VT, nvar_nlp, ncon, bs) + + try + MadNLP.@notice(batch_solver.logger, "MadIPM batch solve ($bs problems)\n") + initialize!(batch_solver) + mpc!(batch_solver) + catch e + for i in 1:bs + if ws.status[i] == MadNLP.REGULAR + ws.status[i] = MadNLP.INTERNAL_ERROR + end + end + batch_solver.opt.rethrow_error && rethrow(e) + finally + bcnt = batch_solver.batch_cnt + t_end = time() + bcnt.total_time .= t_end .- bcnt.start_time[] + update_solution!(stats, batch_solver) + for i in 1:bs + MadNLP.@notice(batch_solver.logger, "Problem $i: $(MadNLP.get_status_output(ws.status[i], batch_solver.opt))") + end + end + + return stats +end + +function madipm_batch(bnlp::NLPModels.AbstractBatchNLPModel; kwargs...) + batch_solver = UniformBatchMPCSolver(bnlp; kwargs...) + return solve!(batch_solver) +end + +function IPMOptions( + bnlp::NLPModels.AbstractBatchNLPModel{T}; + kkt_system = MadNLP.SparseKKTSystem, + linear_solver = MadNLP.LDLSolver, + tol = T(1e-8), +) where T + return IPMOptions( + tol = tol, + kkt_system = kkt_system, + linear_solver = linear_solver, + ) +end + +function MadNLP.print_iter(batch_solver::AbstractBatchMPCSolver) + logger = batch_solver.logger + MadNLP.get_level(logger) > MadNLP.INFO && return + ws = batch_solver.workspace + bcnt = batch_solver.batch_cnt + na = batch_solver.kkt.active_batch_size[] + bs = batch_solver.batch_size + k = maximum(bcnt.k) + + mod(k, 10) == 0 && MadNLP.@info(logger, @sprintf( + " iter active max_inf_pr max_inf_du max_inf_compl max_alpha_p")) + MadNLP.@info(logger, @sprintf( + "%4i %3i/%3i %6.2e %6.2e %7.2e %6.2e", + k, na, bs, + maximum(ws.inf_pr), maximum(ws.inf_du), + maximum(ws.inf_compl), maximum(ws.alpha_p), + )) + return +end diff --git a/src/batch/madnlp/callback.jl b/src/batch/madnlp/callback.jl new file mode 100644 index 00000000..145f41eb --- /dev/null +++ b/src/batch/madnlp/callback.jl @@ -0,0 +1,228 @@ +struct UniformBatchCallback{ + T, + VT<:AbstractVector{T}, + MT<:AbstractMatrix{T}, + VI<:AbstractVector{Int}, + BM<:NLPModels.AbstractBatchNLPModel, + FH<:MadNLP.AbstractFixedVariableTreatment, + EH<:MadNLP.AbstractEqualityTreatment, +} <: MadNLP.AbstractCallback{T, VT, FH} + nlp::BM + batch_size::Int + + nvar::Int # per-instance nvar (after fixed variable removal) + ncon::Int + nnzj::Int # per-instance nnzj (after fixed variable removal) + nnzh::Int # per-instance nnzh (after fixed variable removal) + + # Per-instance + con_buffer::MT # ncon × batch_size + jac_buffer::MT # nnzj × batch_size + grad_buffer::MT # nvar_nlp × batch_size + hess_buffer::MT # nnzh × batch_size + + # Shared + jac_I::VI + jac_J::VI + hess_I::VI + hess_J::VI + + # Per-instance + obj_scale::MT # 1 × batch_size per-instance objective scale + obj_sign::MT # 1 × batch_size ±1.0 per instance + con_scale::MT # ncon × batch_size + jac_scale::MT # nnzj × batch_size + + # Shared + fixed_handler::FH + equality_handler::EH + ind_eq::VI + ind_ineq::VI + ind_fixed::VI + ind_lb::VI + ind_ub::VI + ind_llb::VI + ind_uub::VI +end + +function MadNLP.create_sparse_fixed_handler( + ::Type{MadNLP.MakeParameter}, + bnlp::NLPModels.AbstractBatchNLPModel{T}, + jac_I, jac_J, hess_I, hess_J, hess_buffer, +) where T + n = NLPModels.get_nvar(bnlp) + lvar = view(bnlp.meta.lvar, :, 1) + uvar = view(bnlp.meta.uvar, :, 1) + nnzj = NLPModels.get_nnzj(bnlp) + nnzh = NLPModels.get_nnzh(bnlp) + + bs = NLPModels.get_nbatch(bnlp) + x_full = similar(lvar, n * bs) + g_full = similar(lvar, n * bs) + + isfixed = (lvar .== uvar) + isfree = (lvar .< uvar) + + fixed = findall(isfixed) + nfixed = length(fixed) + + if nfixed == 0 + return MadNLP.NoFixedVariables(), n, nnzj, nnzh + end + + free = findall(isfree) + nx = length(free) + map_full_to_free = similar(jac_I, n); fill!(map_full_to_free, -1) + map_full_to_free[free] .= 1:nx + + ind_jac_free = findall(@view(isfree[jac_J])) + ind_hess_free = findall(@view(isfree[hess_I]) .&& @view(isfree[hess_J])) + + nnzh = length(ind_hess_free) + Hi, Hj = similar(hess_I, nnzh), similar(hess_J, nnzh) + copyto!(Hi, map_full_to_free[hess_I[ind_hess_free]]) + copyto!(Hj, map_full_to_free[hess_J[ind_hess_free]]) + resize!(hess_I, nnzh) + resize!(hess_J, nnzh) + copyto!(hess_I, Hi) + copyto!(hess_J, Hj) + + nnzj = length(ind_jac_free) + Ji, Jj = similar(jac_I, nnzj), similar(jac_J, nnzj) + copyto!(Ji, jac_I[ind_jac_free]) + copyto!(Jj, map_full_to_free[jac_J[ind_jac_free]]) + resize!(jac_I, nnzj) + resize!(jac_J, nnzj) + copyto!(jac_I, Ji) + copyto!(jac_J, Jj) + + fixed_handler = MadNLP.MakeParameter( + free, + fixed, + ind_jac_free, + ind_hess_free, + Ref(NaN), + x_full, + g_full, + ) + + return fixed_handler, nx, nnzj, nnzh +end + +function MadNLP.create_callback( + ::Type{UniformBatchCallback{T,VT,MT,VI}}, + bnlp::NLPModels.AbstractBatchNLPModel{T}; + fixed_variable_treatment=MadNLP.MakeParameter, + equality_treatment=MadNLP.EnforceEquality, +) where {T,VT,MT,VI} + bmeta = bnlp.meta + batch_size = bmeta.nbatch + + n = bmeta.nvar + m = bmeta.ncon + nnzj = bmeta.nnzj + nnzh = bmeta.nnzh + + x0 = NLPModels.get_x0(bnlp) + + jac_I = similar(x0, Int, nnzj) + jac_J = similar(x0, Int, nnzj) + hess_I = similar(x0, Int, nnzh) + hess_J = similar(x0, Int, nnzh) + + obj_scale = fill!(similar(x0, 1, batch_size), one(T)) + con_scale = fill!(similar(x0, m, batch_size), one(T)) + con_buffer = fill!(similar(x0, m, batch_size), zero(T)) + jac_buffer = fill!(similar(x0, nnzj, batch_size), zero(T)) + hess_buffer = fill!(similar(x0, nnzh, batch_size), zero(T)) + + if nnzj > 0 + NLPModels.jac_structure!(bnlp, jac_I, jac_J) + end + if nnzh > 0 + NLPModels.hess_structure!(bnlp, hess_I, hess_J) + end + + # FIXME: assert structure? + lvar = view(bmeta.lvar, :, 1) + uvar = view(bmeta.uvar, :, 1) + lcon = view(bmeta.lcon, :, 1) + ucon = view(bmeta.ucon, :, 1) + + fixed_handler, nvar, nnzj, nnzh = MadNLP.create_sparse_fixed_handler( + fixed_variable_treatment, + bnlp, + jac_I, + jac_J, + hess_I, + hess_J, + nothing, # hess_buffer not used + ) + equality_handler = equality_treatment() + + # Allocate with reduced sizes (after fixed var removal) + jac_scale = similar(x0, nnzj, batch_size); fill!(jac_scale, one(T)) + grad_buffer = fill!(similar(x0, nvar, batch_size), zero(T)) + + # Get fixed variables + ind_fixed = findall(lvar .== uvar) + if length(ind_fixed) > 0 && fixed_variable_treatment == MadNLP.MakeParameter + ind_free = findall(lvar .< uvar) + # Remove fixed variables from problem's formulation + lvar = lvar[ind_free] + uvar = uvar[ind_free] + end + + indexes = MadNLP._parse_indexes(lvar, uvar, lcon, ucon, equality_treatment) + + return UniformBatchCallback{T, VT, MT, VI, typeof(bnlp), typeof(fixed_handler), typeof(equality_handler)}( + bnlp, + batch_size, + nvar, + m, + nnzj, + nnzh, + con_buffer, + jac_buffer, + grad_buffer, + hess_buffer, + jac_I, + jac_J, + hess_I, + hess_J, + obj_scale, + fill!(similar(x0, 1, batch_size), bmeta.minimize ? one(T) : -one(T)), # obj_sign + con_scale, + jac_scale, + fixed_handler, + equality_handler, + indexes.ind_eq, + indexes.ind_ineq, + ind_fixed, + indexes.ind_lb, + indexes.ind_ub, + indexes.ind_llb, + indexes.ind_uub, + ) +end + + + +function MadNLP._jac_sparsity_wrapper!(bcb::UniformBatchCallback, I::AbstractVector, J::AbstractVector) + copyto!(I, bcb.jac_I) + copyto!(J, bcb.jac_J) + return +end + +function MadNLP._hess_sparsity_wrapper!(bcb::UniformBatchCallback, I::AbstractVector, J::AbstractVector) + copyto!(I, bcb.hess_I) + copyto!(J, bcb.hess_J) + return +end + +function MadNLP.build_hessian_structure(bcb::UniformBatchCallback, ::Type{<:MadNLP.ExactHessian}) + hess_I = MadNLP.create_array(bcb, Int32, bcb.nnzh) + hess_J = MadNLP.create_array(bcb, Int32, bcb.nnzh) + MadNLP._hess_sparsity_wrapper!(bcb, hess_I, hess_J) + return hess_I, hess_J +end \ No newline at end of file diff --git a/src/batch/madnlp/initialization.jl b/src/batch/madnlp/initialization.jl new file mode 100644 index 00000000..456661df --- /dev/null +++ b/src/batch/madnlp/initialization.jl @@ -0,0 +1,104 @@ +function MadNLP.set_initial_bounds!(xl::AbstractMatrix{T}, xu::AbstractMatrix{T}, tol) where T + if tol > zero(T) + xl .= xl .- max.(one(T), abs.(xl)) .* tol + xu .= xu .+ max.(one(T), abs.(xu)) .* tol + end +end + +function MadNLP.initialize!( + bcb::UniformBatchCallback{T}, + x, + xl, + xu, + y, + rhs, + ind_ineq; + tol=1e-8, + bound_push=1e-2, + bound_fac=1e-2, +) where T + x0 = MadNLP.variable(x) + lvar = MadNLP.variable(xl) + uvar = MadNLP.variable(xu) + + x0 .= MadNLP.get_x0(bcb) + lvar .= MadNLP.get_lvar(bcb) + uvar .= MadNLP.get_uvar(bcb) + y .= MadNLP.get_y0(bcb) + lcon = copy(MadNLP.get_lcon(bcb)) + ucon = copy(MadNLP.get_ucon(bcb)) + + MadNLP._treat_equality_initialize!(bcb.equality_handler, lcon, ucon, tol) + MadNLP._treat_fixed_variable_initialize!(bcb, x0, lvar, uvar) # this is no-op, recall bcb <: MadNLP.AbstractCallback. will need to be adjust for batchdensecallback + + MadNLP.set_initial_bounds!(lvar, uvar, tol) + x0 .= MadNLP._initialize_variables!.(x0, lvar, uvar, bound_push, bound_fac) + + x_full = MadNLP._update_x!(bcb, x0) + MadNLP._eval_cons_wrapper!(bcb, x_full, bcb.con_buffer) + + MadNLP.slack(xl) .= view(lcon, ind_ineq, :) + MadNLP.slack(xu) .= view(ucon, ind_ineq, :) + rhs .= (lcon .== ucon) .* lcon + copyto!(MadNLP.slack(x), view(bcb.con_buffer, ind_ineq, :)) + + MadNLP.set_initial_bounds!(MadNLP.slack(xl), MadNLP.slack(xu), tol) + MadNLP.slack(x) .= MadNLP._initialize_variables!.(MadNLP.slack(x), MadNLP.slack(xl), MadNLP.slack(xu), bound_push, bound_fac) + + return +end + +function MadNLP.set_con_scale_sparse!( + con_scale::MT, + jac_I, + jac_buffer, + max_gradient, +) where {T,MT<:AbstractMatrix{T}} + fill!(con_scale, one(T)) + MadNLP._set_con_scale_sparse!(con_scale, jac_I, jac_buffer) + con_scale .= min.(one(T), max_gradient ./ con_scale) + return con_scale +end + +function MadNLP._set_con_scale_sparse!(con_scale::MT, jac_I, jac_buffer) where {T,MT<:AbstractMatrix{T}} + nnzj = length(jac_I) + bs = size(jac_buffer, 2) + @inbounds for k in 1:nnzj + row = jac_I[k] + for j in 1:bs + con_scale[row, j] = max(con_scale[row, j], abs(jac_buffer[k, j])) + end + end + return con_scale +end + +function MadNLP.set_jac_scale_sparse!(jac_scale::MT, con_scale, jac_I) where {T,MT<:AbstractMatrix{T}} + return copyto!(jac_scale, @view(con_scale[jac_I, :])) +end + +function MadNLP.set_obj_scale!(obj_scale, F::MT, max_gradient) where {T,MT<:AbstractMatrix{T}} + return obj_scale .= min.(one(T), max_gradient ./ maximum(abs, F; dims=1)) +end + +function MadNLP.set_scaling!( + cb::UniformBatchCallback, + x, xl, xu, y, rhs, ind_ineq, nlp_scaling_max_gradient, +) + x0 = MadNLP.variable(x) + x_full = MadNLP._update_x!(cb, x0) + + jac_free = MadNLP._eval_jac_wrapper!(cb, x_full, cb.jac_buffer) + MadNLP.set_con_scale_sparse!(cb.con_scale, cb.jac_I, jac_free, nlp_scaling_max_gradient) + MadNLP.set_jac_scale_sparse!(cb.jac_scale, cb.con_scale, cb.jac_I) + + MadNLP._eval_grad_f_wrapper!(cb, x_full, cb.grad_buffer) + MadNLP.set_obj_scale!(cb.obj_scale, cb.grad_buffer, nlp_scaling_max_gradient) + + con_scale_slk = @view(cb.con_scale[ind_ineq, :]) + y ./= cb.con_scale + rhs .*= cb.con_scale + MadNLP.slack(x) .*= con_scale_slk + MadNLP.slack(xl) .*= con_scale_slk + MadNLP.slack(xu) .*= con_scale_slk + return +end diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl new file mode 100644 index 00000000..809d9fb0 --- /dev/null +++ b/src/batch/madnlp/kernels.jl @@ -0,0 +1,17 @@ +_adjust_bound_lb(x_lr::T, xl_r, c1, c2) where T = + x_lr - xl_r < c1 ? xl_r - c2 * max(one(T), abs(x_lr)) : xl_r +_adjust_bound_ub(x_ur::T, xu_r, c1, c2) where T = + xu_r - x_ur < c1 ? xu_r + c2 * max(one(T), abs(x_ur)) : xu_r + +function MadNLP.adjust_boundary!( + x_lr::AbstractMatrix{T}, + xl_r::AbstractMatrix{T}, + x_ur::AbstractMatrix{T}, + xu_r::AbstractMatrix{T}, + mu, +) where T + c1 = eps(T) .* mu + c2 = T(eps(T)^(3/4)) + xl_r .= _adjust_bound_lb.(x_lr, xl_r, c1, c2) + xu_r .= _adjust_bound_ub.(x_ur, xu_r, c1, c2) +end diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl new file mode 100644 index 00000000..b8d916f4 --- /dev/null +++ b/src/batch/madnlp/kkt.jl @@ -0,0 +1,377 @@ +abstract type AbstractBatchKKTSystem{T} end + +"""Batched version of SparseKKTSystem""" +struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} + nzVals::MT # (aug_mat_length × batch_size) COO nonzero values + aug_I::VI32 # shared row indices + aug_J::VI32 # shared column indices + batch_solver::LS # batched linear solver + rhs_buffer::VT # contiguous (n+m)*batch_size for batch solve + batch_size::Int + aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values + batch_csc_map::VI # flattened COO→CSC map for all instances + n_tot::Int # n + n_slack (total primal variables) + m::Int # number of constraints + nnzh::Int # number of Hessian nonzeros + # Diagonal and bound data (for _kktmul!) + reg::MT # (n_tot × batch_size) primal regularization + l_diag::MT # (nlb × batch_size) lower bound diagonals + u_diag::MT # (nub × batch_size) upper bound diagonals + l_lower::MT # (nlb × batch_size) lower bound multipliers + u_lower::MT # (nub × batch_size) upper bound multipliers + # Hessian scatter (for mul!) + hess_scatter::SMT # (n_tot × n_hess_sym) scatter matrix + hess_nz_map::VI # nzVals row indices (with sym duplication) + hess_var_map::VI # variable indices for x multiplication + hess_buffer::MT # (n_hess_sym × batch_size) workspace + # J^T scatter (for jtprod! and mul!) + jt_scatter::SMT # (n_tot × n_jac_aug) scatter: S[var_idx, k] = 1 + jt_nz_map::VI # nzVals row indices for Jacobian entries + jt_con_map::VI # maps each Jac nonzero to its constraint index + jt_buffer::MT # (n_jac_aug × batch_size) buffer for jtprod + # J scatter (for mul!) + j_scatter::SMT # (m × n_jac_aug) scatter: S[con_idx, k] = 1 + j_var_map::VI # variable indices for J entries + j_buffer::MT # (n_jac_aug × batch_size) buffer for jprod + # Workspace for mul! (GPU needs full matrices, not SubArray views) + _mul_w_primal::MT # (n_tot × batch_size) + _mul_w_dual::MT # (m × batch_size) + # Batch tracking + batch_map::Vector{Int} # original index → active position (0 if inactive) + batch_map_rev::Vector{Int} # active position → original index + active_batch_size::Base.RefValue{Int} +end + +pr_diag(bkkt::SparseUniformBatchKKTSystem) = view(bkkt.nzVals, 1:bkkt.n_tot, :) +function du_diag(bkkt::SparseUniformBatchKKTSystem) + du_off = size(bkkt.nzVals, 1) - bkkt.m + return view(bkkt.nzVals, du_off+1:du_off+bkkt.m, :) +end + +function MadNLP.create_kkt_system( + ::Type{MadNLP.SparseKKTSystem}, + bcb::UniformBatchCallback{T, VT}, + uniformbatch_linear_solver = LoopedBatchLinearSolver; + opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), +) where {T, VT} + batch_size = bcb.batch_size + + n_slack = length(bcb.ind_ineq) + n = bcb.nvar + m = bcb.ncon + jac_sparsity_I = MadNLP.create_array(bcb, Int32, bcb.nnzj) + jac_sparsity_J = MadNLP.create_array(bcb, Int32, bcb.nnzj) + MadNLP._jac_sparsity_wrapper!(bcb, jac_sparsity_I, jac_sparsity_J) + + hess_sparsity_I, hess_sparsity_J = MadNLP.build_hessian_structure(bcb, MadNLP.ExactHessian) + + nlb = length(bcb.ind_lb) + nub = length(bcb.ind_ub) + + MadNLP.force_lower_triangular!(hess_sparsity_I, hess_sparsity_J) + + ind_ineq = bcb.ind_ineq + + n_slack = length(ind_ineq) + n_jac = length(jac_sparsity_I) + n_hess = length(hess_sparsity_I) + n_tot = n + n_slack + + aug_vec_length = n_tot+m + aug_mat_length = n_tot+m+n_hess+n_jac+n_slack + + I = MadNLP.create_array(bcb, Int32, aug_mat_length) + J = MadNLP.create_array(bcb, Int32, aug_mat_length) + nzVals = similar(bcb.con_buffer, aug_mat_length, batch_size) + fill!(nzVals, zero(T)) + V = _madnlp_unsafe_column_wrap(nzVals, aug_mat_length, 1, VT) + + offset = n_tot+n_jac+n_slack+n_hess+m + + I[1:n_tot] .= 1:n_tot + I[n_tot+1:n_tot+n_hess] = hess_sparsity_I + I[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= (jac_sparsity_I.+n_tot) + I[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= ind_ineq .+ n_tot + I[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) + + J[1:n_tot] .= 1:n_tot + J[n_tot+1:n_tot+n_hess] = hess_sparsity_J + J[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= jac_sparsity_J + J[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= (n+1:n+n_slack) + J[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) + + aug_raw = MadNLP.SparseMatrixCOO(aug_vec_length, aug_vec_length, I, J, V) + aug_com, aug_csc_map = MadNLP.coo_to_csc(aug_raw) + + nnz_csc = SparseArrays.nnz(aug_com) + aug_com_nzvals = similar(nzVals, nnz_csc, batch_size) + fill!(aug_com_nzvals, zero(T)) + + csc_offsets = similar(aug_csc_map, 1, batch_size) + csc_offsets .= (0:batch_size-1)' .* nnz_csc + batch_csc_map = vec(aug_csc_map .+ csc_offsets) + + batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, aug_vec_length; opt=opt_linear_solver) + + rhs_buffer = VT(undef, aug_vec_length * batch_size) + fill!(rhs_buffer, zero(T)) + + jac_range = n_tot+n_hess+1:n_tot+n_hess+n_jac+n_slack + + hess_scatter, hess_nz_map, hess_var_map, hess_buffer = _build_hess_scatter( + I, J, n_tot, n_hess, nzVals, aug_csc_map, batch_size, + ) + jt_scatter, jt_nz_map, jt_con_map, jt_buffer = _build_scatter( + I, J, jac_range, n_tot, nzVals, aug_csc_map, batch_size, + ) + j_scatter, _, j_var_map, j_buffer = _build_jac_scatter( + I, J, jac_range, n_tot, m, nzVals, aug_csc_map, batch_size, + ) + + reg = similar(nzVals, n_tot, batch_size) + l_diag = similar(nzVals, nlb, batch_size) + u_diag = similar(nzVals, nub, batch_size) + l_lower = similar(nzVals, nlb, batch_size) + u_lower = similar(nzVals, nub, batch_size) + + _mul_w_primal = similar(nzVals, n_tot, batch_size) + _mul_w_dual = similar(nzVals, m, batch_size) + + batch_map = collect(1:batch_size) + batch_map_rev = collect(1:batch_size) + active_batch_size = Ref(batch_size) + + LS = typeof(batch_ls) + MT = typeof(nzVals) + VI = typeof(aug_csc_map) + VI32 = typeof(I) + SMT = typeof(jt_scatter) + return SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT}( + nzVals, I, J, batch_ls, rhs_buffer, batch_size, + aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, + reg, l_diag, u_diag, l_lower, u_lower, + hess_scatter, hess_nz_map, hess_var_map, hess_buffer, + jt_scatter, jt_nz_map, jt_con_map, jt_buffer, + j_scatter, j_var_map, j_buffer, + _mul_w_primal, _mul_w_dual, + batch_map, batch_map_rev, active_batch_size, + ) +end + +function update_active_set!(bkkt::SparseUniformBatchKKTSystem, status::Vector{MadNLP.Status}) + active_pos = 0 + for i in 1:bkkt.batch_size + if status[i] == MadNLP.REGULAR + active_pos += 1 + bkkt.batch_map[i] = active_pos + bkkt.batch_map_rev[active_pos] = i + else + bkkt.batch_map[i] = 0 + end + end + for j in (active_pos + 1):bkkt.batch_size + bkkt.batch_map_rev[j] = 0 + end + bkkt.active_batch_size[] = active_pos +end + +function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) + na = bkkt.active_batch_size[] + nzvals = bkkt.aug_com_nzvals + @inbounds for j in 1:na + i = bkkt.batch_map_rev[j] + i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) + end + _active_factorize!(bkkt.batch_solver, na) + return +end + +function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractVector) where T + na = bkkt.active_batch_size[] + bs = bkkt.batch_size + n = length(rhs) ÷ bs + + rhs_mat = reshape(rhs, n, bs) + @inbounds for j in 1:na + i = bkkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) + end + _active_solve!(bkkt.batch_solver, rhs, na, n) + + @inbounds for j in na:-1:1 + i = bkkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) + end + + @inbounds for i in 1:bs + bkkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) + end + return rhs +end + +function MadNLP.reduce_rhs!(bkkt::SparseUniformBatchKKTSystem, d::BatchUnreducedKKTVector) + MadNLP.reduce_rhs!( + xp_lr(d), MadNLP.dual_lb(d), bkkt.l_diag, + xp_ur(d), MadNLP.dual_ub(d), bkkt.u_diag, + ) + return +end + +function MadNLP.finish_aug_solve!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + zl, zu = batch_solver.zl, batch_solver.zu + dzl = MadNLP.dual_lb(d) + dzu = MadNLP.dual_ub(d) + dzl .= (.-dzl .+ lower(zl) .* xp_lr(d)) ./ bkkt.l_diag + dzu .= (dzu .- upper(zu) .* xp_ur(d)) ./ bkkt.u_diag + return +end + +function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + n_pd = d.n + d.m + bs = bkkt.batch_size + na = bkkt.active_batch_size[] + + MadNLP.reduce_rhs!(bkkt, d) + + if na == 1 + # only one active: skip packing + orig_col = bkkt.batch_map_rev[1] + _active_solve!(bkkt.batch_solver, d.views[orig_col], 1, n_pd) + else + rhs = bkkt.rhs_buffer + pd_view = MadNLP.primal_dual(d) + copyto!(reshape(rhs, n_pd, bs), pd_view) + MadNLP.solve_linear_system!(bkkt, rhs) + copyto!(pd_view, reshape(rhs, n_pd, bs)) + end + + MadNLP.finish_aug_solve!(bkkt, batch_solver) + return +end + +function MadNLP.build_kkt!(bkkt::SparseUniformBatchKKTSystem) + MadNLP._transfer!(vec(bkkt.aug_com_nzvals), vec(bkkt.nzVals), bkkt.batch_csc_map) + return +end + +function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) + MadNLP.@trace(batch_solver.logger, "Factorization started.") + MadNLP.build_kkt!(batch_solver.kkt) + batch_solver.batch_cnt.linear_solver_time[] += @elapsed MadNLP.factorize_kkt!(batch_solver.kkt) + return +end + +function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) + bkkt.jt_buffer .= view(bkkt.nzVals, bkkt.jt_nz_map, :) .* MadNLP.full(y)[bkkt.jt_con_map, :] + mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) + return res +end + +function MadNLP.jtprod!(jacl::BatchVector, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) + return MadNLP.jtprod!(MadNLP.full(jacl), bkkt, y) +end + +function MadNLP.eval_jac_wrapper!( + batch_solver::AbstractBatchMPCSolver, + bkkt::SparseUniformBatchKKTSystem, +) + bcb = batch_solver.bcb + ws = batch_solver.workspace + nzVals = bkkt.nzVals + n_tot = bkkt.n_tot + nnzj = bcb.nnzj + n_slack = length(bcb.ind_ineq) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + jac_free = MadNLP._eval_jac_wrapper!(bcb, ws.bx, bcb.jac_buffer) + + jac_offset = n_tot + bkkt.nnzh + view(nzVals, jac_offset+1:jac_offset+nnzj, :) .= jac_free + + if n_slack > 0 + view(nzVals, jac_offset+nnzj+1:jac_offset+nnzj+n_slack, :) .= -one(eltype(nzVals)) + end + return +end + +function MadNLP.eval_lag_hess_wrapper!( + batch_solver::AbstractBatchMPCSolver, + bkkt::SparseUniformBatchKKTSystem, +) + bcb = batch_solver.bcb + ws = batch_solver.workspace + nzVals = bkkt.nzVals + n_tot = bkkt.n_tot + nnzh = bkkt.nnzh + + if nnzh > 0 + hess = view(nzVals, n_tot+1:n_tot+nnzh, :) + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess) + end + return +end + +function MadNLP.initialize!(bkkt::SparseUniformBatchKKTSystem{T}) where T + pr_diag(bkkt) .= one(T) + if bkkt.nnzh > 0 + n_tot = bkkt.n_tot + view(bkkt.nzVals, n_tot+1:n_tot+bkkt.nnzh, :) .= zero(T) + end + du_diag(bkkt) .= zero(T) + + fill!(bkkt.reg, zero(T)) + fill!(bkkt.l_diag, one(T)) + fill!(bkkt.u_diag, one(T)) + fill!(bkkt.l_lower, zero(T)) + fill!(bkkt.u_lower, zero(T)) + + fill!(bkkt.aug_com_nzvals, zero(T)) + return +end + +@inbounds function _kktmul!( + w::BatchUnreducedKKTVector, + x::BatchUnreducedKKTVector, + reg, du_diag_val, l_lower, u_lower, l_diag, u_diag, + alpha, beta, +) + MadNLP.primal(w) .+= alpha .* reg .* MadNLP.primal(x) + MadNLP.dual(w) .+= alpha .* du_diag_val .* MadNLP.dual(x) + xp_lr(w) .-= alpha .* MadNLP.dual_lb(x) + xp_ur(w) .+= alpha .* MadNLP.dual_ub(x) + MadNLP.dual_lb(w) .= beta .* MadNLP.dual_lb(w) .+ alpha .* (xp_lr(x) .* l_lower .- MadNLP.dual_lb(x) .* l_diag) + MadNLP.dual_ub(w) .= beta .* MadNLP.dual_ub(w) .+ alpha .* (xp_ur(x) .* u_lower .+ MadNLP.dual_ub(x) .* u_diag) + return +end + +function LinearAlgebra.mul!( + w::BatchUnreducedKKTVector{T}, + bkkt::SparseUniformBatchKKTSystem{T}, + x::BatchUnreducedKKTVector{T}, + alpha = one(T), + beta = zero(T), +) where T + nzV = bkkt.nzVals + wp = bkkt._mul_w_primal + wd = bkkt._mul_w_dual + + # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) + bkkt.hess_buffer .= view(nzV, bkkt.hess_nz_map, :) .* MadNLP.primal(x)[bkkt.hess_var_map, :] + mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) + MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp + + # mul!(primal(w), jac_com', dual(x), alpha, one(T)) + bkkt.jt_buffer .= view(nzV, bkkt.jt_nz_map, :) .* MadNLP.dual(x)[bkkt.jt_con_map, :] + mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) + MadNLP.primal(w) .+= alpha .* wp + + # mul!(dual(w), jac_com, primal(x), alpha, beta) + bkkt.j_buffer .= view(nzV, bkkt.jt_nz_map, :) .* MadNLP.primal(x)[bkkt.j_var_map, :] + mul!(wd, bkkt.j_scatter, bkkt.j_buffer) + MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd + _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) + return w +end diff --git a/src/batch/madnlp/linear_solver.jl b/src/batch/madnlp/linear_solver.jl new file mode 100644 index 00000000..19730c4e --- /dev/null +++ b/src/batch/madnlp/linear_solver.jl @@ -0,0 +1,54 @@ +# FIXME: threads/polyester version + +struct LoopedBatchLinearSolver{T, LS<:MadNLP.AbstractLinearSolver{T}} <: MadNLP.AbstractLinearSolver{T} + solvers::Vector{LS} + batch_size::Int +end + +function LoopedBatchLinearSolver(solvers::Vector{LS}) where {T, LS<:MadNLP.AbstractLinearSolver{T}} + return LoopedBatchLinearSolver{T, LS}(solvers, length(solvers)) +end + +@kwdef mutable struct LoopedBatchLinearSolverOptions <: MadNLP.AbstractOptions + looped_linear_solver::Type = MadNLP.MumpsSolver +end + +MadNLP.default_options(::Type{LoopedBatchLinearSolver}) = LoopedBatchLinearSolverOptions() + +function LoopedBatchLinearSolver( + aug_com, + nzvals_mat::AbstractMatrix{T}, + n::Int; + opt::LoopedBatchLinearSolverOptions = LoopedBatchLinearSolverOptions(), +) where T + linear_solver = opt.looped_linear_solver + per_instance_opt = MadNLP.default_options(linear_solver) + batch_size = size(nzvals_mat, 2) + nnz_csc = size(nzvals_mat, 1) + VT = typeof(similar(nzvals_mat, T, 0)) + individual_solvers = map(1:batch_size) do i + nzval_i = _madnlp_unsafe_column_wrap(nzvals_mat, nnz_csc, (i - 1) * nnz_csc + 1, VT) + csc_i = _csc_with_nzval(aug_com, nzval_i, n) + linear_solver(csc_i; opt=per_instance_opt) + end + LoopedBatchLinearSolver(individual_solvers) +end + +function is_factorized(batch_linear_solver::LoopedBatchLinearSolver) + return all(is_factorized(s) for s in batch_linear_solver.solvers) +end + +function _active_factorize!(s::LoopedBatchLinearSolver, na::Int) + for j in 1:na + MadNLP.factorize!(s.solvers[j]) + end + return +end + +function _active_solve!(s::LoopedBatchLinearSolver, rhs::AbstractVector, na::Int, n::Int) + for j in 1:na + xj = MadNLP._madnlp_unsafe_wrap(rhs, n, (j-1)*n + 1) + MadNLP.solve_linear_system!(s.solvers[j], xj) + end + return +end diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl new file mode 100644 index 00000000..8ab9bf21 --- /dev/null +++ b/src/batch/madnlp/nlpmodels.jl @@ -0,0 +1,202 @@ +const _DefaultFH = Union{MadNLP.NoFixedVariables, MadNLP.RelaxBound} + +function MadNLP._update_x!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + return x +end + +function MadNLP._update_x!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + fh = bcb.fixed_handler + nvar_nlp = bcb.nlp.meta.nvar + bs = bcb.batch_size + BX = reshape(fh.x_full, nvar_nlp, bs) + view(BX, fh.fixed, :) .= view(bcb.nlp.meta.lvar, fh.fixed, :) + view(BX, fh.free, :) .= x + return BX +end + +function MadNLP.unpack_x!( + X_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x::BatchPrimalVector, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + X_full .= MadNLP.variable(x) +end + +function MadNLP.unpack_x!( + X_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x::BatchPrimalVector, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + fh = bcb.fixed_handler + X_full[fh.free, :] .= MadNLP.variable(x) + X_full[fh.fixed, :] .= view(bcb.nlp.meta.lvar, fh.fixed, :) +end + +function MadNLP.unpack_z!( + Z_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, z_free, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + Z_full .= z_free ./ bcb.obj_scale +end + +function MadNLP.unpack_z!( + Z_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, z_free, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + Z_full[bcb.fixed_handler.free, :] .= z_free ./ bcb.obj_scale +end + +function MadNLP.get_x0(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + view(bcb.nlp.meta.x0, bcb.fixed_handler.free, :) +end + +MadNLP.get_y0(bcb::UniformBatchCallback) = bcb.nlp.meta.y0 + +MadNLP.get_lvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} = bcb.nlp.meta.lvar +function MadNLP.get_lvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + view(bcb.nlp.meta.lvar, bcb.fixed_handler.free, :) +end + +MadNLP.get_uvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} = bcb.nlp.meta.uvar +function MadNLP.get_uvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + view(bcb.nlp.meta.uvar, bcb.fixed_handler.free, :) +end + +MadNLP.get_lcon(bcb::UniformBatchCallback) = bcb.nlp.meta.lcon +MadNLP.get_ucon(bcb::UniformBatchCallback) = bcb.nlp.meta.ucon + +function MadNLP.unpack_y!(y_full, bcb::UniformBatchCallback, y) + y_full .= y .* bcb.con_scale .* (bcb.obj_sign ./ bcb.obj_scale) +end + +MadNLP.unpack_obj(bcb::UniformBatchCallback, obj_val) = vec(bcb.obj_sign .* obj_val ./ bcb.obj_scale) + +function MadNLP.unpack_cons!(c_full, bcb::UniformBatchCallback, c, rhs, ind_ineq, slack) + c_full .= c ./ bcb.con_scale .+ rhs + if length(ind_ineq) > 0 + view(c_full, ind_ineq, :) .+= slack + end +end + +function MadNLP._eval_f_wrapper(bcb::UniformBatchCallback, bx::AbstractMatrix, bf::AbstractVector) + NLPModels.obj!(bcb.nlp, bx, bf) + bf .*= vec(bcb.obj_scale) + return bf +end + +function MadNLP._eval_cons_wrapper!(bcb::UniformBatchCallback, bx::AbstractMatrix, bc_mat::AbstractMatrix) + NLPModels.cons!(bcb.nlp, bx, bc_mat) + bc_mat .*= bcb.con_scale + return bc_mat +end + +function MadNLP._eval_grad_f_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, bg::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + NLPModels.grad!(bcb.nlp, bx, bg) + bg .*= bcb.obj_scale + return bg +end + +function MadNLP._eval_grad_f_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, bg::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + fh = bcb.fixed_handler + nvar_nlp = bcb.nlp.meta.nvar + bs = bcb.batch_size + GF = reshape(fh.g_full, nvar_nlp, bs) + NLPModels.grad!(bcb.nlp, bx, GF) + view(bg, 1:bcb.nvar, :) .= view(GF, fh.free, :) .* bcb.obj_scale + return bg +end + +function MadNLP._eval_jac_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, jac_buffer::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + NLPModels.jac_coord!(bcb.nlp, bx, jac_buffer) + jac_buffer .*= bcb.jac_scale + return jac_buffer +end + +function MadNLP._eval_jac_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, jac_buffer::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + NLPModels.jac_coord!(bcb.nlp, bx, jac_buffer) + jac_free = view(jac_buffer, bcb.fixed_handler.ind_jac_free, :) + jac_free .*= bcb.jac_scale + return jac_free +end + +function MadNLP._eval_lag_hess_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, + bx::AbstractMatrix, + y_mat::AbstractMatrix, + bv::AbstractMatrix, + hess::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} + bv .= y_mat .* bcb.con_scale + NLPModels.hess_coord!(bcb.nlp, bx, bv, vec(bcb.obj_scale), hess) + return +end + +function MadNLP._eval_lag_hess_wrapper!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, + bx::AbstractMatrix, + y_mat::AbstractMatrix, + bv::AbstractMatrix, + hess::AbstractMatrix, +) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + bv .= y_mat .* bcb.con_scale + NLPModels.hess_coord!(bcb.nlp, bx, bv, vec(bcb.obj_scale), bcb.hess_buffer) + hess .= view(bcb.hess_buffer, bcb.fixed_handler.ind_hess_free, :) + return +end + +function MadNLP.eval_f_wrapper(solver::AbstractBatchMPCSolver, bx::AbstractMatrix) + ws = solver.workspace + bcb = solver.bcb + bcnt = solver.batch_cnt + + t = @elapsed begin + MadNLP._eval_f_wrapper(bcb, bx, ws.bf) + ws.bf .*= vec(bcb.obj_sign) + vec(ws.obj_val) .= ws.bf + end + bcnt.eval_function_time[] += t + bcnt.obj_cnt[] += 1 + return +end + +function MadNLP.eval_cons_wrapper!(solver::AbstractBatchMPCSolver, bx::AbstractMatrix) + ws = solver.workspace + bcb = solver.bcb + bcnt = solver.batch_cnt + ind_ineq = bcb.ind_ineq + ns = length(ind_ineq) + + t = @elapsed begin + MadNLP._eval_cons_wrapper!(bcb, bx, MadNLP.full(solver.c)) + if ns > 0 + view(MadNLP.full(solver.c), ind_ineq, :) .-= MadNLP.slack(solver.x) + end + MadNLP.full(solver.c) .-= MadNLP.full(solver.rhs) + end + bcnt.eval_function_time[] += t + bcnt.con_cnt[] += 1 + return +end + +function MadNLP.eval_grad_f_wrapper!(solver::AbstractBatchMPCSolver, bx::AbstractMatrix) + ws = solver.workspace + bcb = solver.bcb + bcnt = solver.batch_cnt + nvar = bcb.nvar + + t = @elapsed begin + MadNLP._eval_grad_f_wrapper!(bcb, bx, ws.bg) + BG = view(ws.bg, 1:nvar, :) + BG .*= bcb.obj_sign + copyto!(MadNLP.variable(solver.f), BG) + end + bcnt.eval_function_time[] += t + bcnt.obj_grad_cnt[] += 1 + return +end diff --git a/src/batch/madnlp/rhs.jl b/src/batch/madnlp/rhs.jl new file mode 100644 index 00000000..72f0be8b --- /dev/null +++ b/src/batch/madnlp/rhs.jl @@ -0,0 +1,71 @@ +struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} + values::MT + views::Vector{VT} + n::Int + m::Int + nlb::Int + nub::Int + ind_lb::VI + ind_ub::VI +end + +function BatchUnreducedKKTVector( + ::Type{MT}, ::Type{VT}, + n::Int, m::Int, nlb::Int, nub::Int, batch_size::Int, + ind_lb, ind_ub, +) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + total = n + m + nlb + nub + values = MT(undef, total, batch_size) + fill!(values, zero(T)) + + views = Vector{VT}(undef, batch_size) + for i in 1:batch_size + col_start = (i-1) * total + 1 + views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) + end + + return BatchUnreducedKKTVector{T, MT, VT, typeof(ind_lb)}(values, views, n, m, nlb, nub, ind_lb, ind_ub) +end + +MadNLP.full(bv::BatchUnreducedKKTVector) = bv.values +MadNLP.primal(bv::BatchUnreducedKKTVector) = view(bv.values, 1:bv.n, :) +MadNLP.dual(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+1:bv.n+bv.m, :) +MadNLP.primal_dual(bv::BatchUnreducedKKTVector) = view(bv.values, 1:bv.n+bv.m, :) +MadNLP.dual_lb(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+bv.m+1:bv.n+bv.m+bv.nlb, :) +MadNLP.dual_ub(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+bv.m+bv.nlb+1:bv.n+bv.m+bv.nlb+bv.nub, :) +xp_lr(bv::BatchUnreducedKKTVector) = view(bv.values, bv.ind_lb, :) +xp_ur(bv::BatchUnreducedKKTVector) = view(bv.values, bv.ind_ub, :) + +struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} + values::MT + views::Vector{VT} + nx::Int + ns::Int + ind_lb::VI + ind_ub::VI +end + +function BatchPrimalVector( + ::Type{MT}, ::Type{VT}, + nx::Int, ns::Int, batch_size::Int, + ind_lb, ind_ub, +) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + total = nx + ns + values = MT(undef, total, batch_size) + fill!(values, zero(T)) + + views = Vector{VT}(undef, batch_size) + for i in 1:batch_size + col_start = (i-1) * total + 1 + views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) + end + + return BatchPrimalVector{T, MT, VT, typeof(ind_lb)}(values, views, nx, ns, ind_lb, ind_ub) +end + +MadNLP.variable(bpv::BatchPrimalVector) = view(bpv.values, 1:bpv.nx, :) +MadNLP.slack(bpv::BatchPrimalVector) = view(bpv.values, bpv.nx+1:bpv.nx+bpv.ns, :) +lower(bpv::BatchPrimalVector) = view(bpv.values, bpv.ind_lb, :) +upper(bpv::BatchPrimalVector) = view(bpv.values, bpv.ind_ub, :) +MadNLP.full(bpv::BatchPrimalVector) = bpv.values +MadNLP.primal(bpv::BatchPrimalVector) = bpv.values diff --git a/src/batch/structure.jl b/src/batch/structure.jl new file mode 100644 index 00000000..84960ce9 --- /dev/null +++ b/src/batch/structure.jl @@ -0,0 +1,207 @@ +struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI<:AbstractMatrix{Int32}} + alpha_xl::MT + alpha_xu::MT + alpha_zl::MT + alpha_zu::MT + alpha_p::MT + alpha_d::MT + idx_xl::MI + idx_xu::MI + idx_zl::MI + idx_zu::MI + tau::MT + + mu_batch::MT + mu_curr::MT + mu_affine::MT + sum_lb::MT + sum_ub::MT + + obj_val::MT + norm_b::MT + norm_c::MT + inf_pr::MT + inf_du::MT + inf_compl::MT + best_complementarity::MT + dual_obj::MT + status::Vector{MadNLP.Status} + term_converged::Vector{Bool} + term_infeasible::Vector{Bool} + term_diverging::Vector{Bool} + + active_mask::MT + + bx::MT + bf::VT + bg::MT + bv::MT +end + +function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch_size::Int; + nvar_nlp::Int=0, ncon::Int=0) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + _proto = MT(undef, 1, batch_size) + MI = typeof(similar(_proto, Int32)) + return UniformBatchWorkspace{T, VT, MT, MI}( + MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_xl, alpha_xu + MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_zl, alpha_zu + MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_p, alpha_d + similar(_proto, Int32), similar(_proto, Int32), # idx_xl, idx_xu + similar(_proto, Int32), similar(_proto, Int32), # idx_zl, idx_zu + MT(undef, 1, batch_size), # tau + MT(undef, 1, batch_size), MT(undef, 1, batch_size), # mu_batch, mu_curr + MT(undef, 1, batch_size), MT(undef, 1, batch_size), # mu_affine, sum_lb + MT(undef, 1, batch_size), # sum_ub + MT(undef, 1, batch_size), # obj_val + MT(undef, 1, batch_size), # norm_b + MT(undef, 1, batch_size), # norm_c + MT(undef, 1, batch_size), # inf_pr + MT(undef, 1, batch_size), # inf_du + MT(undef, 1, batch_size), # inf_compl + MT(undef, 1, batch_size), # best_complementarity + MT(undef, 1, batch_size), # dual_obj + fill(MadNLP.INITIAL, batch_size), # status + falses(batch_size), # term_converged + falses(batch_size), # term_infeasible + falses(batch_size), # term_diverging + fill!(MT(undef, 1, batch_size), one(T)), # active_mask + MT(undef, nvar_nlp, batch_size), # bx + VT(undef, batch_size), # bf + MT(undef, nvar_nlp, batch_size), # bg + MT(undef, ncon, batch_size), # bv + ) +end + +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, MI} <: AbstractBatchMPCSolver{T} + batch_size::Int + + d::BatchUnreducedKKTVector{T, MT, VT} + p::BatchUnreducedKKTVector{T, MT, VT} + _w1::BatchUnreducedKKTVector{T, MT, VT} + _w2::BatchUnreducedKKTVector{T, MT, VT} + + x::BatchPrimalVector{T, MT, VT} + xl::BatchPrimalVector{T, MT, VT} + xu::BatchPrimalVector{T, MT, VT} + zl::BatchPrimalVector{T, MT, VT} + zu::BatchPrimalVector{T, MT, VT} + f::BatchPrimalVector{T, MT, VT} + + y::BatchVector{T, MT, VT} + c::BatchVector{T, MT, VT} + jacl::BatchVector{T, MT, VT} + rhs::BatchVector{T, MT, VT} + correction_lb::BatchVector{T, MT, VT} + correction_ub::BatchVector{T, MT, VT} + + workspace::UniformBatchWorkspace{T, VT, MT, MI} + + opt::IPMOptions + batch_cnt::BatchCounters + logger::MadNLP.MadNLPLogger + kkt::AbstractBatchKKTSystem{T} + + del_w::MT + del_c::MT + + nlp::BM + bcb::BCB +end + +_get_ind_lb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_lb +_get_ind_ub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_ub +_get_ind_llb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_llb +_get_ind_uub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_uub + + +""" + UniformBatchMPCSolver(bnlp::AbstractBatchNLPModel; linear_solver, kwargs...) + +Construct a batch solver from a `AbstractBatchNLPModel`. +""" +function UniformBatchMPCSolver( + bnlp::NLPModels.AbstractBatchNLPModel{T}; + MT = typeof(similar(NLPModels.get_x0(bnlp), T, 0, 0)), + VT = typeof(similar(NLPModels.get_x0(bnlp), T, 0)), + VI = typeof(similar(NLPModels.get_x0(bnlp), Int, 0)), + uniformbatch_linear_solver = LoopedBatchLinearSolver, + kwargs..., +) where {T} + bmeta = bnlp.meta + batch_size = bmeta.nbatch + @assert batch_size > 0 "Need at least one instance in batch" + + nvar_nlp = bmeta.nvar + + opt_batch_ls = MadNLP.default_options(uniformbatch_linear_solver) + remaining_kwargs = MadNLP.set_options!(opt_batch_ls, kwargs) + + options = load_options(bnlp; remaining_kwargs...) + ipm_opt = options.interior_point + logger = options.logger + + batch_cnt = BatchCounters(batch_size) + bcb = MadNLP.create_callback( + UniformBatchCallback{T,VT,MT,VI}, + bnlp; + fixed_variable_treatment=ipm_opt.fixed_variable_treatment, + equality_treatment=ipm_opt.equality_treatment, + ) + + ind_lb = bcb.ind_lb + ind_ub = bcb.ind_ub + + ns = length(bcb.ind_ineq) + nx = bcb.nvar + n = nx + ns + m = bcb.ncon + nlb = length(ind_lb) + nub = length(ind_ub) + + batch_kkts = MadNLP.create_kkt_system( + ipm_opt.kkt_system, + bcb, + uniformbatch_linear_solver; + opt_linear_solver = opt_batch_ls, + ) + + batch_x = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + batch_xl = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + batch_xu = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + batch_zl = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + batch_zu = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + batch_f = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) + + batch_d = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) + batch_p = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) + batch_w1 = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) + batch_w2 = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) + + batch_correction_lb = BatchVector(MT, VT, nlb, batch_size) + batch_correction_ub = BatchVector(MT, VT, nub, batch_size) + batch_jacl = BatchVector(MT, VT, n, batch_size) + batch_y = BatchVector(MT, VT, m, batch_size) + batch_c = BatchVector(MT, VT, m, batch_size) + batch_rhs = BatchVector(MT, VT, m, batch_size) + + workspace = UniformBatchWorkspace(MT, VT, nlb, nub, batch_size; + nvar_nlp=nvar_nlp, ncon=m) + + batch_del_w = fill!(MT(undef, 1, batch_size), zero(T)) + batch_del_c = fill!(MT(undef, 1, batch_size), zero(T)) + + MI = typeof(workspace.idx_xl) + return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb), MI}( + batch_size, + batch_d, batch_p, batch_w1, batch_w2, + batch_x, batch_xl, batch_xu, batch_zl, batch_zu, batch_f, + batch_y, batch_c, batch_jacl, batch_rhs, + batch_correction_lb, batch_correction_ub, + workspace, + ipm_opt, batch_cnt, logger, + batch_kkts, + batch_del_w, batch_del_c, + bnlp, + bcb, + ) +end diff --git a/src/batch/utils.jl b/src/batch/utils.jl new file mode 100644 index 00000000..db6aa059 --- /dev/null +++ b/src/batch/utils.jl @@ -0,0 +1,182 @@ +abstract type AbstractBatchMPCSolver{T} end + +function _madnlp_unsafe_column_wrap(mat::MT, n, shift, ::Type{VT}) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + return unsafe_wrap(VT, pointer(mat, shift), n) +end + +function _csc_with_nzval(A::SparseArrays.SparseMatrixCSC, nzval, n) + return SparseArrays.SparseMatrixCSC(n, n, SparseArrays.getcolptr(A), SparseArrays.rowvals(A), nzval) +end + +function zero_inactive_step!(batch_solver::AbstractBatchMPCSolver{T}) where T + ws = batch_solver.workspace + ws.alpha_p .*= ws.active_mask + ws.alpha_d .*= ws.active_mask +end + +function _coo_to_scatter( + coo_I, nrows::Int, n_entries::Int, + proto_I, nzVals::AbstractMatrix{T}, batch_size::Int, +) where T + if n_entries == 0 + scatter = SparseArrays.sparse(Int32[], Int32[], T[], nrows, 0) + buffer = similar(nzVals, 0, batch_size) + return scatter, buffer + end + coo_J = similar(proto_I, n_entries) + coo_J .= Int32(1):Int32(n_entries) + coo_V = similar(nzVals, n_entries) + fill!(coo_V, one(T)) + scatter, _ = MadNLP.coo_to_csc( + MadNLP.SparseMatrixCOO(nrows, n_entries, coo_I, coo_J, coo_V), + ) + fill!(_nzval(scatter), one(T)) + buffer = similar(nzVals, n_entries, batch_size) + fill!(buffer, zero(T)) + return scatter, buffer +end + +function _build_scatter( + aug_I, aug_J, jac_range, n_tot, + nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, +) where T + n_jac = length(jac_range) + coo_I = similar(aug_I, n_jac) + coo_I .= aug_J[jac_range] + scatter, buffer = _coo_to_scatter(coo_I, n_tot, n_jac, aug_I, nzVals, batch_size) + nz_map = similar(aug_csc_map, n_jac) + nz_map .= jac_range + con_map = similar(aug_csc_map, n_jac) + con_map .= aug_I[jac_range] .- Int32(n_tot) + return scatter, nz_map, con_map, buffer +end + +function _build_jac_scatter( + aug_I, aug_J, jac_range, n_tot, m, + nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, +) where T + n_jac = length(jac_range) + coo_I = similar(aug_I, n_jac) + coo_I .= aug_I[jac_range] .- Int32(n_tot) + scatter, buffer = _coo_to_scatter(coo_I, m, n_jac, aug_I, nzVals, batch_size) + nz_map = similar(aug_csc_map, n_jac) + nz_map .= jac_range + var_map = similar(aug_csc_map, n_jac) + var_map .= aug_J[jac_range] + return scatter, nz_map, var_map, buffer +end + +function _build_hess_scatter( + aug_I, aug_J, n_tot, n_hess, + nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, +) where T + if n_hess == 0 + scatter, buffer = _coo_to_scatter(similar(aug_I, 0), n_tot, 0, aug_I, nzVals, batch_size) + nz_map = similar(aug_csc_map, 0) + var_map = similar(aug_csc_map, 0) + return scatter, nz_map, var_map, buffer + end + + hess_range = n_tot+1:n_tot+n_hess + hess_I = aug_I[hess_range] + hess_J = aug_J[hess_range] + + offdiag_idx = findall(hess_I .!= hess_J) + n_hess_sym = n_hess + length(offdiag_idx) + + scatter_rows = similar(aug_I, n_hess_sym) + scatter_rows[1:n_hess] .= hess_I + scatter_rows[n_hess+1:end] .= hess_J[offdiag_idx] + + scatter, buffer = _coo_to_scatter(scatter_rows, n_tot, n_hess_sym, aug_I, nzVals, batch_size) + + nz_map = similar(aug_csc_map, n_hess_sym) + nz_map[1:n_hess] .= hess_range + nz_map[n_hess+1:end] .= n_tot .+ offdiag_idx + + var_map = similar(aug_csc_map, n_hess_sym) + var_map[1:n_hess] .= hess_J + var_map[n_hess+1:end] .= hess_I[offdiag_idx] + + return scatter, nz_map, var_map, buffer +end + +struct BatchVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + values::MT + views::Vector{VT} +end + +MadNLP.full(bv::BatchVector) = bv.values + +function BatchVector( + ::Type{MT}, ::Type{VT}, + len::Int, batch_size::Int, +) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + values = MT(undef, len, batch_size) + fill!(values, zero(T)) + + views = Vector{VT}(undef, batch_size) + for i in 1:batch_size + views[i] = _madnlp_unsafe_column_wrap(values, len, (i-1) * len + 1, VT) + end + + return BatchVector{T, MT, VT}(values, views) +end + +mutable struct BatchExecutionStats{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}} + status::Vector{MadNLP.Status} # (bs,) + solution::MT # (nvar_nlp, bs) + objective::VT # (bs,) + constraints::MT # (ncon, bs) + dual_feas::VT # (bs,) + primal_feas::VT # (bs,) + multipliers::MT # (ncon, bs) + multipliers_L::MT # (nvar_nlp, bs) + multipliers_U::MT # (nvar_nlp, bs) + iter::Vector{Int} # (bs,) + total_time::Vector{Float64} # (bs,) +end + +function BatchExecutionStats(::Type{MT}, ::Type{VT}, nvar_nlp::Int, ncon::Int, batch_size::Int) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} + return BatchExecutionStats{T, VT, MT}( + fill(MadNLP.INITIAL, batch_size), + MT(undef, nvar_nlp, batch_size), + VT(undef, batch_size), + MT(undef, ncon, batch_size), + VT(undef, batch_size), + VT(undef, batch_size), + MT(undef, ncon, batch_size), + MT(undef, nvar_nlp, batch_size), + MT(undef, nvar_nlp, batch_size), + zeros(Int, batch_size), + zeros(Float64, batch_size), + ) +end + +function Base.getindex(stats::BatchExecutionStats, i::Int) + return ( + status = stats.status[i], + solution = view(stats.solution, :, i), + objective = stats.objective[i], + constraints = view(stats.constraints, :, i), + dual_feas = stats.dual_feas[i], + primal_feas = stats.primal_feas[i], + multipliers = view(stats.multipliers, :, i), + multipliers_L = view(stats.multipliers_L, :, i), + multipliers_U = view(stats.multipliers_U, :, i), + iter = stats.iter[i], + total_time = stats.total_time[i], + ) +end + +struct BatchCounters + k::Vector{Int} # per-instance iteration count + start_time::Base.RefValue{Float64} + total_time::Vector{Float64} # per-instance total solve time + linear_solver_time::Base.RefValue{Float64} + eval_function_time::Base.RefValue{Float64} + obj_cnt::Base.RefValue{Int} + obj_grad_cnt::Base.RefValue{Int} + con_cnt::Base.RefValue{Int} +end +BatchCounters(batch_size::Int) = BatchCounters(zeros(Int, batch_size), Ref(0.0), zeros(Float64, batch_size), Ref(0.0), Ref(0.0), Ref(0), Ref(0), Ref(0)) From 3503f8cf1caf90fe17b0d019239c3d655b610127 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 11:32:09 -0400 Subject: [PATCH 02/51] typo --- src/batch/madipm/kernels.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 69ac1653..f1c7b785 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -131,7 +131,7 @@ function get_complementarity_measure!(solver::AbstractBatchMPCSolver) sum!(ws.sum_lb, _scratch_lb) _scratch_ub = MadNLP.dual_ub(solver._w2) - @. _scratch_ub = (x_ur - xu_r) * zu_r + @. _scratch_ub = (xu_r - x_ur) * zu_r sum!(ws.sum_ub, _scratch_ub) @. ws.mu_curr = (ws.sum_lb + ws.sum_ub) / (nlb + nub) From 848d6231108dc1ebdcac703bf09d849f20e4dd4d Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 12:55:27 -0400 Subject: [PATCH 03/51] more typos --- src/batch/madipm/kernels.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index f1c7b785..941496ef 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -175,7 +175,7 @@ function update_barrier!(::Mehrotra, solver::AbstractBatchMPCSolver, mu_affine) ws = solver.workspace T = eltype(ws.mu_curr) - has_inequalities = (length(_get_ind_llb(solver)) + length(_get_ind_uub(solver))) > 0 + has_inequalities = (solver.d.nlb + solver.d.nub) > 0 get_complementarity_measure!(solver) From e6da03c4f77b43a616934e183f3c279e04574348 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 15:52:24 -0400 Subject: [PATCH 04/51] closer to seq --- src/batch/madipm/kernels.jl | 50 ++++++++++++------------------------ src/batch/madipm/solver.jl | 27 ++++++++++++++----- src/batch/madnlp/callback.jl | 5 ++++ src/batch/madnlp/kkt.jl | 5 ++-- 4 files changed, 44 insertions(+), 43 deletions(-) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 941496ef..1f58ab2f 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -126,14 +126,14 @@ function get_complementarity_measure!(solver::AbstractBatchMPCSolver) x_ur = upper(solver.x) zu_r = upper(solver.zu) - _scratch_lb = MadNLP.dual_lb(solver._w2) - @. _scratch_lb = (x_lr - xl_r) * zl_r - sum!(ws.sum_lb, _scratch_lb) - - _scratch_ub = MadNLP.dual_ub(solver._w2) - @. _scratch_ub = (xu_r - x_ur) * zu_r - sum!(ws.sum_ub, _scratch_ub) - + ws.sum_lb .= mapreduce( + (x, xl, z) -> (x - xl) * z, +, x_lr, xl_r, zl_r; + dims=1, init=zero(T), + ) + ws.sum_ub .= mapreduce( + (xu, x, z) -> (xu - x) * z, +, xu_r, x_ur, zu_r; + dims=1, init=zero(T), + ) @. ws.mu_curr = (ws.sum_lb + ws.sum_ub) / (nlb + nub) return ws.mu_curr end @@ -164,7 +164,7 @@ function get_affine_complementarity_measure!(solver::AbstractBatchMPCSolver, alp sum!(ws.sum_lb, _scratch_lb) _scratch_ub = MadNLP.dual_ub(solver._w2) - @. _scratch_ub = (xu_r - x_ur - alpha_p * dx_ur) * (zu_r + alpha_d * dzub) + @. _scratch_ub = (xu_r - (x_ur + alpha_p * dx_ur)) * (zu_r + alpha_d * dzub) sum!(ws.sum_ub, _scratch_ub) @. ws.mu_affine = (ws.sum_lb + ws.sum_ub) / (nlb + nub) @@ -194,26 +194,18 @@ function get_fraction_to_boundary_step!(batch_solver::AbstractBatchMPCSolver) zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d nlb, nub = d.nlb, d.nub T = eltype(ws.alpha_p) - inf_val = T(Inf) + # can't use mapreduce since tau is (1, bs), not (nlb, bs) if nlb > 0 - _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) - _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) # (nlb, bs) + _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) # (nlb, bs) + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) # (nlb, bs) - map!( - (dx, xl, x) -> dx < 0 ? (xl - x) / dx : inf_val, - _scratch_lb, _dx_lr, _xl_r, _x_lr, - ) + @. _scratch_lb = ifelse(_dx_lr < 0, (-_x_lr + _xl_r) * ws.tau / _dx_lr, T(Inf)) minimum!(ws.alpha_xl, _scratch_lb) - @. ws.alpha_xl *= ws.tau - map!( - (dz, z) -> dz < 0 ? -z / dz : inf_val, - _scratch_lb, _dzlb, _zl_r, - ) + @. _scratch_lb = ifelse(_dzlb < 0, (-_zl_r) * ws.tau / _dzlb, T(Inf)) minimum!(ws.alpha_zl, _scratch_lb) - @. ws.alpha_zl *= ws.tau else fill!(ws.alpha_xl, one(T)) fill!(ws.alpha_zl, one(T)) @@ -224,19 +216,11 @@ function get_fraction_to_boundary_step!(batch_solver::AbstractBatchMPCSolver) _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - map!( - (dx, xu, x) -> dx > 0 ? (xu - x) / dx : inf_val, - _scratch_ub, _dx_ur, _xu_r, _x_ur, - ) + @. _scratch_ub = ifelse(_dx_ur > 0, (-_x_ur + _xu_r) * ws.tau / _dx_ur, T(Inf)) minimum!(ws.alpha_xu, _scratch_ub) - @. ws.alpha_xu *= ws.tau - map!( - (dz, z) -> (dz < 0) & (z + dz < 0) ? -z / dz : inf_val, - _scratch_ub, _dzub, _zu_r, - ) + @. _scratch_ub = ifelse((_dzub < 0) & (_zu_r + _dzub < 0), (-_zu_r) * ws.tau / _dzub, T(Inf)) minimum!(ws.alpha_zu, _scratch_ub) - @. ws.alpha_zu *= ws.tau else fill!(ws.alpha_xu, one(T)) fill!(ws.alpha_zu, one(T)) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 18664fdd..56020459 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -118,14 +118,20 @@ function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T fill!(μ, zero(T)) if nlb_init > 0 _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = (xl - lb) * zl + @. _scratch_lb = xl * zl sum!(ws.sum_lb, _scratch_lb) μ .+= ws.sum_lb + @. _scratch_lb = lb * zl + sum!(ws.sum_lb, _scratch_lb) + μ .-= ws.sum_lb end if nub_init > 0 _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = (ub - xu) * zu + @. _scratch_ub = ub * zu sum!(ws.sum_ub, _scratch_ub) + @. _scratch_ub = xu * zu + sum!(ws.sum_lb, _scratch_ub) + ws.sum_ub .-= ws.sum_lb μ .+= ws.sum_ub end @@ -284,7 +290,9 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w # inf_compl[i] = get_optimality_gap / max(1, norm_c[i]) if nlb > 0 - x_lr = lower(x); xl_r = lower(xl); zl_r = lower(zl) + x_lr = lower(x) + xl_r = lower(xl) + zl_r = lower(zl) _scratch_lb = MadNLP.dual_lb(batch_solver._w2) @. _scratch_lb = abs(x_lr - xl_r) * zl_r maximum!(ws.sum_lb, _scratch_lb) @@ -292,7 +300,9 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w fill!(ws.sum_lb, zero(T)) end if nub > 0 - xu_r = upper(xu); x_ur = upper(x); zu_r = upper(zu) + xu_r = upper(xu) + x_ur = upper(x) + zu_r = upper(zu) _scratch_ub = MadNLP.dual_ub(batch_solver._w2) @. _scratch_ub = abs(xu_r - x_ur) * zu_r maximum!(ws.sum_ub, _scratch_ub) @@ -458,13 +468,16 @@ function apply_step!(batch_solver::AbstractBatchMPCSolver) zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d batch_size = batch_solver.batch_size nlb, nub = d.nlb, d.nub - n, m = d.n, d.m # x += alpha_p * dx - MadNLP.full(x) .+= ws.alpha_p .* MadNLP.primal(d) + _x = MadNLP.full(x) + _dx = MadNLP.primal(d) + @. _x = muladd(ws.alpha_p, _dx, _x) # y += alpha_d * d_dual - MadNLP.full(batch_solver.y) .+= ws.alpha_d .* MadNLP.dual(d) + _y = MadNLP.full(batch_solver.y) + _dy = MadNLP.dual(d) + @. _y = muladd(ws.alpha_d, _dy, _y) # zl_r += alpha_d * dzl, zu_r += alpha_d * dzu if nlb > 0 diff --git a/src/batch/madnlp/callback.jl b/src/batch/madnlp/callback.jl index 145f41eb..4d2dde1f 100644 --- a/src/batch/madnlp/callback.jl +++ b/src/batch/madnlp/callback.jl @@ -208,6 +208,11 @@ end +function MadNLP.n_variables(bcb::UniformBatchCallback{T, VT, MT, VI, BM, FH}) where {T, VT, MT, VI, BM, FH<:MadNLP.MakeParameter} + return length(bcb.fixed_handler.free) +end +MadNLP.n_variables(bcb::UniformBatchCallback) = bcb.nvar + function MadNLP._jac_sparsity_wrapper!(bcb::UniformBatchCallback, I::AbstractVector, J::AbstractVector) copyto!(I, bcb.jac_I) copyto!(J, bcb.jac_J) diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl index b8d916f4..4da25c75 100644 --- a/src/batch/madnlp/kkt.jl +++ b/src/batch/madnlp/kkt.jl @@ -219,11 +219,10 @@ end function MadNLP.finish_aug_solve!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) d = batch_solver.d - zl, zu = batch_solver.zl, batch_solver.zu dzl = MadNLP.dual_lb(d) dzu = MadNLP.dual_ub(d) - dzl .= (.-dzl .+ lower(zl) .* xp_lr(d)) ./ bkkt.l_diag - dzu .= (dzu .- upper(zu) .* xp_ur(d)) ./ bkkt.u_diag + dzl .= (.-dzl .+ bkkt.l_lower .* xp_lr(d)) ./ bkkt.l_diag + dzu .= (dzu .- bkkt.u_lower .* xp_ur(d)) ./ bkkt.u_diag return end From deb8ecded77d55b2d9b8b5fa9c5dcee47222b222 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 20:26:26 -0400 Subject: [PATCH 05/51] contiguous bx buffer --- src/batch/madipm/solver.jl | 4 +++- src/batch/madnlp/initialization.jl | 12 ++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 56020459..30e81110 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -204,7 +204,8 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T batch_solver.xu, MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), - bcb.ind_ineq; + bcb.ind_ineq, + ws.bx; tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac, @@ -221,6 +222,7 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T MadNLP.full(batch_solver.rhs), bcb.ind_ineq, T(opt.nlp_scaling_max_gradient), + ws.bx, ) end diff --git a/src/batch/madnlp/initialization.jl b/src/batch/madnlp/initialization.jl index 456661df..c086c975 100644 --- a/src/batch/madnlp/initialization.jl +++ b/src/batch/madnlp/initialization.jl @@ -12,7 +12,8 @@ function MadNLP.initialize!( xu, y, rhs, - ind_ineq; + ind_ineq, + bx_buffer; tol=1e-8, bound_push=1e-2, bound_fac=1e-2, @@ -35,7 +36,8 @@ function MadNLP.initialize!( x0 .= MadNLP._initialize_variables!.(x0, lvar, uvar, bound_push, bound_fac) x_full = MadNLP._update_x!(bcb, x0) - MadNLP._eval_cons_wrapper!(bcb, x_full, bcb.con_buffer) + copyto!(bx_buffer, x_full) + MadNLP._eval_cons_wrapper!(bcb, bx_buffer, bcb.con_buffer) MadNLP.slack(xl) .= view(lcon, ind_ineq, :) MadNLP.slack(xu) .= view(ucon, ind_ineq, :) @@ -83,15 +85,17 @@ end function MadNLP.set_scaling!( cb::UniformBatchCallback, x, xl, xu, y, rhs, ind_ineq, nlp_scaling_max_gradient, + bx_buffer, ) x0 = MadNLP.variable(x) x_full = MadNLP._update_x!(cb, x0) + copyto!(bx_buffer, x_full) - jac_free = MadNLP._eval_jac_wrapper!(cb, x_full, cb.jac_buffer) + jac_free = MadNLP._eval_jac_wrapper!(cb, bx_buffer, cb.jac_buffer) MadNLP.set_con_scale_sparse!(cb.con_scale, cb.jac_I, jac_free, nlp_scaling_max_gradient) MadNLP.set_jac_scale_sparse!(cb.jac_scale, cb.con_scale, cb.jac_I) - MadNLP._eval_grad_f_wrapper!(cb, x_full, cb.grad_buffer) + MadNLP._eval_grad_f_wrapper!(cb, bx_buffer, cb.grad_buffer) MadNLP.set_obj_scale!(cb.obj_scale, cb.grad_buffer, nlp_scaling_max_gradient) con_scale_slk = @view(cb.con_scale[ind_ineq, :]) From 45f9c93aa2b5008bf2fe12bc278a41234b82b395 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 20:58:32 -0400 Subject: [PATCH 06/51] it seems we do need atomix --- Project.toml | 7 +++++-- ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 5 +++-- ext/MadIPMCUDAExt/cuda_wrapper.jl | 3 +-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index afae98fb..8b446aeb 100644 --- a/Project.toml +++ b/Project.toml @@ -13,6 +13,7 @@ QuadraticModels = "f468eda6-eac5-11e8-05a5-ff9e497bcd19" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" [weakdeps] +Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -23,10 +24,11 @@ MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" QuadraticModels = {rev = "mk/rhsbatch", url = "https://github.com/klamike/QuadraticModels.jl.git"} [extensions] -MadIPMCUDAExt = ["CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] +MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] MadIPMMathOptInterfaceExt = "MathOptInterface" [compat] +Atomix = "1" CUDA = "5.4.0" CUDSS = "0.6" KernelAbstractions = "0.9" @@ -44,6 +46,7 @@ Test = "1.10" julia = "1.10" [extras] +Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -53,4 +56,4 @@ MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["CUDA", "CUDSS", "KernelAbstractions", "MadNLPTests", "MadNLPGPU", "MathOptInterface", "Test"] +test = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPTests", "MadNLPGPU", "MathOptInterface", "Test"] diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index 0c8ca828..bbea26f9 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -8,13 +8,15 @@ using CUDA using CUDA.CUSPARSE using CUDSS using KernelAbstractions +import Atomix import QuadraticModels: SparseMatrixCOO +import LinearAlgebra: BlasFloat import MadIPM import MadNLP +include("operators.jl") include("cuda_wrapper.jl") include("cuda_batch_kernels.jl") -include("operators.jl") function MadIPM._csc_with_nzval(A::CUSPARSE.CuSparseMatrixCSC, nzval, n) return CUSPARSE.CuSparseMatrixCSC(A.colPtr, A.rowVal, nzval, (n, n)) @@ -145,4 +147,3 @@ function Base.convert(::Type{QuadraticModel{T, S}}, qp::QuadraticModel{T}) where end end - diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index b4ba5e08..513db00f 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -5,8 +5,7 @@ import MadNLP: getptr @kernel function _transfer_to_map!(dest, to_map, src) k = @index(Global, Linear) @inbounds begin - # TODO: do we need Atomix? - dest[to_map[k]] += src[k] + Atomix.@atomic dest[to_map[k]] += src[k] end end From e6aa6d47c43f127bc451c68505987508896213e1 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Sun, 8 Mar 2026 21:07:48 -0400 Subject: [PATCH 07/51] batch nlp with madipmoperator --- ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 1 + ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 191 ++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+) create mode 100644 ext/MadIPMCUDAExt/cuda_batch_nlp.jl diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index bbea26f9..07d870cb 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -17,6 +17,7 @@ import MadNLP include("operators.jl") include("cuda_wrapper.jl") include("cuda_batch_kernels.jl") +include("cuda_batch_nlp.jl") function MadIPM._csc_with_nzval(A::CUSPARSE.CuSparseMatrixCSC, nzval, n) return CUSPARSE.CuSparseMatrixCSC(A.colPtr, A.rowVal, nzval, (n, n)) diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl new file mode 100644 index 00000000..f57d1ad5 --- /dev/null +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -0,0 +1,191 @@ +import QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData + +function NLPModels.obj!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, + bx::AbstractMatrix{T}, bf::AbstractVector{T}, +) where {T, S, M1 <: MadIPMOperator, M2, MT} + if !bqp.meta.islp + mul!(bqp._HX, bqp.data.H, bx) + bf .= bqp.data.c0 .+ vec(sum(bqp.c_batch .* bx, dims=1)) .+ T(0.5) .* vec(sum(bx .* bqp._HX, dims=1)) + else + bf .= bqp.data.c0 .+ vec(sum(bqp.c_batch .* bx, dims=1)) + end + return bf +end + +function NLPModels.grad!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, + bx::AbstractMatrix{T}, bg::AbstractMatrix{T}, +) where {T, S, M1 <: MadIPMOperator, M2, MT} + if !bqp.meta.islp + mul!(bg, bqp.data.H, bx) + bg .+= bqp.c_batch + else + copyto!(bg, bqp.c_batch) + end + return bg +end + +function NLPModels.cons!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, + bx::AbstractMatrix{T}, bc::AbstractMatrix{T}, +) where {T, S, M1, M2 <: MadIPMOperator, MT} + mul!(bc, bqp.data.A, bx) + return bc +end + +function NLPModels.jac_structure!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, + jrows::AbstractVector{<:Integer}, + jcols::AbstractVector{<:Integer}, +) where {T, S, M1, M2 <: MadIPMOperator} + fill_structure!(bqp.data.A.A, jrows, jcols) + return jrows, jcols +end + +function NLPModels.hess_structure!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, + hrows::AbstractVector{<:Integer}, + hcols::AbstractVector{<:Integer}, +) where {T, S, M1 <: MadIPMOperator, M2} + fill_structure!(bqp.data.H.A, hrows, hcols) + return hrows, hcols +end + +function NLPModels.jac_coord!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, + bx::AbstractMatrix, + bjvals::AbstractMatrix, +) where {T, S, M1, M2 <: MadIPMOperator} + bjvals .= bqp.data.A.A.nzVal + return bjvals +end + +function NLPModels.hess_coord!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, + bx::AbstractMatrix, + by::AbstractMatrix, + bobj_weight::AbstractVector, + bhvals::AbstractMatrix, +) where {T, S, M1 <: MadIPMOperator, M2} + H = bqp.data.H.A + nnzh = nnz(H) + nnzh == 0 && return bhvals + bhvals .= H.nzVal .* bobj_weight' + return bhvals +end + +function NLPModels.hprod!( + bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, + bx::AbstractMatrix{T}, by::AbstractMatrix{T}, bv::AbstractMatrix{T}, + bobj_weight::AbstractVector{T}, bHv::AbstractMatrix{T}, +) where {T, S, M1 <: MadIPMOperator, M2, MT} + mul!(bHv, bqp.data.H, bv) + bHv .*= bobj_weight' + return bHv +end + +function _expand_symmetric_coo(H::SparseMatrixCOO{Tv, Ti}) where {Tv, Ti} + rows, cols, vals = H.rows, H.cols, H.vals + m, n = size(H) + offdiag = findall(i -> rows[i] != cols[i], 1:length(rows)) + new_rows = vcat(rows, cols[offdiag]) + new_cols = vcat(cols, rows[offdiag]) + new_vals = vcat(vals, vals[offdiag]) + return SparseMatrixCOO(m, n, new_rows, new_cols, new_vals) +end + +function Base.convert(::Type{ObjRHSBatchQuadraticModel{T, S}}, bnlp::ObjRHSBatchQuadraticModel{T}) where {T, S<:CuArray} + nbatch = bnlp.meta.nbatch + nvar = bnlp.meta.nvar + ncon = bnlp.meta.ncon + + H_orig_csr = CUSPARSE.CuSparseMatrixCSR(bnlp.data.H) + + H_full = _expand_symmetric_coo(bnlp.data.H) + H_full_csr = CUSPARSE.CuSparseMatrixCSR(H_full) + + A_csr = CUSPARSE.CuSparseMatrixCSR(bnlp.data.A) + + H_op = MadIPMOperator(H_full_csr; symmetric=false, spmm_ncols=nbatch) + H_op.A = H_orig_csr + A_op = MadIPMOperator(A_csr; symmetric=false, spmm_ncols=nbatch) + + c_gpu = CuVector{T}(bnlp.data.c) + v_gpu = CuVector{T}(bnlp.data.v) + data_gpu = QPData(bnlp.data.c0, c_gpu, v_gpu, H_op, A_op) + + c_batch_gpu = CuMatrix{T}(bnlp.c_batch) + _HX_gpu = CUDA.zeros(T, nvar, nbatch) + _AX_gpu = CUDA.zeros(T, ncon, nbatch) + + VT = typeof(c_gpu) + MT = typeof(c_batch_gpu) + + meta_gpu = NLPModels.BatchNLPModelMeta{T, MT}( + nbatch, nvar; + x0 = CuMatrix{T}(bnlp.meta.x0), + lvar = CuMatrix{T}(bnlp.meta.lvar), + uvar = CuMatrix{T}(bnlp.meta.uvar), + ncon = ncon, + lcon = CuMatrix{T}(bnlp.meta.lcon), + ucon = CuMatrix{T}(bnlp.meta.ucon), + nnzj = bnlp.meta.nnzj, + nnzh = bnlp.meta.nnzh, + islp = bnlp.meta.islp, + ) + + return ObjRHSBatchQuadraticModel{T, VT, typeof(H_op), typeof(A_op), MT}( + meta_gpu, data_gpu, c_batch_gpu, _HX_gpu, _AX_gpu, + ) +end + +function Base.convert(::Type{BatchQuadraticModel{T, MT}}, bnlp::BatchQuadraticModel{T}) where {T, MT<:CuMatrix} + nbatch = bnlp.meta.nbatch + nvar = bnlp.meta.nvar + ncon = bnlp.meta.ncon + + c_batch_gpu = MT(bnlp.c_batch) + c0_batch_gpu = CuVector{T}(bnlp.c0_batch) + H_nzvals_gpu = MT(bnlp.H_nzvals) + A_nzvals_gpu = MT(bnlp.A_nzvals) + + hess_rows_gpu = CuVector{Int}(bnlp.hess_rows) + hess_cols_gpu = CuVector{Int}(bnlp.hess_cols) + + _jac_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._jac_scatter) + _jact_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._jact_scatter) + _hess_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._hess_scatter) + + _hess_sym_gather_cols_gpu = CuVector{Int}(bnlp._hess_sym_gather_cols) + _hess_sym_nzidx_gpu = CuVector{Int}(bnlp._hess_sym_nzidx) + + _HX_gpu = CUDA.zeros(T, nvar, nbatch) + + VT = typeof(c0_batch_gpu) + VI = typeof(hess_rows_gpu) + SpMT_J = typeof(_jac_scatter_gpu) + SpMT_H = typeof(_hess_scatter_gpu) + + meta_gpu = NLPModels.BatchNLPModelMeta{T, MT}( + nbatch, nvar; + x0 = MT(bnlp.meta.x0), + lvar = MT(bnlp.meta.lvar), + uvar = MT(bnlp.meta.uvar), + ncon = ncon, + lcon = MT(bnlp.meta.lcon), + ucon = MT(bnlp.meta.ucon), + nnzj = bnlp.meta.nnzj, + nnzh = bnlp.meta.nnzh, + islp = bnlp.meta.islp, + ) + + return BatchQuadraticModel{T, MT, SpMT_J, SpMT_H, VT, VI}( + meta_gpu, + c_batch_gpu, c0_batch_gpu, H_nzvals_gpu, A_nzvals_gpu, + hess_rows_gpu, hess_cols_gpu, + _jac_scatter_gpu, _jact_scatter_gpu, _hess_scatter_gpu, + _hess_sym_gather_cols_gpu, _hess_sym_nzidx_gpu, + _HX_gpu, + ) +end From c31eefa30e14c72192bb938e2f78bf834b57a669 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 09:26:35 -0400 Subject: [PATCH 08/51] support cuda --- Project.toml | 5 +- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 106 ++++-------------------- ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 6 +- ext/MadIPMCUDAExt/cuda_wrapper.jl | 1 - src/batch/madipm/kernels.jl | 52 +++++------- src/batch/madnlp/kkt.jl | 16 ++-- 6 files changed, 55 insertions(+), 131 deletions(-) diff --git a/Project.toml b/Project.toml index 8b446aeb..cfcfdd75 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "MadIPM" uuid = "4406d55e-1099-4ca6-ac6a-6cf051b10024" -authors = ["François Pacaud ", "Alexis Montoison "] version = "0.2.0" +authors = ["François Pacaud ", "Alexis Montoison "] [deps] LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" @@ -21,7 +21,8 @@ MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [sources] -QuadraticModels = {rev = "mk/rhsbatch", url = "https://github.com/klamike/QuadraticModels.jl.git"} +NLPModels = {path = "../NLPModels.jl"} +QuadraticModels = {path = "../QuadraticModels.jl"} [extensions] MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index 6018974f..0744d27a 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -1,15 +1,17 @@ -@kernel function _batch_set_con_scale_sparse_kernel!( - con_scale, - @Const(ptr), - @Const(inds), - @Const(jac_buffer), -) - (index, j) = @index(Global, NTuple) - @inbounds begin - rng = ptr[index]:ptr[index+1]-1 - for k in rng - (row, i) = inds[k] - con_scale[row, j] = max(con_scale[row, j], abs(jac_buffer[i, j])) +@kernel function _set_con_scale_kernel!(con_scale, @Const(jac_I), @Const(jac_buffer)) + k = @index(Global, Linear) + row = jac_I[k] + bs = size(jac_buffer, 2) + @inbounds for j in 1:bs + val = abs(jac_buffer[k, j]) + # CAS loop for atomic max (CUDA.atomic_max! only supports integers) + old = con_scale[row, j] + while val > old + result = Atomix.@atomicreplace con_scale[row, j] old => val + old = result.old + if result.success + break + end end end end @@ -19,72 +21,15 @@ function MadNLP._set_con_scale_sparse!( jac_I::CuVector{<:Integer}, jac_buffer::CuMatrix{T}, ) where T - ind_jac = CuVector{Int}(1:length(jac_I)) - inds = map((i, j) -> (i, j), jac_I, ind_jac) - !isempty(inds) && sort!(inds) - ptr = getptr(inds; by = ((x1, x2), (y1, y2)) -> x1 != y1) - if length(ptr) > 1 + nnzj = length(jac_I) + if nnzj > 0 backend = CUDABackend() - _batch_set_con_scale_sparse_kernel!(backend)( - con_scale, - ptr, - inds, - jac_buffer; - ndrange = (length(ptr) - 1, size(con_scale, 2)), - ) + _set_con_scale_kernel!(backend)(con_scale, jac_I, jac_buffer; ndrange=nnzj) KernelAbstractions.synchronize(backend) end return con_scale end -@kernel function _block_argmin_kernel!(out_val, out_idx, @Const(parent_data), offset, nrows) - tid = @index(Local, Linear) - j = @index(Group, Linear) - gs = @groupsize()[1] - T = eltype(out_val) - - sval = @localmem T (64,) - sidx = @localmem Int32 (64,) - - # Strided scan - local_min = T(Inf) - local_idx = Int32(0) - @inbounds begin - i = Int32(tid) - while i <= nrows - v = parent_data[offset + i, j] - if v < local_min - local_min = v - local_idx = i - end - i += Int32(gs) - end - sval[tid] = local_min - sidx[tid] = local_idx - end - @synchronize() - - # Tree reduction - @inbounds begin - stride = Int32(gs) >> Int32(1) - while stride > Int32(0) - if Int32(tid) <= stride - if sval[tid + stride] < sval[tid] - sval[tid] = sval[tid + stride] - sidx[tid] = sidx[tid + stride] - end - end - @synchronize() - stride >>= Int32(1) - end - - if tid == 1 - out_val[1, j] = sval[1] - out_idx[1, j] = sidx[1] - end - end -end - @kernel function _mehrotra_correction_kernel!( alpha_p, alpha_d, @Const(mu), @@ -137,23 +82,6 @@ end @inbounds alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) end -const _MEHROTRA_BLOCK = 64 - -function MadIPM._argmin_columns!( - out_val::CuMatrix{T}, out_idx::CuMatrix{Int32}, - parent_data::CuMatrix{T}, offset::Int, nrows::Int; threads_per_column = _MEHROTRA_BLOCK -) where T - ncols = size(out_val, 2) - if ncols > 0 && nrows > 0 - backend = CUDABackend() - _block_argmin_kernel!(backend, threads_per_column)( - out_val, out_idx, parent_data, Int32(offset), Int32(nrows); - ndrange = threads_per_column * ncols, - ) - KernelAbstractions.synchronize(backend) - end -end - function MadIPM._mehrotra_correct_steps!( alpha_p::CuMatrix{T}, alpha_d::CuMatrix{T}, mu, val_xl, idx_xl, val_xu, idx_xu, diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl index f57d1ad5..e0ce8b5e 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -161,6 +161,10 @@ function Base.convert(::Type{BatchQuadraticModel{T, MT}}, bnlp::BatchQuadraticMo _hess_sym_nzidx_gpu = CuVector{Int}(bnlp._hess_sym_nzidx) _HX_gpu = CUDA.zeros(T, nvar, nbatch) + nnzj = bnlp.meta.nnzj + sym_nnzh = size(bnlp._hess_buffer, 1) + _jac_buffer_gpu = CUDA.zeros(T, nnzj, nbatch) + _hess_buffer_gpu = CUDA.zeros(T, sym_nnzh, nbatch) VT = typeof(c0_batch_gpu) VI = typeof(hess_rows_gpu) @@ -186,6 +190,6 @@ function Base.convert(::Type{BatchQuadraticModel{T, MT}}, bnlp::BatchQuadraticMo hess_rows_gpu, hess_cols_gpu, _jac_scatter_gpu, _jact_scatter_gpu, _hess_scatter_gpu, _hess_sym_gather_cols_gpu, _hess_sym_nzidx_gpu, - _HX_gpu, + _HX_gpu, _jac_buffer_gpu, _hess_buffer_gpu, ) end diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 513db00f..34dd7f7c 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -1,6 +1,5 @@ using MadNLPGPU import LinearAlgebra: BlasFloat -import MadNLP: getptr @kernel function _transfer_to_map!(dest, to_map, src) k = @index(Global, Linear) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 1f58ab2f..72af7c0e 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -244,25 +244,6 @@ function update_step!(rule::Union{ConservativeStep, AdaptiveStep}, batch_solver: return end -function _argmin_columns!( - out_val::AbstractMatrix{T}, out_idx::AbstractMatrix{Int32}, - parent_data::AbstractMatrix{T}, offset::Int, nrows::Int, -) where T - @inbounds for j in axes(out_val, 2) - min_val = T(Inf) - min_idx = Int32(0) - for i in 1:nrows - v = parent_data[offset + i, j] - if v < min_val - min_val = v - min_idx = Int32(i) - end - end - out_val[1, j] = min_val - out_idx[1, j] = min_idx - end -end - function _mehrotra_correct_steps!( alpha_p, alpha_d, mu, val_xl, idx_xl, val_xu, idx_xu, @@ -281,11 +262,11 @@ function _mehrotra_correct_steps!( if max_ap < one(T) i_xl = idx_xl[1, j] i_xu = idx_xu[1, j] - if val_xl[1, j] <= val_xu[1, j] && i_xl > Int32(0) + if val_xl[1, j] <= val_xu[1, j] && i_xl > 0 idx = ind_lb[i_xl] zl_stepped = zl_vals[idx, j] + max_ad * d_vals[dlb_off + i_xl, j] corrected_p = (x_vals[idx, j] - xl_vals[idx, j] - mu_j / zl_stepped) / (-d_vals[idx, j]) - elseif i_xu > Int32(0) + elseif i_xu > 0 idx = ind_ub[i_xu] zu_stepped = zu_vals[idx, j] + max_ad * d_vals[dub_off + i_xu, j] corrected_p = (xu_vals[idx, j] - x_vals[idx, j] - mu_j / zu_stepped) / d_vals[idx, j] @@ -298,11 +279,11 @@ function _mehrotra_correct_steps!( if max_ad < one(T) i_zl = idx_zl[1, j] i_zu = idx_zu[1, j] - if val_zl[1, j] <= val_zu[1, j] && i_zl > Int32(0) + if val_zl[1, j] <= val_zu[1, j] && i_zl > 0 idx = ind_lb[i_zl] x_gap = x_vals[idx, j] + max_ap * d_vals[idx, j] - xl_vals[idx, j] corrected_d = -(zl_vals[idx, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] - elseif i_zu > Int32(0) + elseif i_zu > 0 idx = ind_ub[i_zu] x_gap = xu_vals[idx, j] - x_vals[idx, j] - max_ap * d_vals[idx, j] corrected_d = -(zu_vals[idx, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] @@ -328,7 +309,6 @@ function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPC mu_full = ws.mu_curr @. mu_full = ws.mu_affine / gamma_a - w2_vals = batch_solver._w2.values dlb_off = d.n + d.m dub_off = d.n + d.m + d.nlb @@ -338,13 +318,17 @@ function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPC _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) map!((dx, xl, x) -> dx < 0 ? (xl - x) / dx : T(Inf), _scratch_lb, _dx_lr, _xl_r, _x_lr) - _argmin_columns!(ws.alpha_xl, ws.idx_xl, w2_vals, dlb_off, nlb) + _vals, _inds = findmin(_scratch_lb; dims=1) + copyto!(ws.alpha_xl, _vals) + ws.idx_xl .= getindex.(_inds, 1) map!((dz, z) -> dz < 0 ? -z / dz : T(Inf), _scratch_lb, _dzlb, _zl_r) - _argmin_columns!(ws.alpha_zl, ws.idx_zl, w2_vals, dlb_off, nlb) + _vals, _inds = findmin(_scratch_lb; dims=1) + copyto!(ws.alpha_zl, _vals) + ws.idx_zl .= getindex.(_inds, 1) else - fill!(ws.alpha_xl, one(T)); fill!(ws.idx_xl, Int32(0)) - fill!(ws.alpha_zl, one(T)); fill!(ws.idx_zl, Int32(0)) + fill!(ws.alpha_xl, one(T)); fill!(ws.idx_xl, 0) + fill!(ws.alpha_zl, one(T)); fill!(ws.idx_zl, 0) end if nub > 0 @@ -353,13 +337,17 @@ function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPC _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) map!((dx, xu, x) -> dx > 0 ? (xu - x) / dx : T(Inf), _scratch_ub, _dx_ur, _xu_r, _x_ur) - _argmin_columns!(ws.alpha_xu, ws.idx_xu, w2_vals, dub_off, nub) + _vals, _inds = findmin(_scratch_ub; dims=1) + copyto!(ws.alpha_xu, _vals) + ws.idx_xu .= getindex.(_inds, 1) map!((dz, z) -> (dz < 0) & (z + dz < 0) ? -z / dz : T(Inf), _scratch_ub, _dzub, _zu_r) - _argmin_columns!(ws.alpha_zu, ws.idx_zu, w2_vals, dub_off, nub) + _vals, _inds = findmin(_scratch_ub; dims=1) + copyto!(ws.alpha_zu, _vals) + ws.idx_zu .= getindex.(_inds, 1) else - fill!(ws.alpha_xu, one(T)); fill!(ws.idx_xu, Int32(0)) - fill!(ws.alpha_zu, one(T)); fill!(ws.idx_zu, Int32(0)) + fill!(ws.alpha_xu, one(T)); fill!(ws.idx_xu, 0) + fill!(ws.alpha_zu, one(T)); fill!(ws.idx_zu, 0) end _mehrotra_correct_steps!( diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl index 4da25c75..3ad30a39 100644 --- a/src/batch/madnlp/kkt.jl +++ b/src/batch/madnlp/kkt.jl @@ -27,7 +27,8 @@ struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatc # J^T scatter (for jtprod! and mul!) jt_scatter::SMT # (n_tot × n_jac_aug) scatter: S[var_idx, k] = 1 jt_nz_map::VI # nzVals row indices for Jacobian entries - jt_con_map::VI # maps each Jac nonzero to its constraint index + jt_con_map::VI # maps each Jac nonzero to its constraint index (1:m) + jt_con_map_full::VI # jt_con_map offset by n_tot (for indexing into full KKT vector) jt_buffer::MT # (n_jac_aug × batch_size) buffer for jtprod # J scatter (for mul!) j_scatter::SMT # (m × n_jac_aug) scatter: S[con_idx, k] = 1 @@ -124,6 +125,8 @@ function MadNLP.create_kkt_system( jt_scatter, jt_nz_map, jt_con_map, jt_buffer = _build_scatter( I, J, jac_range, n_tot, nzVals, aug_csc_map, batch_size, ) + jt_con_map_full = similar(jt_con_map) + jt_con_map_full .= jt_con_map .+ Int32(n_tot) j_scatter, _, j_var_map, j_buffer = _build_jac_scatter( I, J, jac_range, n_tot, m, nzVals, aug_csc_map, batch_size, ) @@ -151,7 +154,7 @@ function MadNLP.create_kkt_system( aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, reg, l_diag, u_diag, l_lower, u_lower, hess_scatter, hess_nz_map, hess_var_map, hess_buffer, - jt_scatter, jt_nz_map, jt_con_map, jt_buffer, + jt_scatter, jt_nz_map, jt_con_map, jt_con_map_full, jt_buffer, j_scatter, j_var_map, j_buffer, _mul_w_primal, _mul_w_dual, batch_map, batch_map_rev, active_batch_size, @@ -263,7 +266,7 @@ function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) end function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) - bkkt.jt_buffer .= view(bkkt.nzVals, bkkt.jt_nz_map, :) .* MadNLP.full(y)[bkkt.jt_con_map, :] + @views bkkt.jt_buffer .= bkkt.nzVals[bkkt.jt_nz_map, :] .* MadNLP.full(y)[bkkt.jt_con_map, :] mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) return res end @@ -358,17 +361,18 @@ function LinearAlgebra.mul!( wd = bkkt._mul_w_dual # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) - bkkt.hess_buffer .= view(nzV, bkkt.hess_nz_map, :) .* MadNLP.primal(x)[bkkt.hess_var_map, :] + xv = MadNLP.full(x) + @views bkkt.hess_buffer .= nzV[bkkt.hess_nz_map, :] .* xv[bkkt.hess_var_map, :] mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp # mul!(primal(w), jac_com', dual(x), alpha, one(T)) - bkkt.jt_buffer .= view(nzV, bkkt.jt_nz_map, :) .* MadNLP.dual(x)[bkkt.jt_con_map, :] + @views bkkt.jt_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.jt_con_map_full, :] mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) MadNLP.primal(w) .+= alpha .* wp # mul!(dual(w), jac_com, primal(x), alpha, beta) - bkkt.j_buffer .= view(nzV, bkkt.jt_nz_map, :) .* MadNLP.primal(x)[bkkt.j_var_map, :] + @views bkkt.j_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.j_var_map, :] mul!(wd, bkkt.j_scatter, bkkt.j_buffer) MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) From f16f44eeaaf8341d066779c591e9e7f23d8b467f Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 10:08:26 -0400 Subject: [PATCH 09/51] fixmes --- ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 2 +- src/batch/madipm/solver.jl | 1 + src/batch/madnlp/kkt.jl | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl index e0ce8b5e..4b440099 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -85,7 +85,7 @@ function NLPModels.hprod!( return bHv end -function _expand_symmetric_coo(H::SparseMatrixCOO{Tv, Ti}) where {Tv, Ti} +function _expand_symmetric_coo(H::SparseMatrixCOO{Tv, Ti}) where {Tv, Ti} # FIXME: allocates rows, cols, vals = H.rows, H.cols, H.vals m, n = size(H) offdiag = findall(i -> rows[i] != cols[i], 1:length(rows)) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 30e81110..fba172c4 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -381,6 +381,7 @@ function solve_system!( view(MadNLP.full(p), :, i) .= zero(T) end end + # FIXME: per-instance reduction? norm_w = norm(MadNLP.full(w), Inf) norm_p = norm(MadNLP.full(p), Inf) diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl index 3ad30a39..b6aedd70 100644 --- a/src/batch/madnlp/kkt.jl +++ b/src/batch/madnlp/kkt.jl @@ -181,7 +181,7 @@ end function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) na = bkkt.active_batch_size[] nzvals = bkkt.aug_com_nzvals - @inbounds for j in 1:na + @inbounds for j in 1:na # FIXME: refactor to avoid `na` launches i = bkkt.batch_map_rev[j] i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) end From 94f3de26466b86db0035aefd7c21f6a2ed0e4478 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 10:08:56 -0400 Subject: [PATCH 10/51] 20,000 tests --- test/batch_kernels/runtests.jl | 12 + test/batch_kernels/setup.jl | 1186 ++++++++++++++++++++ test/batch_kernels/test_batch_multi.jl | 56 + test/batch_kernels/test_complementarity.jl | 53 + test/batch_kernels/test_iteration_match.jl | 444 ++++++++ test/batch_kernels/test_kkt.jl | 86 ++ test/batch_kernels/test_regularization.jl | 61 + test/batch_kernels/test_rhs.jl | 53 + test/batch_kernels/test_solver_steps.jl | 100 ++ test/batch_kernels/test_step.jl | 64 ++ test/batch_kernels/test_termination.jl | 43 + test/runtests.jl | 10 +- test/test_batch.jl | 126 +++ test/test_batch_gpu.jl | 271 +++++ 14 files changed, 2562 insertions(+), 3 deletions(-) create mode 100644 test/batch_kernels/runtests.jl create mode 100644 test/batch_kernels/setup.jl create mode 100644 test/batch_kernels/test_batch_multi.jl create mode 100644 test/batch_kernels/test_complementarity.jl create mode 100644 test/batch_kernels/test_iteration_match.jl create mode 100644 test/batch_kernels/test_kkt.jl create mode 100644 test/batch_kernels/test_regularization.jl create mode 100644 test/batch_kernels/test_rhs.jl create mode 100644 test/batch_kernels/test_solver_steps.jl create mode 100644 test/batch_kernels/test_step.jl create mode 100644 test/batch_kernels/test_termination.jl create mode 100644 test/test_batch.jl create mode 100644 test/test_batch_gpu.jl diff --git a/test/batch_kernels/runtests.jl b/test/batch_kernels/runtests.jl new file mode 100644 index 00000000..e1c9645c --- /dev/null +++ b/test/batch_kernels/runtests.jl @@ -0,0 +1,12 @@ +@testset "Batch kernel tests" begin + include("setup.jl") + include("test_rhs.jl") + include("test_kkt.jl") + include("test_complementarity.jl") + include("test_step.jl") + include("test_regularization.jl") + include("test_solver_steps.jl") + include("test_batch_multi.jl") + include("test_termination.jl") + include("test_iteration_match.jl") +end diff --git a/test/batch_kernels/setup.jl b/test/batch_kernels/setup.jl new file mode 100644 index 00000000..bf1b79c5 --- /dev/null +++ b/test/batch_kernels/setup.jl @@ -0,0 +1,1186 @@ +using Test +using LinearAlgebra +using SparseArrays +using MadNLP +using MadIPM +using QuadraticModels +using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel + +# ────────────────────────────────────────────────────────────── +# Test problems +# ────────────────────────────────────────────────────────────── + +# LP: n=2, m=1, no upper bounds (nub=0) +function _setup_simple_lp() + c = ones(2) + Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + QuadraticModel(c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + c0=0.0, x0=ones(2), name="simpleLP") +end + +# QP: n=4, m=2, with finite upper bounds (nub>0) +function _setup_small_qp() + n, m = 4, 2 + c = [1.0, -2.0, 0.5, 1.0] + Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 3.0, 1.5] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 3, 4]; Avals = [1.0, 1.0, 1.0, 1.0] + QuadraticModel(c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[2.0, 1.5], + lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n)) +end + +# QP with only upper bounds (nlb=0, nub>0) +function _setup_upper_only_qp() + QuadraticModel( + [1.0, -1.0], + [1, 2], [1, 2], [1.0, 1.0], + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[-Inf, -Inf], uvar=[5.0, 5.0], + x0=[2.5, 2.5], + ) +end + +# QP with doubly-bounded variables (both lvar and uvar finite) +# Regression test for has_inequalities bug where ind_llb/ind_uub are empty +# but nlb+nub > 0 (variables have bound multipliers on both sides) +function _setup_doubly_bounded_qp() + QuadraticModel( + [1.0, -1.0], + [1, 2], [1, 2], [1.0, 1.0], + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[5.0, 5.0], + x0=[2.5, 2.5], + ) +end + +# Free-variable QP (nlb=0, nub=0, no bound multipliers) +function _setup_free_qp() + QuadraticModel( + [1.0, -1.0], + [1, 2], [1, 2], [2.0, 2.0], + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[-Inf, -Inf], uvar=[Inf, Inf], + x0=[0.5, 0.5], + ) +end + +# QP with large coefficients that trigger non-unit scaling (obj_scale, con_scale < 1) +function _setup_scaled_qp() + QuadraticModel( + [500.0, -300.0, 400.0], + [1, 2, 3], [1, 2, 3], [1.0, 1.0, 1.0], + Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[200.0, 150.0, 100.0], + lcon=[1.0], ucon=[1.0], + lvar=zeros(3), uvar=fill(Inf, 3), + x0=ones(3), + ) +end + +# QP with large coefficients + inequality constraints (exercises slack scaling) +function _setup_scaled_ineq_qp() + QuadraticModel( + [500.0, -300.0, 400.0], + [1, 2, 3], [1, 2, 3], [1.0, 1.0, 1.0], + Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[200.0, 150.0, 100.0], + lcon=[0.5], ucon=[2.0], + lvar=zeros(3), uvar=fill(Inf, 3), + x0=ones(3), + ) +end + +# QP with a fixed variable (lvar[1]==uvar[1]) exercising MakeParameter +function _setup_fixed_var_qp() + QuadraticModel( + [1.0, -1.0, 0.5], + [1, 2, 3], [1, 2, 3], [2.0, 1.0, 1.0], + Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[1.0, 1.0, 1.0], + lcon=[3.0], ucon=[3.0], + lvar=[2.0, 0.0, 0.0], uvar=[2.0, Inf, Inf], + x0=[2.0, 0.5, 0.5], + ) +end + +# QP with all inequality constraints (ns == m, full slack initialization) +function _setup_all_ineq_qp() + QuadraticModel( + [1.0, -2.0, 0.5], + [1, 2, 3], [1, 2, 3], [2.0, 1.0, 1.5], + Arows=[1, 1, 2, 2], Acols=[1, 2, 2, 3], Avals=[1.0, 1.0, 1.0, 1.0], + lcon=[0.0, 0.0], ucon=[3.0, 3.0], + lvar=zeros(3), uvar=fill(Inf, 3), + x0=ones(3), + ) +end + +# QP with mixed bound types (lower-only, upper-only, doubly-bounded, free) + inequality +function _setup_mixed_bounds_qp() + QuadraticModel( + [1.0, -1.0, 0.5, -0.5], + [1, 2, 3, 4], [1, 2, 3, 4], [2.0, 1.0, 1.5, 1.0], + Arows=[1, 1, 1, 1, 2, 2], Acols=[1, 2, 3, 4, 1, 3], + Avals=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + lcon=[1.0, 0.5], ucon=[3.0, 1.5], + lvar=[0.0, -Inf, 0.0, -Inf], uvar=[Inf, 5.0, 10.0, Inf], + x0=[1.0, 2.0, 5.0, 0.0], + ) +end + +# QP with non-diagonal (dense lower-triangular) Hessian +# H = [4 .; 2 3] (lower triangular of symmetric PD matrix) +function _setup_dense_hess_qp() + QuadraticModel( + [1.0, -1.0], + [1, 2, 2], [1, 1, 2], [4.0, 2.0, 3.0], + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=[0.5, 0.5], + ) +end + +# Larger QP with non-diagonal Hessian + inequality + mixed bounds +# H = [5 . .; 1 4 .; 0 2 3] (sparse lower-triangular, off-diagonal entries) +function _setup_dense_hess_mixed_qp() + QuadraticModel( + [1.0, -2.0, 0.5], + [1, 2, 2, 3, 3], [1, 1, 2, 2, 3], [5.0, 1.0, 4.0, 2.0, 3.0], + Arows=[1, 1, 1, 2, 2], Acols=[1, 2, 3, 1, 3], + Avals=[1.0, 1.0, 1.0, 1.0, 1.0], + lcon=[1.0, 0.5], ucon=[3.0, 1.5], + lvar=[0.0, -Inf, 0.0], uvar=[Inf, 5.0, 10.0], + x0=[1.0, 2.0, 1.0], + ) +end + +# All test problems with descriptions +const ALL_TEST_PROBLEMS = [ + ("LP (nlb>0, nub=0)", _setup_simple_lp), + ("QP (nlb>0, nub>0)", _setup_small_qp), + ("QP (nlb=0, nub>0)", _setup_upper_only_qp), + ("QP doubly-bounded", _setup_doubly_bounded_qp), + ("QP free vars (nlb=0, nub=0)", _setup_free_qp), + ("QP scaled", _setup_scaled_qp), + ("QP scaled+ineq", _setup_scaled_ineq_qp), + ("QP fixed var", _setup_fixed_var_qp), + ("QP all-ineq", _setup_all_ineq_qp), + ("QP mixed bounds", _setup_mixed_bounds_qp), + ("QP dense Hessian", _setup_dense_hess_qp), + ("QP dense Hessian+mixed", _setup_dense_hess_mixed_qp), +] + +# ────────────────────────────────────────────────────────────── +# Paired test problems for batch_size=2 (same structure, different data) +# ────────────────────────────────────────────────────────────── + +function _paired_lower_only() + Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp_a = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + qp_b = QuadraticModel([2.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.5, 0.5]) + return (qp_a, qp_b) +end + +function _paired_doubly_bounded() + Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [5.0, 5.0] + qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) + qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) + return (qp_a, qp_b) +end + +function _paired_scaled() + Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [1.0, 1.0, 1.0] + Arows = [1, 1, 1]; Acols = [1, 2, 3]; Avals = [200.0, 150.0, 100.0] + lvar = zeros(3); uvar = fill(Inf, 3) + qp_a = QuadraticModel([500.0, -300.0, 400.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=ones(3)) + qp_b = QuadraticModel([100.0, -200.0, 150.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5, 0.5]) + return (qp_a, qp_b) +end + +function _paired_all_ineq() + Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 1.5] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] + lvar = zeros(3); uvar = fill(Inf, 3) + qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.0, 0.0], ucon=[3.0, 3.0], lvar=lvar, uvar=uvar, x0=ones(3)) + qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5, 0.5], ucon=[4.0, 4.0], lvar=lvar, uvar=uvar, x0=[0.5, 1.5, 1.0]) + return (qp_a, qp_b) +end + +function _paired_mixed_bounds() + Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 1.5, 1.0] + Arows = [1, 1, 1, 1, 2, 2]; Acols = [1, 2, 3, 4, 1, 3] + Avals = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + lvar = [0.0, -Inf, 0.0, -Inf]; uvar = [Inf, 5.0, 10.0, Inf] + qp_a = QuadraticModel([1.0, -1.0, 0.5, -0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, + x0=[1.0, 2.0, 5.0, 0.0]) + qp_b = QuadraticModel([2.0, -2.0, 1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, + x0=[0.5, 3.0, 4.0, 0.5]) + return (qp_a, qp_b) +end + +function _paired_dense_hess() + Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + return (qp_a, qp_b) +end + +function _paired_dense_hess_mixed() + Hrows = [1, 2, 2, 3, 3]; Hcols = [1, 1, 2, 2, 3]; Hvals = [5.0, 1.0, 4.0, 2.0, 3.0] + Arows = [1, 1, 1, 2, 2]; Acols = [1, 2, 3, 1, 3] + Avals = [1.0, 1.0, 1.0, 1.0, 1.0] + lvar = [0.0, -Inf, 0.0]; uvar = [Inf, 5.0, 10.0] + qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, + x0=[1.0, 2.0, 1.0]) + qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, + x0=[0.5, 1.5, 0.5]) + return (qp_a, qp_b) +end + +const PAIRED_PROBLEMS = [ + ("paired lower-only LP", _paired_lower_only), + ("paired doubly-bounded QP", _paired_doubly_bounded), + ("paired scaled QP", _paired_scaled), + ("paired all-ineq QP", _paired_all_ineq), + ("paired mixed-bounds QP", _paired_mixed_bounds), + ("paired dense-hess QP", _paired_dense_hess), + ("paired dense-hess+mixed QP", _paired_dense_hess_mixed), +] + +# Staggered convergence: easy QP (converges fast) + harder QP (converges slow) +# Same Hessian, Jacobian (ObjRHSBatch shares these). Differ only in c, lcon/ucon, x0. +function _paired_staggered() + Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + # Easy QP: symmetric objective, tight constraint → fast convergence + qp_easy = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + # Harder QP: asymmetric objective, looser constraint → more iterations + qp_hard = QuadraticModel([0.01, -0.99], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[5.0, 5.0]) + return (qp_easy, qp_hard) +end + +# Quad problems for batch_size=4 tests (all lower-bounded LPs, same structure) +function _quad_lower_only() + Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp1 = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + qp2 = QuadraticModel([2.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.5, 0.5]) + qp3 = QuadraticModel([0.5, 1.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.5, 1.0]) + qp4 = QuadraticModel([-1.0, 3.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[0.3, 0.2]) + return (qp1, qp2, qp3, qp4) +end + +# Quad problems for batch_size=4 with doubly-bounded QPs +function _quad_doubly_bounded() + Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [5.0, 5.0] + qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) + qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) + qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + qp4 = QuadraticModel([2.0, -2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[4.0], ucon=[4.0], lvar=lvar, uvar=uvar, x0=[2.0, 2.0]) + return (qp1, qp2, qp3, qp4) +end + +# Quad problems for batch_size=4 with non-diagonal Hessian QPs +function _quad_dense_hess() + Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.75, 0.75]) + qp4 = QuadraticModel([2.0, -2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[0.3, 0.2]) + return (qp1, qp2, qp3, qp4) +end + +# ────────────────────────────────────────────────────────────── +# Build initialized sequential solver +# ────────────────────────────────────────────────────────────── +function build_seq(qp; kwargs...) + solver = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, kwargs...) + opt = solver.opt + + MadNLP.initialize!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, solver.ind_ineq; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(solver.jacl, 0.0) + if opt.scaling + MadNLP.set_scaling!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, + solver.ind_ineq, Float64(opt.nlp_scaling_max_gradient)) + end + MadNLP.initialize!(solver.kkt) + MadIPM.init_regularization!(solver, opt.regularization) + + solver.obj_val = MadNLP.eval_f_wrapper(solver, solver.x) + MadNLP.eval_jac_wrapper!(solver, solver.kkt, solver.x) + MadNLP.eval_grad_f_wrapper!(solver, solver.f, solver.x) + MadNLP.eval_cons_wrapper!(solver, solver.c, solver.x) + MadNLP.eval_lag_hess_wrapper!(solver, solver.kkt, solver.x, solver.y) + solver.norm_b = norm(solver.rhs, Inf) + solver.norm_c = norm(MadNLP.primal(solver.f), Inf) + + MadIPM.init_starting_point!(solver) + solver.mu = opt.mu_init + solver.cnt.start_time = time() + solver.best_complementarity = typemax(Float64) + solver.status = MadNLP.REGULAR + MadNLP.jtprod!(solver.jacl, solver.kkt, solver.y) + return solver +end + +# ────────────────────────────────────────────────────────────── +# Build initialized batch solver (bs=1) +# ────────────────────────────────────────────────────────────── +function build_batch(qp; kwargs...) + bnlp = ObjRHSBatchQuadraticModel([qp]) + batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) + ws = batch_solver.workspace + bcb = batch_solver.bcb + opt = batch_solver.opt + + MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + batch_solver.workspace.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(batch_solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + Float64(opt.nlp_scaling_max_gradient), + batch_solver.workspace.bx) + end + MadNLP.initialize!(batch_solver.kkt) + MadIPM.init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + MadIPM.init_starting_point!(batch_solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) + batch_solver.batch_cnt.start_time[] = time() + fill!(batch_solver.batch_cnt.k, 0) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return batch_solver +end + +# ────────────────────────────────────────────────────────────── +# Build initialized batch solver with batch_size > 1 +# ────────────────────────────────────────────────────────────── +function build_batch_n(qp, n::Int) + bnlp = ObjRHSBatchQuadraticModel([qp for _ in 1:n]) + batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR) + ws = batch_solver.workspace + bcb = batch_solver.bcb + opt = batch_solver.opt + + MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + batch_solver.workspace.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(batch_solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + Float64(opt.nlp_scaling_max_gradient), + batch_solver.workspace.bx) + end + MadNLP.initialize!(batch_solver.kkt) + MadIPM.init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + MadIPM.init_starting_point!(batch_solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) + batch_solver.batch_cnt.start_time[] = time() + fill!(batch_solver.batch_cnt.k, 0) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return batch_solver +end + +# ────────────────────────────────────────────────────────────── +# Comparison helper: max absolute difference +# ────────────────────────────────────────────────────────────── +# Safe comparison that handles Inf values (e.g. xu with uvar=Inf) +function cmp(a, b) + d = 0.0 + for (ai, bi) in zip(a, b) + if isfinite(ai) && isfinite(bi) + d = max(d, abs(ai - bi) / max(abs(ai), abs(bi), 1.0)) + elseif ai !== bi + return Inf + end + end + return d +end + +# Extract column 1 from a batch matrix/vector +col1(x::AbstractMatrix) = view(x, :, 1) +col1(x::AbstractVector) = x # already a vector (scalar workspace) + +# ────────────────────────────────────────────────────────────── +# Run first factorize_system to get post-factorization state +# ────────────────────────────────────────────────────────────── +function do_first_factorize!(seq, bat) + # Sequential + MadIPM.update_regularization!(seq, seq.opt.regularization) + MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) + MadNLP.build_kkt!(seq.kkt) + MadNLP.factorize_kkt!(seq.kkt) + # Batch + MadIPM.update_regularization!(bat, bat.opt.regularization) + MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) + MadNLP.build_kkt!(bat.kkt) + MadNLP.factorize_kkt!(bat.kkt) +end + +# ────────────────────────────────────────────────────────────── +# Build initialized batch solver from multiple (potentially different) QPs +# ────────────────────────────────────────────────────────────── +function build_batch_from_qps(qps::Vector; kwargs...) + bnlp = ObjRHSBatchQuadraticModel(qps) + batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) + ws = batch_solver.workspace + bcb = batch_solver.bcb + opt = batch_solver.opt + + MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + batch_solver.workspace.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(batch_solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + Float64(opt.nlp_scaling_max_gradient), + batch_solver.workspace.bx) + end + MadNLP.initialize!(batch_solver.kkt) + MadIPM.init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + MadIPM.init_starting_point!(batch_solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) + batch_solver.batch_cnt.start_time[] = time() + fill!(batch_solver.batch_cnt.k, 0) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return batch_solver +end + +# ────────────────────────────────────────────────────────────── +# Build initialized batch solver using BatchQuadraticModel (bs=1) +# ────────────────────────────────────────────────────────────── +function build_fullbatch(qp; kwargs...) + bnlp = BatchQuadraticModel([qp]) + batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) + ws = batch_solver.workspace + bcb = batch_solver.bcb + opt = batch_solver.opt + + MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + batch_solver.workspace.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(batch_solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + Float64(opt.nlp_scaling_max_gradient), + batch_solver.workspace.bx) + end + MadNLP.initialize!(batch_solver.kkt) + MadIPM.init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + MadIPM.init_starting_point!(batch_solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) + batch_solver.batch_cnt.start_time[] = time() + fill!(batch_solver.batch_cnt.k, 0) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return batch_solver +end + +# ────────────────────────────────────────────────────────────── +# Build initialized batch solver from multiple QPs using BatchQuadraticModel +# ────────────────────────────────────────────────────────────── +function build_fullbatch_from_qps(qps::Vector; kwargs...) + bnlp = BatchQuadraticModel(qps) + batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) + ws = batch_solver.workspace + bcb = batch_solver.bcb + opt = batch_solver.opt + + MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + batch_solver.workspace.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(batch_solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, + MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, + Float64(opt.nlp_scaling_max_gradient), + batch_solver.workspace.bx) + end + MadNLP.initialize!(batch_solver.kkt) + MadIPM.init_regularization!(batch_solver, opt.regularization) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP.eval_f_wrapper(batch_solver, ws.bx) + MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) + MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) + MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + + MadIPM.init_starting_point!(batch_solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) + batch_solver.batch_cnt.start_time[] = time() + fill!(batch_solver.batch_cnt.k, 0) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return batch_solver +end + +# ────────────────────────────────────────────────────────────── +# Paired problems for BatchQuadraticModel: DIFFERENT H and A values +# (same sparsity pattern, same bound pattern, different everything else) +# ────────────────────────────────────────────────────────────── + +# Diagonal Hessian + different A values, equality constraint +function _fullbatch_paired_diagonal() + Hrows = [1, 2]; Hcols = [1, 2] + Arows = [1, 1]; Acols = [1, 2] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [2.0, 1.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [3.0, 2.0]; + Arows=Arows, Acols=Acols, Avals=[2.0, 0.5], + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + return (qp_a, qp_b) +end + +# Off-diagonal Hessian + different H and A values, equality constraint +function _fullbatch_paired_dense_hess() + Hrows = [1, 2, 2]; Hcols = [1, 1, 2] + Arows = [1, 1]; Acols = [1, 2] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; + Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + return (qp_a, qp_b) +end + +# Different H and A values, doubly-bounded variables +function _fullbatch_paired_doubly_bounded() + Hrows = [1, 2]; Hcols = [1, 2] + Arows = [1, 1]; Acols = [1, 2] + lvar = [0.0, 0.0]; uvar = [5.0, 5.0] + qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [2.0, 1.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) + qp_b = QuadraticModel([-0.5, 1.5], Hrows, Hcols, [3.0, 2.5]; + Arows=Arows, Acols=Acols, Avals=[2.0, 1.0], + lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) + return (qp_a, qp_b) +end + +# Different H and A values, inequality constraints (exercises different scaling per instance) +function _fullbatch_paired_inequality() + Hrows = [1, 2, 3]; Hcols = [1, 2, 3] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3] + lvar = zeros(3); uvar = fill(Inf, 3) + qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, [2.0, 1.0, 1.5]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0, 1.0, 1.0], + lcon=[0.0, 0.0], ucon=[3.0, 3.0], lvar=lvar, uvar=uvar, x0=ones(3)) + qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, [3.0, 2.0, 1.0]; + Arows=Arows, Acols=Acols, Avals=[2.0, 0.5, 1.5, 0.5], + lcon=[0.5, 0.5], ucon=[4.0, 4.0], lvar=lvar, uvar=uvar, x0=[0.5, 1.5, 1.0]) + return (qp_a, qp_b) +end + +# Mixed bounds + different H and A values + off-diagonal Hessian +function _fullbatch_paired_mixed() + Hrows = [1, 2, 2, 3, 3]; Hcols = [1, 1, 2, 2, 3] + Arows = [1, 1, 1, 2, 2]; Acols = [1, 2, 3, 1, 3] + lvar = [0.0, -Inf, 0.0]; uvar = [Inf, 5.0, 10.0] + qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, [5.0, 1.0, 4.0, 2.0, 3.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0, 1.0, 1.0, 1.0], + lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, + x0=[1.0, 2.0, 1.0]) + qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, [7.0, 2.0, 5.0, 1.0, 4.0]; + Arows=Arows, Acols=Acols, Avals=[2.0, 0.5, 1.0, 1.5, 0.5], + lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, + x0=[0.5, 1.5, 0.5]) + return (qp_a, qp_b) +end + +# Large coefficients triggering scaling + different H and A values +function _fullbatch_paired_scaled() + Hrows = [1, 2, 3]; Hcols = [1, 2, 3] + Arows = [1, 1, 1]; Acols = [1, 2, 3] + lvar = zeros(3); uvar = fill(Inf, 3) + qp_a = QuadraticModel([500.0, -300.0, 400.0], Hrows, Hcols, [1.0, 1.0, 1.0]; + Arows=Arows, Acols=Acols, Avals=[200.0, 150.0, 100.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=ones(3)) + qp_b = QuadraticModel([100.0, -200.0, 150.0], Hrows, Hcols, [2.0, 3.0, 1.5]; + Arows=Arows, Acols=Acols, Avals=[300.0, 50.0, 50.0], + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5, 0.5]) + return (qp_a, qp_b) +end + +const FULLBATCH_PAIRED_PROBLEMS = [ + ("fullbatch diagonal", _fullbatch_paired_diagonal), + ("fullbatch dense Hessian", _fullbatch_paired_dense_hess), + ("fullbatch doubly-bounded", _fullbatch_paired_doubly_bounded), + ("fullbatch inequality", _fullbatch_paired_inequality), + ("fullbatch mixed bounds", _fullbatch_paired_mixed), + ("fullbatch scaled", _fullbatch_paired_scaled), +] + +# ────────────────────────────────────────────────────────────── +# Extract column n from a batch matrix/vector +# ────────────────────────────────────────────────────────────── +coln(x::AbstractMatrix, n::Int) = view(x, :, n) +coln(x::AbstractVector, ::Int) = x + +# ────────────────────────────────────────────────────────────── +# Assertion helpers: compare sequential solver state against batch column +# ────────────────────────────────────────────────────────────── + +function assert_termination_match(seq, bat, col; tol=1e-10) + @test abs(seq.inf_pr - bat.workspace.inf_pr[col]) < tol + @test abs(seq.inf_du - bat.workspace.inf_du[col]) < tol + @test abs(seq.inf_compl - bat.workspace.inf_compl[col]) < tol + @test seq.status == bat.workspace.status[col] + # dual_obj: used in infeasibility detection + seq_dobj = MadIPM.dual_objective(seq) + @test abs(seq_dobj - bat.workspace.dual_obj[col]) < tol +end + +function assert_regularization_match(seq, bat, col; tol=1e-10) + @test abs(seq.del_w - bat.del_w[col]) < tol + @test abs(seq.del_c - bat.del_c[col]) < tol +end + +function assert_prediction_match(seq, bat, col; tol=1e-10) + @test cmp(MadNLP.full(seq.d), coln(MadNLP.full(bat.d), col)) < tol + @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol + @test cmp(seq.correction_lb, coln(MadNLP.full(bat.correction_lb), col)) < tol + @test cmp(seq.correction_ub, coln(MadNLP.full(bat.correction_ub), col)) < tol +end + +function assert_direction_match(seq, bat, col; tol=1e-10) + @test cmp(MadNLP.full(seq.d), coln(MadNLP.full(bat.d), col)) < tol +end + +function assert_step_match(seq, bat, col; tol=1e-10) + @test abs(seq.alpha_p - bat.workspace.alpha_p[col]) < tol + @test abs(seq.alpha_d - bat.workspace.alpha_d[col]) < tol +end + +function assert_iterate_match(seq, bat, col; tol=1e-10) + @test cmp(MadNLP.full(seq.x), coln(MadNLP.full(bat.x), col)) < tol + @test cmp(seq.y, coln(MadNLP.full(bat.y), col)) < tol + @test cmp(MadNLP.full(seq.zl), coln(MadNLP.full(bat.zl), col)) < tol + @test cmp(MadNLP.full(seq.zu), coln(MadNLP.full(bat.zu), col)) < tol + @test cmp(MadNLP.full(seq.xl), coln(MadNLP.full(bat.xl), col)) < tol + @test cmp(MadNLP.full(seq.xu), coln(MadNLP.full(bat.xu), col)) < tol +end + +function assert_model_match(seq, bat, col; tol=1e-10) + @test abs(seq.obj_val - bat.workspace.obj_val[col]) < tol + @test cmp(MadNLP.primal(seq.f), coln(MadNLP.primal(bat.f), col)) < tol + @test cmp(seq.c, coln(MadNLP.full(bat.c), col)) < tol + @test cmp(seq.jacl, coln(MadNLP.full(bat.jacl), col)) < tol +end + +function assert_kkt_diagonals_match(seq, bat, col; tol=1e-10) + skkt = seq.kkt + bkkt = bat.kkt + # l_diag, u_diag, l_lower, u_lower + if length(skkt.l_diag) > 0 + @test cmp(skkt.l_diag, coln(bkkt.l_diag, col)) < tol + @test cmp(skkt.l_lower, coln(bkkt.l_lower, col)) < tol + end + if length(skkt.u_diag) > 0 + @test cmp(skkt.u_diag, coln(bkkt.u_diag, col)) < tol + @test cmp(skkt.u_lower, coln(bkkt.u_lower, col)) < tol + end + # reg (primal regularization) + @test cmp(skkt.reg, coln(bkkt.reg, col)) < tol + # pr_diag (includes reg + bound contributions) + @test cmp(skkt.pr_diag, coln(MadIPM.pr_diag(bkkt), col)) < tol + # du_diag (dual regularization) + @test cmp(skkt.du_diag, coln(MadIPM.du_diag(bkkt), col)) < tol +end + +function assert_init_match(seq, bat, col; tol=1e-10) + # Full state after init_starting_point! + initialization + @test cmp(MadNLP.full(seq.x), coln(MadNLP.full(bat.x), col)) < tol + @test cmp(seq.y, coln(MadNLP.full(bat.y), col)) < tol + @test cmp(MadNLP.full(seq.zl), coln(MadNLP.full(bat.zl), col)) < tol + @test cmp(MadNLP.full(seq.zu), coln(MadNLP.full(bat.zu), col)) < tol + @test cmp(MadNLP.full(seq.xl), coln(MadNLP.full(bat.xl), col)) < tol + @test cmp(MadNLP.full(seq.xu), coln(MadNLP.full(bat.xu), col)) < tol + @test abs(seq.obj_val - bat.workspace.obj_val[col]) < tol + @test cmp(MadNLP.primal(seq.f), coln(MadNLP.primal(bat.f), col)) < tol + @test cmp(seq.c, coln(MadNLP.full(bat.c), col)) < tol + @test cmp(seq.jacl, coln(MadNLP.full(bat.jacl), col)) < tol + # Regularization state + @test abs(seq.del_w - bat.del_w[col]) < tol + @test abs(seq.del_c - bat.del_c[col]) < tol + # Barrier parameter + @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol + # Normalization constants + @test abs(seq.norm_b - bat.workspace.norm_b[col]) < tol + @test abs(seq.norm_c - bat.workspace.norm_c[col]) < tol +end + +function assert_correction_match(seq, bat, col; tol=1e-10) + @test cmp(seq.correction_lb, coln(MadNLP.full(bat.correction_lb), col)) < tol + @test cmp(seq.correction_ub, coln(MadNLP.full(bat.correction_ub), col)) < tol +end + +function assert_kkt_matrix_match(seq, bat, col; tol=1e-10) + skkt = seq.kkt + bkkt = bat.kkt + @test cmp(SparseArrays.nonzeros(skkt.aug_com), coln(bkkt.aug_com_nzvals, col)) < tol +end + +function assert_barrier_match(seq, bat, col; mu_affine_seq=nothing, tol=1e-10) + # mu_curr: complementarity measure (set in update_barrier! → get_complementarity_measure) + @test abs(seq.mu_curr - bat.workspace.mu_curr[col]) < tol + # mu_batch (= sigma * mu_curr): final barrier parameter + @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol + # mu_affine: affine complementarity measure (if provided) + if mu_affine_seq !== nothing + @test abs(mu_affine_seq - bat.workspace.mu_affine[col]) < tol + end +end + +function assert_tau_match(seq_tau, bat, col; tol=1e-10) + @test abs(seq_tau - bat.workspace.tau[col]) < tol +end + +function assert_rhs_match(seq, bat, col; tol=1e-10) + @test cmp(MadNLP.full(seq.p), coln(MadNLP.full(bat.p), col)) < tol +end + +function assert_reduce_rhs_match(seq, bat, col; tol=1e-10) + # After reduce_rhs!, the primal-dual part of d has been modified. + # The dual_lb and dual_ub parts contain the original RHS values + # that will be used in finish_aug_solve! + sd = MadNLP.full(seq.d) + bd = coln(MadNLP.full(bat.d), col) + @test cmp(sd, bd) < tol +end + +function assert_finish_aug_solve_match(seq, bat, col; tol=1e-10) + # After finish_aug_solve!, dzl and dzu are computed from the linear solve result + @test cmp(MadNLP.dual_lb(seq.d), coln(MadNLP.dual_lb(bat.d), col)) < tol + @test cmp(MadNLP.dual_ub(seq.d), coln(MadNLP.dual_ub(bat.d), col)) < tol +end + +# ────────────────────────────────────────────────────────────── +# Patching: copy sequential solver state into a batch column +# (test-only; prevents FP noise from accumulating across iterations) +# ────────────────────────────────────────────────────────────── + +function patch_batch_col_from_seq!(seq, bat, col) + ws = bat.workspace + + # Iterate: x (includes slacks), xl, xu, y, zl, zu + coln(MadNLP.full(bat.x), col) .= MadNLP.full(seq.x) + coln(MadNLP.full(bat.xl), col) .= MadNLP.full(seq.xl) + coln(MadNLP.full(bat.xu), col) .= MadNLP.full(seq.xu) + coln(MadNLP.full(bat.y), col) .= seq.y + coln(MadNLP.full(bat.zl), col) .= MadNLP.full(seq.zl) + coln(MadNLP.full(bat.zu), col) .= MadNLP.full(seq.zu) + + # Model evaluations: gradient (primal part), constraints, jacl, obj + coln(MadNLP.primal(bat.f), col) .= MadNLP.primal(seq.f) + coln(MadNLP.full(bat.c), col) .= seq.c + coln(MadNLP.full(bat.jacl), col) .= seq.jacl + ws.obj_val[col] = seq.obj_val + + # Barrier / termination scalars + ws.mu_batch[col] = seq.mu + ws.best_complementarity[col] = seq.best_complementarity +end + +# ────────────────────────────────────────────────────────────── +# Iteration harness: step sequential + batch solvers in lockstep +# ────────────────────────────────────────────────────────────── + +function run_iterations_bs1!(seq, bat, n_iters; tol=1e-10) + ws = bat.workspace + for iter in 1:n_iters + MadIPM.update_termination_criteria!(seq) + MadIPM.update_termination_criteria!(bat) + assert_termination_match(seq, bat, 1; tol) + + MadIPM.is_done(seq) && break + MadIPM.update_active_set!(bat.kkt, ws.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + + # Factorize + MadIPM.factorize_system!(seq) + MadIPM.factorize_system!(bat) + assert_regularization_match(seq, bat, 1; tol) + assert_kkt_diagonals_match(seq, bat, 1; tol) + assert_kkt_matrix_match(seq, bat, 1; tol) + + # Prediction step (decomposed for finer-grained checks) + # 1. Set predictive RHS + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + assert_rhs_match(seq, bat, 1; tol) + # 2. Solve system (includes reduce_rhs! → linear solve → finish_aug_solve!) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + assert_direction_match(seq, bat, 1; tol) + assert_finish_aug_solve_match(seq, bat, 1; tol) + # 3. Affine step sizes (tau=1) + barrier update + alpha_aff_p, alpha_aff_d = MadIPM.get_fraction_to_boundary_step(seq, 1.0) + fill!(ws.tau, one(eltype(ws.tau))) + MadIPM.get_fraction_to_boundary_step!(bat) + MadIPM.zero_inactive_step!(bat) + mu_affine_seq = MadIPM.get_affine_complementarity_measure(seq, alpha_aff_p, alpha_aff_d) + MadIPM.get_affine_complementarity_measure!(bat, ws.alpha_p, ws.alpha_d) + MadIPM.get_correction!(seq, seq.correction_lb, seq.correction_ub) + MadIPM.get_correction!(bat, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub)) + assert_correction_match(seq, bat, 1; tol) + seq.mu_curr = MadIPM.update_barrier!(seq.opt.barrier_update, seq, mu_affine_seq) + MadIPM.update_barrier!(bat.opt.barrier_update, bat, ws.mu_affine) + assert_prediction_match(seq, bat, 1; tol) + assert_barrier_match(seq, bat, 1; mu_affine_seq, tol) + + # Mehrotra correction + MadIPM.set_correction_rhs!(seq, seq.kkt, seq.mu, seq.correction_lb, seq.correction_ub, seq.ind_lb, seq.ind_ub) + MadIPM.set_correction_rhs!(bat, bat.kkt, ws.mu_batch, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub), nothing, nothing) + assert_rhs_match(seq, bat, 1; tol) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + assert_direction_match(seq, bat, 1; tol) + assert_finish_aug_solve_match(seq, bat, 1; tol) + + # Update step (decomposed to check tau) + MadIPM.update_step!(seq.opt.step_rule, seq) + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.zero_inactive_step!(bat) + # Check tau: compute sequential tau from step rule + seq_tau = if seq.opt.step_rule isa MadIPM.ConservativeStep + seq.opt.step_rule.tau + elseif seq.opt.step_rule isa MadIPM.AdaptiveStep + max(1 - seq.mu, seq.opt.step_rule.tau_min) + else + 1.0 # MehrotraAdaptiveStep uses tau=1.0 internally + end + assert_tau_match(seq_tau, bat, 1; tol) + assert_step_match(seq, bat, 1; tol) + + # Apply step + MadIPM.apply_step!(seq) + MadIPM.apply_step!(bat) + assert_iterate_match(seq, bat, 1; tol) + + # Evaluate model + MadIPM.evaluate_model!(seq) + MadIPM.evaluate_model!(bat) + assert_model_match(seq, bat, 1; tol) + end + @test seq.status == bat.workspace.status[1] + @test seq.status == MadNLP.SOLVE_SUCCEEDED +end + +function run_iterations_bs2!(seq1, seq2, bat, n_iters; tol=1e-10, patch=false) + seq1_done = false + seq2_done = false + + # Patch initial state to eliminate init FP seed differences + if patch + patch_batch_col_from_seq!(seq1, bat, 1) + patch_batch_col_from_seq!(seq2, bat, 2) + end + + for iter in 1:n_iters + if !seq1_done; MadIPM.update_termination_criteria!(seq1); end + if !seq2_done; MadIPM.update_termination_criteria!(seq2); end + MadIPM.update_termination_criteria!(bat) + + if !seq1_done; assert_termination_match(seq1, bat, 1; tol); end + if !seq2_done; assert_termination_match(seq2, bat, 2; tol); end + + seq1_done = seq1_done || MadIPM.is_done(seq1) + seq2_done = seq2_done || MadIPM.is_done(seq2) + + MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + + (seq1_done && seq2_done) && break + + # Factorize + if !seq1_done; MadIPM.factorize_system!(seq1); end + if !seq2_done; MadIPM.factorize_system!(seq2); end + MadIPM.factorize_system!(bat) + if !seq1_done + assert_regularization_match(seq1, bat, 1; tol) + assert_kkt_diagonals_match(seq1, bat, 1; tol) + end + if !seq2_done + assert_regularization_match(seq2, bat, 2; tol) + assert_kkt_diagonals_match(seq2, bat, 2; tol) + end + + # Prediction step + if !seq1_done; MadIPM.prediction_step!(seq1); end + if !seq2_done; MadIPM.prediction_step!(seq2); end + MadIPM.prediction_step!(bat) + if !seq1_done + assert_prediction_match(seq1, bat, 1; tol) + assert_barrier_match(seq1, bat, 1; tol) + end + if !seq2_done + assert_prediction_match(seq2, bat, 2; tol) + assert_barrier_match(seq2, bat, 2; tol) + end + + # Mehrotra correction + if !seq1_done; MadIPM.mehrotra_correction_direction!(seq1); end + if !seq2_done; MadIPM.mehrotra_correction_direction!(seq2); end + MadIPM.mehrotra_correction_direction!(bat) + if !seq1_done; assert_direction_match(seq1, bat, 1; tol); end + if !seq2_done; assert_direction_match(seq2, bat, 2; tol); end + + # Update step + if !seq1_done; MadIPM.update_step!(seq1.opt.step_rule, seq1); end + if !seq2_done; MadIPM.update_step!(seq2.opt.step_rule, seq2); end + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.zero_inactive_step!(bat) + if !seq1_done; assert_step_match(seq1, bat, 1; tol); end + if !seq2_done; assert_step_match(seq2, bat, 2; tol); end + + # Apply step + if !seq1_done; MadIPM.apply_step!(seq1); end + if !seq2_done; MadIPM.apply_step!(seq2); end + MadIPM.apply_step!(bat) + if !seq1_done; assert_iterate_match(seq1, bat, 1; tol); end + if !seq2_done; assert_iterate_match(seq2, bat, 2; tol); end + + # Evaluate model + if !seq1_done; MadIPM.evaluate_model!(seq1); end + if !seq2_done; MadIPM.evaluate_model!(seq2); end + MadIPM.evaluate_model!(bat) + if !seq1_done; assert_model_match(seq1, bat, 1; tol); end + if !seq2_done; assert_model_match(seq2, bat, 2; tol); end + + # Patch batch columns from sequential to prevent FP noise accumulation. + # For QPs the Hessian and Jacobian are constant, so patching iterate + + # model evals is sufficient to get a clean start for the next iteration. + if patch + if !seq1_done; patch_batch_col_from_seq!(seq1, bat, 1); end + if !seq2_done; patch_batch_col_from_seq!(seq2, bat, 2); end + end + end + + @test seq1.status == bat.workspace.status[1] + @test seq2.status == bat.workspace.status[2] +end + +function run_iterations_bsN!(seqs::Vector, bat, n_iters; tol=1e-10) + N = length(seqs) + done = falses(N) + + for iter in 1:n_iters + for i in 1:N + done[i] || MadIPM.update_termination_criteria!(seqs[i]) + end + MadIPM.update_termination_criteria!(bat) + for i in 1:N + done[i] || assert_termination_match(seqs[i], bat, i; tol) + end + + for i in 1:N + done[i] = done[i] || MadIPM.is_done(seqs[i]) + end + + MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + all(done) && break + + # Factorize + for i in 1:N; done[i] || MadIPM.factorize_system!(seqs[i]); end + MadIPM.factorize_system!(bat) + for i in 1:N + if !done[i] + assert_regularization_match(seqs[i], bat, i; tol) + assert_kkt_diagonals_match(seqs[i], bat, i; tol) + end + end + + # Prediction step + for i in 1:N; done[i] || MadIPM.prediction_step!(seqs[i]); end + MadIPM.prediction_step!(bat) + for i in 1:N + if !done[i] + assert_prediction_match(seqs[i], bat, i; tol) + assert_barrier_match(seqs[i], bat, i; tol) + end + end + + # Mehrotra correction + for i in 1:N; done[i] || MadIPM.mehrotra_correction_direction!(seqs[i]); end + MadIPM.mehrotra_correction_direction!(bat) + for i in 1:N; done[i] || assert_direction_match(seqs[i], bat, i; tol); end + + # Update step + for i in 1:N; done[i] || MadIPM.update_step!(seqs[i].opt.step_rule, seqs[i]); end + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.zero_inactive_step!(bat) + for i in 1:N; done[i] || assert_step_match(seqs[i], bat, i; tol); end + + # Apply step + for i in 1:N; done[i] || MadIPM.apply_step!(seqs[i]); end + MadIPM.apply_step!(bat) + for i in 1:N; done[i] || assert_iterate_match(seqs[i], bat, i; tol); end + + # Evaluate model + for i in 1:N; done[i] || MadIPM.evaluate_model!(seqs[i]); end + MadIPM.evaluate_model!(bat) + for i in 1:N; done[i] || assert_model_match(seqs[i], bat, i; tol); end + end + + for i in 1:N + @test seqs[i].status == bat.workspace.status[i] + end +end diff --git a/test/batch_kernels/test_batch_multi.jl b/test/batch_kernels/test_batch_multi.jl new file mode 100644 index 00000000..17149313 --- /dev/null +++ b/test/batch_kernels/test_batch_multi.jl @@ -0,0 +1,56 @@ +@testset "Multi-batch (batch_size > 1)" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + @testset "batch_size=3 consistency" begin + # All instances identical: every column should match the bs=1 result + qp = make_qp() + bat1 = build_batch(qp) + bat3 = build_batch_n(qp, 3) + do_first_factorize!(bat1, bat3) + + # Run one full iteration on both + MadIPM.prediction_step!(bat1) + MadIPM.prediction_step!(bat3) + MadIPM.mehrotra_correction_direction!(bat1) + MadIPM.mehrotra_correction_direction!(bat3) + MadIPM.update_step!(bat1.opt.step_rule, bat1) + MadIPM.update_step!(bat3.opt.step_rule, bat3) + MadIPM.apply_step!(bat1) + MadIPM.apply_step!(bat3) + + # Every column of batch_size=3 should match the single-instance result + for i in 1:3 + @test cmp(col1(MadNLP.full(bat1.x)), view(MadNLP.full(bat3.x), :, i)) < 1e-10 + @test cmp(col1(MadNLP.full(bat1.y)), view(MadNLP.full(bat3.y), :, i)) < 1e-10 + @test cmp(col1(MadNLP.full(bat1.zl)), view(MadNLP.full(bat3.zl), :, i)) < 1e-10 + @test cmp(col1(MadNLP.full(bat1.zu)), view(MadNLP.full(bat3.zu), :, i)) < 1e-10 + end + @test bat3.workspace.alpha_p[1] == bat3.workspace.alpha_p[2] == bat3.workspace.alpha_p[3] + @test bat3.workspace.alpha_d[1] == bat3.workspace.alpha_d[2] == bat3.workspace.alpha_d[3] + end + + @testset "active-set deactivation" begin + qp = make_qp() + bat = build_batch_n(qp, 3) + + # Mark instance 2 as converged + bat.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED + MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + + @test bat.kkt.active_batch_size[] == 2 + @test bat.kkt.batch_map[2] == 0 # deactivated + @test bat.kkt.batch_map[1] != 0 # still active + @test bat.kkt.batch_map[3] != 0 # still active + + # zero_inactive_step! should zero out the deactivated instance + fill!(bat.workspace.active_mask, 1.0) + MadIPM._update_active_mask!(bat) + fill!(bat.workspace.alpha_p, 0.5) + fill!(bat.workspace.alpha_d, 0.5) + MadIPM.zero_inactive_step!(bat) + @test bat.workspace.alpha_p[2] == 0.0 + @test bat.workspace.alpha_d[2] == 0.0 + @test bat.workspace.alpha_p[1] == 0.5 + @test bat.workspace.alpha_d[1] == 0.5 + end + end +end diff --git a/test/batch_kernels/test_complementarity.jl b/test/batch_kernels/test_complementarity.jl new file mode 100644 index 00000000..20cbf671 --- /dev/null +++ b/test/batch_kernels/test_complementarity.jl @@ -0,0 +1,53 @@ +@testset "Barrier / complementarity" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + + @testset "get_complementarity_measure!" begin + seq_mu = MadIPM.get_complementarity_measure(seq) + MadIPM.get_complementarity_measure!(bat) + bat_mu = bat.workspace.mu_curr[1] + @test abs(seq_mu - bat_mu) < 1e-12 + end + + @testset "get_affine_complementarity_measure!" begin + do_first_factorize!(seq, bat) + # Compute affine direction + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + # Compute alpha with tau=1 + seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) + fill!(bat.workspace.tau, 1.0) + MadIPM.get_fraction_to_boundary_step!(bat) + + seq_mu_aff = MadIPM.get_affine_complementarity_measure(seq, seq_ap, seq_ad) + MadIPM.get_affine_complementarity_measure!(bat, bat.workspace.alpha_p, bat.workspace.alpha_d) + bat_mu_aff = bat.workspace.mu_affine[1] + @test abs(seq_mu_aff - bat_mu_aff) < 1e-10 + end + + @testset "update_barrier! (Mehrotra)" begin + do_first_factorize!(seq, bat) + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) + fill!(bat.workspace.tau, 1.0) + MadIPM.get_fraction_to_boundary_step!(bat) + + seq_mu_aff = MadIPM.get_affine_complementarity_measure(seq, seq_ap, seq_ad) + MadIPM.get_affine_complementarity_measure!(bat, bat.workspace.alpha_p, bat.workspace.alpha_d) + + seq.mu_curr = MadIPM.update_barrier!(seq.opt.barrier_update, seq, seq_mu_aff) + MadIPM.update_barrier!(bat.opt.barrier_update, bat, bat.workspace.mu_affine) + + @test abs(seq.mu - bat.workspace.mu_batch[1]) < 1e-10 + end + end +end diff --git a/test/batch_kernels/test_iteration_match.jl b/test/batch_kernels/test_iteration_match.jl new file mode 100644 index 00000000..b898464e --- /dev/null +++ b/test/batch_kernels/test_iteration_match.jl @@ -0,0 +1,444 @@ +@testset "Iteration-level state matching" begin + + # ────────────────────────────────────────────────────────── + # init_starting_point! explicit snapshot: verify post-init state matches + # ────────────────────────────────────────────────────────── + @testset "init snapshot: $label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + assert_init_match(seq, bat, 1) + end + + # ────────────────────────────────────────────────────────── + # batch_size=1: every problem, full iteration trace + # ────────────────────────────────────────────────────────── + @testset "batch_size=1: $label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # batch_size=2: two DIFFERENT problems in one batch + # ────────────────────────────────────────────────────────── + @testset "batch_size=2: $label" for (label, make_paired) in PAIRED_PROBLEMS + qp_a, qp_b = make_paired() + seq1 = build_seq(qp_a) + seq2 = build_seq(qp_b) + bat = build_batch_from_qps([qp_a, qp_b]) + run_iterations_bs2!(seq1, seq2, bat, 50) + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + end + + # ────────────────────────────────────────────────────────── + # Scaling verification: confirm non-unit scales are active + # ────────────────────────────────────────────────────────── + @testset "scaling verification" begin + @testset "$label" for (label, make_qp) in [ + ("scaled_qp", _setup_scaled_qp), + ("scaled_ineq", _setup_scaled_ineq_qp), + ] + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + # Verify obj_scale < 1.0 + @test seq.cb.obj_scale[] < 1.0 + @test bat.bcb.obj_scale[1] < 1.0 + # Verify con_scale has non-unit entries + @test minimum(seq.cb.con_scale) < 1.0 + @test minimum(bat.bcb.con_scale[:, 1]) < 1.0 + # Verify scales match between sequential and batch + @test abs(seq.cb.obj_scale[] - bat.bcb.obj_scale[1]) < 1e-12 + @test cmp(seq.cb.con_scale, bat.bcb.con_scale[:, 1]) < 1e-12 + end + end + + # ────────────────────────────────────────────────────────── + # Fixed variable verification + # ────────────────────────────────────────────────────────── + @testset "fixed variable verification" begin + qp = _setup_fixed_var_qp() + seq = build_seq(qp) + bat = build_batch(qp) + # Verify MakeParameter handler is active for sequential + @test seq.cb.fixed_handler isa MadNLP.MakeParameter + @test bat.bcb.fixed_handler isa MadNLP.MakeParameter + # Sequential removes the fixed var from nvar; batch may keep it + @test MadNLP.n_variables(seq.cb) < 3 # reduced from 3 vars + # Run iterations (also part of ALL_TEST_PROBLEMS bs1 above, but + # this standalone test documents the fixed-variable intent) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # Staggered convergence: one instance converges much earlier + # with explicit zero_inactive_step! and active mask verification + # ────────────────────────────────────────────────────────── + @testset "staggered convergence (bs=2)" begin + qp_easy, qp_hard = _paired_staggered() + seq1 = build_seq(qp_easy) + seq2 = build_seq(qp_hard) + bat = build_batch_from_qps([qp_easy, qp_hard]) + + # Run iterations manually to verify zero_inactive_step! and active mask + ws = bat.workspace + seq1_done = false + seq2_done = false + inactive_verified = false + + for iter in 1:50 + if !seq1_done; MadIPM.update_termination_criteria!(seq1); end + if !seq2_done; MadIPM.update_termination_criteria!(seq2); end + MadIPM.update_termination_criteria!(bat) + + prev_seq1_done = seq1_done + prev_seq2_done = seq2_done + seq1_done = seq1_done || MadIPM.is_done(seq1) + seq2_done = seq2_done || MadIPM.is_done(seq2) + + MadIPM.update_active_set!(bat.kkt, ws.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + + # Verify active mask matches expected state + if seq1_done && !seq2_done + @test bat.kkt.batch_map[1] == 0 # instance 1 inactive + @test bat.kkt.batch_map[2] != 0 # instance 2 still active + @test ws.active_mask[1] == 0.0 + @test ws.active_mask[2] == 1.0 + elseif !seq1_done && seq2_done + @test bat.kkt.batch_map[1] != 0 + @test bat.kkt.batch_map[2] == 0 + @test ws.active_mask[1] == 1.0 + @test ws.active_mask[2] == 0.0 + elseif !seq1_done && !seq2_done + @test bat.kkt.batch_map[1] != 0 + @test bat.kkt.batch_map[2] != 0 + end + + (seq1_done && seq2_done) && break + + # Factorize + if !seq1_done; MadIPM.factorize_system!(seq1); end + if !seq2_done; MadIPM.factorize_system!(seq2); end + MadIPM.factorize_system!(bat) + + # Prediction + Mehrotra + if !seq1_done; MadIPM.prediction_step!(seq1); end + if !seq2_done; MadIPM.prediction_step!(seq2); end + MadIPM.prediction_step!(bat) + if !seq1_done; MadIPM.mehrotra_correction_direction!(seq1); end + if !seq2_done; MadIPM.mehrotra_correction_direction!(seq2); end + MadIPM.mehrotra_correction_direction!(bat) + + # Update step + if !seq1_done; MadIPM.update_step!(seq1.opt.step_rule, seq1); end + if !seq2_done; MadIPM.update_step!(seq2.opt.step_rule, seq2); end + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.zero_inactive_step!(bat) + + # Verify zero_inactive_step! zeroed the inactive instance + if seq1_done && !seq2_done + @test ws.alpha_p[1] == 0.0 + @test ws.alpha_d[1] == 0.0 + @test ws.alpha_p[2] > 0.0 # active instance has nonzero step + inactive_verified = true + elseif !seq1_done && seq2_done + @test ws.alpha_p[2] == 0.0 + @test ws.alpha_d[2] == 0.0 + @test ws.alpha_p[1] > 0.0 + inactive_verified = true + end + if !seq1_done; assert_step_match(seq1, bat, 1); end + if !seq2_done; assert_step_match(seq2, bat, 2); end + + # Apply step + evaluate model + if !seq1_done; MadIPM.apply_step!(seq1); end + if !seq2_done; MadIPM.apply_step!(seq2); end + MadIPM.apply_step!(bat) + if !seq1_done; MadIPM.evaluate_model!(seq1); end + if !seq2_done; MadIPM.evaluate_model!(seq2); end + MadIPM.evaluate_model!(bat) + end + + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + # Verify the easy one converged earlier + @test seq1.cnt.k < seq2.cnt.k + # Verify we actually tested the inactive masking path + @test inactive_verified + end + + # ────────────────────────────────────────────────────────── + # batch_size=4: four DIFFERENT problems in one batch + # ────────────────────────────────────────────────────────── + @testset "batch_size=4: $label" for (label, make_quad) in [ + ("lower-only LP", _quad_lower_only), + ("doubly-bounded QP", _quad_doubly_bounded), + ("dense Hessian QP", _quad_dense_hess), + ] + qp1, qp2, qp3, qp4 = make_quad() + seqs = [build_seq(qp) for qp in [qp1, qp2, qp3, qp4]] + bat = build_batch_from_qps([qp1, qp2, qp3, qp4]) + run_iterations_bsN!(seqs, bat, 50) + for s in seqs + @test s.status == MadNLP.SOLVE_SUCCEEDED + end + end + + # ────────────────────────────────────────────────────────── + # Factorization retry: verify retry loop works for both seq and batch + # ────────────────────────────────────────────────────────── + @testset "factorization retry path" begin + qp = _setup_doubly_bounded_qp() + seq = build_seq(qp) + bat = build_batch(qp) + + # Artificially set del_w/del_c to large values (as if retry triggered) + seq.del_w = 1e-4 + seq.del_c = -1e-4 + bat.del_w .= 1e-4 + bat.del_c .= -1e-4 + + MadIPM.factorize_system!(seq) + MadIPM.factorize_system!(bat) + assert_regularization_match(seq, bat, 1) + assert_kkt_diagonals_match(seq, bat, 1) + assert_kkt_matrix_match(seq, bat, 1) + + MadIPM.prediction_step!(seq) + MadIPM.prediction_step!(bat) + assert_prediction_match(seq, bat, 1) + end + + # ────────────────────────────────────────────────────────── + # KKT mul! verification: batch scatter mul vs sequential sparse mul + # ────────────────────────────────────────────────────────── + @testset "KKT mul! $label" for (label, make_qp) in [ + ("QP (nlb>0, nub>0)", _setup_small_qp), + ("QP doubly-bounded", _setup_doubly_bounded_qp), + ("QP dense Hessian+mixed", _setup_dense_hess_mixed_qp), + ("QP mixed bounds", _setup_mixed_bounds_qp), + ] + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + + # Set up KKT diagonals (needed for mul!) + MadIPM.update_regularization!(seq, seq.opt.regularization) + MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) + MadIPM.update_regularization!(bat, bat.opt.regularization) + MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) + + # Use the current d as input vector (has non-trivial values after init) + # Copy d → p to have a known input + copyto!(MadNLP.full(seq.p), MadNLP.full(seq.d)) + copyto!(MadNLP.full(bat.p), MadNLP.full(bat.d)) + + # Compute mul!(w1, kkt, p) for both + seq_w = seq._w1 + bat_w = bat._w1 + fill!(MadNLP.full(seq_w), 0.0) + fill!(MadNLP.full(bat_w), 0.0) + mul!(seq_w, seq.kkt, seq.p) + mul!(bat_w, bat.kkt, bat.p) + + @test cmp(MadNLP.full(seq_w), coln(MadNLP.full(bat_w), 1)) < 1e-10 + end + + # ────────────────────────────────────────────────────────── + # Step rule variants (bs=1, with _setup_small_qp) + # ────────────────────────────────────────────────────────── + @testset "step_rule=$label" for (label, step_rule) in [ + ("ConservativeStep", MadIPM.ConservativeStep(0.995)), + ("AdaptiveStep", MadIPM.AdaptiveStep(0.99)), + ("MehrotraAdaptiveStep", MadIPM.MehrotraAdaptiveStep(0.99)), + ] + qp = _setup_small_qp() + seq = build_seq(qp; step_rule=step_rule) + bat = build_batch(qp; step_rule=step_rule) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # Step rule variants with diverse problem types + # ────────────────────────────────────────────────────────── + @testset "step_rule=$sr_label on $prob_label" for (sr_label, step_rule) in [ + ("ConservativeStep", MadIPM.ConservativeStep(0.995)), + ("AdaptiveStep", MadIPM.AdaptiveStep(0.99)), + ("MehrotraAdaptiveStep", MadIPM.MehrotraAdaptiveStep(0.99)), + ], (prob_label, make_qp) in [ + ("free vars", _setup_free_qp), + ("all-ineq", _setup_all_ineq_qp), + ("doubly-bounded", _setup_doubly_bounded_qp), + ("dense Hessian+mixed", _setup_dense_hess_mixed_qp), + ] + qp = make_qp() + seq = build_seq(qp; step_rule=step_rule) + bat = build_batch(qp; step_rule=step_rule) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # Regularization variants (bs=1, with _setup_small_qp) + # ────────────────────────────────────────────────────────── + @testset "regularization=$label" for (label, make_reg) in [ + ("NoRegularization", () -> MadIPM.NoRegularization()), + ("FixedRegularization", () -> MadIPM.FixedRegularization(1e-10, 1e-10)), + ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), + ] + qp = _setup_small_qp() + seq = build_seq(qp; regularization=make_reg()) + bat = build_batch(qp; regularization=make_reg()) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # Regularization variants with diverse problem types + # ────────────────────────────────────────────────────────── + @testset "regularization=$reg_label on $prob_label" for (reg_label, make_reg) in [ + ("NoRegularization", () -> MadIPM.NoRegularization()), + ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), + ], (prob_label, make_qp) in [ + ("free vars", _setup_free_qp), + ("all-ineq", _setup_all_ineq_qp), + ("doubly-bounded", _setup_doubly_bounded_qp), + ("dense Hessian+mixed", _setup_dense_hess_mixed_qp), + ] + qp = make_qp() + seq = build_seq(qp; regularization=make_reg()) + bat = build_batch(qp; regularization=make_reg()) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # Step rule + regularization variants with bs=2 + # ────────────────────────────────────────────────────────── + @testset "options bs=2: $label" for (label, make_opts) in [ + ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), + ("Mehrotra+AdaptiveReg", () -> (step_rule=MadIPM.MehrotraAdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), + ] + qp_a, qp_b = _paired_doubly_bounded() + # Each solver needs its own copy of mutable options (e.g. AdaptiveRegularization) + opts1 = make_opts() + opts2 = make_opts() + opts_bat = make_opts() + seq1 = build_seq(qp_a; opts1...) + seq2 = build_seq(qp_b; opts2...) + bat = build_batch_from_qps([qp_a, qp_b]; opts_bat...) + run_iterations_bs2!(seq1, seq2, bat, 50) + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + end + + # ────────────────────────────────────────────────────────── + # Combined step+reg options with diverse bs=2 paired problems + # ────────────────────────────────────────────────────────── + @testset "options bs=2: $opt_label on $prob_label" for (opt_label, make_opts) in [ + ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), + ("Adaptive+AdaptiveReg", () -> (step_rule=MadIPM.AdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), + ], (prob_label, make_paired) in [ + ("dense-hess", _paired_dense_hess), + ("all-ineq", _paired_all_ineq), + ("mixed-bounds", _paired_mixed_bounds), + ] + qp_a, qp_b = make_paired() + opts1 = make_opts() + opts2 = make_opts() + opts_bat = make_opts() + seq1 = build_seq(qp_a; opts1...) + seq2 = build_seq(qp_b; opts2...) + bat = build_batch_from_qps([qp_a, qp_b]; opts_bat...) + run_iterations_bs2!(seq1, seq2, bat, 50) + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + end + + # ────────────────────────────────────────────────────────── + # init_regularization! explicit comparison + # ────────────────────────────────────────────────────────── + @testset "init_regularization! $label" for (label, make_reg) in [ + ("NoRegularization", () -> MadIPM.NoRegularization()), + ("FixedRegularization", () -> MadIPM.FixedRegularization(1e-10, 1e-10)), + ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), + ] + qp = _setup_small_qp() + seq = build_seq(qp; regularization=make_reg()) + bat = build_batch(qp; regularization=make_reg()) + # After build, init_regularization! has been called. + # del_w and del_c should match. + assert_regularization_match(seq, bat, 1) + end + + # ══════════════════════════════════════════════════════════ + # BatchQuadraticModel tests (different H/A values per instance) + # ══════════════════════════════════════════════════════════ + + # ────────────────────────────────────────────────────────── + # BatchQuadraticModel bs=1: verify same results as ObjRHSBatch + # ────────────────────────────────────────────────────────── + @testset "fullbatch bs=1: $label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_fullbatch(qp) + assert_init_match(seq, bat, 1) + run_iterations_bs1!(seq, bat, 50) + end + + # ────────────────────────────────────────────────────────── + # BatchQuadraticModel bs=2: different H AND A values + # ────────────────────────────────────────────────────────── + @testset "fullbatch bs=2: $label" for (label, make_paired) in FULLBATCH_PAIRED_PROBLEMS + qp_a, qp_b = make_paired() + seq1 = build_seq(qp_a) + seq2 = build_seq(qp_b) + bat = build_fullbatch_from_qps([qp_a, qp_b]) + run_iterations_bs2!(seq1, seq2, bat, 50) + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + end + + # ────────────────────────────────────────────────────────── + # BatchQuadraticModel scaling: verify per-instance scaling with different A + # ────────────────────────────────────────────────────────── + @testset "fullbatch scaling verification" begin + qp_a, qp_b = _fullbatch_paired_scaled() + seq1 = build_seq(qp_a) + seq2 = build_seq(qp_b) + bat = build_fullbatch_from_qps([qp_a, qp_b]) + # Verify non-unit scales + @test seq1.cb.obj_scale[] < 1.0 + @test seq2.cb.obj_scale[] < 1.0 + @test bat.bcb.obj_scale[1] < 1.0 + @test bat.bcb.obj_scale[2] < 1.0 + # Verify per-instance scales match sequential + @test abs(seq1.cb.obj_scale[] - bat.bcb.obj_scale[1]) < 1e-12 + @test abs(seq2.cb.obj_scale[] - bat.bcb.obj_scale[2]) < 1e-12 + @test cmp(seq1.cb.con_scale, bat.bcb.con_scale[:, 1]) < 1e-12 + @test cmp(seq2.cb.con_scale, bat.bcb.con_scale[:, 2]) < 1e-12 + # Verify that the two instances have DIFFERENT scales (since A values differ) + @test bat.bcb.con_scale[1, 1] != bat.bcb.con_scale[1, 2] + end + + # ────────────────────────────────────────────────────────── + # BatchQuadraticModel with options variants + # ────────────────────────────────────────────────────────── + @testset "fullbatch options: $label" for (label, make_opts) in [ + ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), + ("Adaptive+AdaptiveReg", () -> (step_rule=MadIPM.AdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), + ] + qp_a, qp_b = _fullbatch_paired_dense_hess() + opts1 = make_opts() + opts2 = make_opts() + opts_bat = make_opts() + seq1 = build_seq(qp_a; opts1...) + seq2 = build_seq(qp_b; opts2...) + bat = build_fullbatch_from_qps([qp_a, qp_b]; opts_bat...) + run_iterations_bs2!(seq1, seq2, bat, 50) + @test seq1.status == MadNLP.SOLVE_SUCCEEDED + @test seq2.status == MadNLP.SOLVE_SUCCEEDED + end +end diff --git a/test/batch_kernels/test_kkt.jl b/test/batch_kernels/test_kkt.jl new file mode 100644 index 00000000..a3f9ff71 --- /dev/null +++ b/test/batch_kernels/test_kkt.jl @@ -0,0 +1,86 @@ +@testset "KKT / augmented system" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + + @testset "set_aug_diagonal_reg!" begin + MadIPM.update_regularization!(seq, seq.opt.regularization) + MadIPM.update_regularization!(bat, bat.opt.regularization) + MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) + MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) + + skkt = seq.kkt + bkkt = bat.kkt + @test cmp(skkt.reg, col1(bkkt.reg)) < 1e-12 + @test cmp(skkt.l_diag, col1(bkkt.l_diag)) < 1e-12 + @test cmp(skkt.u_diag, col1(bkkt.u_diag)) < 1e-12 + @test cmp(skkt.l_lower, col1(bkkt.l_lower)) < 1e-12 + @test cmp(skkt.u_lower, col1(bkkt.u_lower)) < 1e-12 + @test cmp(skkt.pr_diag, col1(MadIPM.pr_diag(bkkt))) < 1e-12 + end + + @testset "build_kkt!" begin + MadIPM.update_regularization!(seq, seq.opt.regularization) + MadIPM.update_regularization!(bat, bat.opt.regularization) + MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) + MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) + MadNLP.build_kkt!(seq.kkt) + MadNLP.build_kkt!(bat.kkt) + + # Compare COO values + seq_V = seq.kkt.aug_raw.V + bat_V = bat.kkt.nzVals[:, 1] + @test cmp(seq_V, bat_V) < 1e-12 + + # Compare CSC nzvals + seq_csc = SparseArrays.nonzeros(seq.kkt.aug_com) + bat_csc = bat.kkt.aug_com_nzvals[:, 1] + @test cmp(seq_csc, bat_csc) < 1e-12 + end + + @testset "factorize + solve" begin + do_first_factorize!(seq, bat) + + # Set same RHS + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 + end + + @testset "mul! (KKT)" begin + do_first_factorize!(seq, bat) + + # Set up input vector x from predictive rhs + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + # w = K * d + w_seq = seq._w1 + w_bat = bat._w1 + fill!(MadNLP.full(w_seq), 0.0) + fill!(MadNLP.full(w_bat), 0.0) + mul!(w_seq, seq.kkt, seq.d) + mul!(w_bat, bat.kkt, bat.d) + @test cmp(MadNLP.full(w_seq), col1(MadNLP.full(w_bat))) < 1e-10 + end + + @testset "jtprod!" begin + # Explicit jtprod! call with initialized solvers + seq_jacl = similar(seq.jacl) + fill!(seq_jacl, 0.0) + MadNLP.jtprod!(seq_jacl, seq.kkt, seq.y) + + bat_jacl = similar(MadNLP.full(bat.jacl)) + fill!(bat_jacl, 0.0) + MadNLP.jtprod!(bat_jacl, bat.kkt, bat.y) + + @test cmp(seq_jacl, col1(bat_jacl)) < 1e-12 + end + end +end diff --git a/test/batch_kernels/test_regularization.jl b/test/batch_kernels/test_regularization.jl new file mode 100644 index 00000000..49cf1a8e --- /dev/null +++ b/test/batch_kernels/test_regularization.jl @@ -0,0 +1,61 @@ +@testset "Regularization" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + @testset "NoRegularization" begin + qp = make_qp() + seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=MadIPM.NoRegularization()) + bat_nlp = ObjRHSBatchQuadraticModel([qp]) + bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=MadIPM.NoRegularization()) + + MadIPM.init_regularization!(seq, MadIPM.NoRegularization()) + MadIPM.init_regularization!(bat, MadIPM.NoRegularization()) + @test seq.del_w == bat.del_w[1] + @test seq.del_c == bat.del_c[1] + + MadIPM.update_regularization!(seq, MadIPM.NoRegularization()) + MadIPM.update_regularization!(bat, MadIPM.NoRegularization()) + @test seq.del_w == bat.del_w[1] + @test seq.del_c == bat.del_c[1] + end + + @testset "FixedRegularization" begin + reg = MadIPM.FixedRegularization(1e-8, -1e-9) + qp = make_qp() + seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=reg) + reg2 = MadIPM.FixedRegularization(1e-8, -1e-9) + bat_nlp = ObjRHSBatchQuadraticModel([qp]) + bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=reg2) + + MadIPM.init_regularization!(seq, reg) + MadIPM.init_regularization!(bat, reg2) + @test seq.del_w == bat.del_w[1] + @test seq.del_c == bat.del_c[1] + + MadIPM.update_regularization!(seq, reg) + MadIPM.update_regularization!(bat, reg2) + @test seq.del_w == bat.del_w[1] + @test seq.del_c == bat.del_c[1] + end + + @testset "AdaptiveRegularization" begin + reg = MadIPM.AdaptiveRegularization(1e-8, -1e-9, 1e-9) + qp = make_qp() + seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=reg) + reg2 = MadIPM.AdaptiveRegularization(1e-8, -1e-9, 1e-9) + bat_nlp = ObjRHSBatchQuadraticModel([qp]) + bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=reg2) + + MadIPM.init_regularization!(seq, reg) + MadIPM.init_regularization!(bat, reg2) + @test seq.del_w == bat.del_w[1] + @test seq.del_c == bat.del_c[1] + + # Multiple updates to test the adaptive decay + for _ in 1:3 + MadIPM.update_regularization!(seq, reg) + MadIPM.update_regularization!(bat, reg2) + @test abs(seq.del_w - bat.del_w[1]) < 1e-15 + @test abs(seq.del_c - bat.del_c[1]) < 1e-15 + end + end + end +end diff --git a/test/batch_kernels/test_rhs.jl b/test/batch_kernels/test_rhs.jl new file mode 100644 index 00000000..57c0b7ef --- /dev/null +++ b/test/batch_kernels/test_rhs.jl @@ -0,0 +1,53 @@ +@testset "RHS setup functions" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + do_first_factorize!(seq, bat) + + @testset "set_initial_primal_rhs!" begin + MadIPM.set_initial_primal_rhs!(seq) + MadIPM.set_initial_primal_rhs!(bat) + @test cmp(MadNLP.full(seq.p), col1(MadNLP.full(bat.p))) < 1e-12 + end + + @testset "set_initial_dual_rhs!" begin + MadIPM.set_initial_dual_rhs!(seq) + MadIPM.set_initial_dual_rhs!(bat) + @test cmp(MadNLP.full(seq.p), col1(MadNLP.full(bat.p))) < 1e-12 + end + + @testset "set_predictive_rhs!" begin + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + @test cmp(MadNLP.primal(seq.p), col1(MadNLP.primal(bat.p))) < 1e-12 + @test cmp(MadNLP.dual(seq.p), col1(MadNLP.dual(bat.p))) < 1e-12 + @test cmp(MadNLP.dual_lb(seq.p), col1(MadNLP.dual_lb(bat.p))) < 1e-12 + @test cmp(MadNLP.dual_ub(seq.p), col1(MadNLP.dual_ub(bat.p))) < 1e-12 + end + + # Need an affine direction to test correction/correction_rhs + @testset "get_correction!" begin + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + MadIPM.get_correction!(seq, seq.correction_lb, seq.correction_ub) + MadIPM.get_correction!(bat, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub)) + @test cmp(seq.correction_lb, col1(MadNLP.full(bat.correction_lb))) < 1e-10 + @test cmp(seq.correction_ub, col1(MadNLP.full(bat.correction_ub))) < 1e-10 + end + + @testset "set_correction_rhs!" begin + # Use the state from get_correction! above + mu_val = seq.mu + MadIPM.set_correction_rhs!(seq, seq.kkt, mu_val, seq.correction_lb, seq.correction_ub, seq.ind_lb, seq.ind_ub) + MadIPM.set_correction_rhs!(bat, bat.kkt, bat.workspace.mu_batch, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub), nothing, nothing) + @test cmp(MadNLP.primal(seq.p), col1(MadNLP.primal(bat.p))) < 1e-10 + @test cmp(MadNLP.dual(seq.p), col1(MadNLP.dual(bat.p))) < 1e-10 + @test cmp(MadNLP.dual_lb(seq.p), col1(MadNLP.dual_lb(bat.p))) < 1e-10 + @test cmp(MadNLP.dual_ub(seq.p), col1(MadNLP.dual_ub(bat.p))) < 1e-10 + end + end +end diff --git a/test/batch_kernels/test_solver_steps.jl b/test/batch_kernels/test_solver_steps.jl new file mode 100644 index 00000000..ae6e4c0d --- /dev/null +++ b/test/batch_kernels/test_solver_steps.jl @@ -0,0 +1,100 @@ +@testset "High-level solver steps" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + @testset "init_starting_point!" begin + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + + @test cmp(MadNLP.full(seq.x), col1(MadNLP.full(bat.x))) < 1e-10 + @test cmp(MadNLP.full(seq.xl), col1(MadNLP.full(bat.xl))) < 1e-10 + @test cmp(MadNLP.full(seq.xu), col1(MadNLP.full(bat.xu))) < 1e-10 + @test cmp(seq.y, col1(MadNLP.full(bat.y))) < 1e-10 + @test cmp(MadNLP.full(seq.zl), col1(MadNLP.full(bat.zl))) < 1e-10 + @test cmp(MadNLP.full(seq.zu), col1(MadNLP.full(bat.zu))) < 1e-10 + end + + @testset "Full first IPM iteration" begin + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + do_first_factorize!(seq, bat) + + # --- affine_direction! --- + MadIPM.affine_direction!(seq) + MadIPM.affine_direction!(bat) + @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 + + # --- prediction_step! (from post-factorize state) --- + # Rebuild fresh solvers since affine_direction! mutated state + seq = build_seq(qp) + bat = build_batch(qp) + do_first_factorize!(seq, bat) + + MadIPM.prediction_step!(seq) + MadIPM.prediction_step!(bat) + + # Check alpha + seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) + @test abs(seq_ap - bat.workspace.alpha_p[1]) < 1e-10 + @test abs(seq_ad - bat.workspace.alpha_d[1]) < 1e-10 + # Check mu + @test abs(seq.mu - bat.workspace.mu_batch[1]) < 1e-10 + # Check corrections + @test cmp(seq.correction_lb, col1(MadNLP.full(bat.correction_lb))) < 1e-10 + @test cmp(seq.correction_ub, col1(MadNLP.full(bat.correction_ub))) < 1e-10 + + # --- mehrotra_correction_direction! --- + MadIPM.mehrotra_correction_direction!(seq) + MadIPM.mehrotra_correction_direction!(bat) + @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 + + # --- update_step! --- + MadIPM.update_step!(seq.opt.step_rule, seq) + MadIPM.update_step!(bat.opt.step_rule, bat) + @test abs(seq.alpha_p - bat.workspace.alpha_p[1]) < 1e-10 + @test abs(seq.alpha_d - bat.workspace.alpha_d[1]) < 1e-10 + + # --- apply_step! --- + MadIPM.apply_step!(seq) + MadIPM.apply_step!(bat) + @test cmp(MadNLP.full(seq.x), col1(MadNLP.full(bat.x))) < 1e-10 + @test cmp(seq.y, col1(MadNLP.full(bat.y))) < 1e-10 + @test cmp(MadNLP.full(seq.zl), col1(MadNLP.full(bat.zl))) < 1e-10 + @test cmp(MadNLP.full(seq.zu), col1(MadNLP.full(bat.zu))) < 1e-10 + + # --- evaluate_model! --- + MadIPM.evaluate_model!(seq) + MadIPM.evaluate_model!(bat) + @test cmp(MadNLP.primal(seq.f), col1(MadNLP.primal(bat.f))) < 1e-10 + @test cmp(seq.c, col1(MadNLP.full(bat.c))) < 1e-10 + @test cmp(seq.jacl, col1(MadNLP.full(bat.jacl))) < 1e-10 + end + + @testset "update_termination_criteria!" begin + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + # Run one full iteration first + do_first_factorize!(seq, bat) + MadIPM.prediction_step!(seq) + MadIPM.prediction_step!(bat) + MadIPM.mehrotra_correction_direction!(seq) + MadIPM.mehrotra_correction_direction!(bat) + MadIPM.update_step!(seq.opt.step_rule, seq) + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.apply_step!(seq) + MadIPM.apply_step!(bat) + MadIPM.evaluate_model!(seq) + MadIPM.evaluate_model!(bat) + + # Now check termination criteria + MadIPM.update_termination_criteria!(seq) + MadIPM.update_termination_criteria!(bat) + + @test abs(seq.inf_pr - bat.workspace.inf_pr[1]) < 1e-10 + @test abs(seq.inf_du - bat.workspace.inf_du[1]) < 1e-10 + @test abs(seq.inf_compl - bat.workspace.inf_compl[1]) < 1e-10 + @test seq.status == bat.workspace.status[1] + end + end +end diff --git a/test/batch_kernels/test_step.jl b/test/batch_kernels/test_step.jl new file mode 100644 index 00000000..b01cdc6a --- /dev/null +++ b/test/batch_kernels/test_step.jl @@ -0,0 +1,64 @@ +@testset "Step size computation" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + @testset "get_fraction_to_boundary_step!" begin + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + do_first_factorize!(seq, bat) + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + + for tau_val in [1.0, 0.99, 0.995] + seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, tau_val) + fill!(bat.workspace.tau, tau_val) + MadIPM.get_fraction_to_boundary_step!(bat) + @test abs(seq_ap - bat.workspace.alpha_p[1]) < 1e-10 + @test abs(seq_ad - bat.workspace.alpha_d[1]) < 1e-10 + end + end + + @testset "set_tau! (ConservativeStep)" begin + qp = make_qp() + bat = build_batch(qp) + rule = MadIPM.ConservativeStep(0.99) + MadIPM.set_tau!(rule, bat) + @test bat.workspace.tau[1] == 0.99 + end + + @testset "set_tau! (AdaptiveStep)" begin + qp = make_qp() + bat = build_batch(qp) + rule = MadIPM.AdaptiveStep(0.99) + MadIPM.set_tau!(rule, bat) + expected_tau = max(1.0 - bat.workspace.mu_batch[1], 0.99) + @test bat.workspace.tau[1] ≈ expected_tau atol=1e-12 + end + + @testset "update_step! ($rule_name)" for (rule_name, make_rule) in [ + ("ConservativeStep", () -> MadIPM.ConservativeStep(0.99)), + ("AdaptiveStep", () -> MadIPM.AdaptiveStep(0.99)), + ("MehrotraAdaptiveStep", () -> MadIPM.MehrotraAdaptiveStep(0.99)), + ] + qp = make_qp() + seq = build_seq(qp) + bat = build_batch(qp) + do_first_factorize!(seq, bat) + + # Standard flow: prediction → correction → update_step + MadIPM.prediction_step!(seq) + MadIPM.prediction_step!(bat) + MadIPM.mehrotra_correction_direction!(seq) + MadIPM.mehrotra_correction_direction!(bat) + + rule = make_rule() + seq.opt.step_rule = rule + bat.opt.step_rule = rule + MadIPM.update_step!(rule, seq) + MadIPM.update_step!(rule, bat) + @test abs(seq.alpha_p - bat.workspace.alpha_p[1]) < 1e-10 + @test abs(seq.alpha_d - bat.workspace.alpha_d[1]) < 1e-10 + end + end +end diff --git a/test/batch_kernels/test_termination.jl b/test/batch_kernels/test_termination.jl new file mode 100644 index 00000000..b842d43a --- /dev/null +++ b/test/batch_kernels/test_termination.jl @@ -0,0 +1,43 @@ +@testset "Termination status branches" begin + @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS + @testset "SOLVE_SUCCEEDED" begin + # Run enough iterations to converge + qp = make_qp() + bat = build_batch(qp) + + # Run several iterations until convergence + for _ in 1:50 + bat.workspace.status[1] != MadNLP.REGULAR && break + MadIPM.update_termination_criteria!(bat) + MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + MadIPM.factorize_system!(bat) + MadIPM.prediction_step!(bat) + MadIPM.mehrotra_correction_direction!(bat) + MadIPM.update_step!(bat.opt.step_rule, bat) + MadIPM.zero_inactive_step!(bat) + MadIPM.apply_step!(bat) + MadIPM.evaluate_model!(bat) + end + @test bat.workspace.status[1] == MadNLP.SOLVE_SUCCEEDED + end + + @testset "MAXIMUM_ITERATIONS_EXCEEDED" begin + qp = make_qp() + bat = build_batch(qp) + bat.opt.max_iter = 0 # no iterations allowed + MadIPM.update_termination_criteria!(bat) + @test bat.workspace.status[1] == MadNLP.MAXIMUM_ITERATIONS_EXCEEDED + end + + @testset "MAXIMUM_WALLTIME_EXCEEDED" begin + qp = make_qp() + bat = build_batch(qp) + bat.opt.max_wall_time = 0.0 # zero walltime + bat.batch_cnt.start_time[] = time() - 1.0 # started 1s ago + MadIPM.update_termination_criteria!(bat) + @test bat.workspace.status[1] == MadNLP.MAXIMUM_WALLTIME_EXCEEDED + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index e68a1dfd..471f5d47 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -210,10 +210,14 @@ end @test sol.solution[3] == 2.0 end -@testset "MathOptInterface" begin - include("MOI_wrapper.jl") -end +# @testset "MathOptInterface" begin +# include("MOI_wrapper.jl") +# end + +include("test_batch.jl") +include("batch_kernels/runtests.jl") if CUDA.functional() include("test_gpu.jl") + include("test_batch_gpu.jl") end diff --git a/test/test_batch.jl b/test/test_batch.jl new file mode 100644 index 00000000..89ed257b --- /dev/null +++ b/test/test_batch.jl @@ -0,0 +1,126 @@ +using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel + +function _make_small_qp() + # Small QP: min 0.5 xᵀHx + cᵀx s.t. lcon ≤ Ax ≤ ucon, lvar ≤ x ≤ uvar + n, m = 4, 2 + c = [1.0, -2.0, 0.5, 1.0] + Hrows = [1, 2, 3, 4] + Hcols = [1, 2, 3, 4] + Hvals = [2.0, 1.0, 3.0, 1.5] + Arows = [1, 1, 2, 2] + Acols = [1, 2, 3, 4] + Avals = [1.0, 1.0, 1.0, 1.0] + return QuadraticModel( + c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[2.0, 1.5], + lvar=zeros(n), uvar=fill(Inf, n), + x0=ones(n), + ) +end + +function _test_batch_lp(; batch_kwargs...) + qp = simple_lp() + ref = MadIPM.madipm(qp; print_level=MadNLP.ERROR) + @test ref.status == MadNLP.SOLVE_SUCCEEDED + + bs = 4 + qps = [simple_lp() for _ in 1:bs] + + @testset "ObjRHSBatch" begin + bnlp = ObjRHSBatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) + for i in 1:bs + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ ref.objective atol=1e-6 + @test si.solution ≈ ref.solution atol=1e-6 + end + end + + @testset "FullBatch" begin + bnlp = BatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) + for i in 1:bs + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ ref.objective atol=1e-6 + @test si.solution ≈ ref.solution atol=1e-6 + end + end +end + +function _test_batch_qp(; batch_kwargs...) + qp = _make_small_qp() + ref = MadIPM.madipm(qp; print_level=MadNLP.ERROR) + @test ref.status == MadNLP.SOLVE_SUCCEEDED + + bs = 3 + qps = [_make_small_qp() for _ in 1:bs] + + @testset "ObjRHSBatch" begin + bnlp = ObjRHSBatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) + for i in 1:bs + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ ref.objective atol=1e-6 + @test si.solution ≈ ref.solution atol=1e-6 + end + end + + @testset "FullBatch" begin + bnlp = BatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) + for i in 1:bs + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ ref.objective atol=1e-6 + @test si.solution ≈ ref.solution atol=1e-6 + end + end +end + +function _test_fullbatch_different_data(; batch_kwargs...) + # QP with different H and A values per instance (same sparsity) + Hrows = [1, 2, 2]; Hcols = [1, 1, 2] + Arows = [1, 1]; Acols = [1, 2] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + + qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; + Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, [3.0, 1.5, 4.0]; + Arows=Arows, Acols=Acols, Avals=[0.5, 2.0], + lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.75, 0.75]) + + qps = [qp1, qp2, qp3] + refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] + for r in refs + @test r.status == MadNLP.SOLVE_SUCCEEDED + end + + bnlp = BatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) + for i in 1:3 + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ refs[i].objective atol=1e-6 + @test si.solution ≈ refs[i].solution atol=1e-6 + end +end + +@testset "Batch solver (CPU)" begin + @testset "Batch LP" begin + _test_batch_lp() + end + @testset "Batch QP" begin + _test_batch_qp() + end + @testset "FullBatch different H/A data" begin + _test_fullbatch_different_data() + end +end diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl new file mode 100644 index 00000000..067b1c44 --- /dev/null +++ b/test/test_batch_gpu.jl @@ -0,0 +1,271 @@ +using KernelAbstractions +using MadNLPGPU +using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel + +# ============================================================ +# Test problem constructors for GPU batch tests +# ============================================================ + +function _gpu_small_qp() + n, m = 4, 2 + c = [1.0, -2.0, 0.5, 1.0] + Hrows = [1, 2, 3, 4] + Hcols = [1, 2, 3, 4] + Hvals = [2.0, 1.0, 3.0, 1.5] + Arows = [1, 1, 2, 2] + Acols = [1, 2, 3, 4] + Avals = [1.0, 1.0, 1.0, 1.0] + return QuadraticModel( + c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[2.0, 1.5], + lvar=zeros(n), uvar=fill(Inf, n), + x0=ones(n), + ) +end + +function _gpu_doubly_bounded_qp() + n, m = 3, 2 + c = [1.0, -1.0, 0.5] + Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 3.0] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] + return QuadraticModel( + c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[1.0, 0.5], + lvar=[0.0, 0.0, 0.0], uvar=[5.0, 5.0, 5.0], + x0=ones(n), + ) +end + +function _gpu_dense_hess_qp() + Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + return QuadraticModel( + [1.0, -1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=[0.5, 0.5], + ) +end + +# ============================================================ +# Helper: solve batch on GPU and compare with CPU reference +# ============================================================ + +function _test_gpu_batch(qps; atol=1e-6, batch_kwargs...) + bs = length(qps) + + # CPU reference: solve each QP independently + refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] + for r in refs + @test r.status == MadNLP.SOLVE_SUCCEEDED + end + + # Build CPU batch model, convert to GPU, solve + cpu_bnlp = ObjRHSBatchQuadraticModel(qps) + gpu_bnlp = convert(ObjRHSBatchQuadraticModel{Float64, CuVector{Float64}}, cpu_bnlp) + stats = try + MadIPM.madipm_batch( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + batch_kwargs..., + ) + catch e + @error "madipm_batch failed" exception=(e, catch_backtrace()) + rethrow(e) + end + + CUDA.@allowscalar for i in 1:bs + si = stats[i] + if si.status != MadNLP.SOLVE_SUCCEEDED + @error "Instance $i failed" status=si.status objective=si.objective + end + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ refs[i].objective atol=atol + @test Array(si.solution) ≈ refs[i].solution atol=atol + end +end + +# ============================================================ +# Tests +# ============================================================ + +@testset "Batch solver (CUDA)" begin + + # ---------------------------------------------------------- + # Identical instances (sanity check) + # ---------------------------------------------------------- + @testset "Identical LP (bs=4)" begin + _test_gpu_batch([simple_lp() for _ in 1:4]; atol=1e-5) + end + + @testset "Identical QP (bs=3)" begin + _test_gpu_batch([_gpu_small_qp() for _ in 1:3]) + end + + @testset "Identical doubly-bounded QP (bs=2)" begin + _test_gpu_batch([_gpu_doubly_bounded_qp() for _ in 1:2]) + end + + @testset "Identical dense-Hessian QP (bs=2)" begin + _test_gpu_batch([_gpu_dense_hess_qp() for _ in 1:2]) + end + + # ---------------------------------------------------------- + # Different-data instances via ObjRHSBatch + # ---------------------------------------------------------- + @testset "Different LP data (bs=3)" begin + qp1 = QuadraticModel( + [1.0, 1.0], Int[], Int[], Float64[]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=ones(2), + ) + qp2 = QuadraticModel( + [2.0, 0.5], Int[], Int[], Float64[]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[2.0], ucon=[2.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=ones(2), + ) + qp3 = QuadraticModel( + [0.5, 3.0], Int[], Int[], Float64[]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[0.5], ucon=[0.5], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=ones(2), + ) + _test_gpu_batch([qp1, qp2, qp3]; atol=1e-5) + end + + @testset "Different QP data (bs=2)" begin + n, m = 4, 2 + Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 3.0, 1.5] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 3, 4]; Avals = [1.0, 1.0, 1.0, 1.0] + qp1 = QuadraticModel( + [1.0, -2.0, 0.5, 1.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[2.0, 1.5], + lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n), + ) + qp2 = QuadraticModel( + [-1.0, 1.0, -0.5, 2.0], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5, 1.0], ucon=[1.5, 2.0], + lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n), + ) + _test_gpu_batch([qp1, qp2]) + end + + @testset "Different doubly-bounded QP data (bs=2)" begin + n, m = 3, 2 + Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 3.0] + Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] + qp1 = QuadraticModel( + [1.0, -1.0, 0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[1.0, 0.5], ucon=[1.0, 0.5], + lvar=[0.0, 0.0, 0.0], uvar=[5.0, 5.0, 5.0], + x0=ones(n), + ) + qp2 = QuadraticModel( + [-1.0, 2.0, -0.5], Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[0.5, 1.0], ucon=[0.5, 1.0], + lvar=[0.0, 0.0, 0.0], uvar=[3.0, 3.0, 3.0], + x0=ones(n), + ) + _test_gpu_batch([qp1, qp2]) + end + + @testset "Different dense-Hessian QP data (bs=4)" begin + Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] + Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] + make_qp(c, rhs) = QuadraticModel( + c, Hrows, Hcols, Hvals; + Arows=Arows, Acols=Acols, Avals=Avals, + lcon=[rhs], ucon=[rhs], + lvar=[0.0, 0.0], uvar=[Inf, Inf], + x0=[0.5, 0.5], + ) + _test_gpu_batch([ + make_qp([1.0, -1.0], 1.0), + make_qp([-1.0, 2.0], 2.0), + make_qp([0.5, 0.5], 0.5), + make_qp([2.0, -2.0], 1.5), + ]) + end + + # ---------------------------------------------------------- + # Batch size variations + # ---------------------------------------------------------- + @testset "batch_size=1" begin + _test_gpu_batch([simple_lp()]; atol=1e-5) + end + + @testset "batch_size=8" begin + _test_gpu_batch([_gpu_small_qp() for _ in 1:8]) + end + + # ---------------------------------------------------------- + # BatchQuadraticModel (different H/A values per instance) + # ---------------------------------------------------------- + @testset "FullBatch different H/A (bs=2)" begin + Hrows = [1, 2, 2]; Hcols = [1, 1, 2] + Arows = [1, 1]; Acols = [1, 2] + lvar = [0.0, 0.0]; uvar = [Inf, Inf] + + qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) + qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; + Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], + lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) + + qps = [qp1, qp2] + refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] + for r in refs; @test r.status == MadNLP.SOLVE_SUCCEEDED; end + + cpu_bnlp = BatchQuadraticModel(qps) + gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) + stats = MadIPM.madipm_batch( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + ) + CUDA.@allowscalar for i in 1:2 + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ refs[i].objective atol=1e-6 + @test Array(si.solution) ≈ refs[i].solution atol=1e-6 + end + end + + @testset "FullBatch identical QP (bs=3)" begin + qps = [_gpu_small_qp() for _ in 1:3] + ref = MadIPM.madipm(qps[1]; print_level=MadNLP.ERROR) + @test ref.status == MadNLP.SOLVE_SUCCEEDED + + cpu_bnlp = BatchQuadraticModel(qps) + gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) + stats = MadIPM.madipm_batch( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + ) + CUDA.@allowscalar for i in 1:3 + si = stats[i] + @test si.status == MadNLP.SOLVE_SUCCEEDED + @test si.objective ≈ ref.objective atol=1e-6 + @test Array(si.solution) ≈ ref.solution atol=1e-6 + end + end + +end From 2212ec48bc0ff1b9374777e516bf9923393e2932 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 15:03:29 -0400 Subject: [PATCH 11/51] refactor batch kkt --- src/MadIPM.jl | 2 +- src/batch/KKT/KKT.jl | 19 ++ src/batch/KKT/Sparse/augmented.jl | 363 ++++++++++++++++++++++++++++++ src/batch/madipm/kernels.jl | 2 +- src/batch/madnlp/kkt.jl | 361 +++++++++++++++++++++++++++++ 5 files changed, 745 insertions(+), 2 deletions(-) create mode 100644 src/batch/KKT/KKT.jl create mode 100644 src/batch/KKT/Sparse/augmented.jl diff --git a/src/MadIPM.jl b/src/MadIPM.jl index ec839a1e..49088b19 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -19,7 +19,7 @@ include("solver.jl") include("batch/utils.jl") include("batch/madnlp/rhs.jl") include("batch/madnlp/callback.jl") -include("batch/madnlp/kkt.jl") +include("batch/KKT/KKT.jl") include("batch/structure.jl") include("batch/madnlp/linear_solver.jl") include("batch/madnlp/kernels.jl") diff --git a/src/batch/KKT/KKT.jl b/src/batch/KKT/KKT.jl new file mode 100644 index 00000000..84ff805a --- /dev/null +++ b/src/batch/KKT/KKT.jl @@ -0,0 +1,19 @@ +abstract type AbstractBatchKKTSystem{T} end + +@inbounds function _kktmul!( + w::BatchUnreducedKKTVector, + x::BatchUnreducedKKTVector, + reg, du_diag_val, l_lower, u_lower, l_diag, u_diag, + alpha, beta, +) + MadNLP.primal(w) .+= alpha .* reg .* MadNLP.primal(x) + MadNLP.dual(w) .+= alpha .* du_diag_val .* MadNLP.dual(x) + xp_lr(w) .-= alpha .* MadNLP.dual_lb(x) + xp_ur(w) .+= alpha .* MadNLP.dual_ub(x) + MadNLP.dual_lb(w) .= beta .* MadNLP.dual_lb(w) .+ alpha .* (xp_lr(x) .* l_lower .- MadNLP.dual_lb(x) .* l_diag) + MadNLP.dual_ub(w) .= beta .* MadNLP.dual_ub(w) .+ alpha .* (xp_ur(x) .* u_lower .+ MadNLP.dual_ub(x) .* u_diag) + return +end + +include("Sparse/augmented.jl") +include("Sparse/normal.jl") \ No newline at end of file diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl new file mode 100644 index 00000000..491be534 --- /dev/null +++ b/src/batch/KKT/Sparse/augmented.jl @@ -0,0 +1,363 @@ +"""Batched version of SparseKKTSystem""" +struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} + nzVals::MT # (aug_mat_length × batch_size) COO nonzero values + aug_I::VI32 # shared row indices + aug_J::VI32 # shared column indices + batch_solver::LS # batched linear solver + rhs_buffer::VT # contiguous (n+m)*batch_size for batch solve + batch_size::Int + aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values + batch_csc_map::VI # flattened COO→CSC map for all instances + n_tot::Int # n + n_slack (total primal variables) + m::Int # number of constraints + nnzh::Int # number of Hessian nonzeros + # Diagonal and bound data (for _kktmul!) + reg::MT # (n_tot × batch_size) primal regularization + l_diag::MT # (nlb × batch_size) lower bound diagonals + u_diag::MT # (nub × batch_size) upper bound diagonals + l_lower::MT # (nlb × batch_size) lower bound multipliers + u_lower::MT # (nub × batch_size) upper bound multipliers + # Hessian scatter (for mul!) + hess_scatter::SMT # (n_tot × n_hess_sym) scatter matrix + hess_nz_map::VI # nzVals row indices (with sym duplication) + hess_var_map::VI # variable indices for x multiplication + hess_buffer::MT # (n_hess_sym × batch_size) workspace + # J^T scatter (for jtprod! and mul!) + jt_scatter::SMT # (n_tot × n_jac_aug) scatter: S[var_idx, k] = 1 + jt_nz_map::VI # nzVals row indices for Jacobian entries + jt_con_map::VI # maps each Jac nonzero to its constraint index (1:m) + jt_con_map_full::VI # jt_con_map offset by n_tot (for indexing into full KKT vector) + jt_buffer::MT # (n_jac_aug × batch_size) buffer for jtprod + # J scatter (for mul!) + j_scatter::SMT # (m × n_jac_aug) scatter: S[con_idx, k] = 1 + j_var_map::VI # variable indices for J entries + j_buffer::MT # (n_jac_aug × batch_size) buffer for jprod + # Workspace for mul! (GPU needs full matrices, not SubArray views) + _mul_w_primal::MT # (n_tot × batch_size) + _mul_w_dual::MT # (m × batch_size) + # Batch tracking + batch_map::Vector{Int} # original index → active position (0 if inactive) + batch_map_rev::Vector{Int} # active position → original index + active_batch_size::Base.RefValue{Int} +end + +pr_diag(bkkt::SparseUniformBatchKKTSystem) = view(bkkt.nzVals, 1:bkkt.n_tot, :) +function du_diag(bkkt::SparseUniformBatchKKTSystem) + du_off = size(bkkt.nzVals, 1) - bkkt.m + return view(bkkt.nzVals, du_off+1:du_off+bkkt.m, :) +end + +function MadNLP.create_kkt_system( + ::Type{MadNLP.SparseKKTSystem}, + bcb::UniformBatchCallback{T, VT}, + uniformbatch_linear_solver = LoopedBatchLinearSolver; + opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), +) where {T, VT} + batch_size = bcb.batch_size + + n_slack = length(bcb.ind_ineq) + n = bcb.nvar + m = bcb.ncon + jac_sparsity_I = MadNLP.create_array(bcb, Int32, bcb.nnzj) + jac_sparsity_J = MadNLP.create_array(bcb, Int32, bcb.nnzj) + MadNLP._jac_sparsity_wrapper!(bcb, jac_sparsity_I, jac_sparsity_J) + + hess_sparsity_I, hess_sparsity_J = MadNLP.build_hessian_structure(bcb, MadNLP.ExactHessian) + + nlb = length(bcb.ind_lb) + nub = length(bcb.ind_ub) + + MadNLP.force_lower_triangular!(hess_sparsity_I, hess_sparsity_J) + + ind_ineq = bcb.ind_ineq + + n_slack = length(ind_ineq) + n_jac = length(jac_sparsity_I) + n_hess = length(hess_sparsity_I) + n_tot = n + n_slack + + aug_vec_length = n_tot+m + aug_mat_length = n_tot+m+n_hess+n_jac+n_slack + + I = MadNLP.create_array(bcb, Int32, aug_mat_length) + J = MadNLP.create_array(bcb, Int32, aug_mat_length) + nzVals = similar(bcb.con_buffer, aug_mat_length, batch_size) + fill!(nzVals, zero(T)) + V = _madnlp_unsafe_column_wrap(nzVals, aug_mat_length, 1, VT) + + offset = n_tot+n_jac+n_slack+n_hess+m + + I[1:n_tot] .= 1:n_tot + I[n_tot+1:n_tot+n_hess] = hess_sparsity_I + I[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= (jac_sparsity_I.+n_tot) + I[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= ind_ineq .+ n_tot + I[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) + + J[1:n_tot] .= 1:n_tot + J[n_tot+1:n_tot+n_hess] = hess_sparsity_J + J[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= jac_sparsity_J + J[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= (n+1:n+n_slack) + J[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) + + aug_raw = MadNLP.SparseMatrixCOO(aug_vec_length, aug_vec_length, I, J, V) + aug_com, aug_csc_map = MadNLP.coo_to_csc(aug_raw) + + nnz_csc = SparseArrays.nnz(aug_com) + aug_com_nzvals = similar(nzVals, nnz_csc, batch_size) + fill!(aug_com_nzvals, zero(T)) + + csc_offsets = similar(aug_csc_map, 1, batch_size) + csc_offsets .= (0:batch_size-1)' .* nnz_csc + batch_csc_map = vec(aug_csc_map .+ csc_offsets) + + batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, aug_vec_length; opt=opt_linear_solver) + + rhs_buffer = VT(undef, aug_vec_length * batch_size) + fill!(rhs_buffer, zero(T)) + + jac_range = n_tot+n_hess+1:n_tot+n_hess+n_jac+n_slack + + hess_scatter, hess_nz_map, hess_var_map, hess_buffer = _build_hess_scatter( + I, J, n_tot, n_hess, nzVals, aug_csc_map, batch_size, + ) + jt_scatter, jt_nz_map, jt_con_map, jt_buffer = _build_scatter( + I, J, jac_range, n_tot, nzVals, aug_csc_map, batch_size, + ) + jt_con_map_full = similar(jt_con_map) + jt_con_map_full .= jt_con_map .+ Int32(n_tot) + j_scatter, _, j_var_map, j_buffer = _build_jac_scatter( + I, J, jac_range, n_tot, m, nzVals, aug_csc_map, batch_size, + ) + + reg = similar(nzVals, n_tot, batch_size) + l_diag = similar(nzVals, nlb, batch_size) + u_diag = similar(nzVals, nub, batch_size) + l_lower = similar(nzVals, nlb, batch_size) + u_lower = similar(nzVals, nub, batch_size) + + _mul_w_primal = similar(nzVals, n_tot, batch_size) + _mul_w_dual = similar(nzVals, m, batch_size) + + batch_map = collect(1:batch_size) + batch_map_rev = collect(1:batch_size) + active_batch_size = Ref(batch_size) + + LS = typeof(batch_ls) + MT = typeof(nzVals) + VI = typeof(aug_csc_map) + VI32 = typeof(I) + SMT = typeof(jt_scatter) + return SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT}( + nzVals, I, J, batch_ls, rhs_buffer, batch_size, + aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, + reg, l_diag, u_diag, l_lower, u_lower, + hess_scatter, hess_nz_map, hess_var_map, hess_buffer, + jt_scatter, jt_nz_map, jt_con_map, jt_con_map_full, jt_buffer, + j_scatter, j_var_map, j_buffer, + _mul_w_primal, _mul_w_dual, + batch_map, batch_map_rev, active_batch_size, + ) +end + +function update_active_set!(bkkt::SparseUniformBatchKKTSystem, status::Vector{MadNLP.Status}) + active_pos = 0 + for i in 1:bkkt.batch_size + if status[i] == MadNLP.REGULAR + active_pos += 1 + bkkt.batch_map[i] = active_pos + bkkt.batch_map_rev[active_pos] = i + else + bkkt.batch_map[i] = 0 + end + end + for j in (active_pos + 1):bkkt.batch_size + bkkt.batch_map_rev[j] = 0 + end + bkkt.active_batch_size[] = active_pos +end + +function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) + na = bkkt.active_batch_size[] + nzvals = bkkt.aug_com_nzvals + @inbounds for j in 1:na # FIXME: refactor to avoid `na` launches + i = bkkt.batch_map_rev[j] + i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) + end + _active_factorize!(bkkt.batch_solver, na) + return +end + +function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractVector) where T + na = bkkt.active_batch_size[] + bs = bkkt.batch_size + n = length(rhs) ÷ bs + + rhs_mat = reshape(rhs, n, bs) + @inbounds for j in 1:na + i = bkkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) + end + _active_solve!(bkkt.batch_solver, rhs, na, n) + + @inbounds for j in na:-1:1 + i = bkkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) + end + + @inbounds for i in 1:bs + bkkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) + end + return rhs +end + +function MadNLP.reduce_rhs!(bkkt::SparseUniformBatchKKTSystem, d::BatchUnreducedKKTVector) + MadNLP.reduce_rhs!( + xp_lr(d), MadNLP.dual_lb(d), bkkt.l_diag, + xp_ur(d), MadNLP.dual_ub(d), bkkt.u_diag, + ) + return +end + +function MadNLP.finish_aug_solve!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + dzl = MadNLP.dual_lb(d) + dzu = MadNLP.dual_ub(d) + dzl .= (.-dzl .+ bkkt.l_lower .* xp_lr(d)) ./ bkkt.l_diag + dzu .= (dzu .- bkkt.u_lower .* xp_ur(d)) ./ bkkt.u_diag + return +end + +function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + n_pd = d.n + d.m + bs = bkkt.batch_size + na = bkkt.active_batch_size[] + + MadNLP.reduce_rhs!(bkkt, d) + + if na == 1 + # only one active: skip packing + orig_col = bkkt.batch_map_rev[1] + _active_solve!(bkkt.batch_solver, d.views[orig_col], 1, n_pd) + else + rhs = bkkt.rhs_buffer + pd_view = MadNLP.primal_dual(d) + copyto!(reshape(rhs, n_pd, bs), pd_view) + MadNLP.solve_linear_system!(bkkt, rhs) + copyto!(pd_view, reshape(rhs, n_pd, bs)) + end + + MadNLP.finish_aug_solve!(bkkt, batch_solver) + return +end + +function MadNLP.build_kkt!(bkkt::SparseUniformBatchKKTSystem) + MadNLP._transfer!(vec(bkkt.aug_com_nzvals), vec(bkkt.nzVals), bkkt.batch_csc_map) + return +end + +function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) + MadNLP.@trace(batch_solver.logger, "Factorization started.") + MadNLP.build_kkt!(batch_solver.kkt) + batch_solver.batch_cnt.linear_solver_time[] += @elapsed MadNLP.factorize_kkt!(batch_solver.kkt) + return +end + +function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) + @views bkkt.jt_buffer .= bkkt.nzVals[bkkt.jt_nz_map, :] .* MadNLP.full(y)[bkkt.jt_con_map, :] + mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) + return res +end + +function MadNLP.jtprod!(jacl::BatchVector, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) + return MadNLP.jtprod!(MadNLP.full(jacl), bkkt, y) +end + +function MadNLP.eval_jac_wrapper!( + batch_solver::AbstractBatchMPCSolver, + bkkt::SparseUniformBatchKKTSystem, +) + bcb = batch_solver.bcb + ws = batch_solver.workspace + nzVals = bkkt.nzVals + n_tot = bkkt.n_tot + nnzj = bcb.nnzj + n_slack = length(bcb.ind_ineq) + + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + jac_free = MadNLP._eval_jac_wrapper!(bcb, ws.bx, bcb.jac_buffer) + + jac_offset = n_tot + bkkt.nnzh + view(nzVals, jac_offset+1:jac_offset+nnzj, :) .= jac_free + + if n_slack > 0 + view(nzVals, jac_offset+nnzj+1:jac_offset+nnzj+n_slack, :) .= -one(eltype(nzVals)) + end + return +end + +function MadNLP.eval_lag_hess_wrapper!( + batch_solver::AbstractBatchMPCSolver, + bkkt::SparseUniformBatchKKTSystem, +) + bcb = batch_solver.bcb + ws = batch_solver.workspace + nzVals = bkkt.nzVals + n_tot = bkkt.n_tot + nnzh = bkkt.nnzh + + if nnzh > 0 + hess = view(nzVals, n_tot+1:n_tot+nnzh, :) + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess) + end + return +end + +function MadNLP.initialize!(bkkt::SparseUniformBatchKKTSystem{T}) where T + pr_diag(bkkt) .= one(T) + if bkkt.nnzh > 0 + n_tot = bkkt.n_tot + view(bkkt.nzVals, n_tot+1:n_tot+bkkt.nnzh, :) .= zero(T) + end + du_diag(bkkt) .= zero(T) + + fill!(bkkt.reg, zero(T)) + fill!(bkkt.l_diag, one(T)) + fill!(bkkt.u_diag, one(T)) + fill!(bkkt.l_lower, zero(T)) + fill!(bkkt.u_lower, zero(T)) + + fill!(bkkt.aug_com_nzvals, zero(T)) + return +end + +function LinearAlgebra.mul!( + w::BatchUnreducedKKTVector{T}, + bkkt::SparseUniformBatchKKTSystem{T}, + x::BatchUnreducedKKTVector{T}, + alpha = one(T), + beta = zero(T), +) where T + nzV = bkkt.nzVals + wp = bkkt._mul_w_primal + wd = bkkt._mul_w_dual + + # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) + xv = MadNLP.full(x) + @views bkkt.hess_buffer .= nzV[bkkt.hess_nz_map, :] .* xv[bkkt.hess_var_map, :] + mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) + MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp + + # mul!(primal(w), jac_com', dual(x), alpha, one(T)) + @views bkkt.jt_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.jt_con_map_full, :] + mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) + MadNLP.primal(w) .+= alpha .* wp + + # mul!(dual(w), jac_com, primal(x), alpha, beta) + @views bkkt.j_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.j_var_map, :] + mul!(wd, bkkt.j_scatter, bkkt.j_buffer) + MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd + _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) + return w +end diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 72af7c0e..44986946 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -94,7 +94,7 @@ function set_aug_diagonal_reg!(kkt, solver::AbstractBatchMPCSolver) kkt.reg .= solver.del_w du_diag(kkt) .= solver.del_c - + kkt.l_diag .= xl_r .- x_lr kkt.u_diag .= x_ur .- xu_r diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl index b6aedd70..1bb206f7 100644 --- a/src/batch/madnlp/kkt.jl +++ b/src/batch/madnlp/kkt.jl @@ -43,6 +43,53 @@ struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatc active_batch_size::Base.RefValue{Int} end +"""Batched normal-equations KKT for LPs with shared A (ObjRHSBatch).""" +struct NormalRHSBatchKKTSystem{T, LS, VT, MT, VI, VI32, CSC_T, OP_T, OPT_T} <: AbstractBatchKKTSystem{T} + # Condensed system C = A Σ⁻¹ Aᵀ (m × m, lower triangular CSC) + aug_com_colptr::VI32 # shared colptr for C + aug_com_rowval::VI32 # shared rowval for C + aug_com_nzvals::MT # (nnz_C × batch_size) per-instance values + + # Shared Jacobian + A_coo::MadNLP.SparseMatrixCOO # COO (with slacks) + AT_csc::CSC_T # Aᵀ in CSC format + A_csr_map::VI # COO → CSC transfer map + jac::VT # view into A_coo.V (first nnzj entries) + + # SpMM operators + AT_op::OP_T # Aᵀ * X operator (non-transpose) + A_op::OPT_T # A * X operator (transpose) + + # Batch linear solver (on m × m system) + batch_solver::LS + rhs_buffer::VT # contiguous m * batch_size for solve + batch_size::Int + + # Per-instance diagonals + pr_diag::MT # (n_tot × batch_size) — full Σ + du_diag::MT # (m × batch_size) — dual regularization + reg::MT # (n_tot × batch_size) — primal regularization + l_diag::MT; u_diag::MT # (nlb/nub × batch_size) + l_lower::MT; u_lower::MT # (nlb/nub × batch_size) + + # Solve buffers + buffer_n::MT # (n_tot × batch_size) + buffer_m::MT # (m × batch_size) + + # Info + ind_ineq::VI + n_tot::Int # n + n_slack + m::Int + + # Active set tracking + batch_map::Vector{Int} + batch_map_rev::Vector{Int} + active_batch_size::Base.RefValue{Int} +end + +pr_diag(kkt::NormalRHSBatchKKTSystem) = kkt.pr_diag +du_diag(kkt::NormalRHSBatchKKTSystem) = kkt.du_diag + pr_diag(bkkt::SparseUniformBatchKKTSystem) = view(bkkt.nzVals, 1:bkkt.n_tot, :) function du_diag(bkkt::SparseUniformBatchKKTSystem) du_off = size(bkkt.nzVals, 1) - bkkt.m @@ -316,6 +363,28 @@ function MadNLP.eval_lag_hess_wrapper!( return end +function MadNLP.eval_jac_wrapper!( + batch_solver::AbstractBatchMPCSolver, + kkt::NormalRHSBatchKKTSystem, +) + bcb = batch_solver.bcb + ws = batch_solver.workspace + MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) + MadNLP._eval_jac_wrapper!(bcb, ws.bx, bcb.jac_buffer) + # A is shared across batch — take column 1 (all identical) + kkt.jac .= view(bcb.jac_buffer, :, 1) + MadNLP.compress_jacobian!(kkt) + return +end + +function MadNLP.eval_lag_hess_wrapper!( + ::AbstractBatchMPCSolver, + ::NormalRHSBatchKKTSystem, +) + # No-op for LP (nnzh == 0) + return +end + function MadNLP.initialize!(bkkt::SparseUniformBatchKKTSystem{T}) where T pr_diag(bkkt) .= one(T) if bkkt.nnzh > 0 @@ -378,3 +447,295 @@ function LinearAlgebra.mul!( _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) return w end + +# =================================================================== +# NormalRHSBatchKKTSystem methods +# =================================================================== + +function MadNLP.create_kkt_system( + ::Type{NormalKKTSystem}, + bcb::UniformBatchCallback{T, VT}, + uniformbatch_linear_solver = LoopedBatchLinearSolver; + opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), +) where {T, VT} + batch_size = bcb.batch_size + + n = bcb.nvar + m = bcb.ncon + ind_ineq = bcb.ind_ineq + n_slack = length(ind_ineq) + nlb = length(bcb.ind_lb) + nub = length(bcb.ind_ub) + n_tot = n + n_slack + + @assert bcb.nnzh == 0 "NormalRHSBatchKKTSystem supports only LPs (nnzh must be 0)" + + # Build Jacobian with slack variables in COO + nnzj = bcb.nnzj + jac_sparsity_I = MadNLP.create_array(bcb, Int32, nnzj) + jac_sparsity_J = MadNLP.create_array(bcb, Int32, nnzj) + MadNLP._jac_sparsity_wrapper!(bcb, jac_sparsity_I, jac_sparsity_J) + + I = MadNLP.create_array(bcb, Int32, nnzj + n_slack) + J = MadNLP.create_array(bcb, Int32, nnzj + n_slack) + V = VT(undef, nnzj + n_slack) + I[1:nnzj] .= jac_sparsity_I + J[1:nnzj] .= jac_sparsity_J + I[nnzj+1:nnzj+n_slack] .= ind_ineq + J[nnzj+1:nnzj+n_slack] .= (n+1:n+n_slack) + A_coo = MadNLP.SparseMatrixCOO(m, n_tot, I, J, V) + jac = MadNLP._madnlp_unsafe_wrap(V, nnzj, 1) + + # Fill values with continuous range to get COO→CSC mapping + A_coo.V .= 1:(nnzj + n_slack) + Ap, Aj, Ax = coo_to_csr(A_coo) + A_csr_map = convert.(Int, Ax) + + # Store transposed matrix Aᵀ in CSC format + CSC = sparse_csc_format(VT) + AT_csc = CSC <: SparseArrays.SparseMatrixCSC ? CSC(n_tot, m, Ap, Aj, Ax) : CSC(Ap, Aj, Ax, (n_tot, m)) + + # Build C = AΣ⁻¹Aᵀ sparsity structure + if CSC <: SparseArrays.SparseMatrixCSC + AAp, AAj = build_normal_system(m, n_tot, Ap, Aj) + else + AAp_h, AAj_h = build_normal_system(m, n_tot, Vector(Ap), Vector(Aj)) + VII = typeof(Ap) + AAp = VII(AAp_h) + AAj = VII(AAj_h) + end + nnz_C = length(AAj) + + # Allocate batch nzvals for C + proto = VT(undef, 0) + MT = typeof(similar(proto, T, 0, 0)) + aug_com_nzvals = MT(undef, nnz_C, batch_size) + fill!(aug_com_nzvals, zero(T)) + + # Create template CSC for solver + AAx = VT(undef, nnz_C) + aug_com = CSC <: SparseArrays.SparseMatrixCSC ? CSC(m, m, AAp, AAj, AAx) : CSC(AAp, AAj, AAx, (m, m)) + + # Create batch linear solver on m × m system + batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, m; opt=opt_linear_solver) + + # SpMM operators + AT_op = _wrap_spmm_operator(AT_csc, batch_size, 'N') + A_op = _wrap_spmm_operator(AT_csc, batch_size, 'T') + + rhs_buffer = VT(undef, m * batch_size) + fill!(rhs_buffer, zero(T)) + + # Per-instance diagonals and buffers + pr_diag_mat = MT(undef, n_tot, batch_size) + du_diag_mat = MT(undef, m, batch_size) + reg_mat = MT(undef, n_tot, batch_size) + l_diag = MT(undef, nlb, batch_size) + u_diag = MT(undef, nub, batch_size) + l_lower = MT(undef, nlb, batch_size) + u_lower = MT(undef, nub, batch_size) + buffer_n = MT(undef, n_tot, batch_size) + buffer_m = MT(undef, m, batch_size) + + batch_map = collect(1:batch_size) + batch_map_rev = collect(1:batch_size) + active_batch_size = Ref(batch_size) + + VI = typeof(A_csr_map) + VI32 = typeof(I) + LS = typeof(batch_ls) + CSC_T = typeof(AT_csc) + OP_T = typeof(AT_op) + OPT_T = typeof(A_op) + + fill!(jac, zero(T)) + + return NormalRHSBatchKKTSystem{T, LS, VT, MT, VI, VI32, CSC_T, OP_T, OPT_T}( + AAp, AAj, aug_com_nzvals, + A_coo, AT_csc, A_csr_map, jac, + AT_op, A_op, + batch_ls, rhs_buffer, batch_size, + pr_diag_mat, du_diag_mat, reg_mat, + l_diag, u_diag, l_lower, u_lower, + buffer_n, buffer_m, + ind_ineq, n_tot, m, + batch_map, batch_map_rev, active_batch_size, + ) +end + +function update_active_set!(kkt::NormalRHSBatchKKTSystem, status::Vector{MadNLP.Status}) + active_pos = 0 + for i in 1:kkt.batch_size + if status[i] == MadNLP.REGULAR + active_pos += 1 + kkt.batch_map[i] = active_pos + kkt.batch_map_rev[active_pos] = i + else + kkt.batch_map[i] = 0 + end + end + for j in (active_pos + 1):kkt.batch_size + kkt.batch_map_rev[j] = 0 + end + kkt.active_batch_size[] = active_pos +end + +function MadNLP.initialize!(kkt::NormalRHSBatchKKTSystem{T}) where T + fill!(kkt.pr_diag, one(T)) + fill!(kkt.du_diag, zero(T)) + fill!(kkt.reg, zero(T)) + fill!(kkt.l_diag, one(T)) + fill!(kkt.u_diag, one(T)) + fill!(kkt.l_lower, zero(T)) + fill!(kkt.u_lower, zero(T)) + fill!(kkt.buffer_n, zero(T)) + fill!(kkt.buffer_m, zero(T)) + fill!(kkt.aug_com_nzvals, zero(T)) + return +end + +function MadNLP.is_inertia_correct(kkt::NormalRHSBatchKKTSystem, num_pos, num_zero, num_neg) + return (num_zero == 0) && (num_pos == kkt.m) +end + +function MadNLP.compress_jacobian!(kkt::NormalRHSBatchKKTSystem) + n_slack = length(kkt.ind_ineq) + kkt.A_coo.V[end-n_slack+1:end] .= -1.0 + fill!(_nzval(kkt.AT_csc), 0.0) + for i in eachindex(kkt.A_csr_map) + _nzval(kkt.AT_csc)[i] = kkt.A_coo.V[kkt.A_csr_map[i]] + end + return +end + +function MadNLP.build_kkt!(kkt::NormalRHSBatchKKTSystem) + D = kkt.buffer_n + D .= 1.0 ./ kkt.pr_diag + + Ap = _colptr(kkt.AT_csc) + Aj = _rowval(kkt.AT_csc) + Ax = _nzval(kkt.AT_csc) + Cp = kkt.aug_com_colptr + Cj = kkt.aug_com_rowval + + assemble_normal_system_batch!(kkt.m, kkt.n_tot, Ap, Aj, Ax, Cp, Cj, kkt.aug_com_nzvals, D) + return +end + +function MadNLP.factorize_kkt!(kkt::NormalRHSBatchKKTSystem) + na = kkt.active_batch_size[] + nzvals = kkt.aug_com_nzvals + @inbounds for j in 1:na + i = kkt.batch_map_rev[j] + i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) + end + _active_factorize!(kkt.batch_solver, na) + return +end + +function _solve_condensed!(kkt::NormalRHSBatchKKTSystem{T}, r2::AbstractMatrix{T}) where T + na = kkt.active_batch_size[] + bs = kkt.batch_size + m = kkt.m + + rhs = kkt.rhs_buffer + rhs_mat = reshape(rhs, m, bs) + copyto!(rhs_mat, r2) + # Pack active columns + @inbounds for j in 1:na + i = kkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) + end + _active_solve!(kkt.batch_solver, rhs, na, m) + # Unpack + @inbounds for j in na:-1:1 + i = kkt.batch_map_rev[j] + i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) + end + @inbounds for i in 1:bs + kkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) + end + copyto!(r2, rhs_mat) +end + +function MadNLP.reduce_rhs!(kkt::NormalRHSBatchKKTSystem, d::BatchUnreducedKKTVector) + MadNLP.reduce_rhs!( + xp_lr(d), MadNLP.dual_lb(d), kkt.l_diag, + xp_ur(d), MadNLP.dual_ub(d), kkt.u_diag, + ) + return +end + +function MadNLP.finish_aug_solve!(kkt::NormalRHSBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + dzl = MadNLP.dual_lb(d) + dzu = MadNLP.dual_ub(d) + dzl .= (.-dzl .+ kkt.l_lower .* xp_lr(d)) ./ kkt.l_diag + dzu .= (dzu .- kkt.u_lower .* xp_ur(d)) ./ kkt.u_diag + return +end + +function MadNLP.solve_kkt!(kkt::NormalRHSBatchKKTSystem{T}, batch_solver::AbstractBatchMPCSolver) where T + d = batch_solver.d + MadNLP.reduce_rhs!(kkt, d) + + r1 = kkt.buffer_n # (n_tot, bs) + r2 = kkt.buffer_m # (m, bs) + Σ = kkt.pr_diag + + wx = MadNLP.primal(d) + wy = MadNLP.dual(d) + + # Condensed RHS: r2 = A Σ⁻¹ r₁ - r₂ + r1 .= wx ./ Σ + r2 .= wy + mul!(r2, kkt.A_op, r1, one(T), -one(T)) # A * (wx/Σ) - wy + + # Solve m×m system + _solve_condensed!(kkt, r2) + + # Back-substitute: Δx = Σ⁻¹(wx - Aᵀ Δy) + wy .= r2 # Δy → copy to d + r1 .= wx # original r₁ + mul!(r1, kkt.AT_op, r2, -one(T), one(T)) # wx - Aᵀ Δy (use r2, not SubArray wy) + wx .= r1 ./ Σ # Σ⁻¹(wx - Aᵀ Δy) + + MadNLP.finish_aug_solve!(kkt, batch_solver) + return +end + +function MadNLP.jtprod!(res::AbstractMatrix, kkt::NormalRHSBatchKKTSystem, y::BatchVector) + copyto!(kkt.buffer_m, MadNLP.full(y)) + mul!(kkt.buffer_n, kkt.AT_op, kkt.buffer_m) + copyto!(res, kkt.buffer_n) + return res +end + +function MadNLP.jtprod!(jacl::BatchVector, kkt::NormalRHSBatchKKTSystem, y::BatchVector) + return MadNLP.jtprod!(MadNLP.full(jacl), kkt, y) +end + +function LinearAlgebra.mul!( + w::BatchUnreducedKKTVector{T}, + kkt::NormalRHSBatchKKTSystem{T}, + v::BatchUnreducedKKTVector{T}, + alpha = one(T), + beta = zero(T), +) where T + wp = kkt.buffer_n + wd = kkt.buffer_m + + # A * vx: copy primal(v) → wp, compute A * wp → wd, apply to dual(w) + copyto!(wp, MadNLP.primal(v)) + mul!(wd, kkt.A_op, wp) + MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd + + # Aᵀ * vy: copy dual(v) → wd, compute Aᵀ * wd → wp, apply to primal(w) + copyto!(wd, MadNLP.dual(v)) + mul!(wp, kkt.AT_op, wd) + MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp + + # Diagonal + bound contributions + _kktmul!(w, v, kkt.reg, kkt.du_diag, kkt.l_lower, kkt.u_lower, kkt.l_diag, kkt.u_diag, alpha, beta) + return w +end From 9c7d0e862bba6469bbf918a1fded2db6e1fdfa0f Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 15:15:51 -0400 Subject: [PATCH 12/51] rm old kkt.jl --- src/batch/madnlp/kkt.jl | 741 ---------------------------------------- 1 file changed, 741 deletions(-) delete mode 100644 src/batch/madnlp/kkt.jl diff --git a/src/batch/madnlp/kkt.jl b/src/batch/madnlp/kkt.jl deleted file mode 100644 index 1bb206f7..00000000 --- a/src/batch/madnlp/kkt.jl +++ /dev/null @@ -1,741 +0,0 @@ -abstract type AbstractBatchKKTSystem{T} end - -"""Batched version of SparseKKTSystem""" -struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} - nzVals::MT # (aug_mat_length × batch_size) COO nonzero values - aug_I::VI32 # shared row indices - aug_J::VI32 # shared column indices - batch_solver::LS # batched linear solver - rhs_buffer::VT # contiguous (n+m)*batch_size for batch solve - batch_size::Int - aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values - batch_csc_map::VI # flattened COO→CSC map for all instances - n_tot::Int # n + n_slack (total primal variables) - m::Int # number of constraints - nnzh::Int # number of Hessian nonzeros - # Diagonal and bound data (for _kktmul!) - reg::MT # (n_tot × batch_size) primal regularization - l_diag::MT # (nlb × batch_size) lower bound diagonals - u_diag::MT # (nub × batch_size) upper bound diagonals - l_lower::MT # (nlb × batch_size) lower bound multipliers - u_lower::MT # (nub × batch_size) upper bound multipliers - # Hessian scatter (for mul!) - hess_scatter::SMT # (n_tot × n_hess_sym) scatter matrix - hess_nz_map::VI # nzVals row indices (with sym duplication) - hess_var_map::VI # variable indices for x multiplication - hess_buffer::MT # (n_hess_sym × batch_size) workspace - # J^T scatter (for jtprod! and mul!) - jt_scatter::SMT # (n_tot × n_jac_aug) scatter: S[var_idx, k] = 1 - jt_nz_map::VI # nzVals row indices for Jacobian entries - jt_con_map::VI # maps each Jac nonzero to its constraint index (1:m) - jt_con_map_full::VI # jt_con_map offset by n_tot (for indexing into full KKT vector) - jt_buffer::MT # (n_jac_aug × batch_size) buffer for jtprod - # J scatter (for mul!) - j_scatter::SMT # (m × n_jac_aug) scatter: S[con_idx, k] = 1 - j_var_map::VI # variable indices for J entries - j_buffer::MT # (n_jac_aug × batch_size) buffer for jprod - # Workspace for mul! (GPU needs full matrices, not SubArray views) - _mul_w_primal::MT # (n_tot × batch_size) - _mul_w_dual::MT # (m × batch_size) - # Batch tracking - batch_map::Vector{Int} # original index → active position (0 if inactive) - batch_map_rev::Vector{Int} # active position → original index - active_batch_size::Base.RefValue{Int} -end - -"""Batched normal-equations KKT for LPs with shared A (ObjRHSBatch).""" -struct NormalRHSBatchKKTSystem{T, LS, VT, MT, VI, VI32, CSC_T, OP_T, OPT_T} <: AbstractBatchKKTSystem{T} - # Condensed system C = A Σ⁻¹ Aᵀ (m × m, lower triangular CSC) - aug_com_colptr::VI32 # shared colptr for C - aug_com_rowval::VI32 # shared rowval for C - aug_com_nzvals::MT # (nnz_C × batch_size) per-instance values - - # Shared Jacobian - A_coo::MadNLP.SparseMatrixCOO # COO (with slacks) - AT_csc::CSC_T # Aᵀ in CSC format - A_csr_map::VI # COO → CSC transfer map - jac::VT # view into A_coo.V (first nnzj entries) - - # SpMM operators - AT_op::OP_T # Aᵀ * X operator (non-transpose) - A_op::OPT_T # A * X operator (transpose) - - # Batch linear solver (on m × m system) - batch_solver::LS - rhs_buffer::VT # contiguous m * batch_size for solve - batch_size::Int - - # Per-instance diagonals - pr_diag::MT # (n_tot × batch_size) — full Σ - du_diag::MT # (m × batch_size) — dual regularization - reg::MT # (n_tot × batch_size) — primal regularization - l_diag::MT; u_diag::MT # (nlb/nub × batch_size) - l_lower::MT; u_lower::MT # (nlb/nub × batch_size) - - # Solve buffers - buffer_n::MT # (n_tot × batch_size) - buffer_m::MT # (m × batch_size) - - # Info - ind_ineq::VI - n_tot::Int # n + n_slack - m::Int - - # Active set tracking - batch_map::Vector{Int} - batch_map_rev::Vector{Int} - active_batch_size::Base.RefValue{Int} -end - -pr_diag(kkt::NormalRHSBatchKKTSystem) = kkt.pr_diag -du_diag(kkt::NormalRHSBatchKKTSystem) = kkt.du_diag - -pr_diag(bkkt::SparseUniformBatchKKTSystem) = view(bkkt.nzVals, 1:bkkt.n_tot, :) -function du_diag(bkkt::SparseUniformBatchKKTSystem) - du_off = size(bkkt.nzVals, 1) - bkkt.m - return view(bkkt.nzVals, du_off+1:du_off+bkkt.m, :) -end - -function MadNLP.create_kkt_system( - ::Type{MadNLP.SparseKKTSystem}, - bcb::UniformBatchCallback{T, VT}, - uniformbatch_linear_solver = LoopedBatchLinearSolver; - opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), -) where {T, VT} - batch_size = bcb.batch_size - - n_slack = length(bcb.ind_ineq) - n = bcb.nvar - m = bcb.ncon - jac_sparsity_I = MadNLP.create_array(bcb, Int32, bcb.nnzj) - jac_sparsity_J = MadNLP.create_array(bcb, Int32, bcb.nnzj) - MadNLP._jac_sparsity_wrapper!(bcb, jac_sparsity_I, jac_sparsity_J) - - hess_sparsity_I, hess_sparsity_J = MadNLP.build_hessian_structure(bcb, MadNLP.ExactHessian) - - nlb = length(bcb.ind_lb) - nub = length(bcb.ind_ub) - - MadNLP.force_lower_triangular!(hess_sparsity_I, hess_sparsity_J) - - ind_ineq = bcb.ind_ineq - - n_slack = length(ind_ineq) - n_jac = length(jac_sparsity_I) - n_hess = length(hess_sparsity_I) - n_tot = n + n_slack - - aug_vec_length = n_tot+m - aug_mat_length = n_tot+m+n_hess+n_jac+n_slack - - I = MadNLP.create_array(bcb, Int32, aug_mat_length) - J = MadNLP.create_array(bcb, Int32, aug_mat_length) - nzVals = similar(bcb.con_buffer, aug_mat_length, batch_size) - fill!(nzVals, zero(T)) - V = _madnlp_unsafe_column_wrap(nzVals, aug_mat_length, 1, VT) - - offset = n_tot+n_jac+n_slack+n_hess+m - - I[1:n_tot] .= 1:n_tot - I[n_tot+1:n_tot+n_hess] = hess_sparsity_I - I[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= (jac_sparsity_I.+n_tot) - I[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= ind_ineq .+ n_tot - I[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) - - J[1:n_tot] .= 1:n_tot - J[n_tot+1:n_tot+n_hess] = hess_sparsity_J - J[n_tot+n_hess+1:n_tot+n_hess+n_jac] .= jac_sparsity_J - J[n_tot+n_hess+n_jac+1:n_tot+n_hess+n_jac+n_slack] .= (n+1:n+n_slack) - J[n_tot+n_hess+n_jac+n_slack+1:offset] .= (n_tot+1:n_tot+m) - - aug_raw = MadNLP.SparseMatrixCOO(aug_vec_length, aug_vec_length, I, J, V) - aug_com, aug_csc_map = MadNLP.coo_to_csc(aug_raw) - - nnz_csc = SparseArrays.nnz(aug_com) - aug_com_nzvals = similar(nzVals, nnz_csc, batch_size) - fill!(aug_com_nzvals, zero(T)) - - csc_offsets = similar(aug_csc_map, 1, batch_size) - csc_offsets .= (0:batch_size-1)' .* nnz_csc - batch_csc_map = vec(aug_csc_map .+ csc_offsets) - - batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, aug_vec_length; opt=opt_linear_solver) - - rhs_buffer = VT(undef, aug_vec_length * batch_size) - fill!(rhs_buffer, zero(T)) - - jac_range = n_tot+n_hess+1:n_tot+n_hess+n_jac+n_slack - - hess_scatter, hess_nz_map, hess_var_map, hess_buffer = _build_hess_scatter( - I, J, n_tot, n_hess, nzVals, aug_csc_map, batch_size, - ) - jt_scatter, jt_nz_map, jt_con_map, jt_buffer = _build_scatter( - I, J, jac_range, n_tot, nzVals, aug_csc_map, batch_size, - ) - jt_con_map_full = similar(jt_con_map) - jt_con_map_full .= jt_con_map .+ Int32(n_tot) - j_scatter, _, j_var_map, j_buffer = _build_jac_scatter( - I, J, jac_range, n_tot, m, nzVals, aug_csc_map, batch_size, - ) - - reg = similar(nzVals, n_tot, batch_size) - l_diag = similar(nzVals, nlb, batch_size) - u_diag = similar(nzVals, nub, batch_size) - l_lower = similar(nzVals, nlb, batch_size) - u_lower = similar(nzVals, nub, batch_size) - - _mul_w_primal = similar(nzVals, n_tot, batch_size) - _mul_w_dual = similar(nzVals, m, batch_size) - - batch_map = collect(1:batch_size) - batch_map_rev = collect(1:batch_size) - active_batch_size = Ref(batch_size) - - LS = typeof(batch_ls) - MT = typeof(nzVals) - VI = typeof(aug_csc_map) - VI32 = typeof(I) - SMT = typeof(jt_scatter) - return SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT}( - nzVals, I, J, batch_ls, rhs_buffer, batch_size, - aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, - reg, l_diag, u_diag, l_lower, u_lower, - hess_scatter, hess_nz_map, hess_var_map, hess_buffer, - jt_scatter, jt_nz_map, jt_con_map, jt_con_map_full, jt_buffer, - j_scatter, j_var_map, j_buffer, - _mul_w_primal, _mul_w_dual, - batch_map, batch_map_rev, active_batch_size, - ) -end - -function update_active_set!(bkkt::SparseUniformBatchKKTSystem, status::Vector{MadNLP.Status}) - active_pos = 0 - for i in 1:bkkt.batch_size - if status[i] == MadNLP.REGULAR - active_pos += 1 - bkkt.batch_map[i] = active_pos - bkkt.batch_map_rev[active_pos] = i - else - bkkt.batch_map[i] = 0 - end - end - for j in (active_pos + 1):bkkt.batch_size - bkkt.batch_map_rev[j] = 0 - end - bkkt.active_batch_size[] = active_pos -end - -function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) - na = bkkt.active_batch_size[] - nzvals = bkkt.aug_com_nzvals - @inbounds for j in 1:na # FIXME: refactor to avoid `na` launches - i = bkkt.batch_map_rev[j] - i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) - end - _active_factorize!(bkkt.batch_solver, na) - return -end - -function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractVector) where T - na = bkkt.active_batch_size[] - bs = bkkt.batch_size - n = length(rhs) ÷ bs - - rhs_mat = reshape(rhs, n, bs) - @inbounds for j in 1:na - i = bkkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) - end - _active_solve!(bkkt.batch_solver, rhs, na, n) - - @inbounds for j in na:-1:1 - i = bkkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) - end - - @inbounds for i in 1:bs - bkkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) - end - return rhs -end - -function MadNLP.reduce_rhs!(bkkt::SparseUniformBatchKKTSystem, d::BatchUnreducedKKTVector) - MadNLP.reduce_rhs!( - xp_lr(d), MadNLP.dual_lb(d), bkkt.l_diag, - xp_ur(d), MadNLP.dual_ub(d), bkkt.u_diag, - ) - return -end - -function MadNLP.finish_aug_solve!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) - d = batch_solver.d - dzl = MadNLP.dual_lb(d) - dzu = MadNLP.dual_ub(d) - dzl .= (.-dzl .+ bkkt.l_lower .* xp_lr(d)) ./ bkkt.l_diag - dzu .= (dzu .- bkkt.u_lower .* xp_ur(d)) ./ bkkt.u_diag - return -end - -function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) - d = batch_solver.d - n_pd = d.n + d.m - bs = bkkt.batch_size - na = bkkt.active_batch_size[] - - MadNLP.reduce_rhs!(bkkt, d) - - if na == 1 - # only one active: skip packing - orig_col = bkkt.batch_map_rev[1] - _active_solve!(bkkt.batch_solver, d.views[orig_col], 1, n_pd) - else - rhs = bkkt.rhs_buffer - pd_view = MadNLP.primal_dual(d) - copyto!(reshape(rhs, n_pd, bs), pd_view) - MadNLP.solve_linear_system!(bkkt, rhs) - copyto!(pd_view, reshape(rhs, n_pd, bs)) - end - - MadNLP.finish_aug_solve!(bkkt, batch_solver) - return -end - -function MadNLP.build_kkt!(bkkt::SparseUniformBatchKKTSystem) - MadNLP._transfer!(vec(bkkt.aug_com_nzvals), vec(bkkt.nzVals), bkkt.batch_csc_map) - return -end - -function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) - MadNLP.@trace(batch_solver.logger, "Factorization started.") - MadNLP.build_kkt!(batch_solver.kkt) - batch_solver.batch_cnt.linear_solver_time[] += @elapsed MadNLP.factorize_kkt!(batch_solver.kkt) - return -end - -function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) - @views bkkt.jt_buffer .= bkkt.nzVals[bkkt.jt_nz_map, :] .* MadNLP.full(y)[bkkt.jt_con_map, :] - mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) - return res -end - -function MadNLP.jtprod!(jacl::BatchVector, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) - return MadNLP.jtprod!(MadNLP.full(jacl), bkkt, y) -end - -function MadNLP.eval_jac_wrapper!( - batch_solver::AbstractBatchMPCSolver, - bkkt::SparseUniformBatchKKTSystem, -) - bcb = batch_solver.bcb - ws = batch_solver.workspace - nzVals = bkkt.nzVals - n_tot = bkkt.n_tot - nnzj = bcb.nnzj - n_slack = length(bcb.ind_ineq) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - jac_free = MadNLP._eval_jac_wrapper!(bcb, ws.bx, bcb.jac_buffer) - - jac_offset = n_tot + bkkt.nnzh - view(nzVals, jac_offset+1:jac_offset+nnzj, :) .= jac_free - - if n_slack > 0 - view(nzVals, jac_offset+nnzj+1:jac_offset+nnzj+n_slack, :) .= -one(eltype(nzVals)) - end - return -end - -function MadNLP.eval_lag_hess_wrapper!( - batch_solver::AbstractBatchMPCSolver, - bkkt::SparseUniformBatchKKTSystem, -) - bcb = batch_solver.bcb - ws = batch_solver.workspace - nzVals = bkkt.nzVals - n_tot = bkkt.n_tot - nnzh = bkkt.nnzh - - if nnzh > 0 - hess = view(nzVals, n_tot+1:n_tot+nnzh, :) - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess) - end - return -end - -function MadNLP.eval_jac_wrapper!( - batch_solver::AbstractBatchMPCSolver, - kkt::NormalRHSBatchKKTSystem, -) - bcb = batch_solver.bcb - ws = batch_solver.workspace - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP._eval_jac_wrapper!(bcb, ws.bx, bcb.jac_buffer) - # A is shared across batch — take column 1 (all identical) - kkt.jac .= view(bcb.jac_buffer, :, 1) - MadNLP.compress_jacobian!(kkt) - return -end - -function MadNLP.eval_lag_hess_wrapper!( - ::AbstractBatchMPCSolver, - ::NormalRHSBatchKKTSystem, -) - # No-op for LP (nnzh == 0) - return -end - -function MadNLP.initialize!(bkkt::SparseUniformBatchKKTSystem{T}) where T - pr_diag(bkkt) .= one(T) - if bkkt.nnzh > 0 - n_tot = bkkt.n_tot - view(bkkt.nzVals, n_tot+1:n_tot+bkkt.nnzh, :) .= zero(T) - end - du_diag(bkkt) .= zero(T) - - fill!(bkkt.reg, zero(T)) - fill!(bkkt.l_diag, one(T)) - fill!(bkkt.u_diag, one(T)) - fill!(bkkt.l_lower, zero(T)) - fill!(bkkt.u_lower, zero(T)) - - fill!(bkkt.aug_com_nzvals, zero(T)) - return -end - -@inbounds function _kktmul!( - w::BatchUnreducedKKTVector, - x::BatchUnreducedKKTVector, - reg, du_diag_val, l_lower, u_lower, l_diag, u_diag, - alpha, beta, -) - MadNLP.primal(w) .+= alpha .* reg .* MadNLP.primal(x) - MadNLP.dual(w) .+= alpha .* du_diag_val .* MadNLP.dual(x) - xp_lr(w) .-= alpha .* MadNLP.dual_lb(x) - xp_ur(w) .+= alpha .* MadNLP.dual_ub(x) - MadNLP.dual_lb(w) .= beta .* MadNLP.dual_lb(w) .+ alpha .* (xp_lr(x) .* l_lower .- MadNLP.dual_lb(x) .* l_diag) - MadNLP.dual_ub(w) .= beta .* MadNLP.dual_ub(w) .+ alpha .* (xp_ur(x) .* u_lower .+ MadNLP.dual_ub(x) .* u_diag) - return -end - -function LinearAlgebra.mul!( - w::BatchUnreducedKKTVector{T}, - bkkt::SparseUniformBatchKKTSystem{T}, - x::BatchUnreducedKKTVector{T}, - alpha = one(T), - beta = zero(T), -) where T - nzV = bkkt.nzVals - wp = bkkt._mul_w_primal - wd = bkkt._mul_w_dual - - # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) - xv = MadNLP.full(x) - @views bkkt.hess_buffer .= nzV[bkkt.hess_nz_map, :] .* xv[bkkt.hess_var_map, :] - mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) - MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp - - # mul!(primal(w), jac_com', dual(x), alpha, one(T)) - @views bkkt.jt_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.jt_con_map_full, :] - mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) - MadNLP.primal(w) .+= alpha .* wp - - # mul!(dual(w), jac_com, primal(x), alpha, beta) - @views bkkt.j_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.j_var_map, :] - mul!(wd, bkkt.j_scatter, bkkt.j_buffer) - MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd - _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) - return w -end - -# =================================================================== -# NormalRHSBatchKKTSystem methods -# =================================================================== - -function MadNLP.create_kkt_system( - ::Type{NormalKKTSystem}, - bcb::UniformBatchCallback{T, VT}, - uniformbatch_linear_solver = LoopedBatchLinearSolver; - opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), -) where {T, VT} - batch_size = bcb.batch_size - - n = bcb.nvar - m = bcb.ncon - ind_ineq = bcb.ind_ineq - n_slack = length(ind_ineq) - nlb = length(bcb.ind_lb) - nub = length(bcb.ind_ub) - n_tot = n + n_slack - - @assert bcb.nnzh == 0 "NormalRHSBatchKKTSystem supports only LPs (nnzh must be 0)" - - # Build Jacobian with slack variables in COO - nnzj = bcb.nnzj - jac_sparsity_I = MadNLP.create_array(bcb, Int32, nnzj) - jac_sparsity_J = MadNLP.create_array(bcb, Int32, nnzj) - MadNLP._jac_sparsity_wrapper!(bcb, jac_sparsity_I, jac_sparsity_J) - - I = MadNLP.create_array(bcb, Int32, nnzj + n_slack) - J = MadNLP.create_array(bcb, Int32, nnzj + n_slack) - V = VT(undef, nnzj + n_slack) - I[1:nnzj] .= jac_sparsity_I - J[1:nnzj] .= jac_sparsity_J - I[nnzj+1:nnzj+n_slack] .= ind_ineq - J[nnzj+1:nnzj+n_slack] .= (n+1:n+n_slack) - A_coo = MadNLP.SparseMatrixCOO(m, n_tot, I, J, V) - jac = MadNLP._madnlp_unsafe_wrap(V, nnzj, 1) - - # Fill values with continuous range to get COO→CSC mapping - A_coo.V .= 1:(nnzj + n_slack) - Ap, Aj, Ax = coo_to_csr(A_coo) - A_csr_map = convert.(Int, Ax) - - # Store transposed matrix Aᵀ in CSC format - CSC = sparse_csc_format(VT) - AT_csc = CSC <: SparseArrays.SparseMatrixCSC ? CSC(n_tot, m, Ap, Aj, Ax) : CSC(Ap, Aj, Ax, (n_tot, m)) - - # Build C = AΣ⁻¹Aᵀ sparsity structure - if CSC <: SparseArrays.SparseMatrixCSC - AAp, AAj = build_normal_system(m, n_tot, Ap, Aj) - else - AAp_h, AAj_h = build_normal_system(m, n_tot, Vector(Ap), Vector(Aj)) - VII = typeof(Ap) - AAp = VII(AAp_h) - AAj = VII(AAj_h) - end - nnz_C = length(AAj) - - # Allocate batch nzvals for C - proto = VT(undef, 0) - MT = typeof(similar(proto, T, 0, 0)) - aug_com_nzvals = MT(undef, nnz_C, batch_size) - fill!(aug_com_nzvals, zero(T)) - - # Create template CSC for solver - AAx = VT(undef, nnz_C) - aug_com = CSC <: SparseArrays.SparseMatrixCSC ? CSC(m, m, AAp, AAj, AAx) : CSC(AAp, AAj, AAx, (m, m)) - - # Create batch linear solver on m × m system - batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, m; opt=opt_linear_solver) - - # SpMM operators - AT_op = _wrap_spmm_operator(AT_csc, batch_size, 'N') - A_op = _wrap_spmm_operator(AT_csc, batch_size, 'T') - - rhs_buffer = VT(undef, m * batch_size) - fill!(rhs_buffer, zero(T)) - - # Per-instance diagonals and buffers - pr_diag_mat = MT(undef, n_tot, batch_size) - du_diag_mat = MT(undef, m, batch_size) - reg_mat = MT(undef, n_tot, batch_size) - l_diag = MT(undef, nlb, batch_size) - u_diag = MT(undef, nub, batch_size) - l_lower = MT(undef, nlb, batch_size) - u_lower = MT(undef, nub, batch_size) - buffer_n = MT(undef, n_tot, batch_size) - buffer_m = MT(undef, m, batch_size) - - batch_map = collect(1:batch_size) - batch_map_rev = collect(1:batch_size) - active_batch_size = Ref(batch_size) - - VI = typeof(A_csr_map) - VI32 = typeof(I) - LS = typeof(batch_ls) - CSC_T = typeof(AT_csc) - OP_T = typeof(AT_op) - OPT_T = typeof(A_op) - - fill!(jac, zero(T)) - - return NormalRHSBatchKKTSystem{T, LS, VT, MT, VI, VI32, CSC_T, OP_T, OPT_T}( - AAp, AAj, aug_com_nzvals, - A_coo, AT_csc, A_csr_map, jac, - AT_op, A_op, - batch_ls, rhs_buffer, batch_size, - pr_diag_mat, du_diag_mat, reg_mat, - l_diag, u_diag, l_lower, u_lower, - buffer_n, buffer_m, - ind_ineq, n_tot, m, - batch_map, batch_map_rev, active_batch_size, - ) -end - -function update_active_set!(kkt::NormalRHSBatchKKTSystem, status::Vector{MadNLP.Status}) - active_pos = 0 - for i in 1:kkt.batch_size - if status[i] == MadNLP.REGULAR - active_pos += 1 - kkt.batch_map[i] = active_pos - kkt.batch_map_rev[active_pos] = i - else - kkt.batch_map[i] = 0 - end - end - for j in (active_pos + 1):kkt.batch_size - kkt.batch_map_rev[j] = 0 - end - kkt.active_batch_size[] = active_pos -end - -function MadNLP.initialize!(kkt::NormalRHSBatchKKTSystem{T}) where T - fill!(kkt.pr_diag, one(T)) - fill!(kkt.du_diag, zero(T)) - fill!(kkt.reg, zero(T)) - fill!(kkt.l_diag, one(T)) - fill!(kkt.u_diag, one(T)) - fill!(kkt.l_lower, zero(T)) - fill!(kkt.u_lower, zero(T)) - fill!(kkt.buffer_n, zero(T)) - fill!(kkt.buffer_m, zero(T)) - fill!(kkt.aug_com_nzvals, zero(T)) - return -end - -function MadNLP.is_inertia_correct(kkt::NormalRHSBatchKKTSystem, num_pos, num_zero, num_neg) - return (num_zero == 0) && (num_pos == kkt.m) -end - -function MadNLP.compress_jacobian!(kkt::NormalRHSBatchKKTSystem) - n_slack = length(kkt.ind_ineq) - kkt.A_coo.V[end-n_slack+1:end] .= -1.0 - fill!(_nzval(kkt.AT_csc), 0.0) - for i in eachindex(kkt.A_csr_map) - _nzval(kkt.AT_csc)[i] = kkt.A_coo.V[kkt.A_csr_map[i]] - end - return -end - -function MadNLP.build_kkt!(kkt::NormalRHSBatchKKTSystem) - D = kkt.buffer_n - D .= 1.0 ./ kkt.pr_diag - - Ap = _colptr(kkt.AT_csc) - Aj = _rowval(kkt.AT_csc) - Ax = _nzval(kkt.AT_csc) - Cp = kkt.aug_com_colptr - Cj = kkt.aug_com_rowval - - assemble_normal_system_batch!(kkt.m, kkt.n_tot, Ap, Aj, Ax, Cp, Cj, kkt.aug_com_nzvals, D) - return -end - -function MadNLP.factorize_kkt!(kkt::NormalRHSBatchKKTSystem) - na = kkt.active_batch_size[] - nzvals = kkt.aug_com_nzvals - @inbounds for j in 1:na - i = kkt.batch_map_rev[j] - i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) - end - _active_factorize!(kkt.batch_solver, na) - return -end - -function _solve_condensed!(kkt::NormalRHSBatchKKTSystem{T}, r2::AbstractMatrix{T}) where T - na = kkt.active_batch_size[] - bs = kkt.batch_size - m = kkt.m - - rhs = kkt.rhs_buffer - rhs_mat = reshape(rhs, m, bs) - copyto!(rhs_mat, r2) - # Pack active columns - @inbounds for j in 1:na - i = kkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) - end - _active_solve!(kkt.batch_solver, rhs, na, m) - # Unpack - @inbounds for j in na:-1:1 - i = kkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) - end - @inbounds for i in 1:bs - kkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) - end - copyto!(r2, rhs_mat) -end - -function MadNLP.reduce_rhs!(kkt::NormalRHSBatchKKTSystem, d::BatchUnreducedKKTVector) - MadNLP.reduce_rhs!( - xp_lr(d), MadNLP.dual_lb(d), kkt.l_diag, - xp_ur(d), MadNLP.dual_ub(d), kkt.u_diag, - ) - return -end - -function MadNLP.finish_aug_solve!(kkt::NormalRHSBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) - d = batch_solver.d - dzl = MadNLP.dual_lb(d) - dzu = MadNLP.dual_ub(d) - dzl .= (.-dzl .+ kkt.l_lower .* xp_lr(d)) ./ kkt.l_diag - dzu .= (dzu .- kkt.u_lower .* xp_ur(d)) ./ kkt.u_diag - return -end - -function MadNLP.solve_kkt!(kkt::NormalRHSBatchKKTSystem{T}, batch_solver::AbstractBatchMPCSolver) where T - d = batch_solver.d - MadNLP.reduce_rhs!(kkt, d) - - r1 = kkt.buffer_n # (n_tot, bs) - r2 = kkt.buffer_m # (m, bs) - Σ = kkt.pr_diag - - wx = MadNLP.primal(d) - wy = MadNLP.dual(d) - - # Condensed RHS: r2 = A Σ⁻¹ r₁ - r₂ - r1 .= wx ./ Σ - r2 .= wy - mul!(r2, kkt.A_op, r1, one(T), -one(T)) # A * (wx/Σ) - wy - - # Solve m×m system - _solve_condensed!(kkt, r2) - - # Back-substitute: Δx = Σ⁻¹(wx - Aᵀ Δy) - wy .= r2 # Δy → copy to d - r1 .= wx # original r₁ - mul!(r1, kkt.AT_op, r2, -one(T), one(T)) # wx - Aᵀ Δy (use r2, not SubArray wy) - wx .= r1 ./ Σ # Σ⁻¹(wx - Aᵀ Δy) - - MadNLP.finish_aug_solve!(kkt, batch_solver) - return -end - -function MadNLP.jtprod!(res::AbstractMatrix, kkt::NormalRHSBatchKKTSystem, y::BatchVector) - copyto!(kkt.buffer_m, MadNLP.full(y)) - mul!(kkt.buffer_n, kkt.AT_op, kkt.buffer_m) - copyto!(res, kkt.buffer_n) - return res -end - -function MadNLP.jtprod!(jacl::BatchVector, kkt::NormalRHSBatchKKTSystem, y::BatchVector) - return MadNLP.jtprod!(MadNLP.full(jacl), kkt, y) -end - -function LinearAlgebra.mul!( - w::BatchUnreducedKKTVector{T}, - kkt::NormalRHSBatchKKTSystem{T}, - v::BatchUnreducedKKTVector{T}, - alpha = one(T), - beta = zero(T), -) where T - wp = kkt.buffer_n - wd = kkt.buffer_m - - # A * vx: copy primal(v) → wp, compute A * wp → wd, apply to dual(w) - copyto!(wp, MadNLP.primal(v)) - mul!(wd, kkt.A_op, wp) - MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd - - # Aᵀ * vy: copy dual(v) → wd, compute Aᵀ * wd → wp, apply to primal(w) - copyto!(wd, MadNLP.dual(v)) - mul!(wp, kkt.AT_op, wd) - MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp - - # Diagonal + bound contributions - _kktmul!(w, v, kkt.reg, kkt.du_diag, kkt.l_lower, kkt.u_lower, kkt.l_diag, kkt.u_diag, alpha, beta) - return w -end From 97bbcb266a2d5d45addab839f6b9fb5b53245970 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 15:26:36 -0400 Subject: [PATCH 13/51] git sources --- Project.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index cfcfdd75..f727f806 100644 --- a/Project.toml +++ b/Project.toml @@ -21,8 +21,7 @@ MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [sources] -NLPModels = {path = "../NLPModels.jl"} -QuadraticModels = {path = "../QuadraticModels.jl"} +QuadraticModels = {url = "https://github.com/klamike/QuadraticModels.jl.git", rev = "mk/rhsbatch"} [extensions] MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] From a724fefb129809857966d487a126bc65361ddfbe Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 15:34:28 -0400 Subject: [PATCH 14/51] [temp] disable testing on `lts` so `[sources]` table is used --- .github/workflows/action.yml | 2 +- .github/workflows/gpu.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/action.yml b/.github/workflows/action.yml index b2b37a6e..22a77c07 100644 --- a/.github/workflows/action.yml +++ b/.github/workflows/action.yml @@ -14,7 +14,7 @@ jobs: strategy: matrix: os: [ubuntu-22.04] - julia-version: ['lts', '1'] + julia-version: ['1'] julia-arch: [x64] steps: - uses: actions/checkout@v4 diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml index 88277097..2582ebc9 100644 --- a/.github/workflows/gpu.yml +++ b/.github/workflows/gpu.yml @@ -19,7 +19,7 @@ jobs: labels: cuda strategy: matrix: - julia-version: ['lts', '1'] + julia-version: ['1'] julia-arch: [x64] hsl-version: ['2025.7.11'] From 975f40caa916f2935c77ce3dbe5a8eaaf7331b7c Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 15:38:00 -0400 Subject: [PATCH 15/51] no normal yet --- src/batch/KKT/KKT.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/batch/KKT/KKT.jl b/src/batch/KKT/KKT.jl index 84ff805a..2ff64d46 100644 --- a/src/batch/KKT/KKT.jl +++ b/src/batch/KKT/KKT.jl @@ -16,4 +16,3 @@ abstract type AbstractBatchKKTSystem{T} end end include("Sparse/augmented.jl") -include("Sparse/normal.jl") \ No newline at end of file From 01138322e5df430c8729b7f6db8102aeaa3a149b Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 16:14:02 -0400 Subject: [PATCH 16/51] `using`? --- ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 2 +- ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index 07d870cb..7b172df7 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -9,7 +9,7 @@ using CUDA.CUSPARSE using CUDSS using KernelAbstractions import Atomix -import QuadraticModels: SparseMatrixCOO +using QuadraticModels: SparseMatrixCOO import LinearAlgebra: BlasFloat import MadIPM import MadNLP diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl index 4b440099..530fb6ee 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -1,4 +1,4 @@ -import QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData +using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData function NLPModels.obj!( bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, From 1d467b3813decc09ddddd80355721adcee3e5bcc Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 16:14:34 -0400 Subject: [PATCH 17/51] delete 19,000 tests --- test/batch_kernels/runtests.jl | 2 - test/batch_kernels/setup.jl | 848 +-------------------------------- 2 files changed, 1 insertion(+), 849 deletions(-) diff --git a/test/batch_kernels/runtests.jl b/test/batch_kernels/runtests.jl index e1c9645c..4cc433d8 100644 --- a/test/batch_kernels/runtests.jl +++ b/test/batch_kernels/runtests.jl @@ -5,8 +5,6 @@ include("test_complementarity.jl") include("test_step.jl") include("test_regularization.jl") - include("test_solver_steps.jl") include("test_batch_multi.jl") include("test_termination.jl") - include("test_iteration_match.jl") end diff --git a/test/batch_kernels/setup.jl b/test/batch_kernels/setup.jl index bf1b79c5..3f3fa710 100644 --- a/test/batch_kernels/setup.jl +++ b/test/batch_kernels/setup.jl @@ -4,7 +4,7 @@ using SparseArrays using MadNLP using MadIPM using QuadraticModels -using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel +using QuadraticModels: ObjRHSBatchQuadraticModel # ────────────────────────────────────────────────────────────── # Test problems @@ -176,194 +176,6 @@ const ALL_TEST_PROBLEMS = [ ("QP dense Hessian+mixed", _setup_dense_hess_mixed_qp), ] -# ────────────────────────────────────────────────────────────── -# Paired test problems for batch_size=2 (same structure, different data) -# ────────────────────────────────────────────────────────────── - -function _paired_lower_only() - Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp_a = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - qp_b = QuadraticModel([2.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.5, 0.5]) - return (qp_a, qp_b) -end - -function _paired_doubly_bounded() - Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [5.0, 5.0] - qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) - qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) - return (qp_a, qp_b) -end - -function _paired_scaled() - Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [1.0, 1.0, 1.0] - Arows = [1, 1, 1]; Acols = [1, 2, 3]; Avals = [200.0, 150.0, 100.0] - lvar = zeros(3); uvar = fill(Inf, 3) - qp_a = QuadraticModel([500.0, -300.0, 400.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=ones(3)) - qp_b = QuadraticModel([100.0, -200.0, 150.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5, 0.5]) - return (qp_a, qp_b) -end - -function _paired_all_ineq() - Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 1.5] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] - lvar = zeros(3); uvar = fill(Inf, 3) - qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.0, 0.0], ucon=[3.0, 3.0], lvar=lvar, uvar=uvar, x0=ones(3)) - qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5, 0.5], ucon=[4.0, 4.0], lvar=lvar, uvar=uvar, x0=[0.5, 1.5, 1.0]) - return (qp_a, qp_b) -end - -function _paired_mixed_bounds() - Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 1.5, 1.0] - Arows = [1, 1, 1, 1, 2, 2]; Acols = [1, 2, 3, 4, 1, 3] - Avals = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] - lvar = [0.0, -Inf, 0.0, -Inf]; uvar = [Inf, 5.0, 10.0, Inf] - qp_a = QuadraticModel([1.0, -1.0, 0.5, -0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, - x0=[1.0, 2.0, 5.0, 0.0]) - qp_b = QuadraticModel([2.0, -2.0, 1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, - x0=[0.5, 3.0, 4.0, 0.5]) - return (qp_a, qp_b) -end - -function _paired_dense_hess() - Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - return (qp_a, qp_b) -end - -function _paired_dense_hess_mixed() - Hrows = [1, 2, 2, 3, 3]; Hcols = [1, 1, 2, 2, 3]; Hvals = [5.0, 1.0, 4.0, 2.0, 3.0] - Arows = [1, 1, 1, 2, 2]; Acols = [1, 2, 3, 1, 3] - Avals = [1.0, 1.0, 1.0, 1.0, 1.0] - lvar = [0.0, -Inf, 0.0]; uvar = [Inf, 5.0, 10.0] - qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, - x0=[1.0, 2.0, 1.0]) - qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, - x0=[0.5, 1.5, 0.5]) - return (qp_a, qp_b) -end - -const PAIRED_PROBLEMS = [ - ("paired lower-only LP", _paired_lower_only), - ("paired doubly-bounded QP", _paired_doubly_bounded), - ("paired scaled QP", _paired_scaled), - ("paired all-ineq QP", _paired_all_ineq), - ("paired mixed-bounds QP", _paired_mixed_bounds), - ("paired dense-hess QP", _paired_dense_hess), - ("paired dense-hess+mixed QP", _paired_dense_hess_mixed), -] - -# Staggered convergence: easy QP (converges fast) + harder QP (converges slow) -# Same Hessian, Jacobian (ObjRHSBatch shares these). Differ only in c, lcon/ucon, x0. -function _paired_staggered() - Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - # Easy QP: symmetric objective, tight constraint → fast convergence - qp_easy = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - # Harder QP: asymmetric objective, looser constraint → more iterations - qp_hard = QuadraticModel([0.01, -0.99], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[5.0, 5.0]) - return (qp_easy, qp_hard) -end - -# Quad problems for batch_size=4 tests (all lower-bounded LPs, same structure) -function _quad_lower_only() - Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp1 = QuadraticModel([1.0, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - qp2 = QuadraticModel([2.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.5, 0.5]) - qp3 = QuadraticModel([0.5, 1.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.5, 1.0]) - qp4 = QuadraticModel([-1.0, 3.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[0.3, 0.2]) - return (qp1, qp2, qp3, qp4) -end - -# Quad problems for batch_size=4 with doubly-bounded QPs -function _quad_doubly_bounded() - Hrows = [1, 2]; Hcols = [1, 2]; Hvals = [1.0, 1.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [5.0, 5.0] - qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) - qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) - qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - qp4 = QuadraticModel([2.0, -2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[4.0], ucon=[4.0], lvar=lvar, uvar=uvar, x0=[2.0, 2.0]) - return (qp1, qp2, qp3, qp4) -end - -# Quad problems for batch_size=4 with non-diagonal Hessian QPs -function _quad_dense_hess() - Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.75, 0.75]) - qp4 = QuadraticModel([2.0, -2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5], ucon=[0.5], lvar=lvar, uvar=uvar, x0=[0.3, 0.2]) - return (qp1, qp2, qp3, qp4) -end - # ────────────────────────────────────────────────────────────── # Build initialized sequential solver # ────────────────────────────────────────────────────────────── @@ -526,661 +338,3 @@ function do_first_factorize!(seq, bat) MadNLP.factorize_kkt!(bat.kkt) end -# ────────────────────────────────────────────────────────────── -# Build initialized batch solver from multiple (potentially different) QPs -# ────────────────────────────────────────────────────────────── -function build_batch_from_qps(qps::Vector; kwargs...) - bnlp = ObjRHSBatchQuadraticModel(qps) - batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) - ws = batch_solver.workspace - bcb = batch_solver.bcb - opt = batch_solver.opt - - MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - batch_solver.workspace.bx; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(MadNLP.full(batch_solver.jacl), 0.0) - if opt.scaling - MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - Float64(opt.nlp_scaling_max_gradient), - batch_solver.workspace.bx) - end - MadNLP.initialize!(batch_solver.kkt) - MadIPM.init_regularization!(batch_solver, opt.regularization) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP.eval_f_wrapper(batch_solver, ws.bx) - MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) - MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) - MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) - MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) - - MadIPM.init_starting_point!(batch_solver) - fill!(ws.mu_batch, opt.mu_init) - fill!(ws.best_complementarity, typemax(Float64)) - fill!(ws.status, MadNLP.REGULAR) - fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) - fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) - fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) - batch_solver.batch_cnt.start_time[] = time() - fill!(batch_solver.batch_cnt.k, 0) - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return batch_solver -end - -# ────────────────────────────────────────────────────────────── -# Build initialized batch solver using BatchQuadraticModel (bs=1) -# ────────────────────────────────────────────────────────────── -function build_fullbatch(qp; kwargs...) - bnlp = BatchQuadraticModel([qp]) - batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) - ws = batch_solver.workspace - bcb = batch_solver.bcb - opt = batch_solver.opt - - MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - batch_solver.workspace.bx; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(MadNLP.full(batch_solver.jacl), 0.0) - if opt.scaling - MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - Float64(opt.nlp_scaling_max_gradient), - batch_solver.workspace.bx) - end - MadNLP.initialize!(batch_solver.kkt) - MadIPM.init_regularization!(batch_solver, opt.regularization) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP.eval_f_wrapper(batch_solver, ws.bx) - MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) - MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) - MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) - MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) - - MadIPM.init_starting_point!(batch_solver) - fill!(ws.mu_batch, opt.mu_init) - fill!(ws.best_complementarity, typemax(Float64)) - fill!(ws.status, MadNLP.REGULAR) - fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) - fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) - fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) - batch_solver.batch_cnt.start_time[] = time() - fill!(batch_solver.batch_cnt.k, 0) - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return batch_solver -end - -# ────────────────────────────────────────────────────────────── -# Build initialized batch solver from multiple QPs using BatchQuadraticModel -# ────────────────────────────────────────────────────────────── -function build_fullbatch_from_qps(qps::Vector; kwargs...) - bnlp = BatchQuadraticModel(qps) - batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) - ws = batch_solver.workspace - bcb = batch_solver.bcb - opt = batch_solver.opt - - MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - batch_solver.workspace.bx; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(MadNLP.full(batch_solver.jacl), 0.0) - if opt.scaling - MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - Float64(opt.nlp_scaling_max_gradient), - batch_solver.workspace.bx) - end - MadNLP.initialize!(batch_solver.kkt) - MadIPM.init_regularization!(batch_solver, opt.regularization) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP.eval_f_wrapper(batch_solver, ws.bx) - MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) - MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) - MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) - MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) - - MadIPM.init_starting_point!(batch_solver) - fill!(ws.mu_batch, opt.mu_init) - fill!(ws.best_complementarity, typemax(Float64)) - fill!(ws.status, MadNLP.REGULAR) - fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) - fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) - fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) - batch_solver.batch_cnt.start_time[] = time() - fill!(batch_solver.batch_cnt.k, 0) - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return batch_solver -end - -# ────────────────────────────────────────────────────────────── -# Paired problems for BatchQuadraticModel: DIFFERENT H and A values -# (same sparsity pattern, same bound pattern, different everything else) -# ────────────────────────────────────────────────────────────── - -# Diagonal Hessian + different A values, equality constraint -function _fullbatch_paired_diagonal() - Hrows = [1, 2]; Hcols = [1, 2] - Arows = [1, 1]; Acols = [1, 2] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [2.0, 1.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [3.0, 2.0]; - Arows=Arows, Acols=Acols, Avals=[2.0, 0.5], - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - return (qp_a, qp_b) -end - -# Off-diagonal Hessian + different H and A values, equality constraint -function _fullbatch_paired_dense_hess() - Hrows = [1, 2, 2]; Hcols = [1, 1, 2] - Arows = [1, 1]; Acols = [1, 2] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp_b = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; - Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - return (qp_a, qp_b) -end - -# Different H and A values, doubly-bounded variables -function _fullbatch_paired_doubly_bounded() - Hrows = [1, 2]; Hcols = [1, 2] - Arows = [1, 1]; Acols = [1, 2] - lvar = [0.0, 0.0]; uvar = [5.0, 5.0] - qp_a = QuadraticModel([1.0, -1.0], Hrows, Hcols, [2.0, 1.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[2.5, 2.5]) - qp_b = QuadraticModel([-0.5, 1.5], Hrows, Hcols, [3.0, 2.5]; - Arows=Arows, Acols=Acols, Avals=[2.0, 1.0], - lcon=[3.0], ucon=[3.0], lvar=lvar, uvar=uvar, x0=[1.5, 1.5]) - return (qp_a, qp_b) -end - -# Different H and A values, inequality constraints (exercises different scaling per instance) -function _fullbatch_paired_inequality() - Hrows = [1, 2, 3]; Hcols = [1, 2, 3] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3] - lvar = zeros(3); uvar = fill(Inf, 3) - qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, [2.0, 1.0, 1.5]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0, 1.0, 1.0], - lcon=[0.0, 0.0], ucon=[3.0, 3.0], lvar=lvar, uvar=uvar, x0=ones(3)) - qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, [3.0, 2.0, 1.0]; - Arows=Arows, Acols=Acols, Avals=[2.0, 0.5, 1.5, 0.5], - lcon=[0.5, 0.5], ucon=[4.0, 4.0], lvar=lvar, uvar=uvar, x0=[0.5, 1.5, 1.0]) - return (qp_a, qp_b) -end - -# Mixed bounds + different H and A values + off-diagonal Hessian -function _fullbatch_paired_mixed() - Hrows = [1, 2, 2, 3, 3]; Hcols = [1, 1, 2, 2, 3] - Arows = [1, 1, 1, 2, 2]; Acols = [1, 2, 3, 1, 3] - lvar = [0.0, -Inf, 0.0]; uvar = [Inf, 5.0, 10.0] - qp_a = QuadraticModel([1.0, -2.0, 0.5], Hrows, Hcols, [5.0, 1.0, 4.0, 2.0, 3.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0, 1.0, 1.0, 1.0], - lcon=[1.0, 0.5], ucon=[3.0, 1.5], lvar=lvar, uvar=uvar, - x0=[1.0, 2.0, 1.0]) - qp_b = QuadraticModel([2.0, -1.0, 1.0], Hrows, Hcols, [7.0, 2.0, 5.0, 1.0, 4.0]; - Arows=Arows, Acols=Acols, Avals=[2.0, 0.5, 1.0, 1.5, 0.5], - lcon=[2.0, 1.0], ucon=[4.0, 2.0], lvar=lvar, uvar=uvar, - x0=[0.5, 1.5, 0.5]) - return (qp_a, qp_b) -end - -# Large coefficients triggering scaling + different H and A values -function _fullbatch_paired_scaled() - Hrows = [1, 2, 3]; Hcols = [1, 2, 3] - Arows = [1, 1, 1]; Acols = [1, 2, 3] - lvar = zeros(3); uvar = fill(Inf, 3) - qp_a = QuadraticModel([500.0, -300.0, 400.0], Hrows, Hcols, [1.0, 1.0, 1.0]; - Arows=Arows, Acols=Acols, Avals=[200.0, 150.0, 100.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=ones(3)) - qp_b = QuadraticModel([100.0, -200.0, 150.0], Hrows, Hcols, [2.0, 3.0, 1.5]; - Arows=Arows, Acols=Acols, Avals=[300.0, 50.0, 50.0], - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5, 0.5]) - return (qp_a, qp_b) -end - -const FULLBATCH_PAIRED_PROBLEMS = [ - ("fullbatch diagonal", _fullbatch_paired_diagonal), - ("fullbatch dense Hessian", _fullbatch_paired_dense_hess), - ("fullbatch doubly-bounded", _fullbatch_paired_doubly_bounded), - ("fullbatch inequality", _fullbatch_paired_inequality), - ("fullbatch mixed bounds", _fullbatch_paired_mixed), - ("fullbatch scaled", _fullbatch_paired_scaled), -] - -# ────────────────────────────────────────────────────────────── -# Extract column n from a batch matrix/vector -# ────────────────────────────────────────────────────────────── -coln(x::AbstractMatrix, n::Int) = view(x, :, n) -coln(x::AbstractVector, ::Int) = x - -# ────────────────────────────────────────────────────────────── -# Assertion helpers: compare sequential solver state against batch column -# ────────────────────────────────────────────────────────────── - -function assert_termination_match(seq, bat, col; tol=1e-10) - @test abs(seq.inf_pr - bat.workspace.inf_pr[col]) < tol - @test abs(seq.inf_du - bat.workspace.inf_du[col]) < tol - @test abs(seq.inf_compl - bat.workspace.inf_compl[col]) < tol - @test seq.status == bat.workspace.status[col] - # dual_obj: used in infeasibility detection - seq_dobj = MadIPM.dual_objective(seq) - @test abs(seq_dobj - bat.workspace.dual_obj[col]) < tol -end - -function assert_regularization_match(seq, bat, col; tol=1e-10) - @test abs(seq.del_w - bat.del_w[col]) < tol - @test abs(seq.del_c - bat.del_c[col]) < tol -end - -function assert_prediction_match(seq, bat, col; tol=1e-10) - @test cmp(MadNLP.full(seq.d), coln(MadNLP.full(bat.d), col)) < tol - @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol - @test cmp(seq.correction_lb, coln(MadNLP.full(bat.correction_lb), col)) < tol - @test cmp(seq.correction_ub, coln(MadNLP.full(bat.correction_ub), col)) < tol -end - -function assert_direction_match(seq, bat, col; tol=1e-10) - @test cmp(MadNLP.full(seq.d), coln(MadNLP.full(bat.d), col)) < tol -end - -function assert_step_match(seq, bat, col; tol=1e-10) - @test abs(seq.alpha_p - bat.workspace.alpha_p[col]) < tol - @test abs(seq.alpha_d - bat.workspace.alpha_d[col]) < tol -end - -function assert_iterate_match(seq, bat, col; tol=1e-10) - @test cmp(MadNLP.full(seq.x), coln(MadNLP.full(bat.x), col)) < tol - @test cmp(seq.y, coln(MadNLP.full(bat.y), col)) < tol - @test cmp(MadNLP.full(seq.zl), coln(MadNLP.full(bat.zl), col)) < tol - @test cmp(MadNLP.full(seq.zu), coln(MadNLP.full(bat.zu), col)) < tol - @test cmp(MadNLP.full(seq.xl), coln(MadNLP.full(bat.xl), col)) < tol - @test cmp(MadNLP.full(seq.xu), coln(MadNLP.full(bat.xu), col)) < tol -end - -function assert_model_match(seq, bat, col; tol=1e-10) - @test abs(seq.obj_val - bat.workspace.obj_val[col]) < tol - @test cmp(MadNLP.primal(seq.f), coln(MadNLP.primal(bat.f), col)) < tol - @test cmp(seq.c, coln(MadNLP.full(bat.c), col)) < tol - @test cmp(seq.jacl, coln(MadNLP.full(bat.jacl), col)) < tol -end - -function assert_kkt_diagonals_match(seq, bat, col; tol=1e-10) - skkt = seq.kkt - bkkt = bat.kkt - # l_diag, u_diag, l_lower, u_lower - if length(skkt.l_diag) > 0 - @test cmp(skkt.l_diag, coln(bkkt.l_diag, col)) < tol - @test cmp(skkt.l_lower, coln(bkkt.l_lower, col)) < tol - end - if length(skkt.u_diag) > 0 - @test cmp(skkt.u_diag, coln(bkkt.u_diag, col)) < tol - @test cmp(skkt.u_lower, coln(bkkt.u_lower, col)) < tol - end - # reg (primal regularization) - @test cmp(skkt.reg, coln(bkkt.reg, col)) < tol - # pr_diag (includes reg + bound contributions) - @test cmp(skkt.pr_diag, coln(MadIPM.pr_diag(bkkt), col)) < tol - # du_diag (dual regularization) - @test cmp(skkt.du_diag, coln(MadIPM.du_diag(bkkt), col)) < tol -end - -function assert_init_match(seq, bat, col; tol=1e-10) - # Full state after init_starting_point! + initialization - @test cmp(MadNLP.full(seq.x), coln(MadNLP.full(bat.x), col)) < tol - @test cmp(seq.y, coln(MadNLP.full(bat.y), col)) < tol - @test cmp(MadNLP.full(seq.zl), coln(MadNLP.full(bat.zl), col)) < tol - @test cmp(MadNLP.full(seq.zu), coln(MadNLP.full(bat.zu), col)) < tol - @test cmp(MadNLP.full(seq.xl), coln(MadNLP.full(bat.xl), col)) < tol - @test cmp(MadNLP.full(seq.xu), coln(MadNLP.full(bat.xu), col)) < tol - @test abs(seq.obj_val - bat.workspace.obj_val[col]) < tol - @test cmp(MadNLP.primal(seq.f), coln(MadNLP.primal(bat.f), col)) < tol - @test cmp(seq.c, coln(MadNLP.full(bat.c), col)) < tol - @test cmp(seq.jacl, coln(MadNLP.full(bat.jacl), col)) < tol - # Regularization state - @test abs(seq.del_w - bat.del_w[col]) < tol - @test abs(seq.del_c - bat.del_c[col]) < tol - # Barrier parameter - @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol - # Normalization constants - @test abs(seq.norm_b - bat.workspace.norm_b[col]) < tol - @test abs(seq.norm_c - bat.workspace.norm_c[col]) < tol -end - -function assert_correction_match(seq, bat, col; tol=1e-10) - @test cmp(seq.correction_lb, coln(MadNLP.full(bat.correction_lb), col)) < tol - @test cmp(seq.correction_ub, coln(MadNLP.full(bat.correction_ub), col)) < tol -end - -function assert_kkt_matrix_match(seq, bat, col; tol=1e-10) - skkt = seq.kkt - bkkt = bat.kkt - @test cmp(SparseArrays.nonzeros(skkt.aug_com), coln(bkkt.aug_com_nzvals, col)) < tol -end - -function assert_barrier_match(seq, bat, col; mu_affine_seq=nothing, tol=1e-10) - # mu_curr: complementarity measure (set in update_barrier! → get_complementarity_measure) - @test abs(seq.mu_curr - bat.workspace.mu_curr[col]) < tol - # mu_batch (= sigma * mu_curr): final barrier parameter - @test abs(seq.mu - bat.workspace.mu_batch[col]) < tol - # mu_affine: affine complementarity measure (if provided) - if mu_affine_seq !== nothing - @test abs(mu_affine_seq - bat.workspace.mu_affine[col]) < tol - end -end - -function assert_tau_match(seq_tau, bat, col; tol=1e-10) - @test abs(seq_tau - bat.workspace.tau[col]) < tol -end - -function assert_rhs_match(seq, bat, col; tol=1e-10) - @test cmp(MadNLP.full(seq.p), coln(MadNLP.full(bat.p), col)) < tol -end - -function assert_reduce_rhs_match(seq, bat, col; tol=1e-10) - # After reduce_rhs!, the primal-dual part of d has been modified. - # The dual_lb and dual_ub parts contain the original RHS values - # that will be used in finish_aug_solve! - sd = MadNLP.full(seq.d) - bd = coln(MadNLP.full(bat.d), col) - @test cmp(sd, bd) < tol -end - -function assert_finish_aug_solve_match(seq, bat, col; tol=1e-10) - # After finish_aug_solve!, dzl and dzu are computed from the linear solve result - @test cmp(MadNLP.dual_lb(seq.d), coln(MadNLP.dual_lb(bat.d), col)) < tol - @test cmp(MadNLP.dual_ub(seq.d), coln(MadNLP.dual_ub(bat.d), col)) < tol -end - -# ────────────────────────────────────────────────────────────── -# Patching: copy sequential solver state into a batch column -# (test-only; prevents FP noise from accumulating across iterations) -# ────────────────────────────────────────────────────────────── - -function patch_batch_col_from_seq!(seq, bat, col) - ws = bat.workspace - - # Iterate: x (includes slacks), xl, xu, y, zl, zu - coln(MadNLP.full(bat.x), col) .= MadNLP.full(seq.x) - coln(MadNLP.full(bat.xl), col) .= MadNLP.full(seq.xl) - coln(MadNLP.full(bat.xu), col) .= MadNLP.full(seq.xu) - coln(MadNLP.full(bat.y), col) .= seq.y - coln(MadNLP.full(bat.zl), col) .= MadNLP.full(seq.zl) - coln(MadNLP.full(bat.zu), col) .= MadNLP.full(seq.zu) - - # Model evaluations: gradient (primal part), constraints, jacl, obj - coln(MadNLP.primal(bat.f), col) .= MadNLP.primal(seq.f) - coln(MadNLP.full(bat.c), col) .= seq.c - coln(MadNLP.full(bat.jacl), col) .= seq.jacl - ws.obj_val[col] = seq.obj_val - - # Barrier / termination scalars - ws.mu_batch[col] = seq.mu - ws.best_complementarity[col] = seq.best_complementarity -end - -# ────────────────────────────────────────────────────────────── -# Iteration harness: step sequential + batch solvers in lockstep -# ────────────────────────────────────────────────────────────── - -function run_iterations_bs1!(seq, bat, n_iters; tol=1e-10) - ws = bat.workspace - for iter in 1:n_iters - MadIPM.update_termination_criteria!(seq) - MadIPM.update_termination_criteria!(bat) - assert_termination_match(seq, bat, 1; tol) - - MadIPM.is_done(seq) && break - MadIPM.update_active_set!(bat.kkt, ws.status) - bat.kkt.active_batch_size[] == 0 && break - MadIPM._update_active_mask!(bat) - - # Factorize - MadIPM.factorize_system!(seq) - MadIPM.factorize_system!(bat) - assert_regularization_match(seq, bat, 1; tol) - assert_kkt_diagonals_match(seq, bat, 1; tol) - assert_kkt_matrix_match(seq, bat, 1; tol) - - # Prediction step (decomposed for finer-grained checks) - # 1. Set predictive RHS - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - assert_rhs_match(seq, bat, 1; tol) - # 2. Solve system (includes reduce_rhs! → linear solve → finish_aug_solve!) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - assert_direction_match(seq, bat, 1; tol) - assert_finish_aug_solve_match(seq, bat, 1; tol) - # 3. Affine step sizes (tau=1) + barrier update - alpha_aff_p, alpha_aff_d = MadIPM.get_fraction_to_boundary_step(seq, 1.0) - fill!(ws.tau, one(eltype(ws.tau))) - MadIPM.get_fraction_to_boundary_step!(bat) - MadIPM.zero_inactive_step!(bat) - mu_affine_seq = MadIPM.get_affine_complementarity_measure(seq, alpha_aff_p, alpha_aff_d) - MadIPM.get_affine_complementarity_measure!(bat, ws.alpha_p, ws.alpha_d) - MadIPM.get_correction!(seq, seq.correction_lb, seq.correction_ub) - MadIPM.get_correction!(bat, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub)) - assert_correction_match(seq, bat, 1; tol) - seq.mu_curr = MadIPM.update_barrier!(seq.opt.barrier_update, seq, mu_affine_seq) - MadIPM.update_barrier!(bat.opt.barrier_update, bat, ws.mu_affine) - assert_prediction_match(seq, bat, 1; tol) - assert_barrier_match(seq, bat, 1; mu_affine_seq, tol) - - # Mehrotra correction - MadIPM.set_correction_rhs!(seq, seq.kkt, seq.mu, seq.correction_lb, seq.correction_ub, seq.ind_lb, seq.ind_ub) - MadIPM.set_correction_rhs!(bat, bat.kkt, ws.mu_batch, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub), nothing, nothing) - assert_rhs_match(seq, bat, 1; tol) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - assert_direction_match(seq, bat, 1; tol) - assert_finish_aug_solve_match(seq, bat, 1; tol) - - # Update step (decomposed to check tau) - MadIPM.update_step!(seq.opt.step_rule, seq) - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.zero_inactive_step!(bat) - # Check tau: compute sequential tau from step rule - seq_tau = if seq.opt.step_rule isa MadIPM.ConservativeStep - seq.opt.step_rule.tau - elseif seq.opt.step_rule isa MadIPM.AdaptiveStep - max(1 - seq.mu, seq.opt.step_rule.tau_min) - else - 1.0 # MehrotraAdaptiveStep uses tau=1.0 internally - end - assert_tau_match(seq_tau, bat, 1; tol) - assert_step_match(seq, bat, 1; tol) - - # Apply step - MadIPM.apply_step!(seq) - MadIPM.apply_step!(bat) - assert_iterate_match(seq, bat, 1; tol) - - # Evaluate model - MadIPM.evaluate_model!(seq) - MadIPM.evaluate_model!(bat) - assert_model_match(seq, bat, 1; tol) - end - @test seq.status == bat.workspace.status[1] - @test seq.status == MadNLP.SOLVE_SUCCEEDED -end - -function run_iterations_bs2!(seq1, seq2, bat, n_iters; tol=1e-10, patch=false) - seq1_done = false - seq2_done = false - - # Patch initial state to eliminate init FP seed differences - if patch - patch_batch_col_from_seq!(seq1, bat, 1) - patch_batch_col_from_seq!(seq2, bat, 2) - end - - for iter in 1:n_iters - if !seq1_done; MadIPM.update_termination_criteria!(seq1); end - if !seq2_done; MadIPM.update_termination_criteria!(seq2); end - MadIPM.update_termination_criteria!(bat) - - if !seq1_done; assert_termination_match(seq1, bat, 1; tol); end - if !seq2_done; assert_termination_match(seq2, bat, 2; tol); end - - seq1_done = seq1_done || MadIPM.is_done(seq1) - seq2_done = seq2_done || MadIPM.is_done(seq2) - - MadIPM.update_active_set!(bat.kkt, bat.workspace.status) - bat.kkt.active_batch_size[] == 0 && break - MadIPM._update_active_mask!(bat) - - (seq1_done && seq2_done) && break - - # Factorize - if !seq1_done; MadIPM.factorize_system!(seq1); end - if !seq2_done; MadIPM.factorize_system!(seq2); end - MadIPM.factorize_system!(bat) - if !seq1_done - assert_regularization_match(seq1, bat, 1; tol) - assert_kkt_diagonals_match(seq1, bat, 1; tol) - end - if !seq2_done - assert_regularization_match(seq2, bat, 2; tol) - assert_kkt_diagonals_match(seq2, bat, 2; tol) - end - - # Prediction step - if !seq1_done; MadIPM.prediction_step!(seq1); end - if !seq2_done; MadIPM.prediction_step!(seq2); end - MadIPM.prediction_step!(bat) - if !seq1_done - assert_prediction_match(seq1, bat, 1; tol) - assert_barrier_match(seq1, bat, 1; tol) - end - if !seq2_done - assert_prediction_match(seq2, bat, 2; tol) - assert_barrier_match(seq2, bat, 2; tol) - end - - # Mehrotra correction - if !seq1_done; MadIPM.mehrotra_correction_direction!(seq1); end - if !seq2_done; MadIPM.mehrotra_correction_direction!(seq2); end - MadIPM.mehrotra_correction_direction!(bat) - if !seq1_done; assert_direction_match(seq1, bat, 1; tol); end - if !seq2_done; assert_direction_match(seq2, bat, 2; tol); end - - # Update step - if !seq1_done; MadIPM.update_step!(seq1.opt.step_rule, seq1); end - if !seq2_done; MadIPM.update_step!(seq2.opt.step_rule, seq2); end - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.zero_inactive_step!(bat) - if !seq1_done; assert_step_match(seq1, bat, 1; tol); end - if !seq2_done; assert_step_match(seq2, bat, 2; tol); end - - # Apply step - if !seq1_done; MadIPM.apply_step!(seq1); end - if !seq2_done; MadIPM.apply_step!(seq2); end - MadIPM.apply_step!(bat) - if !seq1_done; assert_iterate_match(seq1, bat, 1; tol); end - if !seq2_done; assert_iterate_match(seq2, bat, 2; tol); end - - # Evaluate model - if !seq1_done; MadIPM.evaluate_model!(seq1); end - if !seq2_done; MadIPM.evaluate_model!(seq2); end - MadIPM.evaluate_model!(bat) - if !seq1_done; assert_model_match(seq1, bat, 1; tol); end - if !seq2_done; assert_model_match(seq2, bat, 2; tol); end - - # Patch batch columns from sequential to prevent FP noise accumulation. - # For QPs the Hessian and Jacobian are constant, so patching iterate + - # model evals is sufficient to get a clean start for the next iteration. - if patch - if !seq1_done; patch_batch_col_from_seq!(seq1, bat, 1); end - if !seq2_done; patch_batch_col_from_seq!(seq2, bat, 2); end - end - end - - @test seq1.status == bat.workspace.status[1] - @test seq2.status == bat.workspace.status[2] -end - -function run_iterations_bsN!(seqs::Vector, bat, n_iters; tol=1e-10) - N = length(seqs) - done = falses(N) - - for iter in 1:n_iters - for i in 1:N - done[i] || MadIPM.update_termination_criteria!(seqs[i]) - end - MadIPM.update_termination_criteria!(bat) - for i in 1:N - done[i] || assert_termination_match(seqs[i], bat, i; tol) - end - - for i in 1:N - done[i] = done[i] || MadIPM.is_done(seqs[i]) - end - - MadIPM.update_active_set!(bat.kkt, bat.workspace.status) - bat.kkt.active_batch_size[] == 0 && break - MadIPM._update_active_mask!(bat) - all(done) && break - - # Factorize - for i in 1:N; done[i] || MadIPM.factorize_system!(seqs[i]); end - MadIPM.factorize_system!(bat) - for i in 1:N - if !done[i] - assert_regularization_match(seqs[i], bat, i; tol) - assert_kkt_diagonals_match(seqs[i], bat, i; tol) - end - end - - # Prediction step - for i in 1:N; done[i] || MadIPM.prediction_step!(seqs[i]); end - MadIPM.prediction_step!(bat) - for i in 1:N - if !done[i] - assert_prediction_match(seqs[i], bat, i; tol) - assert_barrier_match(seqs[i], bat, i; tol) - end - end - - # Mehrotra correction - for i in 1:N; done[i] || MadIPM.mehrotra_correction_direction!(seqs[i]); end - MadIPM.mehrotra_correction_direction!(bat) - for i in 1:N; done[i] || assert_direction_match(seqs[i], bat, i; tol); end - - # Update step - for i in 1:N; done[i] || MadIPM.update_step!(seqs[i].opt.step_rule, seqs[i]); end - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.zero_inactive_step!(bat) - for i in 1:N; done[i] || assert_step_match(seqs[i], bat, i; tol); end - - # Apply step - for i in 1:N; done[i] || MadIPM.apply_step!(seqs[i]); end - MadIPM.apply_step!(bat) - for i in 1:N; done[i] || assert_iterate_match(seqs[i], bat, i; tol); end - - # Evaluate model - for i in 1:N; done[i] || MadIPM.evaluate_model!(seqs[i]); end - MadIPM.evaluate_model!(bat) - for i in 1:N; done[i] || assert_model_match(seqs[i], bat, i; tol); end - end - - for i in 1:N - @test seqs[i].status == bat.workspace.status[i] - end -end From 2c9a9f6d6ea4984b8f35a6bc5022230bcd736c70 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 16:26:30 -0400 Subject: [PATCH 18/51] rm old helper --- src/MadIPM.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MadIPM.jl b/src/MadIPM.jl index 49088b19..7296a515 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -28,7 +28,7 @@ include("batch/madnlp/nlpmodels.jl") include("batch/madipm/kernels.jl") include("batch/madipm/solver.jl") -export MPCSolver, madipm, madipm_batch, madipm_foreach +export MPCSolver, madipm, madipm_batch MadNLP.madsuite(::Val{:madipm}, args...; kwargs...) = madipm(args...; kwargs...) From 2dab05142dbbbd9cdd1ed49e27220f94c9e6bcf7 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 16:28:41 -0400 Subject: [PATCH 19/51] just use `julia-runtest` --- .github/workflows/gpu.yml | 7 +------ ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 2 +- ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 2 +- 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/.github/workflows/gpu.yml b/.github/workflows/gpu.yml index 2582ebc9..17257433 100644 --- a/.github/workflows/gpu.yml +++ b/.github/workflows/gpu.yml @@ -29,10 +29,5 @@ jobs: with: channel: ${{ matrix.julia-version }} - uses: julia-actions/julia-buildpkg@latest - - name: Test MadIPM.jl - shell: julia --color=yes {0} - run: | - using Pkg - Pkg.develop(path=".") - Pkg.test("MadIPM") + - uses: julia-actions/julia-runtest@latest diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index 7b172df7..07d870cb 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -9,7 +9,7 @@ using CUDA.CUSPARSE using CUDSS using KernelAbstractions import Atomix -using QuadraticModels: SparseMatrixCOO +import QuadraticModels: SparseMatrixCOO import LinearAlgebra: BlasFloat import MadIPM import MadNLP diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl index 530fb6ee..4b440099 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -1,4 +1,4 @@ -using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData +import QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData function NLPModels.obj!( bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, From 0ff744b48f62e48d786f304272c62b43311e0adb Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 18:42:06 -0400 Subject: [PATCH 20/51] refactor `update_termination_criteria` --- src/batch/madipm/kernels.jl | 18 ++++++++++ src/batch/madipm/solver.jl | 68 ++++++++----------------------------- src/batch/madnlp/kernels.jl | 30 ++++++++++++++++ 3 files changed, 63 insertions(+), 53 deletions(-) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 44986946..b2f076eb 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -1,3 +1,21 @@ +function dual_objective!(dual_obj, y_vals, rhs_vals, zl_r, xl_r, zu_r, xu_r, + scratch_m, scratch_lb, scratch_ub, sum_lb, sum_ub, nlb, nub) + @. scratch_m = y_vals * rhs_vals + sum!(dual_obj, scratch_m) + dual_obj .*= -one(eltype(dual_obj)) + if nlb > 0 + @. scratch_lb = zl_r * xl_r + sum!(sum_lb, scratch_lb) + dual_obj .+= sum_lb + end + if nub > 0 + @. scratch_ub = zu_r * xu_r + sum!(sum_ub, scratch_ub) + dual_obj .-= sum_ub + end + return dual_obj +end + function set_initial_primal_rhs!(solver::AbstractBatchMPCSolver) p = solver.p fill!(MadNLP.full(p), 0.0) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index fba172c4..ac48e33e 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -273,65 +273,27 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w bs = batch_solver.batch_size nlb, nub = batch_solver.d.nlb, batch_solver.d.nub - f_vals = MadNLP.full(batch_solver.f) - zl_vals = MadNLP.full(zl) - zu_vals = MadNLP.full(zu) - jacl_vals = MadNLP.full(batch_solver.jacl) - y_vals = MadNLP.full(batch_solver.y) - rhs_vals = MadNLP.full(batch_solver.rhs) - - # inf_pr[i] = norm(c[:, i], Inf) / max(1, norm_b[i]) - ws.inf_pr .= maximum(abs, MadNLP.full(batch_solver.c); dims=1) + _scratch_n = MadNLP.primal(batch_solver._w2) + _scratch_m = MadNLP.dual(batch_solver._w2) + _scratch_lb = MadNLP.dual_lb(batch_solver._w2) + _scratch_ub = MadNLP.dual_ub(batch_solver._w2) + + get_inf_pr!(ws.inf_pr, MadNLP.full(batch_solver.c)) @. ws.inf_pr /= max(one(T), ws.norm_b) - # inf_du[i] = max|f-zl+zu+jacl| / max(1, norm_c[i]) - _scratch_n = MadNLP.primal(batch_solver._w2) - @. _scratch_n = abs(f_vals - zl_vals + zu_vals + jacl_vals) - maximum!(ws.inf_du, _scratch_n) + get_inf_du!(ws.inf_du, MadNLP.full(batch_solver.f), MadNLP.full(zl), + MadNLP.full(zu), MadNLP.full(batch_solver.jacl), _scratch_n) @. ws.inf_du /= max(one(T), ws.norm_c) - # inf_compl[i] = get_optimality_gap / max(1, norm_c[i]) - if nlb > 0 - x_lr = lower(x) - xl_r = lower(xl) - zl_r = lower(zl) - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = abs(x_lr - xl_r) * zl_r - maximum!(ws.sum_lb, _scratch_lb) - else - fill!(ws.sum_lb, zero(T)) - end - if nub > 0 - xu_r = upper(xu) - x_ur = upper(x) - zu_r = upper(zu) - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = abs(xu_r - x_ur) * zu_r - maximum!(ws.sum_ub, _scratch_ub) - else - fill!(ws.sum_ub, zero(T)) - end - @. ws.inf_compl = max(ws.sum_lb, ws.sum_ub) / max(one(T), ws.norm_c) + get_inf_compl!(ws.inf_compl, + lower(x), lower(xl), lower(zl), upper(xu), upper(x), upper(zu), + _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) + @. ws.inf_compl /= max(one(T), ws.norm_c) @. ws.best_complementarity = min(ws.best_complementarity, ws.inf_compl) - _scratch_m = MadNLP.dual(batch_solver._w2) - @. _scratch_m = y_vals * rhs_vals - sum!(ws.dual_obj, _scratch_m) - ws.dual_obj .*= -one(T) - if nlb > 0 - zl_r = lower(zl); xl_r = lower(xl) - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = zl_r * xl_r - sum!(ws.sum_lb, _scratch_lb) - ws.dual_obj .+= ws.sum_lb - end - if nub > 0 - zu_r = upper(zu); xu_r = upper(xu) - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = zu_r * xu_r - sum!(ws.sum_ub, _scratch_ub) - ws.dual_obj .-= ws.sum_ub - end + dual_objective!(ws.dual_obj, MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), + lower(zl), lower(xl), upper(zu), upper(xu), + _scratch_m, _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) ds = T(opt.divergence_scale) copyto!(ws.term_converged, diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl index 809d9fb0..74abc781 100644 --- a/src/batch/madnlp/kernels.jl +++ b/src/batch/madnlp/kernels.jl @@ -1,3 +1,33 @@ +function get_inf_pr!(inf_pr, c) + inf_pr .= maximum(abs, c; dims=1) + return inf_pr +end + +function get_inf_du!(inf_du, f_vals, zl_vals, zu_vals, jacl_vals, scratch) + @. scratch = abs(f_vals - zl_vals + zu_vals + jacl_vals) + maximum!(inf_du, scratch) + return inf_du +end + +function get_inf_compl!(inf_compl, x_lr, xl_r, zl_r, xu_r, x_ur, zu_r, + scratch_lb, scratch_ub, sum_lb, sum_ub, nlb, nub) + T = eltype(inf_compl) + if nlb > 0 + @. scratch_lb = abs(x_lr - xl_r) * zl_r + maximum!(sum_lb, scratch_lb) + else + fill!(sum_lb, zero(T)) + end + if nub > 0 + @. scratch_ub = abs(xu_r - x_ur) * zu_r + maximum!(sum_ub, scratch_ub) + else + fill!(sum_ub, zero(T)) + end + @. inf_compl = max(sum_lb, sum_ub) + return inf_compl +end + _adjust_bound_lb(x_lr::T, xl_r, c1, c2) where T = x_lr - xl_r < c1 ? xl_r - c2 * max(one(T), abs(x_lr)) : xl_r _adjust_bound_ub(x_ur::T, xu_r, c1, c2) where T = From 665ba008fdcbf429e3d36f8f88263b25abfaddef Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 23:05:20 -0400 Subject: [PATCH 21/51] improve the performance of `get_inf_compl!`/`update_termination_criteria!` --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 42 +++++++++++++++ src/batch/madipm/solver.jl | 69 +++++++++++++++---------- src/batch/madnlp/kernels.jl | 4 +- src/batch/structure.jl | 29 ++++++----- test/Project.toml | 69 +++++++++++++++++++++++++ test/test_batch.jl | 41 +++++++++++++++ test/test_batch_gpu.jl | 16 ++++++ 7 files changed, 231 insertions(+), 39 deletions(-) create mode 100644 test/Project.toml diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index 0744d27a..facbf1a6 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -82,6 +82,48 @@ end @inbounds alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) end +@kernel function _gather_compl_kernel!( + scratch, @Const(x_vals), @Const(xb_vals), @Const(z_vals), @Const(ind), +) + i, j = @index(Global, NTuple) + @inbounds begin + idx = ind[i] + scratch[i, j] = abs(x_vals[idx, j] - xb_vals[idx, j]) * z_vals[idx, j] + end +end + +function MadIPM.get_inf_compl!( + inf_compl::CuMatrix, x::MadIPM.BatchPrimalVector, xl::MadIPM.BatchPrimalVector, + zl::MadIPM.BatchPrimalVector, xu::MadIPM.BatchPrimalVector, zu::MadIPM.BatchPrimalVector, + scratch_lb::CuMatrix, scratch_ub::CuMatrix, sum_lb, sum_ub, nlb, nub, +) + T = eltype(inf_compl) + bs = size(inf_compl, 2) + backend = CUDABackend() + if nlb > 0 + _gather_compl_kernel!(backend)( + scratch_lb, x.values, xl.values, zl.values, x.ind_lb; + ndrange=(nlb, bs), + ) + KernelAbstractions.synchronize(backend) + maximum!(sum_lb, scratch_lb) + else + fill!(sum_lb, zero(T)) + end + if nub > 0 + _gather_compl_kernel!(backend)( + scratch_ub, xu.values, x.values, zu.values, x.ind_ub; + ndrange=(nub, bs), + ) + KernelAbstractions.synchronize(backend) + maximum!(sum_ub, scratch_ub) + else + fill!(sum_ub, zero(T)) + end + @. inf_compl = max(sum_lb, sum_ub) + return inf_compl +end + function MadIPM._mehrotra_correct_steps!( alpha_p::CuMatrix{T}, alpha_d::CuMatrix{T}, mu, val_xl, idx_xl, val_xu, idx_xu, diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index ac48e33e..19d60174 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -285,8 +285,7 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w MadNLP.full(zu), MadNLP.full(batch_solver.jacl), _scratch_n) @. ws.inf_du /= max(one(T), ws.norm_c) - get_inf_compl!(ws.inf_compl, - lower(x), lower(xl), lower(zl), upper(xu), upper(x), upper(zu), + get_inf_compl!(ws.inf_compl, x, xl, zl, xu, zu, _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) @. ws.inf_compl /= max(one(T), ws.norm_c) @. ws.best_complementarity = min(ws.best_complementarity, ws.inf_compl) @@ -296,23 +295,34 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w _scratch_m, _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) ds = T(opt.divergence_scale) - copyto!(ws.term_converged, - vec(@. max(ws.inf_pr, ws.inf_du, ws.inf_compl) <= opt.tol)) - copyto!(ws.term_infeasible, - vec(@. (ws.inf_compl > opt.divergence_tol * ws.best_complementarity - ) & (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))))) - copyto!(ws.term_diverging, - vec(@. ws.obj_val < -(opt.divergence_tol * max(ds, abs(ws.dual_obj), one(T))))) + tol = T(opt.tol) + div_tol = T(opt.divergence_tol) + Int_SOLVED = Int(MadNLP.SOLVE_SUCCEEDED) + Int_INFEASIBLE = Int(MadNLP.INFEASIBLE_PROBLEM_DETECTED) + Int_DIVERGING = Int(MadNLP.DIVERGING_ITERATES) + Int_REGULAR = Int(MadNLP.REGULAR) + @. ws._term_gpu = ifelse( + max(ws.inf_pr, ws.inf_du, ws.inf_compl) <= tol, + Int_SOLVED, + ifelse( + (ws.inf_compl > div_tol * ws.best_complementarity) & + (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))), + Int_INFEASIBLE, + ifelse( + ws.obj_val < -(div_tol * max(ds, abs(ws.dual_obj), one(T))), + Int_DIVERGING, + Int_REGULAR, + ), + ), + ) + copyto!(ws._term_cpu, vec(ws._term_gpu)) walltime_hit = time() - bcnt.start_time[] >= opt.max_wall_time @inbounds for i in 1:bs ws.status[i] != MadNLP.REGULAR && continue - if ws.term_converged[i] - ws.status[i] = MadNLP.SOLVE_SUCCEEDED - elseif ws.term_infeasible[i] - ws.status[i] = MadNLP.INFEASIBLE_PROBLEM_DETECTED - elseif ws.term_diverging[i] - ws.status[i] = MadNLP.DIVERGING_ITERATES + code = MadNLP.Status(ws._term_cpu[i]) + if code != MadNLP.REGULAR + ws.status[i] = code elseif bcnt.k[i] >= opt.max_iter ws.status[i] = MadNLP.MAXIMUM_ITERATIONS_EXCEEDED elseif walltime_hit @@ -343,17 +353,24 @@ function solve_system!( view(MadNLP.full(p), :, i) .= zero(T) end end - # FIXME: per-instance reduction? - norm_w = norm(MadNLP.full(w), Inf) - norm_p = norm(MadNLP.full(p), Inf) - - residual_ratio = norm_w / max(one(T), norm_p) - MadNLP.@debug( - batch_solver.logger, - @sprintf("Residual after linear solve: %6.2e", residual_ratio), - ) - if isnan(residual_ratio) || (opt.check_residual && (residual_ratio > opt.tol_linear_solve)) - throw(MadNLP.SolveException()) + + ws = batch_solver.workspace + _fw = MadNLP.full(w) + _fw .= abs.(_fw) + maximum!(ws._norm_gpu, _fw) # (1,bs) per-instance norm_w + copyto!(ws._norm_cpu, vec(ws._norm_gpu)) + _fw .= abs.(MadNLP.full(p)) + maximum!(ws._norm_gpu, _fw) # (1,bs) per-instance norm_p + copyto!(ws._norm_cpu2, vec(ws._norm_gpu)) + + @inbounds for i in 1:bs + bkkt.batch_map[i] == 0 && continue + nw = ws._norm_cpu[i] + np = ws._norm_cpu2[i] + ratio = nw / max(one(T), np) + if isnan(ratio) || (opt.check_residual && (ratio > opt.tol_linear_solve)) + ws.status[i] = MadNLP.INTERNAL_ERROR + end end return d end diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl index 74abc781..62acb88f 100644 --- a/src/batch/madnlp/kernels.jl +++ b/src/batch/madnlp/kernels.jl @@ -9,16 +9,18 @@ function get_inf_du!(inf_du, f_vals, zl_vals, zu_vals, jacl_vals, scratch) return inf_du end -function get_inf_compl!(inf_compl, x_lr, xl_r, zl_r, xu_r, x_ur, zu_r, +function get_inf_compl!(inf_compl, x, xl, zl, xu, zu, scratch_lb, scratch_ub, sum_lb, sum_ub, nlb, nub) T = eltype(inf_compl) if nlb > 0 + x_lr = lower(x); xl_r = lower(xl); zl_r = lower(zl) @. scratch_lb = abs(x_lr - xl_r) * zl_r maximum!(sum_lb, scratch_lb) else fill!(sum_lb, zero(T)) end if nub > 0 + xu_r = upper(xu); x_ur = upper(x); zu_r = upper(zu) @. scratch_ub = abs(xu_r - x_ur) * zu_r maximum!(sum_ub, scratch_ub) else diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 84960ce9..14571cea 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -1,4 +1,4 @@ -struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI<:AbstractMatrix{Int32}} +struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI<:AbstractMatrix{Int32}, MI64<:AbstractMatrix{Int64}} alpha_xl::MT alpha_xu::MT alpha_zl::MT @@ -26,9 +26,12 @@ struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI best_complementarity::MT dual_obj::MT status::Vector{MadNLP.Status} - term_converged::Vector{Bool} - term_infeasible::Vector{Bool} - term_diverging::Vector{Bool} + + _term_gpu::MI64 + _term_cpu::Vector{Int64} + _norm_gpu::MT + _norm_cpu::Vector{T} + _norm_cpu2::Vector{T} active_mask::MT @@ -42,7 +45,8 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch nvar_nlp::Int=0, ncon::Int=0) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} _proto = MT(undef, 1, batch_size) MI = typeof(similar(_proto, Int32)) - return UniformBatchWorkspace{T, VT, MT, MI}( + MI64 = typeof(similar(_proto, Int64)) + return UniformBatchWorkspace{T, VT, MT, MI, MI64}( MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_xl, alpha_xu MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_zl, alpha_zu MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_p, alpha_d @@ -61,9 +65,11 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch MT(undef, 1, batch_size), # best_complementarity MT(undef, 1, batch_size), # dual_obj fill(MadNLP.INITIAL, batch_size), # status - falses(batch_size), # term_converged - falses(batch_size), # term_infeasible - falses(batch_size), # term_diverging + similar(_proto, Int64), # _term_gpu + zeros(Int64, batch_size), # _term_cpu + MT(undef, 1, batch_size), # _norm_gpu + zeros(T, batch_size), # _norm_cpu + zeros(T, batch_size), # _norm_cpu2 fill!(MT(undef, 1, batch_size), one(T)), # active_mask MT(undef, nvar_nlp, batch_size), # bx VT(undef, batch_size), # bf @@ -72,7 +78,7 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch ) end -mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, MI} <: AbstractBatchMPCSolver{T} +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPCSolver{T} batch_size::Int d::BatchUnreducedKKTVector{T, MT, VT} @@ -94,7 +100,7 @@ mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, MI} <: AbstractBatc correction_lb::BatchVector{T, MT, VT} correction_ub::BatchVector{T, MT, VT} - workspace::UniformBatchWorkspace{T, VT, MT, MI} + workspace::UniformBatchWorkspace{T, VT, MT} opt::IPMOptions batch_cnt::BatchCounters @@ -190,8 +196,7 @@ function UniformBatchMPCSolver( batch_del_w = fill!(MT(undef, 1, batch_size), zero(T)) batch_del_c = fill!(MT(undef, 1, batch_size), zero(T)) - MI = typeof(workspace.idx_xl) - return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb), MI}( + return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb)}( batch_size, batch_d, batch_p, batch_w1, batch_w2, batch_x, batch_xl, batch_xu, batch_zl, batch_zu, batch_f, diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 00000000..bd0faf2d --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,69 @@ +[deps] +AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ArgTools = "0dad84c5-d112-42e6-8d28-ef12dabb789f" +Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" +Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" +CompilerSupportLibraries_jll = "e66e0078-7015-5450-92f7-15fbd957f2ae" +Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" +ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04" +FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +FileWatching = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +JLLWrappers = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +JuliaSyntaxHighlighting = "ac6e5ff7-fb65-4e79-a425-ec3bc9c03011" +LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" +LibCURL = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" +LibCURL_jll = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" +LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433" +LibGit2_jll = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" +LibSSH2_jll = "29816b5a-b9ab-546f-933c-edad1886dfa8" +Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125" +Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" +METIS_jll = "d00139f3-1899-568f-a2f0-47f597d42d70" +MUMPS_seq_jll = "d7ed1dd3-d0ae-5e8e-bfb4-87a502085b8d" +MadIPM = "4406d55e-1099-4ca6-ac6a-6cf051b10024" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217" +Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" +MozillaCACerts_jll = "14a3606d-f60d-562e-9121-12d972cd8159" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" +NetworkOptions = "ca575930-c2e3-43a9-ace4-1e988b2c1908" +OpenBLAS32_jll = "656ef2d0-ae68-5445-9ca0-591084a874a2" +OpenBLAS_jll = "4536629a-c528-5b80-bd46-f80d51c5b363" +OpenSSL_jll = "458c3c95-2e84-50aa-8efc-19380b2a3a95" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +QuadraticModels = "f468eda6-eac5-11e8-05a5-ff9e497bcd19" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce" +Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" +SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseMatricesCOO = "fa32481b-f100-4b48-8dc8-c62f61b13870" +StyledStrings = "f489334b-da3d-4c2e-b8f0-e476e12c162b" +SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" +SuiteSparse_jll = "bea87d4a-7f5b-5778-9afe-8cc45184846c" +TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" +Tar = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" +UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" +Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" +Zlib_jll = "83775a58-1f1d-513f-b197-d71354ab007a" +libblastrampoline_jll = "8e850b90-86db-534c-a0d3-1478176c7d93" +nghttp2_jll = "8e850ede-7688-5339-a07c-302acd2aaf8d" +p7zip_jll = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" + +[sources] +MadIPM = {path = ".."} +QuadraticModels = {url = "https://github.com/klamike/QuadraticModels.jl.git", rev = "mk/rhsbatch"} diff --git a/test/test_batch.jl b/test/test_batch.jl index 89ed257b..8c3e509f 100644 --- a/test/test_batch.jl +++ b/test/test_batch.jl @@ -123,4 +123,45 @@ end @testset "FullBatch different H/A data" begin _test_fullbatch_different_data() end + @testset "Residual check marks INTERNAL_ERROR (all fail)" begin + # Force residual check failure with tol_linear_solve=0 — should mark + # instances as INTERNAL_ERROR instead of throwing a SolveException. + qps = [simple_lp() for _ in 1:3] + bnlp = ObjRHSBatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; + print_level=MadNLP.ERROR, + check_residual=true, + tol_linear_solve=0.0, + ) + for i in 1:3 + @test stats[i].status == MadNLP.INTERNAL_ERROR + end + end + + @testset "Residual check marks INTERNAL_ERROR (partial)" begin + # Use NaN objective coefficients to produce NaN residuals for one instance. + # The per-instance residual check should mark only that instance as INTERNAL_ERROR. + Hrows = [1, 2]; Hcols = [1, 2] + Arows = [1, 1]; Acols = [1, 2] + + good_qp() = QuadraticModel( + [1.0, 1.0], Hrows, Hcols, [2.0, 2.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], x0=[0.5, 0.5], + ) + bad_qp = QuadraticModel( + [NaN, NaN], Hrows, Hcols, [2.0, 2.0]; + Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[Inf, Inf], x0=[0.5, 0.5], + ) + qps = [good_qp(), bad_qp, good_qp()] + bnlp = BatchQuadraticModel(qps) + stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR) + # Good instances should solve; bad instance should fail gracefully + @test stats[1].status == MadNLP.SOLVE_SUCCEEDED + @test stats[3].status == MadNLP.SOLVE_SUCCEEDED + @test stats[2].status != MadNLP.SOLVE_SUCCEEDED + end end diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl index 067b1c44..d990bc27 100644 --- a/test/test_batch_gpu.jl +++ b/test/test_batch_gpu.jl @@ -247,6 +247,22 @@ end end end + @testset "Residual check marks INTERNAL_ERROR (GPU)" begin + qps = [simple_lp() for _ in 1:3] + cpu_bnlp = ObjRHSBatchQuadraticModel(qps) + gpu_bnlp = convert(ObjRHSBatchQuadraticModel{Float64, CuVector{Float64}}, cpu_bnlp) + stats = MadIPM.madipm_batch(gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + check_residual=true, + tol_linear_solve=0.0, + ) + CUDA.@allowscalar for i in 1:3 + @test stats[i].status == MadNLP.INTERNAL_ERROR + end + end + @testset "FullBatch identical QP (bs=3)" begin qps = [_gpu_small_qp() for _ in 1:3] ref = MadIPM.madipm(qps[1]; print_level=MadNLP.ERROR) From 32edff1c7916a62d8344d88629b9734d49b0d43c Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 9 Mar 2026 23:19:13 -0400 Subject: [PATCH 22/51] precompute views --- src/batch/madnlp/rhs.jl | 52 ++++++++++++++++++++++++++++++----------- 1 file changed, 39 insertions(+), 13 deletions(-) diff --git a/src/batch/madnlp/rhs.jl b/src/batch/madnlp/rhs.jl index 72f0be8b..c9a8af7b 100644 --- a/src/batch/madnlp/rhs.jl +++ b/src/batch/madnlp/rhs.jl @@ -7,6 +7,13 @@ struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, nub::Int ind_lb::VI ind_ub::VI + _primal::SubArray + _dual::SubArray + _primal_dual::SubArray + _dual_lb::SubArray + _dual_ub::SubArray + _xp_lr::SubArray + _xp_ur::SubArray end function BatchUnreducedKKTVector( @@ -24,17 +31,26 @@ function BatchUnreducedKKTVector( views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) end - return BatchUnreducedKKTVector{T, MT, VT, typeof(ind_lb)}(values, views, n, m, nlb, nub, ind_lb, ind_ub) + return BatchUnreducedKKTVector{T, MT, VT, typeof(ind_lb)}( + values, views, n, m, nlb, nub, ind_lb, ind_ub, + view(values, 1:n, :), + view(values, n+1:n+m, :), + view(values, 1:n+m, :), + view(values, n+m+1:n+m+nlb, :), + view(values, n+m+nlb+1:n+m+nlb+nub, :), + view(values, ind_lb, :), + view(values, ind_ub, :), + ) end MadNLP.full(bv::BatchUnreducedKKTVector) = bv.values -MadNLP.primal(bv::BatchUnreducedKKTVector) = view(bv.values, 1:bv.n, :) -MadNLP.dual(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+1:bv.n+bv.m, :) -MadNLP.primal_dual(bv::BatchUnreducedKKTVector) = view(bv.values, 1:bv.n+bv.m, :) -MadNLP.dual_lb(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+bv.m+1:bv.n+bv.m+bv.nlb, :) -MadNLP.dual_ub(bv::BatchUnreducedKKTVector) = view(bv.values, bv.n+bv.m+bv.nlb+1:bv.n+bv.m+bv.nlb+bv.nub, :) -xp_lr(bv::BatchUnreducedKKTVector) = view(bv.values, bv.ind_lb, :) -xp_ur(bv::BatchUnreducedKKTVector) = view(bv.values, bv.ind_ub, :) +MadNLP.primal(bv::BatchUnreducedKKTVector) = bv._primal +MadNLP.dual(bv::BatchUnreducedKKTVector) = bv._dual +MadNLP.primal_dual(bv::BatchUnreducedKKTVector) = bv._primal_dual +MadNLP.dual_lb(bv::BatchUnreducedKKTVector) = bv._dual_lb +MadNLP.dual_ub(bv::BatchUnreducedKKTVector) = bv._dual_ub +xp_lr(bv::BatchUnreducedKKTVector) = bv._xp_lr +xp_ur(bv::BatchUnreducedKKTVector) = bv._xp_ur struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} values::MT @@ -43,6 +59,10 @@ struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} ns::Int ind_lb::VI ind_ub::VI + _variable::SubArray + _slack::SubArray + _lower::SubArray + _upper::SubArray end function BatchPrimalVector( @@ -60,12 +80,18 @@ function BatchPrimalVector( views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) end - return BatchPrimalVector{T, MT, VT, typeof(ind_lb)}(values, views, nx, ns, ind_lb, ind_ub) + return BatchPrimalVector{T, MT, VT, typeof(ind_lb)}( + values, views, nx, ns, ind_lb, ind_ub, + view(values, 1:nx, :), + view(values, nx+1:nx+ns, :), + view(values, ind_lb, :), + view(values, ind_ub, :), + ) end -MadNLP.variable(bpv::BatchPrimalVector) = view(bpv.values, 1:bpv.nx, :) -MadNLP.slack(bpv::BatchPrimalVector) = view(bpv.values, bpv.nx+1:bpv.nx+bpv.ns, :) -lower(bpv::BatchPrimalVector) = view(bpv.values, bpv.ind_lb, :) -upper(bpv::BatchPrimalVector) = view(bpv.values, bpv.ind_ub, :) +MadNLP.variable(bpv::BatchPrimalVector) = bpv._variable +MadNLP.slack(bpv::BatchPrimalVector) = bpv._slack +lower(bpv::BatchPrimalVector) = bpv._lower +upper(bpv::BatchPrimalVector) = bpv._upper MadNLP.full(bpv::BatchPrimalVector) = bpv.values MadNLP.primal(bpv::BatchPrimalVector) = bpv.values From 3a4927f29254a7ad250ae0ffa1f95d38fbdcb5f3 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 10 Mar 2026 00:20:21 -0400 Subject: [PATCH 23/51] remove unused views --- src/batch/madipm/solver.jl | 2 +- src/batch/madnlp/kernels.jl | 5 +++-- src/batch/madnlp/rhs.jl | 13 +++---------- src/batch/structure.jl | 26 +++++++++++++------------- src/batch/utils.jl | 11 ++--------- 5 files changed, 22 insertions(+), 35 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 19d60174..f84b54ac 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -278,7 +278,7 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w _scratch_lb = MadNLP.dual_lb(batch_solver._w2) _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - get_inf_pr!(ws.inf_pr, MadNLP.full(batch_solver.c)) + get_inf_pr!(ws.inf_pr, MadNLP.full(batch_solver.c), _scratch_m) @. ws.inf_pr /= max(one(T), ws.norm_b) get_inf_du!(ws.inf_du, MadNLP.full(batch_solver.f), MadNLP.full(zl), diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl index 62acb88f..a897073d 100644 --- a/src/batch/madnlp/kernels.jl +++ b/src/batch/madnlp/kernels.jl @@ -1,5 +1,6 @@ -function get_inf_pr!(inf_pr, c) - inf_pr .= maximum(abs, c; dims=1) +function get_inf_pr!(inf_pr, c, scratch) + @. scratch = abs(c) + maximum!(inf_pr, scratch) return inf_pr end diff --git a/src/batch/madnlp/rhs.jl b/src/batch/madnlp/rhs.jl index c9a8af7b..9793bab1 100644 --- a/src/batch/madnlp/rhs.jl +++ b/src/batch/madnlp/rhs.jl @@ -52,9 +52,8 @@ MadNLP.dual_ub(bv::BatchUnreducedKKTVector) = bv._dual_ub xp_lr(bv::BatchUnreducedKKTVector) = bv._xp_lr xp_ur(bv::BatchUnreducedKKTVector) = bv._xp_ur -struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} +struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VI} values::MT - views::Vector{VT} nx::Int ns::Int ind_lb::VI @@ -74,14 +73,8 @@ function BatchPrimalVector( values = MT(undef, total, batch_size) fill!(values, zero(T)) - views = Vector{VT}(undef, batch_size) - for i in 1:batch_size - col_start = (i-1) * total + 1 - views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) - end - - return BatchPrimalVector{T, MT, VT, typeof(ind_lb)}( - values, views, nx, ns, ind_lb, ind_ub, + return BatchPrimalVector{T, MT, typeof(ind_lb)}( + values, nx, ns, ind_lb, ind_ub, view(values, 1:nx, :), view(values, nx+1:nx+ns, :), view(values, ind_lb, :), diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 14571cea..57677745 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -86,19 +86,19 @@ mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPC _w1::BatchUnreducedKKTVector{T, MT, VT} _w2::BatchUnreducedKKTVector{T, MT, VT} - x::BatchPrimalVector{T, MT, VT} - xl::BatchPrimalVector{T, MT, VT} - xu::BatchPrimalVector{T, MT, VT} - zl::BatchPrimalVector{T, MT, VT} - zu::BatchPrimalVector{T, MT, VT} - f::BatchPrimalVector{T, MT, VT} - - y::BatchVector{T, MT, VT} - c::BatchVector{T, MT, VT} - jacl::BatchVector{T, MT, VT} - rhs::BatchVector{T, MT, VT} - correction_lb::BatchVector{T, MT, VT} - correction_ub::BatchVector{T, MT, VT} + x::BatchPrimalVector{T, MT} + xl::BatchPrimalVector{T, MT} + xu::BatchPrimalVector{T, MT} + zl::BatchPrimalVector{T, MT} + zu::BatchPrimalVector{T, MT} + f::BatchPrimalVector{T, MT} + + y::BatchVector{T, MT} + c::BatchVector{T, MT} + jacl::BatchVector{T, MT} + rhs::BatchVector{T, MT} + correction_lb::BatchVector{T, MT} + correction_ub::BatchVector{T, MT} workspace::UniformBatchWorkspace{T, VT, MT} diff --git a/src/batch/utils.jl b/src/batch/utils.jl index db6aa059..efe9bfee 100644 --- a/src/batch/utils.jl +++ b/src/batch/utils.jl @@ -101,9 +101,8 @@ function _build_hess_scatter( return scatter, nz_map, var_map, buffer end -struct BatchVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} +struct BatchVector{T, MT<:AbstractMatrix{T}} values::MT - views::Vector{VT} end MadNLP.full(bv::BatchVector) = bv.values @@ -114,13 +113,7 @@ function BatchVector( ) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} values = MT(undef, len, batch_size) fill!(values, zero(T)) - - views = Vector{VT}(undef, batch_size) - for i in 1:batch_size - views[i] = _madnlp_unsafe_column_wrap(values, len, (i-1) * len + 1, VT) - end - - return BatchVector{T, MT, VT}(values, views) + return BatchVector{T, MT}(values) end mutable struct BatchExecutionStats{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}} From c317f1482e89b2e4b7bd9722c3a5f08142298220 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Wed, 11 Mar 2026 23:41:45 -0400 Subject: [PATCH 24/51] prettier status --- src/batch/madipm/solver.jl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index f84b54ac..ded59db6 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -539,8 +539,13 @@ function solve!(batch_solver::AbstractBatchMPCSolver{T}) where T t_end = time() bcnt.total_time .= t_end .- bcnt.start_time[] update_solution!(stats, batch_solver) + status_counts = Dict{MadNLP.SolverStatus, Int}() for i in 1:bs - MadNLP.@notice(batch_solver.logger, "Problem $i: $(MadNLP.get_status_output(ws.status[i], batch_solver.opt))") + s = ws.status[i] + status_counts[s] = get(status_counts, s, 0) + 1 + end + for (s, cnt) in status_counts + MadNLP.@notice(batch_solver.logger, "$(MadNLP.get_status_output(s, batch_solver.opt)): $cnt/$bs") end end From f6c517322158af075057f57a50528d43e76bfd6b Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Wed, 11 Mar 2026 23:45:57 -0400 Subject: [PATCH 25/51] target `ger!` --- ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl index 4b440099..d451fb0c 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl @@ -71,7 +71,7 @@ function NLPModels.hess_coord!( H = bqp.data.H.A nnzh = nnz(H) nnzh == 0 && return bhvals - bhvals .= H.nzVal .* bobj_weight' + mul!(bhvals, H.nzVal, bobj_weight') return bhvals end From 32bbea3bbe25b4115d2a6258c25c2bed9c88e7c8 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 12 Mar 2026 02:28:12 -0400 Subject: [PATCH 26/51] add gather mul fused kernel --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 17 +++++++++++++++++ src/batch/KKT/Sparse/augmented.jl | 13 +++++++++---- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index facbf1a6..bdf4e0c6 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -146,3 +146,20 @@ function MadIPM._mehrotra_correct_steps!( KernelAbstractions.synchronize(backend) end end + +@kernel function _gather_mul_kernel!(out, @Const(A), @Const(nz_map), @Const(B), @Const(val_map)) + i, j = @index(Global, NTuple) + @inbounds out[i, j] = A[nz_map[i], j] * B[val_map[i], j] +end + +function MadIPM._gather_mul!( + out::CuMatrix, A::CuMatrix, nz_map::CuVector, B::CuMatrix, val_map::CuVector, +) + n, bs = size(out) + if n > 0 + backend = CUDABackend() + _gather_mul_kernel!(backend)(out, A, nz_map, B, val_map; ndrange=(n, bs)) + KernelAbstractions.synchronize(backend) + end + return out +end diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 491be534..6f756d66 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -263,8 +263,13 @@ function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) return end +function _gather_mul!(out::AbstractMatrix, A::AbstractMatrix, nz_map, B::AbstractMatrix, val_map) + @views out .= A[nz_map, :] .* B[val_map, :] + return out +end + function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) - @views bkkt.jt_buffer .= bkkt.nzVals[bkkt.jt_nz_map, :] .* MadNLP.full(y)[bkkt.jt_con_map, :] + _gather_mul!(bkkt.jt_buffer, bkkt.nzVals, bkkt.jt_nz_map, MadNLP.full(y), bkkt.jt_con_map) mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) return res end @@ -345,17 +350,17 @@ function LinearAlgebra.mul!( # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) xv = MadNLP.full(x) - @views bkkt.hess_buffer .= nzV[bkkt.hess_nz_map, :] .* xv[bkkt.hess_var_map, :] + _gather_mul!(bkkt.hess_buffer, nzV, bkkt.hess_nz_map, xv, bkkt.hess_var_map) mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp # mul!(primal(w), jac_com', dual(x), alpha, one(T)) - @views bkkt.jt_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.jt_con_map_full, :] + _gather_mul!(bkkt.jt_buffer, nzV, bkkt.jt_nz_map, xv, bkkt.jt_con_map_full) mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) MadNLP.primal(w) .+= alpha .* wp # mul!(dual(w), jac_com, primal(x), alpha, beta) - @views bkkt.j_buffer .= nzV[bkkt.jt_nz_map, :] .* xv[bkkt.j_var_map, :] + _gather_mul!(bkkt.j_buffer, nzV, bkkt.jt_nz_map, xv, bkkt.j_var_map) mul!(wd, bkkt.j_scatter, bkkt.j_buffer) MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) From 180bd6c9b0f95864da3b636bb7f13e027b923d62 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 12 Mar 2026 13:20:37 -0400 Subject: [PATCH 27/51] typo --- src/batch/madipm/solver.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index ded59db6..6552e78d 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -539,7 +539,7 @@ function solve!(batch_solver::AbstractBatchMPCSolver{T}) where T t_end = time() bcnt.total_time .= t_end .- bcnt.start_time[] update_solution!(stats, batch_solver) - status_counts = Dict{MadNLP.SolverStatus, Int}() + status_counts = Dict{MadNLP.Status, Int}() for i in 1:bs s = ws.status[i] status_counts[s] = get(status_counts, s, 0) + 1 From 287f8bc699e4df94b43964570f8956f7e6571310 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Fri, 13 Mar 2026 22:35:38 -0400 Subject: [PATCH 28/51] use `mk/paramnlp` for now --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index f727f806..e1c09de8 100644 --- a/Project.toml +++ b/Project.toml @@ -21,6 +21,7 @@ MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [sources] +NLPModels = {rev = "mk/paramnlp", url = "https://github.com/klamike/NLPModels.jl"} QuadraticModels = {url = "https://github.com/klamike/QuadraticModels.jl.git", rev = "mk/rhsbatch"} [extensions] From 266341e4c98f161ea9ab619acdf69c1cc4e45d3d Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 24 Mar 2026 11:49:08 -0400 Subject: [PATCH 29/51] add MT (batched VT) and VT (batched T) to abstractbatchmpcsolver --- src/batch/KKT/Sparse/augmented.jl | 6 ++---- src/batch/madipm/solver.jl | 4 +--- src/batch/structure.jl | 2 +- src/batch/utils.jl | 2 +- 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 6f756d66..2c9703cf 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -49,10 +49,10 @@ end function MadNLP.create_kkt_system( ::Type{MadNLP.SparseKKTSystem}, - bcb::UniformBatchCallback{T, VT}, + bcb::UniformBatchCallback{T, VT, MT, VI}, uniformbatch_linear_solver = LoopedBatchLinearSolver; opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), -) where {T, VT} +) where {T, VT, MT, VI} batch_size = bcb.batch_size n_slack = length(bcb.ind_ineq) @@ -143,8 +143,6 @@ function MadNLP.create_kkt_system( active_batch_size = Ref(batch_size) LS = typeof(batch_ls) - MT = typeof(nzVals) - VI = typeof(aug_csc_map) VI32 = typeof(I) SMT = typeof(jt_scatter) return SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT}( diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 6552e78d..11205ac6 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -512,15 +512,13 @@ function mpc!(batch_solver::AbstractBatchMPCSolver) end end -function solve!(batch_solver::AbstractBatchMPCSolver{T}) where T +function solve!(batch_solver::AbstractBatchMPCSolver{T, MT, VT}) where {T, MT, VT} ws = batch_solver.workspace bcb = batch_solver.bcb bs = batch_solver.batch_size nvar_nlp = bcb.nlp.meta.nvar ncon = bcb.ncon - VT = typeof(ws.bf) - MT = typeof(MadNLP.full(batch_solver.x)) stats = BatchExecutionStats(MT, VT, nvar_nlp, ncon, bs) try diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 57677745..43506712 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -78,7 +78,7 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch ) end -mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPCSolver{T} +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPCSolver{T, MT, VT} batch_size::Int d::BatchUnreducedKKTVector{T, MT, VT} diff --git a/src/batch/utils.jl b/src/batch/utils.jl index efe9bfee..1e6ea3fe 100644 --- a/src/batch/utils.jl +++ b/src/batch/utils.jl @@ -1,4 +1,4 @@ -abstract type AbstractBatchMPCSolver{T} end +abstract type AbstractBatchMPCSolver{T, MT, VT} end function _madnlp_unsafe_column_wrap(mat::MT, n, shift, ::Type{VT}) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} return unsafe_wrap(VT, pointer(mat, shift), n) From 0f10879224a06c724d840484fd22531a14a71b47 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 24 Mar 2026 11:49:18 -0400 Subject: [PATCH 30/51] get rid of muladd --- src/batch/madipm/solver.jl | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 11205ac6..52369c68 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -446,20 +446,16 @@ end function apply_step!(batch_solver::AbstractBatchMPCSolver) ws = batch_solver.workspace - x, xl, xu = batch_solver.x, batch_solver.xl, batch_solver.xu + x, y, xl, xu = batch_solver.x, batch_solver.y, batch_solver.xl, batch_solver.xu zl, zu, d = batch_solver.zl, batch_solver.zu, batch_solver.d batch_size = batch_solver.batch_size nlb, nub = d.nlb, d.nub # x += alpha_p * dx - _x = MadNLP.full(x) - _dx = MadNLP.primal(d) - @. _x = muladd(ws.alpha_p, _dx, _x) + MadNLP.full(x) .+= ws.alpha_p .* MadNLP.primal(d) # y += alpha_d * d_dual - _y = MadNLP.full(batch_solver.y) - _dy = MadNLP.dual(d) - @. _y = muladd(ws.alpha_d, _dy, _y) + MadNLP.full(y) .+= ws.alpha_d .* MadNLP.dual(d) # zl_r += alpha_d * dzl, zu_r += alpha_d * dzu if nlb > 0 From b4480a663ba06d064fe1619c8e5d85db28e9986e Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 26 Mar 2026 20:26:31 -0400 Subject: [PATCH 31/51] use matrix throughout --- ext/MadIPMCUDAExt/cuda_wrapper.jl | 2 +- src/batch/KKT/Sparse/augmented.jl | 38 ++++++++++++------------------- src/batch/madipm/solver.jl | 15 +++++++++++- src/batch/madnlp/linear_solver.jl | 14 ++++-------- src/batch/madnlp/rhs.jl | 13 +++-------- src/batch/structure.jl | 36 +++++++++++++++++------------ 6 files changed, 58 insertions(+), 60 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 34dd7f7c..4afefa27 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -241,7 +241,7 @@ function MadIPM._active_factorize!(s::MadNLPGPU.CUDSSSolver, na::Int) return end -function MadIPM._active_solve!(s::MadNLPGPU.CUDSSSolver{T}, rhs::CuVector{T}, na::Int, n::Int) where T +function MadIPM._active_solve!(s::MadNLPGPU.CUDSSSolver{T}, rhs::CuMatrix{T}, na::Int, n::Int) where T rhs_active = unsafe_wrap(CuArray{T, 2}, pointer(rhs), (n, na)) CUDSS.cudss_update(s.b_gpu, rhs_active) CUDSS.cudss_update(s.x_gpu, rhs_active) diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 2c9703cf..85ca737d 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -1,10 +1,10 @@ """Batched version of SparseKKTSystem""" -struct SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} +struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} nzVals::MT # (aug_mat_length × batch_size) COO nonzero values aug_I::VI32 # shared row indices aug_J::VI32 # shared column indices batch_solver::LS # batched linear solver - rhs_buffer::VT # contiguous (n+m)*batch_size for batch solve + rhs_buffer::MT # (n+m) × batch_size for batch solve batch_size::Int aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values batch_csc_map::VI # flattened COO→CSC map for all instances @@ -112,7 +112,7 @@ function MadNLP.create_kkt_system( batch_ls = uniformbatch_linear_solver(aug_com, aug_com_nzvals, aug_vec_length; opt=opt_linear_solver) - rhs_buffer = VT(undef, aug_vec_length * batch_size) + rhs_buffer = similar(nzVals, aug_vec_length, batch_size) fill!(rhs_buffer, zero(T)) jac_range = n_tot+n_hess+1:n_tot+n_hess+n_jac+n_slack @@ -145,7 +145,7 @@ function MadNLP.create_kkt_system( LS = typeof(batch_ls) VI32 = typeof(I) SMT = typeof(jt_scatter) - return SparseUniformBatchKKTSystem{T, LS, VT, MT, VI, VI32, SMT}( + return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, SMT}( nzVals, I, J, batch_ls, rhs_buffer, batch_size, aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, reg, l_diag, u_diag, l_lower, u_lower, @@ -185,25 +185,24 @@ function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) return end -function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractVector) where T +function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractMatrix) where T na = bkkt.active_batch_size[] bs = bkkt.batch_size - n = length(rhs) ÷ bs + n = size(rhs, 1) - rhs_mat = reshape(rhs, n, bs) @inbounds for j in 1:na i = bkkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, j) .= view(rhs_mat, :, i)) + i != j && (view(rhs, :, j) .= view(rhs, :, i)) end _active_solve!(bkkt.batch_solver, rhs, na, n) @inbounds for j in na:-1:1 i = bkkt.batch_map_rev[j] - i != j && (view(rhs_mat, :, i) .= view(rhs_mat, :, j)) + i != j && (view(rhs, :, i) .= view(rhs, :, j)) end @inbounds for i in 1:bs - bkkt.batch_map[i] == 0 && (view(rhs_mat, :, i) .= zero(T)) + bkkt.batch_map[i] == 0 && (view(rhs, :, i) .= zero(T)) end return rhs end @@ -227,23 +226,14 @@ end function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) d = batch_solver.d - n_pd = d.n + d.m - bs = bkkt.batch_size - na = bkkt.active_batch_size[] MadNLP.reduce_rhs!(bkkt, d) - if na == 1 - # only one active: skip packing - orig_col = bkkt.batch_map_rev[1] - _active_solve!(bkkt.batch_solver, d.views[orig_col], 1, n_pd) - else - rhs = bkkt.rhs_buffer - pd_view = MadNLP.primal_dual(d) - copyto!(reshape(rhs, n_pd, bs), pd_view) - MadNLP.solve_linear_system!(bkkt, rhs) - copyto!(pd_view, reshape(rhs, n_pd, bs)) - end + rhs = bkkt.rhs_buffer + pd_view = MadNLP.primal_dual(d) + copyto!(rhs, pd_view) + MadNLP.solve_linear_system!(bkkt, rhs) + copyto!(pd_view, rhs) MadNLP.finish_aug_solve!(bkkt, batch_solver) return diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 52369c68..53003cb2 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -318,6 +318,15 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w copyto!(ws._term_cpu, vec(ws._term_gpu)) walltime_hit = time() - bcnt.start_time[] >= opt.max_wall_time + max_iter_hit = walltime_hit ? false : + any(ws.status[i] == MadNLP.REGULAR && bcnt.k[i] >= opt.max_iter for i in 1:bs) + + if !walltime_hit && !max_iter_hit + copyto!(ws._any_nonregular_cpu, ws._any_nonregular_gpu) # TODO: use CuRef + ws._any_nonregular_cpu[1] == Int_REGULAR && return false + end + + copyto!(ws._term_cpu, ws._term_gpu) @inbounds for i in 1:bs ws.status[i] != MadNLP.REGULAR && continue code = MadNLP.Status(ws._term_cpu[i]) @@ -494,7 +503,11 @@ end function _update_active_mask!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace bmap = batch_solver.kkt.batch_map - copyto!(ws.active_mask, reshape(T.(bmap .!= 0), 1, :)) + buf = ws.active_mask_cpu + @inbounds for i in eachindex(bmap) + buf[i] = T(bmap[i] != 0) + end + copyto!(ws.active_mask, buf) end function mpc!(batch_solver::AbstractBatchMPCSolver) diff --git a/src/batch/madnlp/linear_solver.jl b/src/batch/madnlp/linear_solver.jl index 19730c4e..9ce71e01 100644 --- a/src/batch/madnlp/linear_solver.jl +++ b/src/batch/madnlp/linear_solver.jl @@ -1,14 +1,10 @@ # FIXME: threads/polyester version -struct LoopedBatchLinearSolver{T, LS<:MadNLP.AbstractLinearSolver{T}} <: MadNLP.AbstractLinearSolver{T} +struct LoopedBatchLinearSolver{T, VT, LS<:MadNLP.AbstractLinearSolver{T}} <: MadNLP.AbstractLinearSolver{T} solvers::Vector{LS} batch_size::Int end -function LoopedBatchLinearSolver(solvers::Vector{LS}) where {T, LS<:MadNLP.AbstractLinearSolver{T}} - return LoopedBatchLinearSolver{T, LS}(solvers, length(solvers)) -end - @kwdef mutable struct LoopedBatchLinearSolverOptions <: MadNLP.AbstractOptions looped_linear_solver::Type = MadNLP.MumpsSolver end @@ -31,7 +27,7 @@ function LoopedBatchLinearSolver( csc_i = _csc_with_nzval(aug_com, nzval_i, n) linear_solver(csc_i; opt=per_instance_opt) end - LoopedBatchLinearSolver(individual_solvers) + LoopedBatchLinearSolver{T, VT, eltype(individual_solvers)}(individual_solvers, batch_size) end function is_factorized(batch_linear_solver::LoopedBatchLinearSolver) @@ -45,10 +41,10 @@ function _active_factorize!(s::LoopedBatchLinearSolver, na::Int) return end -function _active_solve!(s::LoopedBatchLinearSolver, rhs::AbstractVector, na::Int, n::Int) +function _active_solve!(s::LoopedBatchLinearSolver{T, VT}, rhs::AbstractMatrix{T}, na::Int, n::Int) where {T, VT} for j in 1:na - xj = MadNLP._madnlp_unsafe_wrap(rhs, n, (j-1)*n + 1) - MadNLP.solve_linear_system!(s.solvers[j], xj) + rhs_j = _madnlp_unsafe_column_wrap(rhs, n, (j-1)*n + 1, VT) + MadNLP.solve_linear_system!(s.solvers[j], rhs_j) end return end diff --git a/src/batch/madnlp/rhs.jl b/src/batch/madnlp/rhs.jl index 9793bab1..3bcdd301 100644 --- a/src/batch/madnlp/rhs.jl +++ b/src/batch/madnlp/rhs.jl @@ -1,6 +1,5 @@ -struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}, VI} +struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VI} values::MT - views::Vector{VT} n::Int m::Int nlb::Int @@ -25,14 +24,8 @@ function BatchUnreducedKKTVector( values = MT(undef, total, batch_size) fill!(values, zero(T)) - views = Vector{VT}(undef, batch_size) - for i in 1:batch_size - col_start = (i-1) * total + 1 - views[i] = _madnlp_unsafe_column_wrap(values, total, col_start, VT) - end - - return BatchUnreducedKKTVector{T, MT, VT, typeof(ind_lb)}( - values, views, n, m, nlb, nub, ind_lb, ind_ub, + return BatchUnreducedKKTVector{T, MT, typeof(ind_lb)}( + values, n, m, nlb, nub, ind_lb, ind_ub, view(values, 1:n, :), view(values, n+1:n+m, :), view(values, 1:n+m, :), diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 43506712..b4324587 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -29,11 +29,14 @@ struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI _term_gpu::MI64 _term_cpu::Vector{Int64} - _norm_gpu::MT - _norm_cpu::Vector{T} - _norm_cpu2::Vector{T} + _any_nonregular_gpu::MI64 + _any_nonregular_cpu::Matrix{Int64} + _norm_gpu_w::MT + _norm_gpu_p::MT + _ls_error::MI active_mask::MT + active_mask_cpu::Matrix{T} bx::MT bf::VT @@ -41,8 +44,8 @@ struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI bv::MT end -function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch_size::Int; - nvar_nlp::Int=0, ncon::Int=0) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} +function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, nub::Int, batch_size::Int; + nvar_nlp::Int=0) where {T, MT<:AbstractMatrix{T}, VT<:AbstractVector{T}} _proto = MT(undef, 1, batch_size) MI = typeof(similar(_proto, Int32)) MI64 = typeof(similar(_proto, Int64)) @@ -67,24 +70,27 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, nlb::Int, nub::Int, batch fill(MadNLP.INITIAL, batch_size), # status similar(_proto, Int64), # _term_gpu zeros(Int64, batch_size), # _term_cpu - MT(undef, 1, batch_size), # _norm_gpu - zeros(T, batch_size), # _norm_cpu - zeros(T, batch_size), # _norm_cpu2 + fill!(similar(_proto, Int64, 1, 1), Int64(Int(MadNLP.REGULAR))), # _any_nonregular_gpu + zeros(Int64, 1, 1), # _any_nonregular_cpu + MT(undef, 1, batch_size), # _norm_gpu_w + MT(undef, 1, batch_size), # _norm_gpu_p + fill!(similar(_proto, Int32), zero(Int32)), # _ls_error fill!(MT(undef, 1, batch_size), one(T)), # active_mask + ones(T, 1, batch_size), # active_mask_cpu MT(undef, nvar_nlp, batch_size), # bx VT(undef, batch_size), # bf MT(undef, nvar_nlp, batch_size), # bg - MT(undef, ncon, batch_size), # bv + MT(undef, m, batch_size), # bv ) end mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPCSolver{T, MT, VT} batch_size::Int - d::BatchUnreducedKKTVector{T, MT, VT} - p::BatchUnreducedKKTVector{T, MT, VT} - _w1::BatchUnreducedKKTVector{T, MT, VT} - _w2::BatchUnreducedKKTVector{T, MT, VT} + d::BatchUnreducedKKTVector{T, MT} + p::BatchUnreducedKKTVector{T, MT} + _w1::BatchUnreducedKKTVector{T, MT} + _w2::BatchUnreducedKKTVector{T, MT} x::BatchPrimalVector{T, MT} xl::BatchPrimalVector{T, MT} @@ -190,8 +196,8 @@ function UniformBatchMPCSolver( batch_c = BatchVector(MT, VT, m, batch_size) batch_rhs = BatchVector(MT, VT, m, batch_size) - workspace = UniformBatchWorkspace(MT, VT, nlb, nub, batch_size; - nvar_nlp=nvar_nlp, ncon=m) + workspace = UniformBatchWorkspace(MT, VT, n, m, nlb, nub, batch_size; + nvar_nlp=nvar_nlp) batch_del_w = fill!(MT(undef, 1, batch_size), zero(T)) batch_del_c = fill!(MT(undef, 1, batch_size), zero(T)) From fb9fbca92e2f69466cf8c88760dd8a31f674c327 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 26 Mar 2026 20:30:52 -0400 Subject: [PATCH 32/51] move to BatchQuadraticModels --- Project.toml | 5 +- ext/MadIPMCUDAExt/MadIPMCUDAExt.jl | 2 +- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 17 ++ ext/MadIPMCUDAExt/cuda_batch_nlp.jl | 195 --------------------- ext/MadIPMCUDAExt/cuda_wrapper.jl | 29 --- ext/MadIPMMathOptInterfaceExt/parse_moi.jl | 1 - src/MadIPM.jl | 9 +- src/batch/KKT/Sparse/augmented.jl | 82 ++------- src/batch/utils.jl | 62 +++---- src/utils.jl | 2 - test/Project.toml | 10 +- test/batch_kernels/setup.jl | 2 +- test/test_batch.jl | 2 +- test/test_batch_gpu.jl | 2 +- 14 files changed, 81 insertions(+), 339 deletions(-) delete mode 100644 ext/MadIPMCUDAExt/cuda_batch_nlp.jl diff --git a/Project.toml b/Project.toml index e1c09de8..6aae2f4f 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "0.2.0" authors = ["François Pacaud ", "Alexis Montoison "] [deps] +BatchQuadraticModels = "412afcf0-3ec6-4826-be17-8d792afa05a8" LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" @@ -21,8 +22,7 @@ MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [sources] -NLPModels = {rev = "mk/paramnlp", url = "https://github.com/klamike/NLPModels.jl"} -QuadraticModels = {url = "https://github.com/klamike/QuadraticModels.jl.git", rev = "mk/rhsbatch"} +BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl"} [extensions] MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] @@ -30,6 +30,7 @@ MadIPMMathOptInterfaceExt = "MathOptInterface" [compat] Atomix = "1" +BatchQuadraticModels = "0.1" CUDA = "5.4.0" CUDSS = "0.6" KernelAbstractions = "0.9" diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index 07d870cb..a0569f88 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -3,6 +3,7 @@ module MadIPMCUDAExt using LinearAlgebra using SparseArrays using NLPModels +using BatchQuadraticModels using QuadraticModels using CUDA using CUDA.CUSPARSE @@ -17,7 +18,6 @@ import MadNLP include("operators.jl") include("cuda_wrapper.jl") include("cuda_batch_kernels.jl") -include("cuda_batch_nlp.jl") function MadIPM._csc_with_nzval(A::CUSPARSE.CuSparseMatrixCSC, nzval, n) return CUSPARSE.CuSparseMatrixCSC(A.colPtr, A.rowVal, nzval, (n, n)) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index bdf4e0c6..c1220adf 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -16,6 +16,23 @@ end end +@inline function _atomic_colreduce!(::typeof(+), out, j, value) + Atomix.@atomic out[1, j] += value + return +end + +@inline function _atomic_colreduce!(::typeof(min), out, j, value) + old = out[1, j] + while value < old + result = Atomix.@atomicreplace out[1, j] old => value + old = result.old + if result.success + break + end + end + return +end + function MadNLP._set_con_scale_sparse!( con_scale::CuMatrix{T}, jac_I::CuVector{<:Integer}, diff --git a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl b/ext/MadIPMCUDAExt/cuda_batch_nlp.jl deleted file mode 100644 index d451fb0c..00000000 --- a/ext/MadIPMCUDAExt/cuda_batch_nlp.jl +++ /dev/null @@ -1,195 +0,0 @@ -import QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel, QPData - -function NLPModels.obj!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, - bx::AbstractMatrix{T}, bf::AbstractVector{T}, -) where {T, S, M1 <: MadIPMOperator, M2, MT} - if !bqp.meta.islp - mul!(bqp._HX, bqp.data.H, bx) - bf .= bqp.data.c0 .+ vec(sum(bqp.c_batch .* bx, dims=1)) .+ T(0.5) .* vec(sum(bx .* bqp._HX, dims=1)) - else - bf .= bqp.data.c0 .+ vec(sum(bqp.c_batch .* bx, dims=1)) - end - return bf -end - -function NLPModels.grad!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, - bx::AbstractMatrix{T}, bg::AbstractMatrix{T}, -) where {T, S, M1 <: MadIPMOperator, M2, MT} - if !bqp.meta.islp - mul!(bg, bqp.data.H, bx) - bg .+= bqp.c_batch - else - copyto!(bg, bqp.c_batch) - end - return bg -end - -function NLPModels.cons!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, - bx::AbstractMatrix{T}, bc::AbstractMatrix{T}, -) where {T, S, M1, M2 <: MadIPMOperator, MT} - mul!(bc, bqp.data.A, bx) - return bc -end - -function NLPModels.jac_structure!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, - jrows::AbstractVector{<:Integer}, - jcols::AbstractVector{<:Integer}, -) where {T, S, M1, M2 <: MadIPMOperator} - fill_structure!(bqp.data.A.A, jrows, jcols) - return jrows, jcols -end - -function NLPModels.hess_structure!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, - hrows::AbstractVector{<:Integer}, - hcols::AbstractVector{<:Integer}, -) where {T, S, M1 <: MadIPMOperator, M2} - fill_structure!(bqp.data.H.A, hrows, hcols) - return hrows, hcols -end - -function NLPModels.jac_coord!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, - bx::AbstractMatrix, - bjvals::AbstractMatrix, -) where {T, S, M1, M2 <: MadIPMOperator} - bjvals .= bqp.data.A.A.nzVal - return bjvals -end - -function NLPModels.hess_coord!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2}, - bx::AbstractMatrix, - by::AbstractMatrix, - bobj_weight::AbstractVector, - bhvals::AbstractMatrix, -) where {T, S, M1 <: MadIPMOperator, M2} - H = bqp.data.H.A - nnzh = nnz(H) - nnzh == 0 && return bhvals - mul!(bhvals, H.nzVal, bobj_weight') - return bhvals -end - -function NLPModels.hprod!( - bqp::ObjRHSBatchQuadraticModel{T, S, M1, M2, MT}, - bx::AbstractMatrix{T}, by::AbstractMatrix{T}, bv::AbstractMatrix{T}, - bobj_weight::AbstractVector{T}, bHv::AbstractMatrix{T}, -) where {T, S, M1 <: MadIPMOperator, M2, MT} - mul!(bHv, bqp.data.H, bv) - bHv .*= bobj_weight' - return bHv -end - -function _expand_symmetric_coo(H::SparseMatrixCOO{Tv, Ti}) where {Tv, Ti} # FIXME: allocates - rows, cols, vals = H.rows, H.cols, H.vals - m, n = size(H) - offdiag = findall(i -> rows[i] != cols[i], 1:length(rows)) - new_rows = vcat(rows, cols[offdiag]) - new_cols = vcat(cols, rows[offdiag]) - new_vals = vcat(vals, vals[offdiag]) - return SparseMatrixCOO(m, n, new_rows, new_cols, new_vals) -end - -function Base.convert(::Type{ObjRHSBatchQuadraticModel{T, S}}, bnlp::ObjRHSBatchQuadraticModel{T}) where {T, S<:CuArray} - nbatch = bnlp.meta.nbatch - nvar = bnlp.meta.nvar - ncon = bnlp.meta.ncon - - H_orig_csr = CUSPARSE.CuSparseMatrixCSR(bnlp.data.H) - - H_full = _expand_symmetric_coo(bnlp.data.H) - H_full_csr = CUSPARSE.CuSparseMatrixCSR(H_full) - - A_csr = CUSPARSE.CuSparseMatrixCSR(bnlp.data.A) - - H_op = MadIPMOperator(H_full_csr; symmetric=false, spmm_ncols=nbatch) - H_op.A = H_orig_csr - A_op = MadIPMOperator(A_csr; symmetric=false, spmm_ncols=nbatch) - - c_gpu = CuVector{T}(bnlp.data.c) - v_gpu = CuVector{T}(bnlp.data.v) - data_gpu = QPData(bnlp.data.c0, c_gpu, v_gpu, H_op, A_op) - - c_batch_gpu = CuMatrix{T}(bnlp.c_batch) - _HX_gpu = CUDA.zeros(T, nvar, nbatch) - _AX_gpu = CUDA.zeros(T, ncon, nbatch) - - VT = typeof(c_gpu) - MT = typeof(c_batch_gpu) - - meta_gpu = NLPModels.BatchNLPModelMeta{T, MT}( - nbatch, nvar; - x0 = CuMatrix{T}(bnlp.meta.x0), - lvar = CuMatrix{T}(bnlp.meta.lvar), - uvar = CuMatrix{T}(bnlp.meta.uvar), - ncon = ncon, - lcon = CuMatrix{T}(bnlp.meta.lcon), - ucon = CuMatrix{T}(bnlp.meta.ucon), - nnzj = bnlp.meta.nnzj, - nnzh = bnlp.meta.nnzh, - islp = bnlp.meta.islp, - ) - - return ObjRHSBatchQuadraticModel{T, VT, typeof(H_op), typeof(A_op), MT}( - meta_gpu, data_gpu, c_batch_gpu, _HX_gpu, _AX_gpu, - ) -end - -function Base.convert(::Type{BatchQuadraticModel{T, MT}}, bnlp::BatchQuadraticModel{T}) where {T, MT<:CuMatrix} - nbatch = bnlp.meta.nbatch - nvar = bnlp.meta.nvar - ncon = bnlp.meta.ncon - - c_batch_gpu = MT(bnlp.c_batch) - c0_batch_gpu = CuVector{T}(bnlp.c0_batch) - H_nzvals_gpu = MT(bnlp.H_nzvals) - A_nzvals_gpu = MT(bnlp.A_nzvals) - - hess_rows_gpu = CuVector{Int}(bnlp.hess_rows) - hess_cols_gpu = CuVector{Int}(bnlp.hess_cols) - - _jac_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._jac_scatter) - _jact_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._jact_scatter) - _hess_scatter_gpu = CUSPARSE.CuSparseMatrixCSC(bnlp._hess_scatter) - - _hess_sym_gather_cols_gpu = CuVector{Int}(bnlp._hess_sym_gather_cols) - _hess_sym_nzidx_gpu = CuVector{Int}(bnlp._hess_sym_nzidx) - - _HX_gpu = CUDA.zeros(T, nvar, nbatch) - nnzj = bnlp.meta.nnzj - sym_nnzh = size(bnlp._hess_buffer, 1) - _jac_buffer_gpu = CUDA.zeros(T, nnzj, nbatch) - _hess_buffer_gpu = CUDA.zeros(T, sym_nnzh, nbatch) - - VT = typeof(c0_batch_gpu) - VI = typeof(hess_rows_gpu) - SpMT_J = typeof(_jac_scatter_gpu) - SpMT_H = typeof(_hess_scatter_gpu) - - meta_gpu = NLPModels.BatchNLPModelMeta{T, MT}( - nbatch, nvar; - x0 = MT(bnlp.meta.x0), - lvar = MT(bnlp.meta.lvar), - uvar = MT(bnlp.meta.uvar), - ncon = ncon, - lcon = MT(bnlp.meta.lcon), - ucon = MT(bnlp.meta.ucon), - nnzj = bnlp.meta.nnzj, - nnzh = bnlp.meta.nnzh, - islp = bnlp.meta.islp, - ) - - return BatchQuadraticModel{T, MT, SpMT_J, SpMT_H, VT, VI}( - meta_gpu, - c_batch_gpu, c0_batch_gpu, H_nzvals_gpu, A_nzvals_gpu, - hess_rows_gpu, hess_cols_gpu, - _jac_scatter_gpu, _jact_scatter_gpu, _hess_scatter_gpu, - _hess_sym_gather_cols_gpu, _hess_sym_nzidx_gpu, - _HX_gpu, _jac_buffer_gpu, _hess_buffer_gpu, - ) -end diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 4afefa27..b0770e9c 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -188,35 +188,6 @@ MadIPM._colptr(A::CuSparseMatrixCSC) = A.colPtr MadIPM._rowval(A::CuSparseMatrixCSC) = A.rowVal MadIPM._nzval(A::CuSparseMatrixCSC) = A.nzVal -function MadIPM._coo_to_scatter( - coo_I, nrows::Int, n_entries::Int, - proto_I, nzVals::CuMatrix{T}, batch_size::Int, -) where T - if n_entries == 0 - scatter = CUSPARSE.CuSparseMatrixCSC( - CuVector{Int32}([1]), - CuVector{Int32}(undef, 0), - CuVector{T}(undef, 0), - (nrows, 0), - ) - op = MadIPMOperator(scatter; spmm_ncols=batch_size) - buffer = similar(nzVals, 0, batch_size) - return op, buffer - end - coo_J = similar(proto_I, n_entries) - coo_J .= Int32(1):Int32(n_entries) - coo_V = similar(nzVals, n_entries) - fill!(coo_V, one(T)) - scatter, _ = MadNLP.coo_to_csc( - MadNLP.SparseMatrixCOO(nrows, n_entries, coo_I, coo_J, coo_V), - ) - fill!(MadIPM._nzval(scatter), one(T)) - op = MadIPMOperator(scatter; spmm_ncols=batch_size) - buffer = similar(nzVals, n_entries, batch_size) - fill!(buffer, zero(T)) - return op, buffer -end - # we introduce a new constructor that takes the nzvals as a matrix explicitly function MadNLPGPU.CUDSSSolver( aug_com::CUSPARSE.CuSparseMatrixCSC{T,Cint}, diff --git a/ext/MadIPMMathOptInterfaceExt/parse_moi.jl b/ext/MadIPMMathOptInterfaceExt/parse_moi.jl index 9faab7ab..8051e9bb 100644 --- a/ext/MadIPMMathOptInterfaceExt/parse_moi.jl +++ b/ext/MadIPMMathOptInterfaceExt/parse_moi.jl @@ -170,7 +170,6 @@ function qp_model(moimodel::MOI.ModelLike) data = QuadraticModels.QPData( d, c, - zeros(length(c)), Q, A, ) diff --git a/src/MadIPM.jl b/src/MadIPM.jl index 7296a515..ab589ad6 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -7,7 +7,14 @@ import MadNLP import MadNLP: full, LDLFactorizations import NLPModels import QuadraticModels -import QuadraticModels: SparseMatrixCOO +import QuadraticModels: QPData, QuadraticModel, SparseMatrixCOO + +import BatchQuadraticModels +import BatchQuadraticModels: + ObjRHSBatchQuadraticModel, + BatchSparseOp, + batch_spmv!, + batch_mapreduce! include("utils.jl") include("structure.jl") diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 85ca737d..76d51917 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -1,5 +1,5 @@ """Batched version of SparseKKTSystem""" -struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, SMT} <: AbstractBatchKKTSystem{T} +struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT} <: AbstractBatchKKTSystem{T} nzVals::MT # (aug_mat_length × batch_size) COO nonzero values aug_I::VI32 # shared row indices aug_J::VI32 # shared column indices @@ -17,24 +17,10 @@ struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, SMT} <: AbstractBatchKKT u_diag::MT # (nub × batch_size) upper bound diagonals l_lower::MT # (nlb × batch_size) lower bound multipliers u_lower::MT # (nub × batch_size) upper bound multipliers - # Hessian scatter (for mul!) - hess_scatter::SMT # (n_tot × n_hess_sym) scatter matrix - hess_nz_map::VI # nzVals row indices (with sym duplication) - hess_var_map::VI # variable indices for x multiplication - hess_buffer::MT # (n_hess_sym × batch_size) workspace - # J^T scatter (for jtprod! and mul!) - jt_scatter::SMT # (n_tot × n_jac_aug) scatter: S[var_idx, k] = 1 - jt_nz_map::VI # nzVals row indices for Jacobian entries - jt_con_map::VI # maps each Jac nonzero to its constraint index (1:m) - jt_con_map_full::VI # jt_con_map offset by n_tot (for indexing into full KKT vector) - jt_buffer::MT # (n_jac_aug × batch_size) buffer for jtprod - # J scatter (for mul!) - j_scatter::SMT # (m × n_jac_aug) scatter: S[con_idx, k] = 1 - j_var_map::VI # variable indices for J entries - j_buffer::MT # (n_jac_aug × batch_size) buffer for jprod - # Workspace for mul! (GPU needs full matrices, not SubArray views) - _mul_w_primal::MT # (n_tot × batch_size) - _mul_w_dual::MT # (m × batch_size) + # Operators for batch SpMV (jtprod! and KKT mul!) + hess_op::OPT + jt_op::OPT + j_op::OPT # Batch tracking batch_map::Vector{Int} # original index → active position (0 if inactive) batch_map_rev::Vector{Int} # active position → original index @@ -117,17 +103,9 @@ function MadNLP.create_kkt_system( jac_range = n_tot+n_hess+1:n_tot+n_hess+n_jac+n_slack - hess_scatter, hess_nz_map, hess_var_map, hess_buffer = _build_hess_scatter( - I, J, n_tot, n_hess, nzVals, aug_csc_map, batch_size, - ) - jt_scatter, jt_nz_map, jt_con_map, jt_buffer = _build_scatter( - I, J, jac_range, n_tot, nzVals, aug_csc_map, batch_size, - ) - jt_con_map_full = similar(jt_con_map) - jt_con_map_full .= jt_con_map .+ Int32(n_tot) - j_scatter, _, j_var_map, j_buffer = _build_jac_scatter( - I, J, jac_range, n_tot, m, nzVals, aug_csc_map, batch_size, - ) + hess_op = _build_hess_op(I, J, n_tot, n_hess, nzVals, aug_csc_map) + jt_op = _build_jt_op(I, J, jac_range, n_tot, nzVals, aug_csc_map) + j_op = _build_j_op(I, J, jac_range, n_tot, m, nzVals, aug_csc_map) reg = similar(nzVals, n_tot, batch_size) l_diag = similar(nzVals, nlb, batch_size) @@ -135,24 +113,18 @@ function MadNLP.create_kkt_system( l_lower = similar(nzVals, nlb, batch_size) u_lower = similar(nzVals, nub, batch_size) - _mul_w_primal = similar(nzVals, n_tot, batch_size) - _mul_w_dual = similar(nzVals, m, batch_size) - batch_map = collect(1:batch_size) batch_map_rev = collect(1:batch_size) active_batch_size = Ref(batch_size) LS = typeof(batch_ls) VI32 = typeof(I) - SMT = typeof(jt_scatter) - return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, SMT}( + OPT = typeof(jt_op) + return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT}( nzVals, I, J, batch_ls, rhs_buffer, batch_size, aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, reg, l_diag, u_diag, l_lower, u_lower, - hess_scatter, hess_nz_map, hess_var_map, hess_buffer, - jt_scatter, jt_nz_map, jt_con_map, jt_con_map_full, jt_buffer, - j_scatter, j_var_map, j_buffer, - _mul_w_primal, _mul_w_dual, + hess_op, jt_op, j_op, batch_map, batch_map_rev, active_batch_size, ) end @@ -251,14 +223,8 @@ function MadNLP.factorize_wrapper!(batch_solver::AbstractBatchMPCSolver) return end -function _gather_mul!(out::AbstractMatrix, A::AbstractMatrix, nz_map, B::AbstractMatrix, val_map) - @views out .= A[nz_map, :] .* B[val_map, :] - return out -end - function MadNLP.jtprod!(res::AbstractMatrix, bkkt::SparseUniformBatchKKTSystem, y::BatchVector) - _gather_mul!(bkkt.jt_buffer, bkkt.nzVals, bkkt.jt_nz_map, MadNLP.full(y), bkkt.jt_con_map) - mul!(res, bkkt.jt_scatter, bkkt.jt_buffer) + batch_spmv!(res, bkkt.jt_op, MadNLP.full(y)) return res end @@ -332,25 +298,13 @@ function LinearAlgebra.mul!( alpha = one(T), beta = zero(T), ) where T - nzV = bkkt.nzVals - wp = bkkt._mul_w_primal - wd = bkkt._mul_w_dual - - # mul!(primal(w), Symmetric(hess_com, :L), primal(x), alpha, beta) xv = MadNLP.full(x) - _gather_mul!(bkkt.hess_buffer, nzV, bkkt.hess_nz_map, xv, bkkt.hess_var_map) - mul!(wp, bkkt.hess_scatter, bkkt.hess_buffer) - MadNLP.primal(w) .= beta .* MadNLP.primal(w) .+ alpha .* wp - - # mul!(primal(w), jac_com', dual(x), alpha, one(T)) - _gather_mul!(bkkt.jt_buffer, nzV, bkkt.jt_nz_map, xv, bkkt.jt_con_map_full) - mul!(wp, bkkt.jt_scatter, bkkt.jt_buffer) - MadNLP.primal(w) .+= alpha .* wp - - # mul!(dual(w), jac_com, primal(x), alpha, beta) - _gather_mul!(bkkt.j_buffer, nzV, bkkt.jt_nz_map, xv, bkkt.j_var_map) - mul!(wd, bkkt.j_scatter, bkkt.j_buffer) - MadNLP.dual(w) .= beta .* MadNLP.dual(w) .+ alpha .* wd + # mul!(primal(w), Symmetric(kkt.hess_com, :L), primal(x), alpha, beta) + batch_spmv!(MadNLP.primal(w), bkkt.hess_op, xv, alpha, beta) + # mul!(primal(w), kkt.jac_com', dual(x), alpha, one(T)) + batch_spmv!(MadNLP.primal(w), bkkt.jt_op, xv, alpha, one(T); val_offset=bkkt.n_tot) + # mul!(dual(w), kkt.jac_com, primal(x), alpha, beta) + batch_spmv!(MadNLP.dual(w), bkkt.j_op, xv, alpha, beta) _kktmul!(w, x, bkkt.reg, du_diag(bkkt), bkkt.l_lower, bkkt.u_lower, bkkt.l_diag, bkkt.u_diag, alpha, beta) return w end diff --git a/src/batch/utils.jl b/src/batch/utils.jl index 1e6ea3fe..9ef3f24c 100644 --- a/src/batch/utils.jl +++ b/src/batch/utils.jl @@ -8,73 +8,63 @@ function _csc_with_nzval(A::SparseArrays.SparseMatrixCSC, nzval, n) return SparseArrays.SparseMatrixCSC(n, n, SparseArrays.getcolptr(A), SparseArrays.rowvals(A), nzval) end +batch_maximum!(out::AbstractMatrix, src::AbstractMatrix) = maximum!(out, src) +batch_minimum!(out::AbstractMatrix, src::AbstractMatrix) = minimum!(out, src) +batch_sum!(out::AbstractMatrix, src::AbstractMatrix) = sum!(out, src) + function zero_inactive_step!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace ws.alpha_p .*= ws.active_mask ws.alpha_d .*= ws.active_mask end -function _coo_to_scatter( - coo_I, nrows::Int, n_entries::Int, - proto_I, nzVals::AbstractMatrix{T}, batch_size::Int, -) where T - if n_entries == 0 - scatter = SparseArrays.sparse(Int32[], Int32[], T[], nrows, 0) - buffer = similar(nzVals, 0, batch_size) - return scatter, buffer - end - coo_J = similar(proto_I, n_entries) - coo_J .= Int32(1):Int32(n_entries) - coo_V = similar(nzVals, n_entries) - fill!(coo_V, one(T)) - scatter, _ = MadNLP.coo_to_csc( - MadNLP.SparseMatrixCOO(nrows, n_entries, coo_I, coo_J, coo_V), +function _build_batch_op(nzVals, nz_map, val_map, coo_I, nrows) + rowptr, colidx = BatchQuadraticModels._coo_to_csr(Vector{Int}(coo_I), nrows) + return BatchQuadraticModels._build_storage_op( + nzVals, + rowptr, + Vector{Int}(nz_map), + Vector{Int}(val_map), + colidx, ) - fill!(_nzval(scatter), one(T)) - buffer = similar(nzVals, n_entries, batch_size) - fill!(buffer, zero(T)) - return scatter, buffer end -function _build_scatter( +function _build_jt_op( aug_I, aug_J, jac_range, n_tot, - nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, + nzVals::AbstractMatrix{T}, aug_csc_map, ) where T n_jac = length(jac_range) coo_I = similar(aug_I, n_jac) coo_I .= aug_J[jac_range] - scatter, buffer = _coo_to_scatter(coo_I, n_tot, n_jac, aug_I, nzVals, batch_size) nz_map = similar(aug_csc_map, n_jac) nz_map .= jac_range con_map = similar(aug_csc_map, n_jac) con_map .= aug_I[jac_range] .- Int32(n_tot) - return scatter, nz_map, con_map, buffer + return _build_batch_op(nzVals, nz_map, con_map, coo_I, n_tot) end -function _build_jac_scatter( +function _build_j_op( aug_I, aug_J, jac_range, n_tot, m, - nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, + nzVals::AbstractMatrix{T}, aug_csc_map, ) where T n_jac = length(jac_range) coo_I = similar(aug_I, n_jac) coo_I .= aug_I[jac_range] .- Int32(n_tot) - scatter, buffer = _coo_to_scatter(coo_I, m, n_jac, aug_I, nzVals, batch_size) nz_map = similar(aug_csc_map, n_jac) nz_map .= jac_range var_map = similar(aug_csc_map, n_jac) var_map .= aug_J[jac_range] - return scatter, nz_map, var_map, buffer + return _build_batch_op(nzVals, nz_map, var_map, coo_I, m) end -function _build_hess_scatter( +function _build_hess_op( aug_I, aug_J, n_tot, n_hess, - nzVals::AbstractMatrix{T}, aug_csc_map, batch_size, + nzVals::AbstractMatrix{T}, aug_csc_map, ) where T if n_hess == 0 - scatter, buffer = _coo_to_scatter(similar(aug_I, 0), n_tot, 0, aug_I, nzVals, batch_size) nz_map = similar(aug_csc_map, 0) var_map = similar(aug_csc_map, 0) - return scatter, nz_map, var_map, buffer + return _build_batch_op(nzVals, nz_map, var_map, similar(aug_I, 0), n_tot) end hess_range = n_tot+1:n_tot+n_hess @@ -84,11 +74,9 @@ function _build_hess_scatter( offdiag_idx = findall(hess_I .!= hess_J) n_hess_sym = n_hess + length(offdiag_idx) - scatter_rows = similar(aug_I, n_hess_sym) - scatter_rows[1:n_hess] .= hess_I - scatter_rows[n_hess+1:end] .= hess_J[offdiag_idx] - - scatter, buffer = _coo_to_scatter(scatter_rows, n_tot, n_hess_sym, aug_I, nzVals, batch_size) + coo_rows = similar(aug_I, n_hess_sym) + coo_rows[1:n_hess] .= hess_I + coo_rows[n_hess+1:end] .= hess_J[offdiag_idx] nz_map = similar(aug_csc_map, n_hess_sym) nz_map[1:n_hess] .= hess_range @@ -98,7 +86,7 @@ function _build_hess_scatter( var_map[1:n_hess] .= hess_J var_map[n_hess+1:end] .= hess_I[offdiag_idx] - return scatter, nz_map, var_map, buffer + return _build_batch_op(nzVals, nz_map, var_map, coo_rows, n_tot) end struct BatchVector{T, MT<:AbstractMatrix{T}} diff --git a/src/utils.jl b/src/utils.jl index 84f5e916..d8773830 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -477,7 +477,6 @@ function standard_form_qp(qp::QuadraticModels.QuadraticModel) data = QuadraticModels.QPData( qp.data.c0, [qp.data.c; zeros(ns + nw)], - [qp.data.v; zeros(ns + nw)], Hs, As, ) @@ -502,4 +501,3 @@ function standard_form_qp(qp::QuadraticModels.QuadraticModel) data, ) end - diff --git a/test/Project.toml b/test/Project.toml index bd0faf2d..e2129e29 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,9 +1,11 @@ [deps] AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ArgTools = "0dad84c5-d112-42e6-8d28-ef12dabb789f" Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" +BatchQuadraticModels = "412afcf0-3ec6-4826-be17-8d792afa05a8" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CompilerSupportLibraries_jll = "e66e0078-7015-5450-92f7-15fbd957f2ae" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6" @@ -12,8 +14,8 @@ FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FileWatching = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLLWrappers = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" JuliaSyntaxHighlighting = "ac6e5ff7-fb65-4e79-a425-ec3bc9c03011" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" LibCURL = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" LibCURL_jll = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" @@ -29,9 +31,9 @@ MUMPS_seq_jll = "d7ed1dd3-d0ae-5e8e-bfb4-87a502085b8d" MadIPM = "4406d55e-1099-4ca6-ac6a-6cf051b10024" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" -MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" MadNLPTests = "b52a2a03-04ab-4a5f-9698-6a2deff93217" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" +MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" MozillaCACerts_jll = "14a3606d-f60d-562e-9121-12d972cd8159" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" NLPModelsModifiers = "e01155f1-5c6f-4375-a9d8-616dd036575f" @@ -65,5 +67,5 @@ nghttp2_jll = "8e850ede-7688-5339-a07c-302acd2aaf8d" p7zip_jll = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" [sources] +BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl"} MadIPM = {path = ".."} -QuadraticModels = {url = "https://github.com/klamike/QuadraticModels.jl.git", rev = "mk/rhsbatch"} diff --git a/test/batch_kernels/setup.jl b/test/batch_kernels/setup.jl index 3f3fa710..20f5da74 100644 --- a/test/batch_kernels/setup.jl +++ b/test/batch_kernels/setup.jl @@ -4,7 +4,7 @@ using SparseArrays using MadNLP using MadIPM using QuadraticModels -using QuadraticModels: ObjRHSBatchQuadraticModel +using BatchQuadraticModels: ObjRHSBatchQuadraticModel # ────────────────────────────────────────────────────────────── # Test problems diff --git a/test/test_batch.jl b/test/test_batch.jl index 8c3e509f..ae1ff110 100644 --- a/test/test_batch.jl +++ b/test/test_batch.jl @@ -1,4 +1,4 @@ -using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel +using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel function _make_small_qp() # Small QP: min 0.5 xᵀHx + cᵀx s.t. lcon ≤ Ax ≤ ucon, lvar ≤ x ≤ uvar diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl index d990bc27..1077335f 100644 --- a/test/test_batch_gpu.jl +++ b/test/test_batch_gpu.jl @@ -1,6 +1,6 @@ using KernelAbstractions using MadNLPGPU -using QuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel +using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel # ============================================================ # Test problem constructors for GPU batch tests From d2512465d68064a1d6d4305ab3d2d4a31e2e6500 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 26 Mar 2026 20:31:27 -0400 Subject: [PATCH 33/51] kernels for reduce_rhs!, finish_aug_solve! --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 49 +++++++++++++++++++------ src/batch/KKT/Sparse/augmented.jl | 35 ++++++++++++++---- 2 files changed, 64 insertions(+), 20 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index c1220adf..ecd725c6 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -164,19 +164,44 @@ function MadIPM._mehrotra_correct_steps!( end end -@kernel function _gather_mul_kernel!(out, @Const(A), @Const(nz_map), @Const(B), @Const(val_map)) +@kernel function _reduce_rhs_lb_kernel!(values, @Const(ind_lb), lb_off, @Const(l_diag)) i, j = @index(Global, NTuple) - @inbounds out[i, j] = A[nz_map[i], j] * B[val_map[i], j] + @inbounds values[ind_lb[i], j] -= values[lb_off + i, j] / l_diag[i, j] end -function MadIPM._gather_mul!( - out::CuMatrix, A::CuMatrix, nz_map::CuVector, B::CuMatrix, val_map::CuVector, -) - n, bs = size(out) - if n > 0 - backend = CUDABackend() - _gather_mul_kernel!(backend)(out, A, nz_map, B, val_map; ndrange=(n, bs)) - KernelAbstractions.synchronize(backend) - end - return out +@kernel function _reduce_rhs_ub_kernel!(values, @Const(ind_ub), ub_off, @Const(u_diag)) + i, j = @index(Global, NTuple) + @inbounds values[ind_ub[i], j] -= values[ub_off + i, j] / u_diag[i, j] +end + +function MadIPM._reduce_rhs_batch!(values::CuMatrix, ind_lb, lb_off, l_diag, + ind_ub, ub_off, u_diag) + bs = size(values, 2); backend = CUDABackend() + nlb = length(ind_lb) + nlb > 0 && _reduce_rhs_lb_kernel!(backend)(values, ind_lb, lb_off, l_diag; ndrange=(nlb, bs)) + nub = length(ind_ub) + nub > 0 && _reduce_rhs_ub_kernel!(backend)(values, ind_ub, ub_off, u_diag; ndrange=(nub, bs)) + KernelAbstractions.synchronize(backend) + return +end + +@kernel function _finish_aug_solve_lb_kernel!(values, @Const(ind_lb), lb_off, @Const(l_lower), @Const(l_diag)) + i, j = @index(Global, NTuple) + @inbounds values[lb_off + i, j] = (-values[lb_off + i, j] + l_lower[i, j] * values[ind_lb[i], j]) / l_diag[i, j] +end + +@kernel function _finish_aug_solve_ub_kernel!(values, @Const(ind_ub), ub_off, @Const(u_lower), @Const(u_diag)) + i, j = @index(Global, NTuple) + @inbounds values[ub_off + i, j] = (values[ub_off + i, j] - u_lower[i, j] * values[ind_ub[i], j]) / u_diag[i, j] +end + +function MadIPM._finish_aug_solve_batch!(values::CuMatrix, ind_lb, lb_off, l_lower, l_diag, + ind_ub, ub_off, u_lower, u_diag) + bs = size(values, 2); backend = CUDABackend() + nlb = length(ind_lb) + nlb > 0 && _finish_aug_solve_lb_kernel!(backend)(values, ind_lb, lb_off, l_lower, l_diag; ndrange=(nlb, bs)) + nub = length(ind_ub) + nub > 0 && _finish_aug_solve_ub_kernel!(backend)(values, ind_ub, ub_off, u_lower, u_diag; ndrange=(nub, bs)) + KernelAbstractions.synchronize(backend) + return end diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 76d51917..d3ac4e07 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -179,20 +179,39 @@ function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs:: return rhs end +function _reduce_rhs_batch!(values::AbstractMatrix, ind_lb, lb_off, l_diag, ind_ub, ub_off, u_diag) + nlb = length(ind_lb); nub = length(ind_ub); bs = size(values, 2) + @inbounds for j in 1:bs, i in 1:nlb + values[ind_lb[i], j] -= values[lb_off + i, j] / l_diag[i, j] + end + @inbounds for j in 1:bs, i in 1:nub + values[ind_ub[i], j] -= values[ub_off + i, j] / u_diag[i, j] + end +end + +function _finish_aug_solve_batch!(values::AbstractMatrix, ind_lb, lb_off, l_lower, l_diag, + ind_ub, ub_off, u_lower, u_diag) + nlb = length(ind_lb); nub = length(ind_ub); bs = size(values, 2) + @inbounds for j in 1:bs, i in 1:nlb + values[lb_off + i, j] = (-values[lb_off + i, j] + l_lower[i, j] * values[ind_lb[i], j]) / l_diag[i, j] + end + @inbounds for j in 1:bs, i in 1:nub + values[ub_off + i, j] = (values[ub_off + i, j] - u_lower[i, j] * values[ind_ub[i], j]) / u_diag[i, j] + end +end + function MadNLP.reduce_rhs!(bkkt::SparseUniformBatchKKTSystem, d::BatchUnreducedKKTVector) - MadNLP.reduce_rhs!( - xp_lr(d), MadNLP.dual_lb(d), bkkt.l_diag, - xp_ur(d), MadNLP.dual_ub(d), bkkt.u_diag, - ) + lb_off = d.n + d.m + _reduce_rhs_batch!(d.values, d.ind_lb, lb_off, bkkt.l_diag, + d.ind_ub, lb_off + d.nlb, bkkt.u_diag) return end function MadNLP.finish_aug_solve!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) d = batch_solver.d - dzl = MadNLP.dual_lb(d) - dzu = MadNLP.dual_ub(d) - dzl .= (.-dzl .+ bkkt.l_lower .* xp_lr(d)) ./ bkkt.l_diag - dzu .= (dzu .- bkkt.u_lower .* xp_ur(d)) ./ bkkt.u_diag + lb_off = d.n + d.m + _finish_aug_solve_batch!(d.values, d.ind_lb, lb_off, bkkt.l_lower, bkkt.l_diag, + d.ind_ub, lb_off + d.nlb, bkkt.u_lower, bkkt.u_diag) return end From 78a279464bf942ed46a1fe4e1597c6876ce689c9 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 26 Mar 2026 20:39:58 -0400 Subject: [PATCH 34/51] more kernels --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 385 ++++++++++++++++++------ src/batch/madipm/kernels.jl | 357 +++++++++++++--------- src/batch/madipm/solver.jl | 178 ++++++----- src/batch/madnlp/kernels.jl | 24 +- src/batch/structure.jl | 10 +- test/batch_kernels/test_termination.jl | 11 +- 6 files changed, 606 insertions(+), 359 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index ecd725c6..802b11a5 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -47,121 +47,306 @@ function MadNLP._set_con_scale_sparse!( return con_scale end -@kernel function _mehrotra_correction_kernel!( - alpha_p, alpha_d, - @Const(mu), - @Const(val_xl), @Const(idx_xl), @Const(val_xu), @Const(idx_xu), - @Const(val_zl), @Const(idx_zl), @Const(val_zu), @Const(idx_zu), - @Const(d_vals), @Const(x_vals), @Const(xl_vals), @Const(xu_vals), - @Const(zl_vals), @Const(zu_vals), - @Const(ind_lb), @Const(ind_ub), - dlb_off, dub_off, gamma_f, -) - j = @index(Global, Linear) - T = eltype(alpha_p) +_ftb_primal_lb_kernel!(alpha_out, dx, x, xb, tau, nrows) = begin + bs = size(alpha_out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(alpha_out) + @inbounds if j <= bs + a = T(Inf); τ = tau[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + d = dx[i, j] + d < zero(T) && (a = min(a, (-x[i, j] + xb[i, j]) * τ / d)) + i += blockDim().x * gridDim_reduce + end + a = CUDA.reduce_block(min, a, T(Inf), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(min, alpha_out, j, a) + end + return +end - mu_j = mu[1, j] - max_ap = alpha_p[1, j] - max_ad = alpha_d[1, j] - - # primal step - corrected_p = one(T) - @inbounds if max_ap < one(T) - i_xl = idx_xl[1, j] - i_xu = idx_xu[1, j] - if val_xl[1, j] <= val_xu[1, j] && i_xl > Int32(0) - idx = ind_lb[i_xl] - zl_stepped = zl_vals[idx, j] + max_ad * d_vals[dlb_off + i_xl, j] - corrected_p = (x_vals[idx, j] - xl_vals[idx, j] - mu_j / zl_stepped) / (-d_vals[idx, j]) - elseif i_xu > Int32(0) - idx = ind_ub[i_xu] - zu_stepped = zu_vals[idx, j] + max_ad * d_vals[dub_off + i_xu, j] - corrected_p = (xu_vals[idx, j] - x_vals[idx, j] - mu_j / zu_stepped) / d_vals[idx, j] +_ftb_primal_ub_kernel!(alpha_out, dx, x, xb, tau, nrows) = begin + bs = size(alpha_out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(alpha_out) + @inbounds if j <= bs + a = T(Inf); τ = tau[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + d = dx[i, j] + d > zero(T) && (a = min(a, (-x[i, j] + xb[i, j]) * τ / d)) + i += blockDim().x * gridDim_reduce end + a = CUDA.reduce_block(min, a, T(Inf), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(min, alpha_out, j, a) end - @inbounds alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) - - # dual step - corrected_d = one(T) - @inbounds if max_ad < one(T) - i_zl = idx_zl[1, j] - i_zu = idx_zu[1, j] - if val_zl[1, j] <= val_zu[1, j] && i_zl > Int32(0) - idx = ind_lb[i_zl] - x_gap = x_vals[idx, j] + max_ap * d_vals[idx, j] - xl_vals[idx, j] - corrected_d = -(zl_vals[idx, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] - elseif i_zu > Int32(0) - idx = ind_ub[i_zu] - x_gap = xu_vals[idx, j] - x_vals[idx, j] - max_ap * d_vals[idx, j] - corrected_d = -(zu_vals[idx, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] + return +end + +_ftb_dual_lb_kernel!(alpha_out, dz, z, tau, nrows) = begin + bs = size(alpha_out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(alpha_out) + @inbounds if j <= bs + a = T(Inf); τ = tau[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + d = dz[i, j] + d < zero(T) && (a = min(a, -z[i, j] * τ / d)) + i += blockDim().x * gridDim_reduce end + a = CUDA.reduce_block(min, a, T(Inf), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(min, alpha_out, j, a) end - @inbounds alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) + return end -@kernel function _gather_compl_kernel!( - scratch, @Const(x_vals), @Const(xb_vals), @Const(z_vals), @Const(ind), -) - i, j = @index(Global, NTuple) - @inbounds begin - idx = ind[i] - scratch[i, j] = abs(x_vals[idx, j] - xb_vals[idx, j]) * z_vals[idx, j] +_ftb_dual_ub_kernel!(alpha_out, dz, z, tau, nrows) = begin + bs = size(alpha_out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(alpha_out) + @inbounds if j <= bs + a = T(Inf); τ = tau[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + d = dz[i, j] + (d < zero(T) && z[i, j] + d < zero(T)) && (a = min(a, -z[i, j] * τ / d)) + i += blockDim().x * gridDim_reduce + end + a = CUDA.reduce_block(min, a, T(Inf), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(min, alpha_out, j, a) end + return end -function MadIPM.get_inf_compl!( - inf_compl::CuMatrix, x::MadIPM.BatchPrimalVector, xl::MadIPM.BatchPrimalVector, - zl::MadIPM.BatchPrimalVector, xu::MadIPM.BatchPrimalVector, zu::MadIPM.BatchPrimalVector, - scratch_lb::CuMatrix, scratch_ub::CuMatrix, sum_lb, sum_ub, nlb, nub, -) - T = eltype(inf_compl) - bs = size(inf_compl, 2) - backend = CUDABackend() - if nlb > 0 - _gather_compl_kernel!(backend)( - scratch_lb, x.values, xl.values, zl.values, x.ind_lb; - ndrange=(nlb, bs), - ) - KernelAbstractions.synchronize(backend) - maximum!(sum_lb, scratch_lb) - else - fill!(sum_lb, zero(T)) +function _launch_ftb_kernel!(kernel_fn, alpha_out, nrows, srcs...) + T = eltype(alpha_out) + bs = size(alpha_out, 2) + fill!(alpha_out, T(Inf)) + nrows == 0 && return + kernel = @cuda launch=false kernel_fn(alpha_out, srcs..., nrows) + config = launch_configuration(kernel.fun) + threads = (config.threads ÷ 32) * 32 + reduce_blocks = min(cld(nrows, threads), max(1, cld(config.blocks, bs))) + kernel(alpha_out, srcs..., nrows; threads, blocks = reduce_blocks * bs) +end + +function MadIPM._ftb_primal_lb!(alpha_out::AnyCuMatrix, dx::AnyCuMatrix, x::AnyCuMatrix, xb::AnyCuMatrix, tau::AnyCuMatrix) + _launch_ftb_kernel!(_ftb_primal_lb_kernel!, alpha_out, size(dx, 1), dx, x, xb, tau) +end +function MadIPM._ftb_primal_ub!(alpha_out::AnyCuMatrix, dx::AnyCuMatrix, x::AnyCuMatrix, xb::AnyCuMatrix, tau::AnyCuMatrix) + _launch_ftb_kernel!(_ftb_primal_ub_kernel!, alpha_out, size(dx, 1), dx, x, xb, tau) +end +function MadIPM._ftb_dual_lb!(alpha_out::AnyCuMatrix, dz::AnyCuMatrix, z::AnyCuMatrix, tau::AnyCuMatrix) + _launch_ftb_kernel!(_ftb_dual_lb_kernel!, alpha_out, size(dz, 1), dz, z, tau) +end +function MadIPM._ftb_dual_ub!(alpha_out::AnyCuMatrix, dz::AnyCuMatrix, z::AnyCuMatrix, tau::AnyCuMatrix) + _launch_ftb_kernel!(_ftb_dual_ub_kernel!, alpha_out, size(dz, 1), dz, z, tau) +end + +_affine_compl_lb_kernel!(out, x, xl, z, dx, dz, αp, αd, nrows) = begin + bs = size(out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(out) + @inbounds if j <= bs + s = zero(T); ap = αp[1, j]; ad = αd[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + s += (x[i,j] + ap * dx[i,j] - xl[i,j]) * (z[i,j] + ad * dz[i,j]) + i += blockDim().x * gridDim_reduce + end + s = CUDA.reduce_block(+, s, zero(T), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(+, out, j, s) end - if nub > 0 - _gather_compl_kernel!(backend)( - scratch_ub, xu.values, x.values, zu.values, x.ind_ub; - ndrange=(nub, bs), - ) - KernelAbstractions.synchronize(backend) - maximum!(sum_ub, scratch_ub) - else - fill!(sum_ub, zero(T)) + return +end + +_affine_compl_ub_kernel!(out, xu, x, z, dx, dz, αp, αd, nrows) = begin + bs = size(out, 2) + blockIdx_reduce, j = fldmod1(blockIdx().x, bs) + gridDim_reduce = gridDim().x ÷ bs + T = eltype(out) + @inbounds if j <= bs + s = zero(T); ap = αp[1, j]; ad = αd[1, j] + i = threadIdx().x + (blockIdx_reduce - 1) * blockDim().x + while i <= nrows + s += (xu[i,j] - (x[i,j] + ap * dx[i,j])) * (z[i,j] + ad * dz[i,j]) + i += blockDim().x * gridDim_reduce + end + s = CUDA.reduce_block(+, s, zero(T), Val(true)) + threadIdx().x == 1 && _atomic_colreduce!(+, out, j, s) end - @. inf_compl = max(sum_lb, sum_ub) - return inf_compl + return end -function MadIPM._mehrotra_correct_steps!( - alpha_p::CuMatrix{T}, alpha_d::CuMatrix{T}, mu, - val_xl, idx_xl, val_xu, idx_xu, - val_zl, idx_zl, val_zu, idx_zu, - d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, - ind_lb, ind_ub, dlb_off::Int, dub_off::Int, gamma_f, -) where T - bs = size(alpha_p, 2) - if bs > 0 - backend = CUDABackend() - _mehrotra_correction_kernel!(backend)( - alpha_p, alpha_d, mu, - val_xl, idx_xl, val_xu, idx_xu, - val_zl, idx_zl, val_zu, idx_zu, - d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, - ind_lb, ind_ub, - Int32(dlb_off), Int32(dub_off), gamma_f; - ndrange = bs, - ) - KernelAbstractions.synchronize(backend) +function _launch_reduce_kernel!(kernel_fn, out, nrows, srcs...) + T = eltype(out) + bs = size(out, 2) + fill!(out, zero(T)) + nrows == 0 && return + kernel = @cuda launch=false kernel_fn(out, srcs..., nrows) + config = launch_configuration(kernel.fun) + threads = (config.threads ÷ 32) * 32 + reduce_blocks = min(cld(nrows, threads), max(1, cld(config.blocks, bs))) + kernel(out, srcs..., nrows; threads, blocks = reduce_blocks * bs) +end + +function MadIPM._affine_compl_lb!(out::AnyCuMatrix, x::AnyCuMatrix, xl::AnyCuMatrix, z::AnyCuMatrix, + dx::AnyCuMatrix, dz::AnyCuMatrix, αp::AnyCuMatrix, αd::AnyCuMatrix) + _launch_reduce_kernel!(_affine_compl_lb_kernel!, out, size(x, 1), x, xl, z, dx, dz, αp, αd) +end + +function MadIPM._affine_compl_ub!(out::AnyCuMatrix, xu::AnyCuMatrix, x::AnyCuMatrix, z::AnyCuMatrix, + dx::AnyCuMatrix, dz::AnyCuMatrix, αp::AnyCuMatrix, αd::AnyCuMatrix) + _launch_reduce_kernel!(_affine_compl_ub_kernel!, out, size(x, 1), xu, x, z, dx, dz, αp, αd) +end + +@inline function _warp_argmin(val::T, idx::Int32) where T + offset = Int32(16) + while offset > Int32(0) + other_val = CUDA.shfl_down_sync(0xffffffff, val, offset) + other_idx = CUDA.shfl_down_sync(0xffffffff, idx, offset) + if other_val < val + val = other_val + idx = other_idx + end + offset >>= Int32(1) + end + return val, idx +end + +_mehrotra_step_kernel!( + alpha_p, alpha_d, mu, gamma_f, + dx_lr, x_lr, xl_r, nlb::Int32, dzlb, zl_r, + dx_ur, x_ur, xu_r, nub::Int32, dzub, zu_r, + d_vals, ind_lb, ind_ub, dlb_off::Int32, dub_off::Int32, +) = begin + j = Int32(blockIdx().x) + lane = Int32(threadIdx().x - Int32(1)) # 0:31 + T = eltype(alpha_p) + INF = T(Inf) + + # primal lb + best_xl = INF; i_xl = Int32(0) + @inbounds begin + k = lane + Int32(1) + while k <= nlb + d = dx_lr[k, j] + if d < zero(T) + v = (xl_r[k, j] - x_lr[k, j]) / d + if v < best_xl + best_xl = v; i_xl = k + end + end + k += Int32(32) + end + end + best_xl, i_xl = _warp_argmin(best_xl, i_xl) + + # primal ub + best_xu = INF; i_xu = Int32(0) + @inbounds begin + k = lane + Int32(1) + while k <= nub + d = dx_ur[k, j] + if d > zero(T) + v = (xu_r[k, j] - x_ur[k, j]) / d + if v < best_xu + best_xu = v; i_xu = k + end + end + k += Int32(32) + end + end + best_xu, i_xu = _warp_argmin(best_xu, i_xu) + + # dual lb + best_zl = INF; i_zl = Int32(0) + @inbounds begin + k = lane + Int32(1) + while k <= nlb + d = dzlb[k, j] + if d < zero(T) + v = -zl_r[k, j] / d + if v < best_zl + best_zl = v; i_zl = k + end + end + k += Int32(32) + end + end + best_zl, i_zl = _warp_argmin(best_zl, i_zl) + + # dual ub + best_zu = INF; i_zu = Int32(0) + @inbounds begin + k = lane + Int32(1) + while k <= nub + d = dzub[k, j] + if d < zero(T) && zu_r[k, j] + d < zero(T) + v = -zu_r[k, j] / d + if v < best_zu + best_zu = v; i_zu = k + end + end + k += Int32(32) + end end + best_zu, i_zu = _warp_argmin(best_zu, i_zu) + + # lane 0: compute corrected steps + @inbounds if lane == Int32(0) + mu_j = mu[1, j] + max_ap = alpha_p[1, j] + max_ad = alpha_d[1, j] + + # corrected primal + corrected_p = one(T) + if max_ap < one(T) + if best_xl <= best_xu && i_xl > Int32(0) + zl_stepped = zl_r[i_xl, j] + max_ad * d_vals[dlb_off + i_xl, j] + corrected_p = (x_lr[i_xl, j] - xl_r[i_xl, j] - mu_j / zl_stepped) / (-dx_lr[i_xl, j]) + elseif i_xu > Int32(0) + zu_stepped = zu_r[i_xu, j] + max_ad * d_vals[dub_off + i_xu, j] + corrected_p = (xu_r[i_xu, j] - x_ur[i_xu, j] - mu_j / zu_stepped) / dx_ur[i_xu, j] + end + end + alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) + + # corrected dual + corrected_d = one(T) + if max_ad < one(T) + if best_zl <= best_zu && i_zl > Int32(0) + x_gap = x_lr[i_zl, j] + max_ap * dx_lr[i_zl, j] - xl_r[i_zl, j] + corrected_d = -(zl_r[i_zl, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] + elseif i_zu > Int32(0) + x_gap = xu_r[i_zu, j] - x_ur[i_zu, j] - max_ap * dx_ur[i_zu, j] + corrected_d = -(zu_r[i_zu, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] + end + end + alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) + end + return nothing +end + +function MadIPM._mehrotra_step!( + alpha_p::AnyCuMatrix, alpha_d, mu, gamma_f, + dx_lr, x_lr, xl_r, nlb, dzlb, zl_r, + dx_ur, x_ur, xu_r, nub, dzub, zu_r, + d_vals, ind_lb, ind_ub, dlb_off, dub_off, +) + CUDA.@cuda threads=32 blocks=bs _mehrotra_step_kernel!( + alpha_p, alpha_d, mu, gamma_f, + dx_lr, x_lr, xl_r, Int32(nlb), dzlb, zl_r, + dx_ur, x_ur, xu_r, Int32(nub), dzub, zu_r, + d_vals, ind_lb, ind_ub, Int32(dlb_off), Int32(dub_off), + ) + return end @kernel function _reduce_rhs_lb_kernel!(values, @Const(ind_lb), lb_off, @Const(l_diag)) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index b2f076eb..5a066490 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -1,16 +1,14 @@ function dual_objective!(dual_obj, y_vals, rhs_vals, zl_r, xl_r, zu_r, xu_r, - scratch_m, scratch_lb, scratch_ub, sum_lb, sum_ub, nlb, nub) - @. scratch_m = y_vals * rhs_vals - sum!(dual_obj, scratch_m) - dual_obj .*= -one(eltype(dual_obj)) + sum_lb, sum_ub, nlb, nub) + T = eltype(dual_obj) + batch_mapreduce!(*, +, zero(T), dual_obj, y_vals, rhs_vals) + dual_obj .*= -one(T) if nlb > 0 - @. scratch_lb = zl_r * xl_r - sum!(sum_lb, scratch_lb) + batch_mapreduce!(*, +, zero(T), sum_lb, zl_r, xl_r) dual_obj .+= sum_lb end if nub > 0 - @. scratch_ub = zu_r * xu_r - sum!(sum_ub, scratch_ub) + batch_mapreduce!(*, +, zero(T), sum_ub, zu_r, xu_r) dual_obj .-= sum_ub end return dual_obj @@ -144,14 +142,8 @@ function get_complementarity_measure!(solver::AbstractBatchMPCSolver) x_ur = upper(solver.x) zu_r = upper(solver.zu) - ws.sum_lb .= mapreduce( - (x, xl, z) -> (x - xl) * z, +, x_lr, xl_r, zl_r; - dims=1, init=zero(T), - ) - ws.sum_ub .= mapreduce( - (xu, x, z) -> (xu - x) * z, +, xu_r, x_ur, zu_r; - dims=1, init=zero(T), - ) + batch_mapreduce!((x, xl, z) -> (x - xl) * z, +, zero(T), ws.sum_lb, x_lr, xl_r, zl_r) + batch_mapreduce!((xu, x, z) -> (xu - x) * z, +, zero(T), ws.sum_ub, xu_r, x_ur, zu_r) @. ws.mu_curr = (ws.sum_lb + ws.sum_ub) / (nlb + nub) return ws.mu_curr end @@ -177,14 +169,8 @@ function get_affine_complementarity_measure!(solver::AbstractBatchMPCSolver, alp dzlb = MadNLP.dual_lb(solver.d) dzub = MadNLP.dual_ub(solver.d) - _scratch_lb = MadNLP.dual_lb(solver._w2) - @. _scratch_lb = (x_lr + alpha_p * dx_lr - xl_r) * (zl_r + alpha_d * dzlb) - sum!(ws.sum_lb, _scratch_lb) - - _scratch_ub = MadNLP.dual_ub(solver._w2) - @. _scratch_ub = (xu_r - (x_ur + alpha_p * dx_ur)) * (zu_r + alpha_d * dzub) - sum!(ws.sum_ub, _scratch_ub) - + _affine_compl_lb!(ws.sum_lb, x_lr, xl_r, zl_r, dx_lr, dzlb, alpha_p, alpha_d) + _affine_compl_ub!(ws.sum_ub, xu_r, x_ur, zu_r, dx_ur, dzub, alpha_p, alpha_d) @. ws.mu_affine = (ws.sum_lb + ws.sum_ub) / (nlb + nub) return ws.mu_affine end @@ -213,32 +199,17 @@ function get_fraction_to_boundary_step!(batch_solver::AbstractBatchMPCSolver) nlb, nub = d.nlb, d.nub T = eltype(ws.alpha_p) - # can't use mapreduce since tau is (1, bs), not (nlb, bs) if nlb > 0 - _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) # (nlb, bs) - _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) # (nlb, bs) - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) # (nlb, bs) - - @. _scratch_lb = ifelse(_dx_lr < 0, (-_x_lr + _xl_r) * ws.tau / _dx_lr, T(Inf)) - minimum!(ws.alpha_xl, _scratch_lb) - - @. _scratch_lb = ifelse(_dzlb < 0, (-_zl_r) * ws.tau / _dzlb, T(Inf)) - minimum!(ws.alpha_zl, _scratch_lb) + _ftb_primal_lb!(ws.alpha_xl, xp_lr(d), lower(x), lower(xl), ws.tau) + _ftb_dual_lb!(ws.alpha_zl, MadNLP.dual_lb(d), lower(zl), ws.tau) else fill!(ws.alpha_xl, one(T)) fill!(ws.alpha_zl, one(T)) end if nub > 0 - _dx_ur = xp_ur(d); _xu_r = upper(xu); _x_ur = upper(x) - _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - - @. _scratch_ub = ifelse(_dx_ur > 0, (-_x_ur + _xu_r) * ws.tau / _dx_ur, T(Inf)) - minimum!(ws.alpha_xu, _scratch_ub) - - @. _scratch_ub = ifelse((_dzub < 0) & (_zu_r + _dzub < 0), (-_zu_r) * ws.tau / _dzub, T(Inf)) - minimum!(ws.alpha_zu, _scratch_ub) + _ftb_primal_ub!(ws.alpha_xu, xp_ur(d), upper(x), upper(xu), ws.tau) + _ftb_dual_ub!(ws.alpha_zu, MadNLP.dual_ub(d), upper(zu), ws.tau) else fill!(ws.alpha_xu, one(T)) fill!(ws.alpha_zu, one(T)) @@ -249,6 +220,92 @@ function get_fraction_to_boundary_step!(batch_solver::AbstractBatchMPCSolver) return end +function _ftb_primal_lb!(alpha_out, dx, x, xb, tau) + T = eltype(alpha_out) + n, bs = size(dx) + @inbounds for j in 1:bs + a = T(Inf) + τ = tau[1, j] + for i in 1:n + d = dx[i, j] + d < zero(T) || continue + a = min(a, (-x[i, j] + xb[i, j]) * τ / d) + end + alpha_out[1, j] = a + end +end + +function _ftb_primal_ub!(alpha_out, dx, x, xb, tau) + T = eltype(alpha_out) + n, bs = size(dx) + @inbounds for j in 1:bs + a = T(Inf) + τ = tau[1, j] + for i in 1:n + d = dx[i, j] + d > zero(T) || continue + a = min(a, (-x[i, j] + xb[i, j]) * τ / d) + end + alpha_out[1, j] = a + end +end + +function _ftb_dual_lb!(alpha_out, dz, z, tau) + T = eltype(alpha_out) + n, bs = size(dz) + @inbounds for j in 1:bs + a = T(Inf) + τ = tau[1, j] + for i in 1:n + d = dz[i, j] + d < zero(T) || continue + a = min(a, -z[i, j] * τ / d) + end + alpha_out[1, j] = a + end +end + +function _ftb_dual_ub!(alpha_out, dz, z, tau) + T = eltype(alpha_out) + n, bs = size(dz) + @inbounds for j in 1:bs + a = T(Inf) + τ = tau[1, j] + for i in 1:n + d = dz[i, j] + (d < zero(T) && z[i, j] + d < zero(T)) || continue + a = min(a, -z[i, j] * τ / d) + end + alpha_out[1, j] = a + end +end + +function _affine_compl_lb!(out, x, xl, z, dx, dz, αp, αd) + T = eltype(out) + n, bs = size(x) + @inbounds for j in 1:bs + s = zero(T) + ap = αp[1, j]; ad = αd[1, j] + for i in 1:n + s += (x[i,j] + ap * dx[i,j] - xl[i,j]) * (z[i,j] + ad * dz[i,j]) + end + out[1, j] = s + end +end + +function _affine_compl_ub!(out, xu, x, z, dx, dz, αp, αd) + T = eltype(out) + n, bs = size(x) + @inbounds for j in 1:bs + s = zero(T) + ap = αp[1, j]; ad = αd[1, j] + for i in 1:n + s += (xu[i,j] - (x[i,j] + ap * dx[i,j])) * (z[i,j] + ad * dz[i,j]) + end + out[1, j] = s + end +end + function set_tau!(rule::ConservativeStep, batch_solver::AbstractBatchMPCSolver) fill!(batch_solver.workspace.tau, rule.tau) end @@ -262,53 +319,94 @@ function update_step!(rule::Union{ConservativeStep, AdaptiveStep}, batch_solver: return end -function _mehrotra_correct_steps!( - alpha_p, alpha_d, mu, - val_xl, idx_xl, val_xu, idx_xu, - val_zl, idx_zl, val_zu, idx_zu, - d_vals, x_vals, xl_vals, xu_vals, zl_vals, zu_vals, - ind_lb, ind_ub, dlb_off::Int, dub_off::Int, gamma_f, +function _mehrotra_step!( + alpha_p, alpha_d, mu, gamma_f, + dx_lr, x_lr, xl_r, nlb, dzlb, zl_r, + dx_ur, x_ur, xu_r, nub, dzub, zu_r, + d_vals, ind_lb, ind_ub, dlb_off, dub_off, ) - T = eltype(alpha_p) - @inbounds for j in axes(alpha_p, 2) - mu_j = mu[1, j] - max_ap = alpha_p[1, j] - max_ad = alpha_d[1, j] - - # primal step - corrected_p = one(T) - if max_ap < one(T) - i_xl = idx_xl[1, j] - i_xu = idx_xu[1, j] - if val_xl[1, j] <= val_xu[1, j] && i_xl > 0 - idx = ind_lb[i_xl] - zl_stepped = zl_vals[idx, j] + max_ad * d_vals[dlb_off + i_xl, j] - corrected_p = (x_vals[idx, j] - xl_vals[idx, j] - mu_j / zl_stepped) / (-d_vals[idx, j]) - elseif i_xu > 0 - idx = ind_ub[i_xu] - zu_stepped = zu_vals[idx, j] + max_ad * d_vals[dub_off + i_xu, j] - corrected_p = (xu_vals[idx, j] - x_vals[idx, j] - mu_j / zu_stepped) / d_vals[idx, j] - end + for j in axes(alpha_p, 2) + _mehrotra_step_column!( + j, alpha_p, alpha_d, mu[1, j], gamma_f, + dx_lr, x_lr, xl_r, nlb, dzlb, zl_r, + dx_ur, x_ur, xu_r, nub, dzub, zu_r, + d_vals, ind_lb, ind_ub, dlb_off, dub_off, + ) + end +end + +@inline function _mehrotra_step_column!( + j, alpha_p, alpha_d, mu_j::T, gamma_f::T, + dx_lr, x_lr, xl_r, nlb, dzlb, zl_r, + dx_ur, x_ur, xu_r, nub, dzub, zu_r, + d_vals, ind_lb, ind_ub, dlb_off, dub_off, +) where T + max_ap = alpha_p[1, j] + max_ad = alpha_d[1, j] + + # primal lb + best_xl = T(Inf); i_xl = 0 + @inbounds for i in 1:nlb + d = dx_lr[i, j] + d < zero(T) || continue + v = (xl_r[i, j] - x_lr[i, j]) / d + v < best_xl && (best_xl = v; i_xl = i) + end + # primal ub + best_xu = T(Inf); i_xu = 0 + @inbounds for i in 1:nub + d = dx_ur[i, j] + d > zero(T) || continue + v = (xu_r[i, j] - x_ur[i, j]) / d + v < best_xu && (best_xu = v; i_xu = i) + end + # dual lb + best_zl = T(Inf); i_zl = 0 + @inbounds for i in 1:nlb + d = dzlb[i, j] + d < zero(T) || continue + v = -zl_r[i, j] / d + v < best_zl && (best_zl = v; i_zl = i) + end + # dual ub + best_zu = T(Inf); i_zu = 0 + @inbounds for i in 1:nub + d = dzub[i, j] + (d < zero(T) && zu_r[i, j] + d < zero(T)) || continue + v = -zu_r[i, j] / d + v < best_zu && (best_zu = v; i_zu = i) + end + + # primal step + corrected_p = one(T) + @inbounds if max_ap < one(T) + if best_xl <= best_xu && i_xl > 0 + idx = ind_lb[i_xl] + zl_stepped = zl_r[i_xl, j] + max_ad * d_vals[dlb_off + i_xl, j] + corrected_p = (x_lr[i_xl, j] - xl_r[i_xl, j] - mu_j / zl_stepped) / (-dx_lr[i_xl, j]) + elseif i_xu > 0 + idx = ind_ub[i_xu] + zu_stepped = zu_r[i_xu, j] + max_ad * d_vals[dub_off + i_xu, j] + corrected_p = (xu_r[i_xu, j] - x_ur[i_xu, j] - mu_j / zu_stepped) / dx_ur[i_xu, j] end - alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) - - # dual step - corrected_d = one(T) - if max_ad < one(T) - i_zl = idx_zl[1, j] - i_zu = idx_zu[1, j] - if val_zl[1, j] <= val_zu[1, j] && i_zl > 0 - idx = ind_lb[i_zl] - x_gap = x_vals[idx, j] + max_ap * d_vals[idx, j] - xl_vals[idx, j] - corrected_d = -(zl_vals[idx, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] - elseif i_zu > 0 - idx = ind_ub[i_zu] - x_gap = xu_vals[idx, j] - x_vals[idx, j] - max_ap * d_vals[idx, j] - corrected_d = -(zu_vals[idx, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] - end + end + alpha_p[1, j] = max(corrected_p, gamma_f * max_ap) + + # dual step + corrected_d = one(T) + @inbounds if max_ad < one(T) + if best_zl <= best_zu && i_zl > 0 + idx = ind_lb[i_zl] + x_gap = x_lr[i_zl, j] + max_ap * dx_lr[i_zl, j] - xl_r[i_zl, j] + corrected_d = -(zl_r[i_zl, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] + elseif i_zu > 0 + idx = ind_ub[i_zu] + x_gap = xu_r[i_zu, j] - x_ur[i_zu, j] - max_ap * dx_ur[i_zu, j] + corrected_d = -(zu_r[i_zu, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] end - alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) end + alpha_d[1, j] = max(corrected_d, gamma_f * max_ad) + return end function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPCSolver) @@ -329,79 +427,58 @@ function update_step!(rule::MehrotraAdaptiveStep, batch_solver::AbstractBatchMPC dlb_off = d.n + d.m dub_off = d.n + d.m + d.nlb + bs = batch_solver.batch_size - if nlb > 0 - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - _dx_lr = xp_lr(d); _xl_r = lower(xl); _x_lr = lower(x) - _dzlb = MadNLP.dual_lb(d); _zl_r = lower(zl) - - map!((dx, xl, x) -> dx < 0 ? (xl - x) / dx : T(Inf), _scratch_lb, _dx_lr, _xl_r, _x_lr) - _vals, _inds = findmin(_scratch_lb; dims=1) - copyto!(ws.alpha_xl, _vals) - ws.idx_xl .= getindex.(_inds, 1) - - map!((dz, z) -> dz < 0 ? -z / dz : T(Inf), _scratch_lb, _dzlb, _zl_r) - _vals, _inds = findmin(_scratch_lb; dims=1) - copyto!(ws.alpha_zl, _vals) - ws.idx_zl .= getindex.(_inds, 1) - else - fill!(ws.alpha_xl, one(T)); fill!(ws.idx_xl, 0) - fill!(ws.alpha_zl, one(T)); fill!(ws.idx_zl, 0) - end - - if nub > 0 - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - _dx_ur = xp_ur(d); _xu_r = upper(xu); _x_ur = upper(x) - _dzub = MadNLP.dual_ub(d); _zu_r = upper(zu) - - map!((dx, xu, x) -> dx > 0 ? (xu - x) / dx : T(Inf), _scratch_ub, _dx_ur, _xu_r, _x_ur) - _vals, _inds = findmin(_scratch_ub; dims=1) - copyto!(ws.alpha_xu, _vals) - ws.idx_xu .= getindex.(_inds, 1) - - map!((dz, z) -> (dz < 0) & (z + dz < 0) ? -z / dz : T(Inf), _scratch_ub, _dzub, _zu_r) - _vals, _inds = findmin(_scratch_ub; dims=1) - copyto!(ws.alpha_zu, _vals) - ws.idx_zu .= getindex.(_inds, 1) - else - fill!(ws.alpha_xu, one(T)); fill!(ws.idx_xu, 0) - fill!(ws.alpha_zu, one(T)); fill!(ws.idx_zu, 0) - end - - _mehrotra_correct_steps!( - ws.alpha_p, ws.alpha_d, mu_full, - ws.alpha_xl, ws.idx_xl, ws.alpha_xu, ws.idx_xu, - ws.alpha_zl, ws.idx_zl, ws.alpha_zu, ws.idx_zu, - d.values, x.values, xl.values, xu.values, zl.values, zu.values, - d.ind_lb, d.ind_ub, dlb_off, dub_off, gamma_f, + _mehrotra_step!( + ws.alpha_p, ws.alpha_d, mu_full, gamma_f, + xp_lr(d), lower(x), lower(xl), nlb, MadNLP.dual_lb(d), lower(zl), + xp_ur(d), upper(x), upper(xu), nub, MadNLP.dual_ub(d), upper(zu), + d.values, d.ind_lb, d.ind_ub, dlb_off, dub_off, ) - return end +# FIXME: make it a kernel +function _adjust_boundary_active!(x_lr::AbstractMatrix{T}, xl_r, x_ur, xu_r, mu, mask) where {T} + c2 = eps(T)^(T(3)/T(4)) + c1 = eps(T) + xl_r .= ifelse.( + (mask .!= 0) .& (x_lr .- xl_r .< (c1 .* mu)), + xl_r .- c2 .* max.(one(T), abs.(x_lr)), + xl_r, + ) + xu_r .= ifelse.( + (mask .!= 0) .& (xu_r .- x_ur .< (c1 .* mu)), + xu_r .+ c2 .* max.(one(T), abs.(x_ur)), + xu_r, + ) +end + function init_regularization!(solver::AbstractBatchMPCSolver, ::NoRegularization) fill!(solver.del_w, 1.0) fill!(solver.del_c, 0.0) end -function update_regularization!(solver::AbstractBatchMPCSolver, ::NoRegularization) - fill!(solver.del_w, 0.0) - fill!(solver.del_c, 0.0) +update_regularization!(solver::AbstractBatchMPCSolver, reg) = + update_regularization!(solver, reg, solver.workspace.active_mask) +function update_regularization!(solver::AbstractBatchMPCSolver, ::NoRegularization, mask) + solver.del_w .= ifelse.(mask .== 1, 0.0, solver.del_w) + solver.del_c .= ifelse.(mask .== 1, 0.0, solver.del_c) end function init_regularization!(solver::AbstractBatchMPCSolver, reg::FixedRegularization) fill!(solver.del_w, 1.0) fill!(solver.del_c, reg.delta_d) end -function update_regularization!(solver::AbstractBatchMPCSolver, reg::FixedRegularization) - fill!(solver.del_w, reg.delta_p) - fill!(solver.del_c, reg.delta_d) +function update_regularization!(solver::AbstractBatchMPCSolver, reg::FixedRegularization, mask) + solver.del_w .= ifelse.(mask .== 1, reg.delta_p, solver.del_w) + solver.del_c .= ifelse.(mask .== 1, reg.delta_d, solver.del_c) end function init_regularization!(solver::AbstractBatchMPCSolver, reg::AdaptiveRegularization) fill!(solver.del_w, 1.0) fill!(solver.del_c, reg.delta_d) end -function update_regularization!(solver::AbstractBatchMPCSolver, reg::AdaptiveRegularization) +function update_regularization!(solver::AbstractBatchMPCSolver, reg::AdaptiveRegularization, mask) reg.delta_p = max(reg.delta_p / 10.0, reg.delta_min) reg.delta_d = min(reg.delta_d / 10.0, -reg.delta_min) - fill!(solver.del_w, reg.delta_p) - fill!(solver.del_c, reg.delta_d) + solver.del_w .= ifelse.(mask .== 1, reg.delta_p, solver.del_w) + solver.del_c .= ifelse.(mask .== 1, reg.delta_d, solver.del_c) end diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 53003cb2..25fbb6ab 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -69,22 +69,19 @@ function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace nlb_init, nub_init = batch_solver.d.nlb, batch_solver.d.nub + bs = batch_solver.batch_size _s1 = ws.alpha_xl # (1,bs) scratch _s2 = ws.alpha_xu # (1,bs) scratch # delta_x = max(0, -1.5 * min(xl-lb, 0), -1.5 * min(ub-xu, 0)) if nlb_init > 0 - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = xl - lb - minimum!(_s1, _scratch_lb) - @. _s1 = min(_s1, zero(T)) # clamp: init=0.0 behavior + batch_mapreduce!(-, min, T(Inf), _s1, xl, lb) + @. _s1 = min(_s1, zero(T)) else fill!(_s1, zero(T)) end if nub_init > 0 - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = ub - xu - minimum!(_s2, _scratch_ub) + batch_mapreduce!(-, min, T(Inf), _s2, ub, xu) @. _s2 = min(_s2, zero(T)) else fill!(_s2, zero(T)) @@ -94,13 +91,13 @@ function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T # delta_s = max(0, -1.5 * min(zl, 0), -1.5 * min(zu, 0)) if nlb_init > 0 - minimum!(_s1, zl) + batch_mapreduce!(identity, min, T(Inf), _s1, zl) @. _s1 = min(_s1, zero(T)) else fill!(_s1, zero(T)) end if nub_init > 0 - minimum!(_s2, zu) + batch_mapreduce!(identity, min, T(Inf), _s2, zu) @. _s2 = min(_s2, zero(T)) else fill!(_s2, zero(T)) @@ -117,32 +114,26 @@ function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T μ = ws.mu_affine # (1,bs) scratch fill!(μ, zero(T)) if nlb_init > 0 - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = xl * zl - sum!(ws.sum_lb, _scratch_lb) + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_lb, xl, zl) μ .+= ws.sum_lb - @. _scratch_lb = lb * zl - sum!(ws.sum_lb, _scratch_lb) + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_lb, lb, zl) μ .-= ws.sum_lb end if nub_init > 0 - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = ub * zu - sum!(ws.sum_ub, _scratch_ub) - @. _scratch_ub = xu * zu - sum!(ws.sum_lb, _scratch_ub) + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_ub, ub, zu) + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_lb, xu, zu) ws.sum_ub .-= ws.sum_lb μ .+= ws.sum_ub end # delta_x2 = μ / (2 * (sum(zl) + sum(zu))) if nlb_init > 0 - sum!(ws.sum_lb, zl) + batch_mapreduce!(identity, +, zero(T), ws.sum_lb, zl) else fill!(ws.sum_lb, zero(T)) end if nub_init > 0 - sum!(ws.sum_ub, zu) + batch_mapreduce!(identity, +, zero(T), ws.sum_ub, zu) else fill!(ws.sum_ub, zero(T)) end @@ -151,16 +142,12 @@ function init_starting_point!(batch_solver::AbstractBatchMPCSolver{T}) where T # delta_s2 = μ / (2 * (sum(xl-lb) + sum(ub-xu))) if nlb_init > 0 - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - @. _scratch_lb = xl - lb - sum!(ws.sum_lb, _scratch_lb) + batch_mapreduce!(-, +, zero(T), ws.sum_lb, xl, lb) else fill!(ws.sum_lb, zero(T)) end if nub_init > 0 - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - @. _scratch_ub = ub - xu - sum!(ws.sum_ub, _scratch_ub) + batch_mapreduce!(-, +, zero(T), ws.sum_ub, ub, xu) else fill!(ws.sum_ub, zero(T)) end @@ -236,8 +223,8 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) + batch_mapreduce!(abs, max, typemin(T), ws.norm_b, MadNLP.full(batch_solver.rhs)) + batch_mapreduce!(abs, max, typemin(T), ws.norm_c, MadNLP.full(batch_solver.f)) init_starting_point!(batch_solver) @@ -264,6 +251,36 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T return end +function compute_term_gpu!(ws::UniformBatchWorkspace{T}, opt) where T + ds = T(opt.divergence_scale) + tol = T(opt.tol) + div_tol = T(opt.divergence_tol) + Int_ERROR = Int(MadNLP.INTERNAL_ERROR) + Int_SOLVED = Int(MadNLP.SOLVE_SUCCEEDED) + Int_INFEASIBLE = Int(MadNLP.INFEASIBLE_PROBLEM_DETECTED) + Int_DIVERGING = Int(MadNLP.DIVERGING_ITERATES) + Int_REGULAR = Int(MadNLP.REGULAR) + @. ws._term_gpu = ifelse( + ws._ls_error > zero(Int32), + Int_ERROR, + ifelse( + max(ws.inf_pr, ws.inf_du, ws.inf_compl) <= tol, + Int_SOLVED, + ifelse( + (ws.inf_compl > div_tol * ws.best_complementarity) & + (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))), + Int_INFEASIBLE, + ifelse( + ws.obj_val < -(div_tol * max(ds, abs(ws.dual_obj), one(T))), + Int_DIVERGING, + Int_REGULAR, + ), + ), + ), + ) + minimum!(ws._any_nonregular_gpu, ws._term_gpu) +end + function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace opt = batch_solver.opt @@ -273,56 +290,39 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w bs = batch_solver.batch_size nlb, nub = batch_solver.d.nlb, batch_solver.d.nub - _scratch_n = MadNLP.primal(batch_solver._w2) - _scratch_m = MadNLP.dual(batch_solver._w2) - _scratch_lb = MadNLP.dual_lb(batch_solver._w2) - _scratch_ub = MadNLP.dual_ub(batch_solver._w2) - - get_inf_pr!(ws.inf_pr, MadNLP.full(batch_solver.c), _scratch_m) + get_inf_pr!(ws.inf_pr, MadNLP.full(batch_solver.c)) @. ws.inf_pr /= max(one(T), ws.norm_b) get_inf_du!(ws.inf_du, MadNLP.full(batch_solver.f), MadNLP.full(zl), - MadNLP.full(zu), MadNLP.full(batch_solver.jacl), _scratch_n) + MadNLP.full(zu), MadNLP.full(batch_solver.jacl)) @. ws.inf_du /= max(one(T), ws.norm_c) get_inf_compl!(ws.inf_compl, x, xl, zl, xu, zu, - _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) + ws.sum_lb, ws.sum_ub, nlb, nub) @. ws.inf_compl /= max(one(T), ws.norm_c) @. ws.best_complementarity = min(ws.best_complementarity, ws.inf_compl) dual_objective!(ws.dual_obj, MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), lower(zl), lower(xl), upper(zu), upper(xu), - _scratch_m, _scratch_lb, _scratch_ub, ws.sum_lb, ws.sum_ub, nlb, nub) + ws.sum_lb, ws.sum_ub, nlb, nub) - ds = T(opt.divergence_scale) - tol = T(opt.tol) - div_tol = T(opt.divergence_tol) - Int_SOLVED = Int(MadNLP.SOLVE_SUCCEEDED) - Int_INFEASIBLE = Int(MadNLP.INFEASIBLE_PROBLEM_DETECTED) - Int_DIVERGING = Int(MadNLP.DIVERGING_ITERATES) - Int_REGULAR = Int(MadNLP.REGULAR) - @. ws._term_gpu = ifelse( - max(ws.inf_pr, ws.inf_du, ws.inf_compl) <= tol, - Int_SOLVED, - ifelse( - (ws.inf_compl > div_tol * ws.best_complementarity) & - (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))), - Int_INFEASIBLE, - ifelse( - ws.obj_val < -(div_tol * max(ds, abs(ws.dual_obj), one(T))), - Int_DIVERGING, - Int_REGULAR, - ), - ), - ) - copyto!(ws._term_cpu, vec(ws._term_gpu)) + compute_term_gpu!(ws, opt) + return +end + +function update_termination_status!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + opt = batch_solver.opt + bcnt = batch_solver.batch_cnt + bs = batch_solver.batch_size + Int_REGULAR = Int64(Int(MadNLP.REGULAR)) walltime_hit = time() - bcnt.start_time[] >= opt.max_wall_time max_iter_hit = walltime_hit ? false : any(ws.status[i] == MadNLP.REGULAR && bcnt.k[i] >= opt.max_iter for i in 1:bs) if !walltime_hit && !max_iter_hit - copyto!(ws._any_nonregular_cpu, ws._any_nonregular_gpu) # TODO: use CuRef + copyto!(ws._any_nonregular_cpu, ws._any_nonregular_gpu) ws._any_nonregular_cpu[1] == Int_REGULAR && return false end @@ -338,7 +338,7 @@ function update_termination_criteria!(batch_solver::AbstractBatchMPCSolver{T}) w ws.status[i] = MadNLP.MAXIMUM_WALLTIME_EXCEEDED end end - return + return true end function solve_system!( @@ -346,7 +346,6 @@ function solve_system!( batch_solver::AbstractBatchMPCSolver{T}, p::BatchUnreducedKKTVector{T}, ) where T - opt = batch_solver.opt copyto!(MadNLP.full(d), MadNLP.full(p)) MadNLP.solve_kkt!(batch_solver.kkt, batch_solver) @@ -354,33 +353,20 @@ function solve_system!( copyto!(MadNLP.full(w), MadNLP.full(p)) mul!(w, batch_solver.kkt, d, -one(T), one(T)) - bkkt = batch_solver.kkt - bs = bkkt.batch_size - @inbounds for i in 1:bs - if bkkt.batch_map[i] == 0 - view(MadNLP.full(w), :, i) .= zero(T) - view(MadNLP.full(p), :, i) .= zero(T) - end - end - ws = batch_solver.workspace + MadNLP.full(w) .*= ws.active_mask + MadNLP.full(p) .*= ws.active_mask + + opt = batch_solver.opt + check_res = opt.check_residual + tol_ls = T(opt.tol_linear_solve) _fw = MadNLP.full(w) _fw .= abs.(_fw) - maximum!(ws._norm_gpu, _fw) # (1,bs) per-instance norm_w - copyto!(ws._norm_cpu, vec(ws._norm_gpu)) + batch_maximum!(ws._norm_gpu_w, _fw) # (1,bs) per-instance norm_w _fw .= abs.(MadNLP.full(p)) - maximum!(ws._norm_gpu, _fw) # (1,bs) per-instance norm_p - copyto!(ws._norm_cpu2, vec(ws._norm_gpu)) - - @inbounds for i in 1:bs - bkkt.batch_map[i] == 0 && continue - nw = ws._norm_cpu[i] - np = ws._norm_cpu2[i] - ratio = nw / max(one(T), np) - if isnan(ratio) || (opt.check_residual && (ratio > opt.tol_linear_solve)) - ws.status[i] = MadNLP.INTERNAL_ERROR - end - end + batch_maximum!(ws._norm_gpu_p, _fw) # (1,bs) per-instance norm_p + @. ws._norm_gpu_w /= max(one(T), ws._norm_gpu_p) # ratio in-place + @. ws._ls_error |= isnan(ws._norm_gpu_w) | (check_res & (ws._norm_gpu_w > tol_ls)) return d end @@ -440,14 +426,16 @@ function mehrotra_correction_direction!(solver::AbstractBatchMPCSolver) end function factorize_system!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace update_regularization!(batch_solver, batch_solver.opt.regularization) max_trials = 3 for _ in 1:max_trials set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) MadNLP.factorize_wrapper!(batch_solver) is_factorized(batch_solver.kkt.batch_solver) && break # exit once all are factorized - batch_solver.del_w .*= 100.0 - batch_solver.del_c .*= 100.0 + # FIXME: mask based on is_factorized instead of termination only + batch_solver.del_w .*= 100.0 .* ws.active_mask + batch_solver.del_c .*= 100.0 .* ws.active_mask end return end @@ -474,8 +462,8 @@ function apply_step!(batch_solver::AbstractBatchMPCSolver) upper(zu) .+= ws.alpha_d .* MadNLP.dual_ub(d) end - MadNLP.adjust_boundary!(lower(x), lower(xl), upper(x), upper(xu), ws.mu_batch) - increment_k!(batch_solver) + _adjust_boundary_active!(lower(x), lower(xl), upper(x), upper(xu), ws.mu_batch, ws.active_mask) + increment_k!(batch_solver) # this is CPU work, ends up overlapped return end @@ -491,6 +479,7 @@ function evaluate_model!(batch_solver::AbstractBatchMPCSolver) end function mpc_step!(batch_solver::AbstractBatchMPCSolver) + fill!(batch_solver.workspace._ls_error, zero(Int32)) factorize_system!(batch_solver) prediction_step!(batch_solver) mehrotra_correction_direction!(batch_solver) @@ -514,9 +503,12 @@ function mpc!(batch_solver::AbstractBatchMPCSolver) while true MadNLP.print_iter(batch_solver) update_termination_criteria!(batch_solver) - update_active_set!(batch_solver.kkt, batch_solver.workspace.status) - batch_solver.kkt.active_batch_size[] == 0 && return - _update_active_mask!(batch_solver) + changed = update_termination_status!(batch_solver) + if changed + update_active_set!(batch_solver.kkt, batch_solver.workspace.status) + batch_solver.kkt.active_batch_size[] == 0 && return + _update_active_mask!(batch_solver) + end mpc_step!(batch_solver) end end diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl index a897073d..a04718dc 100644 --- a/src/batch/madnlp/kernels.jl +++ b/src/batch/madnlp/kernels.jl @@ -1,29 +1,25 @@ -function get_inf_pr!(inf_pr, c, scratch) - @. scratch = abs(c) - maximum!(inf_pr, scratch) +function get_inf_pr!(inf_pr, c) + batch_mapreduce!(abs, max, zero(eltype(inf_pr)), inf_pr, c) return inf_pr end -function get_inf_du!(inf_du, f_vals, zl_vals, zu_vals, jacl_vals, scratch) - @. scratch = abs(f_vals - zl_vals + zu_vals + jacl_vals) - maximum!(inf_du, scratch) +function get_inf_du!(inf_du, f_vals, zl_vals, zu_vals, jacl_vals) + batch_mapreduce!((f, zl, zu, jl) -> abs(f - zl + zu + jl), max, zero(eltype(inf_du)), + inf_du, f_vals, zl_vals, zu_vals, jacl_vals) return inf_du end -function get_inf_compl!(inf_compl, x, xl, zl, xu, zu, - scratch_lb, scratch_ub, sum_lb, sum_ub, nlb, nub) +function get_inf_compl!(inf_compl, x, xl, zl, xu, zu, sum_lb, sum_ub, nlb, nub) T = eltype(inf_compl) if nlb > 0 - x_lr = lower(x); xl_r = lower(xl); zl_r = lower(zl) - @. scratch_lb = abs(x_lr - xl_r) * zl_r - maximum!(sum_lb, scratch_lb) + batch_mapreduce!((x, xl, z) -> abs(x - xl) * z, max, zero(T), + sum_lb, lower(x), lower(xl), lower(zl)) else fill!(sum_lb, zero(T)) end if nub > 0 - xu_r = upper(xu); x_ur = upper(x); zu_r = upper(zu) - @. scratch_ub = abs(xu_r - x_ur) * zu_r - maximum!(sum_ub, scratch_ub) + batch_mapreduce!((xu, x, z) -> abs(xu - x) * z, max, zero(T), + sum_ub, upper(xu), upper(x), upper(zu)) else fill!(sum_ub, zero(T)) end diff --git a/src/batch/structure.jl b/src/batch/structure.jl index b4324587..b18bc5ab 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -5,10 +5,6 @@ struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI alpha_zu::MT alpha_p::MT alpha_d::MT - idx_xl::MI - idx_xu::MI - idx_zl::MI - idx_zu::MI tau::MT mu_batch::MT @@ -53,8 +49,6 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_xl, alpha_xu MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_zl, alpha_zu MT(undef, 1, batch_size), MT(undef, 1, batch_size), # alpha_p, alpha_d - similar(_proto, Int32), similar(_proto, Int32), # idx_xl, idx_xu - similar(_proto, Int32), similar(_proto, Int32), # idx_zl, idx_zu MT(undef, 1, batch_size), # tau MT(undef, 1, batch_size), MT(undef, 1, batch_size), # mu_batch, mu_curr MT(undef, 1, batch_size), MT(undef, 1, batch_size), # mu_affine, sum_lb @@ -90,7 +84,6 @@ mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPC d::BatchUnreducedKKTVector{T, MT} p::BatchUnreducedKKTVector{T, MT} _w1::BatchUnreducedKKTVector{T, MT} - _w2::BatchUnreducedKKTVector{T, MT} x::BatchPrimalVector{T, MT} xl::BatchPrimalVector{T, MT} @@ -187,7 +180,6 @@ function UniformBatchMPCSolver( batch_d = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) batch_p = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) batch_w1 = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) - batch_w2 = BatchUnreducedKKTVector(MT, VT, n, m, nlb, nub, batch_size, ind_lb, ind_ub) batch_correction_lb = BatchVector(MT, VT, nlb, batch_size) batch_correction_ub = BatchVector(MT, VT, nub, batch_size) @@ -204,7 +196,7 @@ function UniformBatchMPCSolver( return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb)}( batch_size, - batch_d, batch_p, batch_w1, batch_w2, + batch_d, batch_p, batch_w1, batch_x, batch_xl, batch_xu, batch_zl, batch_zu, batch_f, batch_y, batch_c, batch_jacl, batch_rhs, batch_correction_lb, batch_correction_ub, diff --git a/test/batch_kernels/test_termination.jl b/test/batch_kernels/test_termination.jl index b842d43a..3d9bf3ed 100644 --- a/test/batch_kernels/test_termination.jl +++ b/test/batch_kernels/test_termination.jl @@ -9,9 +9,12 @@ for _ in 1:50 bat.workspace.status[1] != MadNLP.REGULAR && break MadIPM.update_termination_criteria!(bat) - MadIPM.update_active_set!(bat.kkt, bat.workspace.status) - bat.kkt.active_batch_size[] == 0 && break - MadIPM._update_active_mask!(bat) + changed = MadIPM.update_termination_status!(bat) + if changed + MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + bat.kkt.active_batch_size[] == 0 && break + MadIPM._update_active_mask!(bat) + end MadIPM.factorize_system!(bat) MadIPM.prediction_step!(bat) MadIPM.mehrotra_correction_direction!(bat) @@ -28,6 +31,7 @@ bat = build_batch(qp) bat.opt.max_iter = 0 # no iterations allowed MadIPM.update_termination_criteria!(bat) + MadIPM.update_termination_status!(bat) @test bat.workspace.status[1] == MadNLP.MAXIMUM_ITERATIONS_EXCEEDED end @@ -37,6 +41,7 @@ bat.opt.max_wall_time = 0.0 # zero walltime bat.batch_cnt.start_time[] = time() - 1.0 # started 1s ago MadIPM.update_termination_criteria!(bat) + MadIPM.update_termination_status!(bat) @test bat.workspace.status[1] == MadNLP.MAXIMUM_WALLTIME_EXCEEDED end end From 8baa4a0b9648a0e28780f26edb05872effe3f3fb Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Thu, 26 Mar 2026 20:40:06 -0400 Subject: [PATCH 35/51] remove device sync --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 3 --- ext/MadIPMCUDAExt/cuda_wrapper.jl | 4 ---- 2 files changed, 7 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index 802b11a5..fd1eaf7a 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -42,7 +42,6 @@ function MadNLP._set_con_scale_sparse!( if nnzj > 0 backend = CUDABackend() _set_con_scale_kernel!(backend)(con_scale, jac_I, jac_buffer; ndrange=nnzj) - KernelAbstractions.synchronize(backend) end return con_scale end @@ -366,7 +365,6 @@ function MadIPM._reduce_rhs_batch!(values::CuMatrix, ind_lb, lb_off, l_diag, nlb > 0 && _reduce_rhs_lb_kernel!(backend)(values, ind_lb, lb_off, l_diag; ndrange=(nlb, bs)) nub = length(ind_ub) nub > 0 && _reduce_rhs_ub_kernel!(backend)(values, ind_ub, ub_off, u_diag; ndrange=(nub, bs)) - KernelAbstractions.synchronize(backend) return end @@ -387,6 +385,5 @@ function MadIPM._finish_aug_solve_batch!(values::CuMatrix, ind_lb, lb_off, l_low nlb > 0 && _finish_aug_solve_lb_kernel!(backend)(values, ind_lb, lb_off, l_lower, l_diag; ndrange=(nlb, bs)) nub = length(ind_ub) nub > 0 && _finish_aug_solve_ub_kernel!(backend)(values, ind_ub, ub_off, u_lower, u_diag; ndrange=(nub, bs)) - KernelAbstractions.synchronize(backend) return end diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index b0770e9c..6db6c07f 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -21,7 +21,6 @@ function MadNLP._transfer!(dest::CuVector{T}, src::CuVector{T}, map::CuVector{In if length(map) > 0 backend = CUDABackend() _transfer_to_map!(backend)(dest, map, src; ndrange=length(map)) - KernelAbstractions.synchronize(backend) end return end @@ -102,7 +101,6 @@ function MadIPM.assemble_normal_system!( backend = CUDABackend() kernel! = assemble_normal_system_kernel!(backend) kernel!(n_rows, n_cols, Jtp, Jtj, Jtx, Cp, Cj, Cx, Dx, Tv; ndrange = n_rows) - KernelAbstractions.synchronize(backend) end @kernel function count_normal_nnz!(Cp, @Const(Jtp), @Const(Jtj), @Const(n_rows), @Const(n_cols)) @@ -171,7 +169,6 @@ function MadIPM.build_normal_system( Cp = CUDA.ones(Ti, n_rows + 1) kernel1! = count_normal_nnz!(backend) kernel1!(Cp, Jtp, Jtj, n_rows, n_cols; ndrange = n_rows) - KernelAbstractions.synchronize(backend) Cp = cumsum(Cp) nnz_JtJ = CUDA.@allowscalar (Cp[end] - 1) @@ -179,7 +176,6 @@ function MadIPM.build_normal_system( kernel2! = fill_normal_indices!(backend) kernel2!(Cj, Cp, Jtp, Jtj, n_rows, n_cols; ndrange = n_rows) - KernelAbstractions.synchronize(backend) return (Cp, Cj) end From 96e753ddf34db93e25d79a44c37ea36656b2be7e Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Fri, 27 Mar 2026 14:56:23 -0400 Subject: [PATCH 36/51] view system --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 50 ++++++ ext/MadIPMCUDAExt/cuda_wrapper.jl | 7 +- src/MadIPM.jl | 1 + src/batch/KKT/Sparse/augmented.jl | 72 +++----- src/batch/madipm/solver.jl | 52 ++++-- src/batch/madnlp/linear_solver.jl | 29 +++- src/batch/structure.jl | 27 ++- src/batch/views.jl | 172 ++++++++++++++++++++ test/batch_kernels/test_batch_multi.jl | 9 +- test/batch_kernels/test_iteration_match.jl | 14 +- test/batch_kernels/test_termination.jl | 4 +- test/test_batch.jl | 181 +++++++++++++++++++++ test/test_batch_gpu.jl | 32 +++- 13 files changed, 558 insertions(+), 92 deletions(-) create mode 100644 src/batch/views.jl diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index fd1eaf7a..9f3ee0ff 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -16,6 +16,24 @@ end end +@kernel function _gather_batch_view_columns_kernel!(dst, @Const(src), @Const(local_to_root)) + i, j = @index(Global, NTuple) + @inbounds dst[i, j] = src[i, local_to_root[j]] +end + +@kernel function _scatter_batch_view_columns_kernel!(dst, @Const(src), @Const(local_to_root)) + i, j = @index(Global, NTuple) + @inbounds dst[i, local_to_root[j]] = src[i, j] +end + +@kernel function _compact_active_columns_inplace_kernel!(dst, @Const(local_to_root)) + i, j = @index(Global, NTuple) + src_j = local_to_root[j] + @inbounds if src_j != j + dst[i, j] = dst[i, src_j] + end +end + @inline function _atomic_colreduce!(::typeof(+), out, j, value) Atomix.@atomic out[1, j] += value return @@ -33,6 +51,38 @@ end return end +function MadIPM.gather_batch_view_columns!( + dst::CuMatrix{TD}, + src::CuMatrix{TS}, + batch_view::MadIPM.BatchView, +) where {TD, TS} + na = MadIPM.local_batch_size(batch_view) + backend = CUDABackend() + _gather_batch_view_columns_kernel!(backend)(dst, src, MadIPM.local_to_root_dev(batch_view); ndrange=(size(dst, 1), na)) + return dst +end + +function MadIPM.scatter_batch_view_columns!( + dst::CuMatrix{TD}, + src::CuMatrix{TS}, + batch_view::MadIPM.BatchView, +) where {TD, TS} + na = MadIPM.local_batch_size(batch_view) + backend = CUDABackend() + _scatter_batch_view_columns_kernel!(backend)(dst, src, MadIPM.local_to_root_dev(batch_view); ndrange=(size(src, 1), na)) + return dst +end + +function MadIPM.compact_active_columns_inplace!( + dst::CuMatrix{T}, + batch_view::MadIPM.BatchView, +) where T + na = MadIPM.local_batch_size(batch_view) + backend = CUDABackend() + _compact_active_columns_inplace_kernel!(backend)(dst, MadIPM.local_to_root_dev(batch_view); ndrange=(size(dst, 1), na)) + return dst +end + function MadNLP._set_con_scale_sparse!( con_scale::CuMatrix{T}, jac_I::CuVector{<:Integer}, diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 6db6c07f..d1f247a1 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -202,13 +202,16 @@ end MadIPM.is_factorized(::MadNLPGPU.CUDSSSolver) = true -function MadIPM._active_factorize!(s::MadNLPGPU.CUDSSSolver, na::Int) +function MadIPM.factorize_active!(s::MadNLPGPU.CUDSSSolver, active::MadIPM.BatchView) + na = MadIPM.local_batch_size(active) CUDSS.cudss_set(s.inner, "ubatch_size", na) MadNLP.factorize!(s) return end -function MadIPM._active_solve!(s::MadNLPGPU.CUDSSSolver{T}, rhs::CuMatrix{T}, na::Int, n::Int) where T +function MadIPM.solve_active!(s::MadNLPGPU.CUDSSSolver{T}, rhs::CuMatrix{T}, active::MadIPM.BatchView) where T + na = MadIPM.local_batch_size(active) + n = size(rhs, 1) rhs_active = unsafe_wrap(CuArray{T, 2}, pointer(rhs), (n, na)) CUDSS.cudss_update(s.b_gpu, rhs_active) CUDSS.cudss_update(s.x_gpu, rhs_active) diff --git a/src/MadIPM.jl b/src/MadIPM.jl index ab589ad6..21d5a163 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -24,6 +24,7 @@ include("linear_solver.jl") include("solver.jl") include("batch/utils.jl") +include("batch/views.jl") include("batch/madnlp/rhs.jl") include("batch/madnlp/callback.jl") include("batch/KKT/KKT.jl") diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index d3ac4e07..c2071bce 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -1,11 +1,12 @@ """Batched version of SparseKKTSystem""" -struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT} <: AbstractBatchKKTSystem{T} +struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT, BVS} <: AbstractBatchKKTSystem{T} nzVals::MT # (aug_mat_length × batch_size) COO nonzero values aug_I::VI32 # shared row indices aug_J::VI32 # shared column indices batch_solver::LS # batched linear solver rhs_buffer::MT # (n+m) × batch_size for batch solve batch_size::Int + batch_views::BVS aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values batch_csc_map::VI # flattened COO→CSC map for all instances n_tot::Int # n + n_slack (total primal variables) @@ -21,10 +22,6 @@ struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT} <: AbstractBatchKKT hess_op::OPT jt_op::OPT j_op::OPT - # Batch tracking - batch_map::Vector{Int} # original index → active position (0 if inactive) - batch_map_rev::Vector{Int} # active position → original index - active_batch_size::Base.RefValue{Int} end pr_diag(bkkt::SparseUniformBatchKKTSystem) = view(bkkt.nzVals, 1:bkkt.n_tot, :) @@ -38,6 +35,7 @@ function MadNLP.create_kkt_system( bcb::UniformBatchCallback{T, VT, MT, VI}, uniformbatch_linear_solver = LoopedBatchLinearSolver; opt_linear_solver = MadNLP.default_options(uniformbatch_linear_solver), + batch_views, ) where {T, VT, MT, VI} batch_size = bcb.batch_size @@ -113,69 +111,32 @@ function MadNLP.create_kkt_system( l_lower = similar(nzVals, nlb, batch_size) u_lower = similar(nzVals, nub, batch_size) - batch_map = collect(1:batch_size) - batch_map_rev = collect(1:batch_size) - active_batch_size = Ref(batch_size) - LS = typeof(batch_ls) VI32 = typeof(I) OPT = typeof(jt_op) - return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT}( - nzVals, I, J, batch_ls, rhs_buffer, batch_size, + return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT, typeof(batch_views)}( + nzVals, I, J, batch_ls, rhs_buffer, batch_size, batch_views, aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, reg, l_diag, u_diag, l_lower, u_lower, hess_op, jt_op, j_op, - batch_map, batch_map_rev, active_batch_size, ) end -function update_active_set!(bkkt::SparseUniformBatchKKTSystem, status::Vector{MadNLP.Status}) - active_pos = 0 - for i in 1:bkkt.batch_size - if status[i] == MadNLP.REGULAR - active_pos += 1 - bkkt.batch_map[i] = active_pos - bkkt.batch_map_rev[active_pos] = i - else - bkkt.batch_map[i] = 0 - end - end - for j in (active_pos + 1):bkkt.batch_size - bkkt.batch_map_rev[j] = 0 - end - bkkt.active_batch_size[] = active_pos -end - function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) - na = bkkt.active_batch_size[] - nzvals = bkkt.aug_com_nzvals - @inbounds for j in 1:na # FIXME: refactor to avoid `na` launches - i = bkkt.batch_map_rev[j] - i != j && (view(nzvals, :, j) .= view(nzvals, :, i)) + factor_view = active_view(bkkt.batch_views) + if !is_identity_view(factor_view) + compact_active_columns_inplace!(bkkt.aug_com_nzvals, factor_view) end - _active_factorize!(bkkt.batch_solver, na) + factorize_active!(bkkt.batch_solver, factor_view) return end function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractMatrix) where T - na = bkkt.active_batch_size[] - bs = bkkt.batch_size - n = size(rhs, 1) - - @inbounds for j in 1:na - i = bkkt.batch_map_rev[j] - i != j && (view(rhs, :, j) .= view(rhs, :, i)) - end - _active_solve!(bkkt.batch_solver, rhs, na, n) - - @inbounds for j in na:-1:1 - i = bkkt.batch_map_rev[j] - i != j && (view(rhs, :, i) .= view(rhs, :, j)) - end - - @inbounds for i in 1:bs - bkkt.batch_map[i] == 0 && (view(rhs, :, i) .= zero(T)) + active = active_view(bkkt.batch_views) + if !is_identity_view(active) + compact_active_columns_inplace!(rhs, active) end + solve_active!(bkkt.batch_solver, rhs, active) return rhs end @@ -222,9 +183,14 @@ function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::Abst rhs = bkkt.rhs_buffer pd_view = MadNLP.primal_dual(d) + active = active_view(bkkt.batch_views) copyto!(rhs, pd_view) MadNLP.solve_linear_system!(bkkt, rhs) - copyto!(pd_view, rhs) + if is_identity_view(active) + copyto!(pd_view, rhs) + else + scatter_batch_view_columns!(pd_view, rhs, active) + end MadNLP.finish_aug_solve!(bkkt, batch_solver) return diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 25fbb6ab..9c321367 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -231,6 +231,8 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T fill!(ws.mu_batch, opt.mu_init) fill!(ws.best_complementarity, typemax(T)) fill!(ws.status, MadNLP.REGULAR) + reset_active_view!(batch_solver.batch_views) + _update_active_mask!(batch_solver) fill!(ws.inf_pr, zero(T)) fill!(ws.inf_du, zero(T)) fill!(ws.inf_compl, zero(T)) @@ -425,17 +427,42 @@ function mehrotra_correction_direction!(solver::AbstractBatchMPCSolver) return end +function _restrict_to_failed_locals!(batch_solver::AbstractBatchMPCSolver, nfailed::Int) + select_local!(batch_solver.batch_views, batch_solver.batch_views.selected_local_buffer, nfailed) + return active_view(batch_solver.batch_views) +end + +function _bump_failed_regularization!(batch_solver::AbstractBatchMPCSolver{T}, failed_view::BatchView) where T + ws = batch_solver.workspace + fill_batch_view_mask!(ws.active_mask_cpu, failed_view) + @. batch_solver.del_w = ifelse(ws.active_mask_cpu == one(T), T(100) * batch_solver.del_w, batch_solver.del_w) + @. batch_solver.del_c = ifelse(ws.active_mask_cpu == one(T), T(100) * batch_solver.del_c, batch_solver.del_c) + return +end + function factorize_system!(batch_solver::AbstractBatchMPCSolver) ws = batch_solver.workspace + batch_views = batch_solver.batch_views update_regularization!(batch_solver, batch_solver.opt.regularization) max_trials = 3 - for _ in 1:max_trials - set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) - MadNLP.factorize_wrapper!(batch_solver) - is_factorized(batch_solver.kkt.batch_solver) && break # exit once all are factorized - # FIXME: mask based on is_factorized instead of termination only - batch_solver.del_w .*= 100.0 .* ws.active_mask - batch_solver.del_c .*= 100.0 .* ws.active_mask + saved_active = active_view(batch_views) + try + for _ in 1:max_trials + set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) + MadNLP.factorize_wrapper!(batch_solver) + factor_view = active_view(batch_views) + nfailed = failed_factorization_local_count!( + batch_views.selected_local_buffer, + batch_solver.kkt.batch_solver, + factor_view, + ) + nfailed == 0 && break + + failed_view = _restrict_to_failed_locals!(batch_solver, nfailed) + _bump_failed_regularization!(batch_solver, failed_view) + end + finally + restore_state!(batch_views, saved_active) end return end @@ -491,11 +518,8 @@ end function _update_active_mask!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace - bmap = batch_solver.kkt.batch_map buf = ws.active_mask_cpu - @inbounds for i in eachindex(bmap) - buf[i] = T(bmap[i] != 0) - end + fill_batch_view_mask!(buf, active_view(batch_solver.batch_views)) copyto!(ws.active_mask, buf) end @@ -505,8 +529,8 @@ function mpc!(batch_solver::AbstractBatchMPCSolver) update_termination_criteria!(batch_solver) changed = update_termination_status!(batch_solver) if changed - update_active_set!(batch_solver.kkt, batch_solver.workspace.status) - batch_solver.kkt.active_batch_size[] == 0 && return + update_active_set!(batch_solver) + active_batch_size(batch_solver) == 0 && return _update_active_mask!(batch_solver) end mpc_step!(batch_solver) @@ -574,7 +598,7 @@ function MadNLP.print_iter(batch_solver::AbstractBatchMPCSolver) MadNLP.get_level(logger) > MadNLP.INFO && return ws = batch_solver.workspace bcnt = batch_solver.batch_cnt - na = batch_solver.kkt.active_batch_size[] + na = active_batch_size(batch_solver) bs = batch_solver.batch_size k = maximum(bcnt.k) diff --git a/src/batch/madnlp/linear_solver.jl b/src/batch/madnlp/linear_solver.jl index 9ce71e01..f083c8ae 100644 --- a/src/batch/madnlp/linear_solver.jl +++ b/src/batch/madnlp/linear_solver.jl @@ -34,17 +34,36 @@ function is_factorized(batch_linear_solver::LoopedBatchLinearSolver) return all(is_factorized(s) for s in batch_linear_solver.solvers) end -function _active_factorize!(s::LoopedBatchLinearSolver, na::Int) - for j in 1:na +function failed_factorization_local_count!( + failed_local_buffer::Vector{Int32}, + batch_linear_solver::LoopedBatchLinearSolver, + factor_view::BatchView, +) + nfailed = 0 + @inbounds for j in 1:factor_view.n + if !is_factorized(batch_linear_solver.solvers[j]) + nfailed += 1 + failed_local_buffer[nfailed] = j + end + end + return nfailed +end + +function factorize_active!(s::LoopedBatchLinearSolver, factor_view::BatchView) + @inbounds for j in 1:factor_view.n MadNLP.factorize!(s.solvers[j]) end return end -function _active_solve!(s::LoopedBatchLinearSolver{T, VT}, rhs::AbstractMatrix{T}, na::Int, n::Int) where {T, VT} - for j in 1:na - rhs_j = _madnlp_unsafe_column_wrap(rhs, n, (j-1)*n + 1, VT) +function solve_active!(s::LoopedBatchLinearSolver{T, VT}, rhs::AbstractMatrix{T}, active::BatchView) where {T, VT} + na = local_batch_size(active) + na == 0 && return + n = size(rhs, 1) + @inbounds for j in 1:na + rhs_j = _madnlp_unsafe_column_wrap(rhs, n, (j - 1) * n + 1, VT) MadNLP.solve_linear_system!(s.solvers[j], rhs_j) end return end +failed_factorization_local_count!(::Vector{Int32}, ::MadNLP.AbstractLinearSolver, ::BatchView) = 0 diff --git a/src/batch/structure.jl b/src/batch/structure.jl index b18bc5ab..5af001bd 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -78,7 +78,7 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, ) end -mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPCSolver{T, MT, VT} +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, BVS} <: AbstractBatchMPCSolver{T, MT, VT} batch_size::Int d::BatchUnreducedKKTVector{T, MT} @@ -104,6 +104,7 @@ mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB} <: AbstractBatchMPC opt::IPMOptions batch_cnt::BatchCounters logger::MadNLP.MadNLPLogger + batch_views::BVS kkt::AbstractBatchKKTSystem{T} del_w::MT @@ -117,6 +118,24 @@ _get_ind_lb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_lb _get_ind_ub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_ub _get_ind_llb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_llb _get_ind_uub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_uub +active_batch_size(bs::AbstractBatchMPCSolver) = local_batch_size(active_view(bs.batch_views)) + +function update_active_set!(state::BatchViewState, status::Vector{MadNLP.Status}) + nselected = 0 + @inbounds for i in eachindex(status) + if status[i] == MadNLP.REGULAR + nselected += 1 + state.selected_local_buffer[nselected] = i + end + end + if nselected == batch_size_root(root_view(state)) + return reset_active_view!(state) + end + reset_active_view!(state) + return select_local!(state, state.selected_local_buffer, nselected; reset_slots=true) +end + +update_active_set!(bs::AbstractBatchMPCSolver) = update_active_set!(bs.batch_views, bs.workspace.status) """ @@ -163,11 +182,14 @@ function UniformBatchMPCSolver( nlb = length(ind_lb) nub = length(ind_ub) + batch_views = BatchViewState(bcb, batch_size) + batch_kkts = MadNLP.create_kkt_system( ipm_opt.kkt_system, bcb, uniformbatch_linear_solver; opt_linear_solver = opt_batch_ls, + batch_views = batch_views, ) batch_x = BatchPrimalVector(MT, VT, nx, ns, batch_size, ind_lb, ind_ub) @@ -194,7 +216,7 @@ function UniformBatchMPCSolver( batch_del_w = fill!(MT(undef, 1, batch_size), zero(T)) batch_del_c = fill!(MT(undef, 1, batch_size), zero(T)) - return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb)}( + return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb), typeof(batch_views)}( batch_size, batch_d, batch_p, batch_w1, batch_x, batch_xl, batch_xu, batch_zl, batch_zu, batch_f, @@ -202,6 +224,7 @@ function UniformBatchMPCSolver( batch_correction_lb, batch_correction_ub, workspace, ipm_opt, batch_cnt, logger, + batch_views, batch_kkts, batch_del_w, batch_del_c, bnlp, diff --git a/src/batch/views.jl b/src/batch/views.jl new file mode 100644 index 00000000..4f27a195 --- /dev/null +++ b/src/batch/views.jl @@ -0,0 +1,172 @@ +mutable struct BatchView{VI32} + batch_size_root::Int + layer::Int + n::Int + local_to_root::Vector{Int32} + local_to_slot::Vector{Int32} + local_to_root_dev::VI32 +end + +mutable struct BatchViewState{V<:BatchView} + views::Vector{V} + active_layer::Int + selected_local_buffer::Vector{Int32} +end + +batch_size_root(view::BatchView) = view.batch_size_root +view_layer(view::BatchView) = view.layer +local_batch_size(view::BatchView) = view.n +local_to_root_dev(view::BatchView) = view.local_to_root_dev +active_view(state::BatchViewState) = state.views[state.active_layer] +root_view(state::BatchViewState) = state.views[1] + +function is_identity_view(view::BatchView) + view.n == view.batch_size_root || return false + @inbounds for i in 1:view.n + view.local_to_root[i] == i || return false + end + return true +end + +function _sync_local_to_root_dev!(view::BatchView) + copyto!(view.local_to_root_dev, view.local_to_root) + return view +end + +function _init_batch_view!(view::BatchView, batch_size::Int) + view.n = batch_size + @inbounds for i in 1:batch_size + idx = Int32(i) + view.local_to_root[i] = idx + view.local_to_slot[i] = idx + end + return _sync_local_to_root_dev!(view) +end + +function _select_local!(child::BatchView, parent::BatchView, selected_local, nselected::Int, reset_slots::Bool) + child.n = nselected + roots = parent.local_to_root + slots = parent.local_to_slot + @inbounds for j in 1:nselected + parent_j = Int(selected_local[j]) + child.local_to_root[j] = roots[parent_j] + child.local_to_slot[j] = reset_slots ? Int32(j) : slots[parent_j] + end + return _sync_local_to_root_dev!(child) +end + +function reset_active_view!(state::BatchViewState) + state.active_layer = 1 + return state +end + +function restore_state!(state::BatchViewState, old_state::BatchView) + state.active_layer = view_layer(old_state) + return state +end + +function select_local!( + state::BatchViewState, + selected_local::AbstractVector{<:Integer}, + nselected::Int = length(selected_local); + reset_slots::Bool = false, +) + parent = active_view(state) + layer = view_layer(parent) + 1 + @assert layer <= length(state.views) "BatchViewState max_layers exceeded" + old_state = parent + state.active_layer = layer + _select_local!(active_view(state), old_state, selected_local, nselected, reset_slots) + return old_state +end + +function select_local!( + state::BatchViewState, + keep::AbstractVector{Bool}; + reset_slots::Bool = false, +) + current = active_view(state) + @assert length(keep) == current.n + nselected = 0 + @inbounds for i in 1:current.n + if keep[i] + nselected += 1 + state.selected_local_buffer[nselected] = i + end + end + return select_local!(state, state.selected_local_buffer, nselected; reset_slots=reset_slots) +end + +function exclude_local!(state::BatchViewState, exclude_mask::AbstractVector{Bool}) + current = active_view(state) + @assert length(exclude_mask) == current.n + nselected = 0 + @inbounds for i in 1:current.n + if !exclude_mask[i] + nselected += 1 + state.selected_local_buffer[nselected] = i + end + end + return select_local!(state, state.selected_local_buffer, nselected) +end + +function BatchViewState(bcb, batch_size::Int; max_layers::Int = 4) + sample_dev = MadNLP.create_array(bcb, Int32, batch_size) + views = Vector{BatchView{typeof(sample_dev)}}(undef, max_layers) + @inbounds for layer in 1:max_layers + views[layer] = BatchView( + batch_size, + layer, + 0, + Vector{Int32}(undef, batch_size), + Vector{Int32}(undef, batch_size), + layer == 1 ? sample_dev : MadNLP.create_array(bcb, Int32, batch_size), + ) + end + _init_batch_view!(views[1], batch_size) + return BatchViewState(views, 1, Vector{Int32}(undef, batch_size)) +end + +function fill_batch_view_mask!(mask::AbstractMatrix{T}, batch_view::BatchView) where T + @assert size(mask, 1) == 1 + @assert size(mask, 2) == batch_view.batch_size_root + fill!(mask, zero(T)) + @inbounds for j in 1:batch_view.n + mask[1, batch_view.local_to_root[j]] = one(T) + end + return mask +end + +function gather_batch_view_columns!( + dst::AbstractMatrix{TD}, + src::AbstractMatrix{TS}, + batch_view::BatchView, +) where {TD, TS} + roots = batch_view.local_to_root + @inbounds for j in 1:batch_view.n + copyto!(view(dst, :, j), view(src, :, roots[j])) + end + return dst +end + +function compact_active_columns_inplace!(dst::AbstractMatrix{T}, batch_view::BatchView) where T + roots = batch_view.local_to_root + @inbounds for j in 1:batch_view.n + src_j = roots[j] + src_j < j && error("active view must be root-ordered") + src_j != j && copyto!(view(dst, :, j), view(dst, :, src_j)) + end + return dst +end + +function scatter_batch_view_columns!( + dst::AbstractMatrix{TD}, + src::AbstractMatrix{TS}, + batch_view::BatchView, +) where {TD, TS} + roots = batch_view.local_to_root + @inbounds for j in 1:batch_view.n + copyto!(view(dst, :, roots[j]), view(src, :, j)) + end + return dst +end diff --git a/test/batch_kernels/test_batch_multi.jl b/test/batch_kernels/test_batch_multi.jl index 17149313..c69f3b27 100644 --- a/test/batch_kernels/test_batch_multi.jl +++ b/test/batch_kernels/test_batch_multi.jl @@ -34,12 +34,11 @@ # Mark instance 2 as converged bat.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED - MadIPM.update_active_set!(bat.kkt, bat.workspace.status) + MadIPM.update_active_set!(bat) + active = MadIPM.active_view(bat.batch_views) - @test bat.kkt.active_batch_size[] == 2 - @test bat.kkt.batch_map[2] == 0 # deactivated - @test bat.kkt.batch_map[1] != 0 # still active - @test bat.kkt.batch_map[3] != 0 # still active + @test MadIPM.local_batch_size(active) == 2 + @test Int[active.local_to_root[i] for i in 1:active.n] == [1, 3] # zero_inactive_step! should zero out the deactivated instance fill!(bat.workspace.active_mask, 1.0) diff --git a/test/batch_kernels/test_iteration_match.jl b/test/batch_kernels/test_iteration_match.jl index b898464e..a6160323 100644 --- a/test/batch_kernels/test_iteration_match.jl +++ b/test/batch_kernels/test_iteration_match.jl @@ -99,24 +99,22 @@ seq1_done = seq1_done || MadIPM.is_done(seq1) seq2_done = seq2_done || MadIPM.is_done(seq2) - MadIPM.update_active_set!(bat.kkt, ws.status) - bat.kkt.active_batch_size[] == 0 && break + MadIPM.update_active_set!(bat) + active = MadIPM.active_view(bat.batch_views) + MadIPM.local_batch_size(active) == 0 && break MadIPM._update_active_mask!(bat) # Verify active mask matches expected state if seq1_done && !seq2_done - @test bat.kkt.batch_map[1] == 0 # instance 1 inactive - @test bat.kkt.batch_map[2] != 0 # instance 2 still active + @test Int[active.local_to_root[i] for i in 1:active.n] == [2] @test ws.active_mask[1] == 0.0 @test ws.active_mask[2] == 1.0 elseif !seq1_done && seq2_done - @test bat.kkt.batch_map[1] != 0 - @test bat.kkt.batch_map[2] == 0 + @test Int[active.local_to_root[i] for i in 1:active.n] == [1] @test ws.active_mask[1] == 1.0 @test ws.active_mask[2] == 0.0 elseif !seq1_done && !seq2_done - @test bat.kkt.batch_map[1] != 0 - @test bat.kkt.batch_map[2] != 0 + @test Int[active.local_to_root[i] for i in 1:active.n] == [1, 2] end (seq1_done && seq2_done) && break diff --git a/test/batch_kernels/test_termination.jl b/test/batch_kernels/test_termination.jl index 3d9bf3ed..b37bb30c 100644 --- a/test/batch_kernels/test_termination.jl +++ b/test/batch_kernels/test_termination.jl @@ -11,8 +11,8 @@ MadIPM.update_termination_criteria!(bat) changed = MadIPM.update_termination_status!(bat) if changed - MadIPM.update_active_set!(bat.kkt, bat.workspace.status) - bat.kkt.active_batch_size[] == 0 && break + MadIPM.update_active_set!(bat) + MadIPM.local_batch_size(MadIPM.active_view(bat.batch_views)) == 0 && break MadIPM._update_active_mask!(bat) end MadIPM.factorize_system!(bat) diff --git a/test/test_batch.jl b/test/test_batch.jl index ae1ff110..a4800e54 100644 --- a/test/test_batch.jl +++ b/test/test_batch.jl @@ -1,5 +1,72 @@ using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel +_test_local_to_root(view) = Int[view.local_to_root[i] for i in 1:view.n] + +struct RecordingBatchLinearSolver{T, MT} <: MadNLP.AbstractLinearSolver{T} + nzvals_mat::MT + call_counts::Vector{Int} + factorized::Vector{Bool} + fail_marker::T +end + +@kwdef mutable struct RecordingBatchLinearSolverOptions <: MadNLP.AbstractOptions + fail_marker::Float64 = -999.0 +end + +MadNLP.default_options(::Type{RecordingBatchLinearSolver}) = RecordingBatchLinearSolverOptions() + +function RecordingBatchLinearSolver( + aug_com, + nzvals_mat::AbstractMatrix{T}, + n::Int; + opt::RecordingBatchLinearSolverOptions = RecordingBatchLinearSolverOptions(), +) where T + batch_size = size(nzvals_mat, 2) + return RecordingBatchLinearSolver( + nzvals_mat, + zeros(Int, batch_size), + fill(true, batch_size), + T(opt.fail_marker), + ) +end + +function MadIPM.factorize_active!( + s::RecordingBatchLinearSolver{T, MT}, + factor_view::MadIPM.BatchView, +) where {T, MT} + @inbounds for j in 1:MadIPM.local_batch_size(factor_view) + s.call_counts[j] += 1 + col = view(s.nzvals_mat, :, j) + corrupted = any(==(s.fail_marker), col) + s.factorized[j] = !(corrupted && s.call_counts[j] == 1) + end + return +end + +function MadIPM.failed_factorization_local_count!( + failed_local_buffer::Vector{Int32}, + s::RecordingBatchLinearSolver, + factor_view::MadIPM.BatchView, +) + nfailed = 0 + @inbounds for j in 1:MadIPM.local_batch_size(factor_view) + if !s.factorized[j] + nfailed += 1 + failed_local_buffer[nfailed] = j + end + end + return nfailed +end + +MadIPM.is_factorized(s::RecordingBatchLinearSolver) = all(s.factorized) + +MadIPM.solve_active!(s::RecordingBatchLinearSolver, rhs::AbstractMatrix, active::MadIPM.BatchView) = rhs + +function _make_batch_solver(qps; batch_kwargs...) + bnlp = BatchQuadraticModel(qps) + return MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) +end + function _make_small_qp() # Small QP: min 0.5 xᵀHx + cᵀx s.t. lcon ≤ Ax ≤ ucon, lvar ≤ x ≤ uvar n, m = 4, 2 @@ -114,6 +181,59 @@ function _test_fullbatch_different_data(; batch_kwargs...) end @testset "Batch solver (CPU)" begin + @testset "Batch views" begin + solver = _make_batch_solver([simple_lp() for _ in 1:4]) + root = MadIPM.root_view(solver.batch_views) + + @test MadIPM.local_batch_size(root) == 4 + @test MadIPM.batch_size_root(root) == 4 + @test MadIPM.is_identity_view(root) + @test _test_local_to_root(root) == [1, 2, 3, 4] + + saved_root = MadIPM.select_local!(solver.batch_views, [2, 4]) + child = MadIPM.active_view(solver.batch_views) + @test MadIPM.local_batch_size(child) == 2 + @test _test_local_to_root(child) == [2, 4] + @test !MadIPM.is_identity_view(child) + mask = zeros(Float64, 1, 4) + MadIPM.fill_batch_view_mask!(mask, child) + @test mask == [0.0 1.0 0.0 1.0] + + saved_child = MadIPM.select_local!(solver.batch_views, [2]) + grandchild = MadIPM.active_view(solver.batch_views) + @test _test_local_to_root(grandchild) == [4] + MadIPM.restore_state!(solver.batch_views, saved_child) + + MadIPM.select_local!(solver.batch_views, Int[]) + empty_child = MadIPM.active_view(solver.batch_views) + @test MadIPM.local_batch_size(empty_child) == 0 + MadIPM.restore_state!(solver.batch_views, saved_root) + + MadIPM.select_local!(solver.batch_views, [1, 2, 3, 4]) + full_child = MadIPM.active_view(solver.batch_views) + @test MadIPM.is_identity_view(full_child) + MadIPM.restore_state!(solver.batch_views, root) + end + + @testset "Partial active KKT solve preserves rhs" begin + solver = _make_batch_solver([simple_lp() for _ in 1:3]) + MadIPM.initialize!(solver) + status = fill(MadNLP.REGULAR, 3) + status[2] = MadNLP.INTERNAL_ERROR + solver.workspace.status .= status + MadIPM.update_active_set!(solver) + + pd_view = MadNLP.primal_dual(solver.d) + pd_view .= reshape(collect(1.0:length(pd_view)), size(pd_view)) + pd_before = copy(pd_view) + + MadNLP.build_kkt!(solver.kkt) + MadNLP.factorize_kkt!(solver.kkt) + MadNLP.solve_kkt!(solver.kkt, solver) + + @test pd_view[:, 2] == pd_before[:, 2] + end + @testset "Batch LP" begin _test_batch_lp() end @@ -164,4 +284,65 @@ end @test stats[3].status == MadNLP.SOLVE_SUCCEEDED @test stats[2].status != MadNLP.SOLVE_SUCCEEDED end + + @testset "Factorization retry only refactorizes failed instance" begin + solver = _make_batch_solver( + [simple_lp() for _ in 1:3]; + uniformbatch_linear_solver=RecordingBatchLinearSolver, + regularization=MadIPM.FixedRegularization(1.0, -1.0), + ) + MadIPM.initialize!(solver) + fill!(solver.kkt.batch_solver.call_counts, 0) + fill!(solver.kkt.batch_solver.factorized, true) + + corrupt_k = solver.kkt.n_tot + solver.kkt.nnzh + 1 + solver.kkt.nzVals[corrupt_k, 2] = -999.0 + solver.kkt.nzVals[corrupt_k, 3] = -999.0 + # two corruptions -- should see an extra solve in slots 1 and 2 + # and bumped regularizations in instances 2 and 3 + + MadIPM.factorize_system!(solver) + + ls = solver.kkt.batch_solver + @test ls.call_counts == [2, 2, 1] + @test solver.del_w == [1.0 100.0 100.0] + @test solver.del_c == [-1.0 -100.0 -100.0] + end + + @testset "Factorization retry with one inactive instance" begin + solver = _make_batch_solver( + [simple_lp() for _ in 1:3]; + uniformbatch_linear_solver=RecordingBatchLinearSolver, + regularization=MadIPM.FixedRegularization(1.0, -1.0), + ) + MadIPM.initialize!(solver) + fill!(solver.kkt.batch_solver.call_counts, 0) + fill!(solver.kkt.batch_solver.factorized, true) + # reset all call counts to zero + + # terminate an instance, so the last slot should be never used + status = fill(MadNLP.REGULAR, 3) + status[2] = MadNLP.SOLVE_SUCCEEDED + solver.workspace.status .= status + MadIPM.update_active_set!(solver) + MadIPM._update_active_mask!(solver) + + # corrupt instance 3, which is currently at position 2 + corrupt_k = solver.kkt.n_tot + solver.kkt.nnzh + 1 + solver.kkt.nzVals[corrupt_k, 3] = -999.0 + + MadIPM.factorize_system!(solver) + # since instance 3 required regularization bump, + # we expect to do another compaction to move 3 into position 1 + + # two solves in position 1 -- first instance 1, then instance 3's retry + # one solve in position 2 -- instance 3's first try + ls = solver.kkt.batch_solver + @test ls.call_counts == [2, 1, 0] + + # instance 1 and 2's regularization is untouched, + # instance 3's regularization should be bumped + @test solver.del_w == [1.0 1.0 100.0] + @test solver.del_c == [-1.0 -1.0 -100.0] + end end diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl index 1077335f..01700d53 100644 --- a/test/test_batch_gpu.jl +++ b/test/test_batch_gpu.jl @@ -95,6 +95,37 @@ end # ============================================================ @testset "Batch solver (CUDA)" begin + @testset "Batch views gather/scatter" begin + cpu_bnlp = BatchQuadraticModel([simple_lp() for _ in 1:4]) + gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) + solver = MadIPM.UniformBatchMPCSolver( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + ) + root = MadIPM.root_view(solver.batch_views) + saved_root = MadIPM.select_local!(solver.batch_views, [2, 4]) + child = MadIPM.active_view(solver.batch_views) + MadIPM.select_local!(solver.batch_views, [2]) + grandchild = MadIPM.active_view(solver.batch_views) + + src = cu(reshape(collect(1.0:12.0), 3, 4)) + gathered = similar(src, 3, MadIPM.local_batch_size(child)) + MadIPM.gather_batch_view_columns!(gathered, src, child) + @test Array(gathered) == Array(src[:, [2, 4]]) + + gathered_nested = similar(src, 3, MadIPM.local_batch_size(grandchild)) + MadIPM.gather_batch_view_columns!(gathered_nested, src, grandchild) + @test Array(gathered_nested) == Array(src[:, [4]]) + + scattered = CUDA.fill(-1.0, 3, 4) + MadIPM.scatter_batch_view_columns!(scattered, gathered, child) + @test Array(scattered[:, 2]) == Array(src[:, 2]) + @test Array(scattered[:, 4]) == Array(src[:, 4]) + @test Array(scattered[:, 1]) == fill(-1.0, 3) + @test Array(scattered[:, 3]) == fill(-1.0, 3) + MadIPM.restore_state!(solver.batch_views, saved_root) + end # ---------------------------------------------------------- # Identical instances (sanity check) @@ -283,5 +314,4 @@ end @test Array(si.solution) ≈ ref.solution atol=1e-6 end end - end From 7dea242626eecc6f6dcac4ede4442eac6746ddaf Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Fri, 27 Mar 2026 15:09:19 -0400 Subject: [PATCH 37/51] split state initialization into helper --- src/batch/madipm/solver.jl | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 9c321367..e77dbd4f 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -227,7 +227,15 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T batch_mapreduce!(abs, max, typemin(T), ws.norm_c, MadNLP.full(batch_solver.f)) init_starting_point!(batch_solver) + initialize_solver_state!(batch_solver) + MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) + return +end + +function initialize_solver_state!(batch_solver::AbstractBatchMPCSolver{T}) where T + ws = batch_solver.workspace + opt = batch_solver.opt fill!(ws.mu_batch, opt.mu_init) fill!(ws.best_complementarity, typemax(T)) fill!(ws.status, MadNLP.REGULAR) @@ -247,9 +255,6 @@ function initialize!(batch_solver::AbstractBatchMPCSolver{T}) where T batch_solver.batch_cnt.obj_cnt[] = 0 batch_solver.batch_cnt.obj_grad_cnt[] = 0 batch_solver.batch_cnt.con_cnt[] = 0 - - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return end From 52dc15fd80ae4a4e560e7199b6603bd4d65970eb Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Fri, 27 Mar 2026 15:55:20 -0400 Subject: [PATCH 38/51] check structure with optional skip --- src/batch/madnlp/callback.jl | 16 ++++++++++++++-- src/batch/structure.jl | 2 ++ test/test_batch.jl | 25 +++++++++++++++++++++++++ test/test_batch_gpu.jl | 31 +++++++++++++++++++++++++++++++ 4 files changed, 72 insertions(+), 2 deletions(-) diff --git a/src/batch/madnlp/callback.jl b/src/batch/madnlp/callback.jl index 4d2dde1f..90281710 100644 --- a/src/batch/madnlp/callback.jl +++ b/src/batch/madnlp/callback.jl @@ -114,6 +114,7 @@ function MadNLP.create_callback( bnlp::NLPModels.AbstractBatchNLPModel{T}; fixed_variable_treatment=MadNLP.MakeParameter, equality_treatment=MadNLP.EnforceEquality, + check_batch_structure::Bool=true, ) where {T,VT,MT,VI} bmeta = bnlp.meta batch_size = bmeta.nbatch @@ -143,7 +144,18 @@ function MadNLP.create_callback( NLPModels.hess_structure!(bnlp, hess_I, hess_J) end - # FIXME: assert structure? + if check_batch_structure + row_sums = vcat( + sum(bmeta.lvar .== bmeta.uvar; dims=2), + sum(isfinite.(bmeta.lvar); dims=2), + sum(isfinite.(bmeta.uvar); dims=2), + sum(bmeta.lcon .== bmeta.ucon; dims=2), + sum(isfinite.(bmeta.lcon); dims=2), + sum(isfinite.(bmeta.ucon); dims=2), + ) + @assert all((row_sums .== 0) .| (row_sums .== batch_size)) "Batch fixed/bound/equality structure must match across instances" + end + lvar = view(bmeta.lvar, :, 1) uvar = view(bmeta.uvar, :, 1) lcon = view(bmeta.lcon, :, 1) @@ -230,4 +242,4 @@ function MadNLP.build_hessian_structure(bcb::UniformBatchCallback, ::Type{<:MadN hess_J = MadNLP.create_array(bcb, Int32, bcb.nnzh) MadNLP._hess_sparsity_wrapper!(bcb, hess_I, hess_J) return hess_I, hess_J -end \ No newline at end of file +end diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 5af001bd..3ed78aca 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -149,6 +149,7 @@ function UniformBatchMPCSolver( VT = typeof(similar(NLPModels.get_x0(bnlp), T, 0)), VI = typeof(similar(NLPModels.get_x0(bnlp), Int, 0)), uniformbatch_linear_solver = LoopedBatchLinearSolver, + check_batch_structure::Bool = true, kwargs..., ) where {T} bmeta = bnlp.meta @@ -170,6 +171,7 @@ function UniformBatchMPCSolver( bnlp; fixed_variable_treatment=ipm_opt.fixed_variable_treatment, equality_treatment=ipm_opt.equality_treatment, + check_batch_structure=check_batch_structure, ) ind_lb = bcb.ind_lb diff --git a/test/test_batch.jl b/test/test_batch.jl index a4800e54..ae670e61 100644 --- a/test/test_batch.jl +++ b/test/test_batch.jl @@ -215,6 +215,31 @@ end MadIPM.restore_state!(solver.batch_views, root) end + @testset "Batch structure mismatch throws" begin + qp1 = QuadraticModel( + [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[1.0, 1.0], + x0=[0.5, 0.5], + ) + qp2 = QuadraticModel( + [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[0.0, 1.0], + x0=[0.0, 0.5], + ) + bnlp = BatchQuadraticModel([qp1, qp2]) + @test_throws AssertionError MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR) + solver = MadIPM.UniformBatchMPCSolver( + bnlp; + print_level=MadNLP.ERROR, + check_batch_structure=false, + ) + @test solver.batch_size == 2 + end + @testset "Partial active KKT solve preserves rhs" begin solver = _make_batch_solver([simple_lp() for _ in 1:3]) MadIPM.initialize!(solver) diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl index 01700d53..a7bcfb2c 100644 --- a/test/test_batch_gpu.jl +++ b/test/test_batch_gpu.jl @@ -127,6 +127,37 @@ end MadIPM.restore_state!(solver.batch_views, saved_root) end + @testset "Batch structure mismatch throws" begin + qp1 = QuadraticModel( + [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[1.0, 1.0], + x0=[0.5, 0.5], + ) + qp2 = QuadraticModel( + [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; + Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], + lcon=[1.0], ucon=[1.0], + lvar=[0.0, 0.0], uvar=[0.0, 1.0], + x0=[0.0, 0.5], + ) + cpu_bnlp = BatchQuadraticModel([qp1, qp2]) + gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) + @test_throws AssertionError MadIPM.UniformBatchMPCSolver( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + ) + solver = MadIPM.UniformBatchMPCSolver( + gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + check_batch_structure=false, + ) + @test solver.batch_size == 2 + end + # ---------------------------------------------------------- # Identical instances (sanity check) # ---------------------------------------------------------- From acc0e2e4fba55bc933b8059c4a8868e94beb2cfb Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Fri, 27 Mar 2026 17:50:13 -0400 Subject: [PATCH 39/51] use BQM's mapreduce for max, min, sum --- src/batch/utils.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/batch/utils.jl b/src/batch/utils.jl index 9ef3f24c..3345fbd4 100644 --- a/src/batch/utils.jl +++ b/src/batch/utils.jl @@ -8,9 +8,12 @@ function _csc_with_nzval(A::SparseArrays.SparseMatrixCSC, nzval, n) return SparseArrays.SparseMatrixCSC(n, n, SparseArrays.getcolptr(A), SparseArrays.rowvals(A), nzval) end -batch_maximum!(out::AbstractMatrix, src::AbstractMatrix) = maximum!(out, src) -batch_minimum!(out::AbstractMatrix, src::AbstractMatrix) = minimum!(out, src) -batch_sum!(out::AbstractMatrix, src::AbstractMatrix) = sum!(out, src) +batch_maximum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = + batch_mapreduce!(identity, max, typemin(T), out, src) +batch_minimum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = + batch_mapreduce!(identity, min, typemax(T), out, src) +batch_sum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = + batch_mapreduce!(identity, +, zero(T), out, src) function zero_inactive_step!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace From 8a29e71f9703b87726fee73115d7e51892aa72a7 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 13:41:22 -0400 Subject: [PATCH 40/51] remove _DefaultFH --- src/batch/madnlp/nlpmodels.jl | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl index 8ab9bf21..e1d5b444 100644 --- a/src/batch/madnlp/nlpmodels.jl +++ b/src/batch/madnlp/nlpmodels.jl @@ -1,8 +1,6 @@ -const _DefaultFH = Union{MadNLP.NoFixedVariables, MadNLP.RelaxBound} - function MadNLP._update_x!( bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} return x end @@ -20,7 +18,7 @@ end function MadNLP.unpack_x!( X_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x::BatchPrimalVector, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} X_full .= MadNLP.variable(x) end @@ -34,7 +32,7 @@ end function MadNLP.unpack_z!( Z_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, z_free, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} Z_full .= z_free ./ bcb.obj_scale end @@ -50,12 +48,12 @@ end MadNLP.get_y0(bcb::UniformBatchCallback) = bcb.nlp.meta.y0 -MadNLP.get_lvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} = bcb.nlp.meta.lvar +MadNLP.get_lvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH,EH} = bcb.nlp.meta.lvar function MadNLP.get_lvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} view(bcb.nlp.meta.lvar, bcb.fixed_handler.free, :) end -MadNLP.get_uvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} = bcb.nlp.meta.uvar +MadNLP.get_uvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH,EH} = bcb.nlp.meta.uvar function MadNLP.get_uvar(bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} view(bcb.nlp.meta.uvar, bcb.fixed_handler.free, :) end @@ -90,7 +88,7 @@ end function MadNLP._eval_grad_f_wrapper!( bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, bg::AbstractMatrix, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} NLPModels.grad!(bcb.nlp, bx, bg) bg .*= bcb.obj_scale return bg @@ -110,7 +108,7 @@ end function MadNLP._eval_jac_wrapper!( bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, bx::AbstractMatrix, jac_buffer::AbstractMatrix, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} NLPModels.jac_coord!(bcb.nlp, bx, jac_buffer) jac_buffer .*= bcb.jac_scale return jac_buffer @@ -131,7 +129,7 @@ function MadNLP._eval_lag_hess_wrapper!( y_mat::AbstractMatrix, bv::AbstractMatrix, hess::AbstractMatrix, -) where {T,VT,MT,VI,BM,FH<:_DefaultFH,EH} +) where {T,VT,MT,VI,BM,FH,EH} bv .= y_mat .* bcb.con_scale NLPModels.hess_coord!(bcb.nlp, bx, bv, vec(bcb.obj_scale), hess) return From 0b09974ed03b751a93a505a65180d588ce7deba7 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 13:48:56 -0400 Subject: [PATCH 41/51] remove unused test file --- test/batch_kernels/test_iteration_match.jl | 442 --------------------- 1 file changed, 442 deletions(-) delete mode 100644 test/batch_kernels/test_iteration_match.jl diff --git a/test/batch_kernels/test_iteration_match.jl b/test/batch_kernels/test_iteration_match.jl deleted file mode 100644 index a6160323..00000000 --- a/test/batch_kernels/test_iteration_match.jl +++ /dev/null @@ -1,442 +0,0 @@ -@testset "Iteration-level state matching" begin - - # ────────────────────────────────────────────────────────── - # init_starting_point! explicit snapshot: verify post-init state matches - # ────────────────────────────────────────────────────────── - @testset "init snapshot: $label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - assert_init_match(seq, bat, 1) - end - - # ────────────────────────────────────────────────────────── - # batch_size=1: every problem, full iteration trace - # ────────────────────────────────────────────────────────── - @testset "batch_size=1: $label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # batch_size=2: two DIFFERENT problems in one batch - # ────────────────────────────────────────────────────────── - @testset "batch_size=2: $label" for (label, make_paired) in PAIRED_PROBLEMS - qp_a, qp_b = make_paired() - seq1 = build_seq(qp_a) - seq2 = build_seq(qp_b) - bat = build_batch_from_qps([qp_a, qp_b]) - run_iterations_bs2!(seq1, seq2, bat, 50) - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - end - - # ────────────────────────────────────────────────────────── - # Scaling verification: confirm non-unit scales are active - # ────────────────────────────────────────────────────────── - @testset "scaling verification" begin - @testset "$label" for (label, make_qp) in [ - ("scaled_qp", _setup_scaled_qp), - ("scaled_ineq", _setup_scaled_ineq_qp), - ] - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - # Verify obj_scale < 1.0 - @test seq.cb.obj_scale[] < 1.0 - @test bat.bcb.obj_scale[1] < 1.0 - # Verify con_scale has non-unit entries - @test minimum(seq.cb.con_scale) < 1.0 - @test minimum(bat.bcb.con_scale[:, 1]) < 1.0 - # Verify scales match between sequential and batch - @test abs(seq.cb.obj_scale[] - bat.bcb.obj_scale[1]) < 1e-12 - @test cmp(seq.cb.con_scale, bat.bcb.con_scale[:, 1]) < 1e-12 - end - end - - # ────────────────────────────────────────────────────────── - # Fixed variable verification - # ────────────────────────────────────────────────────────── - @testset "fixed variable verification" begin - qp = _setup_fixed_var_qp() - seq = build_seq(qp) - bat = build_batch(qp) - # Verify MakeParameter handler is active for sequential - @test seq.cb.fixed_handler isa MadNLP.MakeParameter - @test bat.bcb.fixed_handler isa MadNLP.MakeParameter - # Sequential removes the fixed var from nvar; batch may keep it - @test MadNLP.n_variables(seq.cb) < 3 # reduced from 3 vars - # Run iterations (also part of ALL_TEST_PROBLEMS bs1 above, but - # this standalone test documents the fixed-variable intent) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # Staggered convergence: one instance converges much earlier - # with explicit zero_inactive_step! and active mask verification - # ────────────────────────────────────────────────────────── - @testset "staggered convergence (bs=2)" begin - qp_easy, qp_hard = _paired_staggered() - seq1 = build_seq(qp_easy) - seq2 = build_seq(qp_hard) - bat = build_batch_from_qps([qp_easy, qp_hard]) - - # Run iterations manually to verify zero_inactive_step! and active mask - ws = bat.workspace - seq1_done = false - seq2_done = false - inactive_verified = false - - for iter in 1:50 - if !seq1_done; MadIPM.update_termination_criteria!(seq1); end - if !seq2_done; MadIPM.update_termination_criteria!(seq2); end - MadIPM.update_termination_criteria!(bat) - - prev_seq1_done = seq1_done - prev_seq2_done = seq2_done - seq1_done = seq1_done || MadIPM.is_done(seq1) - seq2_done = seq2_done || MadIPM.is_done(seq2) - - MadIPM.update_active_set!(bat) - active = MadIPM.active_view(bat.batch_views) - MadIPM.local_batch_size(active) == 0 && break - MadIPM._update_active_mask!(bat) - - # Verify active mask matches expected state - if seq1_done && !seq2_done - @test Int[active.local_to_root[i] for i in 1:active.n] == [2] - @test ws.active_mask[1] == 0.0 - @test ws.active_mask[2] == 1.0 - elseif !seq1_done && seq2_done - @test Int[active.local_to_root[i] for i in 1:active.n] == [1] - @test ws.active_mask[1] == 1.0 - @test ws.active_mask[2] == 0.0 - elseif !seq1_done && !seq2_done - @test Int[active.local_to_root[i] for i in 1:active.n] == [1, 2] - end - - (seq1_done && seq2_done) && break - - # Factorize - if !seq1_done; MadIPM.factorize_system!(seq1); end - if !seq2_done; MadIPM.factorize_system!(seq2); end - MadIPM.factorize_system!(bat) - - # Prediction + Mehrotra - if !seq1_done; MadIPM.prediction_step!(seq1); end - if !seq2_done; MadIPM.prediction_step!(seq2); end - MadIPM.prediction_step!(bat) - if !seq1_done; MadIPM.mehrotra_correction_direction!(seq1); end - if !seq2_done; MadIPM.mehrotra_correction_direction!(seq2); end - MadIPM.mehrotra_correction_direction!(bat) - - # Update step - if !seq1_done; MadIPM.update_step!(seq1.opt.step_rule, seq1); end - if !seq2_done; MadIPM.update_step!(seq2.opt.step_rule, seq2); end - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.zero_inactive_step!(bat) - - # Verify zero_inactive_step! zeroed the inactive instance - if seq1_done && !seq2_done - @test ws.alpha_p[1] == 0.0 - @test ws.alpha_d[1] == 0.0 - @test ws.alpha_p[2] > 0.0 # active instance has nonzero step - inactive_verified = true - elseif !seq1_done && seq2_done - @test ws.alpha_p[2] == 0.0 - @test ws.alpha_d[2] == 0.0 - @test ws.alpha_p[1] > 0.0 - inactive_verified = true - end - if !seq1_done; assert_step_match(seq1, bat, 1); end - if !seq2_done; assert_step_match(seq2, bat, 2); end - - # Apply step + evaluate model - if !seq1_done; MadIPM.apply_step!(seq1); end - if !seq2_done; MadIPM.apply_step!(seq2); end - MadIPM.apply_step!(bat) - if !seq1_done; MadIPM.evaluate_model!(seq1); end - if !seq2_done; MadIPM.evaluate_model!(seq2); end - MadIPM.evaluate_model!(bat) - end - - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - # Verify the easy one converged earlier - @test seq1.cnt.k < seq2.cnt.k - # Verify we actually tested the inactive masking path - @test inactive_verified - end - - # ────────────────────────────────────────────────────────── - # batch_size=4: four DIFFERENT problems in one batch - # ────────────────────────────────────────────────────────── - @testset "batch_size=4: $label" for (label, make_quad) in [ - ("lower-only LP", _quad_lower_only), - ("doubly-bounded QP", _quad_doubly_bounded), - ("dense Hessian QP", _quad_dense_hess), - ] - qp1, qp2, qp3, qp4 = make_quad() - seqs = [build_seq(qp) for qp in [qp1, qp2, qp3, qp4]] - bat = build_batch_from_qps([qp1, qp2, qp3, qp4]) - run_iterations_bsN!(seqs, bat, 50) - for s in seqs - @test s.status == MadNLP.SOLVE_SUCCEEDED - end - end - - # ────────────────────────────────────────────────────────── - # Factorization retry: verify retry loop works for both seq and batch - # ────────────────────────────────────────────────────────── - @testset "factorization retry path" begin - qp = _setup_doubly_bounded_qp() - seq = build_seq(qp) - bat = build_batch(qp) - - # Artificially set del_w/del_c to large values (as if retry triggered) - seq.del_w = 1e-4 - seq.del_c = -1e-4 - bat.del_w .= 1e-4 - bat.del_c .= -1e-4 - - MadIPM.factorize_system!(seq) - MadIPM.factorize_system!(bat) - assert_regularization_match(seq, bat, 1) - assert_kkt_diagonals_match(seq, bat, 1) - assert_kkt_matrix_match(seq, bat, 1) - - MadIPM.prediction_step!(seq) - MadIPM.prediction_step!(bat) - assert_prediction_match(seq, bat, 1) - end - - # ────────────────────────────────────────────────────────── - # KKT mul! verification: batch scatter mul vs sequential sparse mul - # ────────────────────────────────────────────────────────── - @testset "KKT mul! $label" for (label, make_qp) in [ - ("QP (nlb>0, nub>0)", _setup_small_qp), - ("QP doubly-bounded", _setup_doubly_bounded_qp), - ("QP dense Hessian+mixed", _setup_dense_hess_mixed_qp), - ("QP mixed bounds", _setup_mixed_bounds_qp), - ] - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - - # Set up KKT diagonals (needed for mul!) - MadIPM.update_regularization!(seq, seq.opt.regularization) - MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) - MadIPM.update_regularization!(bat, bat.opt.regularization) - MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) - - # Use the current d as input vector (has non-trivial values after init) - # Copy d → p to have a known input - copyto!(MadNLP.full(seq.p), MadNLP.full(seq.d)) - copyto!(MadNLP.full(bat.p), MadNLP.full(bat.d)) - - # Compute mul!(w1, kkt, p) for both - seq_w = seq._w1 - bat_w = bat._w1 - fill!(MadNLP.full(seq_w), 0.0) - fill!(MadNLP.full(bat_w), 0.0) - mul!(seq_w, seq.kkt, seq.p) - mul!(bat_w, bat.kkt, bat.p) - - @test cmp(MadNLP.full(seq_w), coln(MadNLP.full(bat_w), 1)) < 1e-10 - end - - # ────────────────────────────────────────────────────────── - # Step rule variants (bs=1, with _setup_small_qp) - # ────────────────────────────────────────────────────────── - @testset "step_rule=$label" for (label, step_rule) in [ - ("ConservativeStep", MadIPM.ConservativeStep(0.995)), - ("AdaptiveStep", MadIPM.AdaptiveStep(0.99)), - ("MehrotraAdaptiveStep", MadIPM.MehrotraAdaptiveStep(0.99)), - ] - qp = _setup_small_qp() - seq = build_seq(qp; step_rule=step_rule) - bat = build_batch(qp; step_rule=step_rule) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # Step rule variants with diverse problem types - # ────────────────────────────────────────────────────────── - @testset "step_rule=$sr_label on $prob_label" for (sr_label, step_rule) in [ - ("ConservativeStep", MadIPM.ConservativeStep(0.995)), - ("AdaptiveStep", MadIPM.AdaptiveStep(0.99)), - ("MehrotraAdaptiveStep", MadIPM.MehrotraAdaptiveStep(0.99)), - ], (prob_label, make_qp) in [ - ("free vars", _setup_free_qp), - ("all-ineq", _setup_all_ineq_qp), - ("doubly-bounded", _setup_doubly_bounded_qp), - ("dense Hessian+mixed", _setup_dense_hess_mixed_qp), - ] - qp = make_qp() - seq = build_seq(qp; step_rule=step_rule) - bat = build_batch(qp; step_rule=step_rule) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # Regularization variants (bs=1, with _setup_small_qp) - # ────────────────────────────────────────────────────────── - @testset "regularization=$label" for (label, make_reg) in [ - ("NoRegularization", () -> MadIPM.NoRegularization()), - ("FixedRegularization", () -> MadIPM.FixedRegularization(1e-10, 1e-10)), - ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), - ] - qp = _setup_small_qp() - seq = build_seq(qp; regularization=make_reg()) - bat = build_batch(qp; regularization=make_reg()) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # Regularization variants with diverse problem types - # ────────────────────────────────────────────────────────── - @testset "regularization=$reg_label on $prob_label" for (reg_label, make_reg) in [ - ("NoRegularization", () -> MadIPM.NoRegularization()), - ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), - ], (prob_label, make_qp) in [ - ("free vars", _setup_free_qp), - ("all-ineq", _setup_all_ineq_qp), - ("doubly-bounded", _setup_doubly_bounded_qp), - ("dense Hessian+mixed", _setup_dense_hess_mixed_qp), - ] - qp = make_qp() - seq = build_seq(qp; regularization=make_reg()) - bat = build_batch(qp; regularization=make_reg()) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # Step rule + regularization variants with bs=2 - # ────────────────────────────────────────────────────────── - @testset "options bs=2: $label" for (label, make_opts) in [ - ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), - ("Mehrotra+AdaptiveReg", () -> (step_rule=MadIPM.MehrotraAdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), - ] - qp_a, qp_b = _paired_doubly_bounded() - # Each solver needs its own copy of mutable options (e.g. AdaptiveRegularization) - opts1 = make_opts() - opts2 = make_opts() - opts_bat = make_opts() - seq1 = build_seq(qp_a; opts1...) - seq2 = build_seq(qp_b; opts2...) - bat = build_batch_from_qps([qp_a, qp_b]; opts_bat...) - run_iterations_bs2!(seq1, seq2, bat, 50) - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - end - - # ────────────────────────────────────────────────────────── - # Combined step+reg options with diverse bs=2 paired problems - # ────────────────────────────────────────────────────────── - @testset "options bs=2: $opt_label on $prob_label" for (opt_label, make_opts) in [ - ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), - ("Adaptive+AdaptiveReg", () -> (step_rule=MadIPM.AdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), - ], (prob_label, make_paired) in [ - ("dense-hess", _paired_dense_hess), - ("all-ineq", _paired_all_ineq), - ("mixed-bounds", _paired_mixed_bounds), - ] - qp_a, qp_b = make_paired() - opts1 = make_opts() - opts2 = make_opts() - opts_bat = make_opts() - seq1 = build_seq(qp_a; opts1...) - seq2 = build_seq(qp_b; opts2...) - bat = build_batch_from_qps([qp_a, qp_b]; opts_bat...) - run_iterations_bs2!(seq1, seq2, bat, 50) - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - end - - # ────────────────────────────────────────────────────────── - # init_regularization! explicit comparison - # ────────────────────────────────────────────────────────── - @testset "init_regularization! $label" for (label, make_reg) in [ - ("NoRegularization", () -> MadIPM.NoRegularization()), - ("FixedRegularization", () -> MadIPM.FixedRegularization(1e-10, 1e-10)), - ("AdaptiveRegularization", () -> MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12)), - ] - qp = _setup_small_qp() - seq = build_seq(qp; regularization=make_reg()) - bat = build_batch(qp; regularization=make_reg()) - # After build, init_regularization! has been called. - # del_w and del_c should match. - assert_regularization_match(seq, bat, 1) - end - - # ══════════════════════════════════════════════════════════ - # BatchQuadraticModel tests (different H/A values per instance) - # ══════════════════════════════════════════════════════════ - - # ────────────────────────────────────────────────────────── - # BatchQuadraticModel bs=1: verify same results as ObjRHSBatch - # ────────────────────────────────────────────────────────── - @testset "fullbatch bs=1: $label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_fullbatch(qp) - assert_init_match(seq, bat, 1) - run_iterations_bs1!(seq, bat, 50) - end - - # ────────────────────────────────────────────────────────── - # BatchQuadraticModel bs=2: different H AND A values - # ────────────────────────────────────────────────────────── - @testset "fullbatch bs=2: $label" for (label, make_paired) in FULLBATCH_PAIRED_PROBLEMS - qp_a, qp_b = make_paired() - seq1 = build_seq(qp_a) - seq2 = build_seq(qp_b) - bat = build_fullbatch_from_qps([qp_a, qp_b]) - run_iterations_bs2!(seq1, seq2, bat, 50) - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - end - - # ────────────────────────────────────────────────────────── - # BatchQuadraticModel scaling: verify per-instance scaling with different A - # ────────────────────────────────────────────────────────── - @testset "fullbatch scaling verification" begin - qp_a, qp_b = _fullbatch_paired_scaled() - seq1 = build_seq(qp_a) - seq2 = build_seq(qp_b) - bat = build_fullbatch_from_qps([qp_a, qp_b]) - # Verify non-unit scales - @test seq1.cb.obj_scale[] < 1.0 - @test seq2.cb.obj_scale[] < 1.0 - @test bat.bcb.obj_scale[1] < 1.0 - @test bat.bcb.obj_scale[2] < 1.0 - # Verify per-instance scales match sequential - @test abs(seq1.cb.obj_scale[] - bat.bcb.obj_scale[1]) < 1e-12 - @test abs(seq2.cb.obj_scale[] - bat.bcb.obj_scale[2]) < 1e-12 - @test cmp(seq1.cb.con_scale, bat.bcb.con_scale[:, 1]) < 1e-12 - @test cmp(seq2.cb.con_scale, bat.bcb.con_scale[:, 2]) < 1e-12 - # Verify that the two instances have DIFFERENT scales (since A values differ) - @test bat.bcb.con_scale[1, 1] != bat.bcb.con_scale[1, 2] - end - - # ────────────────────────────────────────────────────────── - # BatchQuadraticModel with options variants - # ────────────────────────────────────────────────────────── - @testset "fullbatch options: $label" for (label, make_opts) in [ - ("Conservative+NoReg", () -> (step_rule=MadIPM.ConservativeStep(0.995), regularization=MadIPM.NoRegularization())), - ("Adaptive+AdaptiveReg", () -> (step_rule=MadIPM.AdaptiveStep(0.99), regularization=MadIPM.AdaptiveRegularization(1e-6, 1e-6, 1e-12))), - ] - qp_a, qp_b = _fullbatch_paired_dense_hess() - opts1 = make_opts() - opts2 = make_opts() - opts_bat = make_opts() - seq1 = build_seq(qp_a; opts1...) - seq2 = build_seq(qp_b; opts2...) - bat = build_fullbatch_from_qps([qp_a, qp_b]; opts_bat...) - run_iterations_bs2!(seq1, seq2, bat, 50) - @test seq1.status == MadNLP.SOLVE_SUCCEEDED - @test seq2.status == MadNLP.SOLVE_SUCCEEDED - end -end From 2ff5a98273594164ec1454be0339068b1fc03f4a Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 14:15:52 -0400 Subject: [PATCH 42/51] typos --- src/batch/madipm/solver.jl | 2 +- src/batch/structure.jl | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index e77dbd4f..01d86e49 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -278,7 +278,7 @@ function compute_term_gpu!(ws::UniformBatchWorkspace{T}, opt) where T (ws.dual_obj > max(ds * abs(ws.obj_val), one(T))), Int_INFEASIBLE, ifelse( - ws.obj_val < -(div_tol * max(ds, abs(ws.dual_obj), one(T))), + ws.obj_val < -(div_tol * max(ds * abs(ws.dual_obj), one(T))), Int_DIVERGING, Int_REGULAR, ), diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 3ed78aca..18639aa9 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -24,7 +24,7 @@ struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI status::Vector{MadNLP.Status} _term_gpu::MI64 - _term_cpu::Vector{Int64} + _term_cpu::Matrix{Int64} _any_nonregular_gpu::MI64 _any_nonregular_cpu::Matrix{Int64} _norm_gpu_w::MT @@ -62,8 +62,8 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, MT(undef, 1, batch_size), # best_complementarity MT(undef, 1, batch_size), # dual_obj fill(MadNLP.INITIAL, batch_size), # status - similar(_proto, Int64), # _term_gpu - zeros(Int64, batch_size), # _term_cpu + similar(_proto, Int64, 1, batch_size), # _term_gpu + zeros(Int64, 1, batch_size), # _term_cpu fill!(similar(_proto, Int64, 1, 1), Int64(Int(MadNLP.REGULAR))), # _any_nonregular_gpu zeros(Int64, 1, 1), # _any_nonregular_cpu MT(undef, 1, batch_size), # _norm_gpu_w From ff709448a37822ebdd74471673409baa58e5bcb1 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 17:49:08 -0400 Subject: [PATCH 43/51] use S from nlpmodel --- src/batch/structure.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/batch/structure.jl b/src/batch/structure.jl index 18639aa9..e8302bd3 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -64,8 +64,8 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, fill(MadNLP.INITIAL, batch_size), # status similar(_proto, Int64, 1, batch_size), # _term_gpu zeros(Int64, 1, batch_size), # _term_cpu - fill!(similar(_proto, Int64, 1, 1), Int64(Int(MadNLP.REGULAR))), # _any_nonregular_gpu - zeros(Int64, 1, 1), # _any_nonregular_cpu + fill!(similar(_proto, Int64, 1, 1), Int64(MadNLP.REGULAR)), # _any_nonregular_gpu + fill!(zeros(Int64, 1, 1), Int64(MadNLP.REGULAR)), # _any_nonregular_cpu MT(undef, 1, batch_size), # _norm_gpu_w MT(undef, 1, batch_size), # _norm_gpu_p fill!(similar(_proto, Int32), zero(Int32)), # _ls_error @@ -144,14 +144,13 @@ update_active_set!(bs::AbstractBatchMPCSolver) = update_active_set!(bs.batch_vie Construct a batch solver from a `AbstractBatchNLPModel`. """ function UniformBatchMPCSolver( - bnlp::NLPModels.AbstractBatchNLPModel{T}; - MT = typeof(similar(NLPModels.get_x0(bnlp), T, 0, 0)), + bnlp::NLPModels.AbstractBatchNLPModel{T, MT}; VT = typeof(similar(NLPModels.get_x0(bnlp), T, 0)), VI = typeof(similar(NLPModels.get_x0(bnlp), Int, 0)), uniformbatch_linear_solver = LoopedBatchLinearSolver, check_batch_structure::Bool = true, kwargs..., -) where {T} +) where {T, MT} bmeta = bnlp.meta batch_size = bmeta.nbatch @assert batch_size > 0 "Need at least one instance in batch" From 6860103270f31a37ab0b524cb66c99a552e54792 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 21:21:42 -0400 Subject: [PATCH 44/51] don't modify inactive kkt --- src/batch/madipm/kernels.jl | 50 ++++++++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 5a066490..077c420a 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -100,31 +100,53 @@ function get_correction!( return end -function set_aug_diagonal_reg!(kkt, solver::AbstractBatchMPCSolver) +function _set_aug_diagonal_reg_unmasked!(kkt, solver::AbstractBatchMPCSolver) + kkt.reg .= solver.del_w + du_diag(kkt) .= solver.del_c + kkt.l_diag .= lower(solver.xl) .- lower(solver.x) + kkt.u_diag .= upper(solver.x) .- upper(solver.xu) + kkt.l_lower .= lower(solver.zl) + kkt.u_lower .= upper(solver.zu) + pr_diag(kkt) .= kkt.reg + pr_diag_lb = view(kkt.nzVals, _get_ind_lb(solver), :) + pr_diag_ub = view(kkt.nzVals, _get_ind_ub(solver), :) + pr_diag_lb .-= kkt.l_lower ./ kkt.l_diag + pr_diag_ub .-= kkt.u_lower ./ kkt.u_diag + return +end + +function _set_aug_diagonal_reg_masked!(kkt, solver::AbstractBatchMPCSolver) xl_r = lower(solver.xl) x_lr = lower(solver.x) zl_r = lower(solver.zl) xu_r = upper(solver.xu) x_ur = upper(solver.x) zu_r = upper(solver.zu) - - kkt.reg .= solver.del_w - du_diag(kkt) .= solver.del_c - - kkt.l_diag .= xl_r .- x_lr - kkt.u_diag .= x_ur .- xu_r - - kkt.l_lower .= zl_r - kkt.u_lower .= zu_r - - pr_diag(kkt) .= kkt.reg + mask = solver.workspace.active_mask + _du = du_diag(kkt) + _pr = pr_diag(kkt) + @. kkt.reg = ifelse(mask == 1, solver.del_w, kkt.reg) + @. _du = ifelse(mask == 1, solver.del_c, _du) + @. kkt.l_diag = ifelse(mask == 1, xl_r - x_lr, kkt.l_diag) + @. kkt.u_diag = ifelse(mask == 1, x_ur - xu_r, kkt.u_diag) + @. kkt.l_lower = ifelse(mask == 1, zl_r, kkt.l_lower) + @. kkt.u_lower = ifelse(mask == 1, zu_r, kkt.u_lower) + @. _pr = ifelse(mask == 1, kkt.reg, _pr) pr_diag_lb = view(kkt.nzVals, _get_ind_lb(solver), :) pr_diag_ub = view(kkt.nzVals, _get_ind_ub(solver), :) - pr_diag_lb .-= kkt.l_lower ./ kkt.l_diag - pr_diag_ub .-= kkt.u_lower ./ kkt.u_diag + @. pr_diag_lb = ifelse(mask == 1, pr_diag_lb - kkt.l_lower / kkt.l_diag, pr_diag_lb) + @. pr_diag_ub = ifelse(mask == 1, pr_diag_ub - kkt.u_lower / kkt.u_diag, pr_diag_ub) return end +function set_aug_diagonal_reg!(kkt, solver::AbstractBatchMPCSolver) + if is_identity_view(active_view(solver.batch_views)) + _set_aug_diagonal_reg_unmasked!(kkt, solver) + else + _set_aug_diagonal_reg_masked!(kkt, solver) + end +end + function get_complementarity_measure!(solver::AbstractBatchMPCSolver) ws = solver.workspace nlb, nub = solver.d.nlb, solver.d.nub From 4537ed4e796e682fc57fae94c3625421a24f450a Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 21:36:56 -0400 Subject: [PATCH 45/51] simplify factorize_system since we can't nest between factorize and solve --- src/batch/madipm/solver.jl | 53 +++++++++++++++++++------------------- test/test_batch.jl | 12 +++------ 2 files changed, 29 insertions(+), 36 deletions(-) diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 01d86e49..20684d14 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -432,16 +432,22 @@ function mehrotra_correction_direction!(solver::AbstractBatchMPCSolver) return end -function _restrict_to_failed_locals!(batch_solver::AbstractBatchMPCSolver, nfailed::Int) - select_local!(batch_solver.batch_views, batch_solver.batch_views.selected_local_buffer, nfailed) - return active_view(batch_solver.batch_views) -end - -function _bump_failed_regularization!(batch_solver::AbstractBatchMPCSolver{T}, failed_view::BatchView) where T +function _bump_failed_regularization!(batch_solver::AbstractBatchMPCSolver{T}, failed_locals, nfailed::Int) where T + factor_view = active_view(batch_solver.batch_views) ws = batch_solver.workspace - fill_batch_view_mask!(ws.active_mask_cpu, failed_view) - @. batch_solver.del_w = ifelse(ws.active_mask_cpu == one(T), T(100) * batch_solver.del_w, batch_solver.del_w) - @. batch_solver.del_c = ifelse(ws.active_mask_cpu == one(T), T(100) * batch_solver.del_c, batch_solver.del_c) + # build root-level mask from local failed idx + fill!(ws.active_mask_cpu, zero(T)) + @inbounds for k in 1:nfailed + j = factor_view.local_to_root[failed_locals[k]] + ws.active_mask_cpu[1, j] = one(T) + end + copyto!(ws.active_mask, ws.active_mask_cpu) + mask = ws.active_mask + @. batch_solver.del_w = ifelse(mask == one(T), T(100) * batch_solver.del_w, batch_solver.del_w) + @. batch_solver.del_c = ifelse(mask == one(T), T(100) * batch_solver.del_c, batch_solver.del_c) + # restore active mask + # this is required to not throw away any successful factorization that we need later + _update_active_mask!(batch_solver) return end @@ -450,24 +456,17 @@ function factorize_system!(batch_solver::AbstractBatchMPCSolver) batch_views = batch_solver.batch_views update_regularization!(batch_solver, batch_solver.opt.regularization) max_trials = 3 - saved_active = active_view(batch_views) - try - for _ in 1:max_trials - set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) - MadNLP.factorize_wrapper!(batch_solver) - factor_view = active_view(batch_views) - nfailed = failed_factorization_local_count!( - batch_views.selected_local_buffer, - batch_solver.kkt.batch_solver, - factor_view, - ) - nfailed == 0 && break - - failed_view = _restrict_to_failed_locals!(batch_solver, nfailed) - _bump_failed_regularization!(batch_solver, failed_view) - end - finally - restore_state!(batch_views, saved_active) + factor_view = active_view(batch_views) + failed_locals = batch_views.selected_local_buffer + + for _ in 1:max_trials + set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) + MadNLP.factorize_wrapper!(batch_solver) + nfailed = failed_factorization_local_count!( + failed_locals, batch_solver.kkt.batch_solver, factor_view, + ) + nfailed == 0 && break + _bump_failed_regularization!(batch_solver, failed_locals, nfailed) end return end diff --git a/test/test_batch.jl b/test/test_batch.jl index ae670e61..742e4c7e 100644 --- a/test/test_batch.jl +++ b/test/test_batch.jl @@ -329,7 +329,7 @@ end MadIPM.factorize_system!(solver) ls = solver.kkt.batch_solver - @test ls.call_counts == [2, 2, 1] + @test ls.call_counts == [2, 2, 2] @test solver.del_w == [1.0 100.0 100.0] @test solver.del_c == [-1.0 -100.0 -100.0] end @@ -343,7 +343,6 @@ end MadIPM.initialize!(solver) fill!(solver.kkt.batch_solver.call_counts, 0) fill!(solver.kkt.batch_solver.factorized, true) - # reset all call counts to zero # terminate an instance, so the last slot should be never used status = fill(MadNLP.REGULAR, 3) @@ -352,18 +351,13 @@ end MadIPM.update_active_set!(solver) MadIPM._update_active_mask!(solver) - # corrupt instance 3, which is currently at position 2 + # corrupt instance 3, which is at position 2 after compacting [1,3] corrupt_k = solver.kkt.n_tot + solver.kkt.nnzh + 1 solver.kkt.nzVals[corrupt_k, 3] = -999.0 MadIPM.factorize_system!(solver) - # since instance 3 required regularization bump, - # we expect to do another compaction to move 3 into position 1 - - # two solves in position 1 -- first instance 1, then instance 3's retry - # one solve in position 2 -- instance 3's first try ls = solver.kkt.batch_solver - @test ls.call_counts == [2, 1, 0] + @test ls.call_counts == [2, 2, 0] # instance 1 and 2's regularization is untouched, # instance 3's regularization should be bumped From 1ceeba2e26e6aabaa6b72659da8ac432469cef88 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Mon, 30 Mar 2026 21:51:45 -0400 Subject: [PATCH 46/51] typos --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 1 + src/batch/KKT/Sparse/augmented.jl | 3 ++- src/batch/madnlp/nlpmodels.jl | 10 ++++++---- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index 9f3ee0ff..b18e71fa 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -389,6 +389,7 @@ function MadIPM._mehrotra_step!( dx_ur, x_ur, xu_r, nub, dzub, zu_r, d_vals, ind_lb, ind_ub, dlb_off, dub_off, ) + bs = size(alpha_p, 2) CUDA.@cuda threads=32 blocks=bs _mehrotra_step_kernel!( alpha_p, alpha_d, mu, gamma_f, dx_lr, x_lr, xl_r, Int32(nlb), dzlb, zl_r, diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index c2071bce..0435b27e 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -253,7 +253,8 @@ function MadNLP.eval_lag_hess_wrapper!( if nnzh > 0 hess = view(nzVals, n_tot+1:n_tot+nnzh, :) MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess) + scale = vec(bcb.obj_sign .* bcb.obj_scale) + MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess; obj_weight=scale) end return end diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl index e1d5b444..1d6cfae3 100644 --- a/src/batch/madnlp/nlpmodels.jl +++ b/src/batch/madnlp/nlpmodels.jl @@ -128,10 +128,11 @@ function MadNLP._eval_lag_hess_wrapper!( bx::AbstractMatrix, y_mat::AbstractMatrix, bv::AbstractMatrix, - hess::AbstractMatrix, + hess::AbstractMatrix; + obj_weight::AbstractVector = vec(bcb.obj_scale), ) where {T,VT,MT,VI,BM,FH,EH} bv .= y_mat .* bcb.con_scale - NLPModels.hess_coord!(bcb.nlp, bx, bv, vec(bcb.obj_scale), hess) + NLPModels.hess_coord!(bcb.nlp, bx, bv, obj_weight, hess) return end @@ -140,10 +141,11 @@ function MadNLP._eval_lag_hess_wrapper!( bx::AbstractMatrix, y_mat::AbstractMatrix, bv::AbstractMatrix, - hess::AbstractMatrix, + hess::AbstractMatrix; + obj_weight::AbstractVector = vec(bcb.obj_scale), ) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} bv .= y_mat .* bcb.con_scale - NLPModels.hess_coord!(bcb.nlp, bx, bv, vec(bcb.obj_scale), bcb.hess_buffer) + NLPModels.hess_coord!(bcb.nlp, bx, bv, obj_weight, bcb.hess_buffer) hess .= view(bcb.hess_buffer, bcb.fixed_handler.ind_hess_free, :) return end From fa1f4d5eb8d71bfd1dc76f3ae91b88af6049c9ab Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 31 Mar 2026 01:16:36 -0400 Subject: [PATCH 47/51] rewrite tests --- Project.toml | 2 +- src/MadIPM.jl | 3 +- src/batch/madipm/solver.jl | 7 +- src/batch/madnlp/kernels.jl | 17 - src/batch/madnlp/linear_solver.jl | 8 +- src/batch/utils.jl | 7 - test/batch/fakels.jl | 66 +++ test/batch/gpu.jl | 125 +++++ test/batch/solver.jl | 518 +++++++++++++++++++++ test/batch/views.jl | 112 +++++ test/batch_kernels/runtests.jl | 10 - test/batch_kernels/setup.jl | 340 -------------- test/batch_kernels/test_batch_multi.jl | 55 --- test/batch_kernels/test_complementarity.jl | 53 --- test/batch_kernels/test_kkt.jl | 86 ---- test/batch_kernels/test_regularization.jl | 61 --- test/batch_kernels/test_rhs.jl | 53 --- test/batch_kernels/test_solver_steps.jl | 100 ---- test/batch_kernels/test_step.jl | 64 --- test/batch_kernels/test_termination.jl | 48 -- test/runtests.jl | 6 +- test/test_batch.jl | 367 --------------- test/test_batch_gpu.jl | 348 -------------- 23 files changed, 833 insertions(+), 1623 deletions(-) create mode 100644 test/batch/fakels.jl create mode 100644 test/batch/gpu.jl create mode 100644 test/batch/solver.jl create mode 100644 test/batch/views.jl delete mode 100644 test/batch_kernels/runtests.jl delete mode 100644 test/batch_kernels/setup.jl delete mode 100644 test/batch_kernels/test_batch_multi.jl delete mode 100644 test/batch_kernels/test_complementarity.jl delete mode 100644 test/batch_kernels/test_kkt.jl delete mode 100644 test/batch_kernels/test_regularization.jl delete mode 100644 test/batch_kernels/test_rhs.jl delete mode 100644 test/batch_kernels/test_solver_steps.jl delete mode 100644 test/batch_kernels/test_step.jl delete mode 100644 test/batch_kernels/test_termination.jl delete mode 100644 test/test_batch.jl delete mode 100644 test/test_batch_gpu.jl diff --git a/Project.toml b/Project.toml index 6aae2f4f..49d8cac7 100644 --- a/Project.toml +++ b/Project.toml @@ -22,7 +22,7 @@ MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" [sources] -BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl"} +BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl.git"} [extensions] MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] diff --git a/src/MadIPM.jl b/src/MadIPM.jl index 21d5a163..28517c05 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -14,7 +14,8 @@ import BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchSparseOp, batch_spmv!, - batch_mapreduce! + batch_mapreduce!, + batch_maximum! include("utils.jl") include("structure.jl") diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 20684d14..0ca4e145 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -462,7 +462,7 @@ function factorize_system!(batch_solver::AbstractBatchMPCSolver) for _ in 1:max_trials set_aug_diagonal_reg!(batch_solver.kkt, batch_solver) MadNLP.factorize_wrapper!(batch_solver) - nfailed = failed_factorization_local_count!( + nfailed = is_factorized!( failed_locals, batch_solver.kkt.batch_solver, factor_view, ) nfailed == 0 && break @@ -606,11 +606,12 @@ function MadNLP.print_iter(batch_solver::AbstractBatchMPCSolver) bs = batch_solver.batch_size k = maximum(bcnt.k) + active_str = "$na/$bs" mod(k, 10) == 0 && MadNLP.@info(logger, @sprintf( " iter active max_inf_pr max_inf_du max_inf_compl max_alpha_p")) MadNLP.@info(logger, @sprintf( - "%4i %3i/%3i %6.2e %6.2e %7.2e %6.2e", - k, na, bs, + "%4i ", k) * lpad(active_str, 6) * @sprintf( + " %6.2e %6.2e %7.2e %6.2e", maximum(ws.inf_pr), maximum(ws.inf_du), maximum(ws.inf_compl), maximum(ws.alpha_p), )) diff --git a/src/batch/madnlp/kernels.jl b/src/batch/madnlp/kernels.jl index a04718dc..5164966c 100644 --- a/src/batch/madnlp/kernels.jl +++ b/src/batch/madnlp/kernels.jl @@ -27,20 +27,3 @@ function get_inf_compl!(inf_compl, x, xl, zl, xu, zu, sum_lb, sum_ub, nlb, nub) return inf_compl end -_adjust_bound_lb(x_lr::T, xl_r, c1, c2) where T = - x_lr - xl_r < c1 ? xl_r - c2 * max(one(T), abs(x_lr)) : xl_r -_adjust_bound_ub(x_ur::T, xu_r, c1, c2) where T = - xu_r - x_ur < c1 ? xu_r + c2 * max(one(T), abs(x_ur)) : xu_r - -function MadNLP.adjust_boundary!( - x_lr::AbstractMatrix{T}, - xl_r::AbstractMatrix{T}, - x_ur::AbstractMatrix{T}, - xu_r::AbstractMatrix{T}, - mu, -) where T - c1 = eps(T) .* mu - c2 = T(eps(T)^(3/4)) - xl_r .= _adjust_bound_lb.(x_lr, xl_r, c1, c2) - xu_r .= _adjust_bound_ub.(x_ur, xu_r, c1, c2) -end diff --git a/src/batch/madnlp/linear_solver.jl b/src/batch/madnlp/linear_solver.jl index f083c8ae..208dac11 100644 --- a/src/batch/madnlp/linear_solver.jl +++ b/src/batch/madnlp/linear_solver.jl @@ -30,11 +30,7 @@ function LoopedBatchLinearSolver( LoopedBatchLinearSolver{T, VT, eltype(individual_solvers)}(individual_solvers, batch_size) end -function is_factorized(batch_linear_solver::LoopedBatchLinearSolver) - return all(is_factorized(s) for s in batch_linear_solver.solvers) -end - -function failed_factorization_local_count!( +function is_factorized!( failed_local_buffer::Vector{Int32}, batch_linear_solver::LoopedBatchLinearSolver, factor_view::BatchView, @@ -66,4 +62,4 @@ function solve_active!(s::LoopedBatchLinearSolver{T, VT}, rhs::AbstractMatrix{T} end return end -failed_factorization_local_count!(::Vector{Int32}, ::MadNLP.AbstractLinearSolver, ::BatchView) = 0 +is_factorized!(::Vector{Int32}, ::MadNLP.AbstractLinearSolver, ::BatchView) = 0 diff --git a/src/batch/utils.jl b/src/batch/utils.jl index 3345fbd4..bbabdd49 100644 --- a/src/batch/utils.jl +++ b/src/batch/utils.jl @@ -8,13 +8,6 @@ function _csc_with_nzval(A::SparseArrays.SparseMatrixCSC, nzval, n) return SparseArrays.SparseMatrixCSC(n, n, SparseArrays.getcolptr(A), SparseArrays.rowvals(A), nzval) end -batch_maximum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = - batch_mapreduce!(identity, max, typemin(T), out, src) -batch_minimum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = - batch_mapreduce!(identity, min, typemax(T), out, src) -batch_sum!(out::AbstractMatrix{T}, src::AbstractMatrix{T}) where {T} = - batch_mapreduce!(identity, +, zero(T), out, src) - function zero_inactive_step!(batch_solver::AbstractBatchMPCSolver{T}) where T ws = batch_solver.workspace ws.alpha_p .*= ws.active_mask diff --git a/test/batch/fakels.jl b/test/batch/fakels.jl new file mode 100644 index 00000000..66037f93 --- /dev/null +++ b/test/batch/fakels.jl @@ -0,0 +1,66 @@ +using LDLFactorizations + +mutable struct FailOnDemandLS{T,VT,LS<:MadNLP.AbstractLinearSolver{T}} <: MadNLP.AbstractLinearSolver{T} + solvers::Vector{LS} + batch_size::Int + fail_positions::Set{Int} + fail_remaining::Int +end + +@kwdef mutable struct FailOnDemandLSOptions <: MadNLP.AbstractOptions + looped_linear_solver::Type = MadNLP.LDLSolver +end + +MadNLP.default_options(::Type{FailOnDemandLS}) = FailOnDemandLSOptions() + +function FailOnDemandLS(aug_com, nzvals_mat::AbstractMatrix{T}, n::Int; + opt=FailOnDemandLSOptions()) where T + bs = size(nzvals_mat, 2) + nnz_csc = size(nzvals_mat, 1) + VT = typeof(similar(nzvals_mat, T, 0)) + solvers = map(1:bs) do i + nzval_i = MadIPM._madnlp_unsafe_column_wrap(nzvals_mat, nnz_csc, (i-1)*nnz_csc+1, VT) + csc_i = MadIPM._csc_with_nzval(aug_com, nzval_i, n) + opt.looped_linear_solver(csc_i; opt=MadNLP.default_options(opt.looped_linear_solver)) + end + FailOnDemandLS{T,VT,eltype(solvers)}(solvers, bs, Set{Int}(), 0) +end + +MadIPM.is_factorized(s::FailOnDemandLS) = all(MadIPM.is_factorized(sj) for sj in s.solvers) + +function MadIPM.is_factorized!(buf::Vector{Int32}, s::FailOnDemandLS, v::MadIPM.BatchView) + nf = 0 + @inbounds for j in 1:v.n + if !MadIPM.is_factorized(s.solvers[j]) + nf += 1 + buf[nf] = j + end + end + return nf +end + +function MadIPM.factorize_active!(s::FailOnDemandLS, v::MadIPM.BatchView) + fail = s.fail_remaining > 0 + if fail + s.fail_remaining -= 1 + end + @inbounds for j in 1:v.n + if fail && j in s.fail_positions + nz = s.solvers[j].tril.nzval + saved = nz[1] + nz[1] = 0.0 + MadNLP.factorize!(s.solvers[j]) + nz[1] = saved + else + MadNLP.factorize!(s.solvers[j]) + end + end +end + +function MadIPM.solve_active!(s::FailOnDemandLS{T,VT}, rhs::AbstractMatrix{T}, v::MadIPM.BatchView) where {T,VT} + n = size(rhs, 1) + @inbounds for j in 1:v.n + rhs_j = MadIPM._madnlp_unsafe_column_wrap(rhs, n, (j-1)*n+1, VT) + MadNLP.solve_linear_system!(s.solvers[j], rhs_j) + end +end diff --git a/test/batch/gpu.jl b/test/batch/gpu.jl new file mode 100644 index 00000000..d259a3db --- /dev/null +++ b/test/batch/gpu.jl @@ -0,0 +1,125 @@ +using KernelAbstractions +using MadNLPGPU +using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel + +function _gpu_batch(qps; Model=ObjRHSBatchQuadraticModel, atol=1e-6, batch_kwargs...) + bs = length(qps) + refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] + for r in refs + @test r.status == MadNLP.SOLVE_SUCCEEDED + end + cpu_bnlp = Model(qps) + + # TODO: make this consistent in BQM + CuModel = Model{Float64, Model <: ObjRHSBatchQuadraticModel ? CuVector{Float64} : CuMatrix{Float64}} + gpu_bnlp = convert(CuModel, cpu_bnlp) + stats = MadIPM.madipm_batch(gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + batch_kwargs...) + CUDA.@allowscalar for i in 1:bs + @test stats[i].status == MadNLP.SOLVE_SUCCEEDED + @test stats[i].objective ≈ refs[i].objective atol=atol + @test Array(stats[i].solution) ≈ refs[i].solution atol=atol + end +end + +@testset "Batch solver (CUDA)" begin + +@testset "GPU gather/scatter" begin + cpu_bnlp = BatchQuadraticModel([_lp() for _ in 1:4]) + gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) + solver = MadIPM.UniformBatchMPCSolver(gpu_bnlp; + print_level=MadNLP.ERROR, uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver) + bvs = solver.batch_views + + saved = MadIPM.select_local!(bvs, [2, 4]) + child = MadIPM.active_view(bvs) + + src = cu(reshape(collect(1.0:12.0), 3, 4)) + gathered = similar(src, 3, 2) + MadIPM.gather_batch_view_columns!(gathered, src, child) + @test Array(gathered) == Array(src[:, [2, 4]]) + + scattered = CUDA.fill(-1.0, 3, 4) + MadIPM.scatter_batch_view_columns!(scattered, gathered, child) + @test Array(scattered[:, [2, 4]]) == Array(src[:, [2, 4]]) + @test all(Array(scattered[:, [1, 3]]) .== -1.0) + + MadIPM.restore_state!(bvs, saved) +end + +@testset "ObjRHSBatch LP bs=4" begin + _gpu_batch([_lp() for _ in 1:4]; atol=1e-5) +end + +@testset "ObjRHSBatch QP bs=3" begin + _gpu_batch([_qp() for _ in 1:3]) +end + +@testset "ObjRHSBatch QP doubly bs=2" begin + _gpu_batch([_qp_db() for _ in 1:2]) +end + +@testset "ObjRHSBatch QP dense bs=2" begin + _gpu_batch([_qp_dense() for _ in 1:2]) +end + +@testset "ObjRHSBatch different data bs=3" begin + _gpu_batch([ + QuadraticModel([1.0,1.0], Int[], Int[], Float64[]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=ones(2)), + QuadraticModel([2.0,0.5], Int[], Int[], Float64[]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[2.0], ucon=[2.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=ones(2)), + QuadraticModel([0.5,3.0], Int[], Int[], Float64[]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[0.5], ucon=[0.5], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=ones(2)), + ]; atol=1e-5) +end + +@testset "bs=1" begin + _gpu_batch([_lp()]; atol=1e-5) +end + +@testset "bs=8" begin + _gpu_batch([_qp() for _ in 1:8]) +end + +@testset "FullBatch identical QP bs=3" begin + _gpu_batch([_qp() for _ in 1:3]; Model=BatchQuadraticModel) +end + +@testset "FullBatch different H/A bs=2" begin + Hrows = [1,2,2]; Hcols = [1,1,2]; Arows = [1,1]; Acols = [1,2] + qps = [ + QuadraticModel([1.0,-1.0], Hrows, Hcols, [4.0,2.0,3.0]; + Arows=Arows, Acols=Acols, Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=[0.5,0.5]), + QuadraticModel([-1.0,2.0], Hrows, Hcols, [6.0,1.0,5.0]; + Arows=Arows, Acols=Acols, Avals=[1.5,0.5], + lcon=[2.0], ucon=[2.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=[1.0,1.0]), + ] + _gpu_batch(qps; Model=BatchQuadraticModel) +end + +@testset "MehrotraAdaptiveStep" begin + _gpu_batch([_qp() for _ in 1:3]; step_rule=MadIPM.MehrotraAdaptiveStep(0.99)) +end + +@testset "residual check INTERNAL_ERROR" begin + cpu_bnlp = ObjRHSBatchQuadraticModel([_lp() for _ in 1:3]) + gpu_bnlp = convert(ObjRHSBatchQuadraticModel{Float64, CuVector{Float64}}, cpu_bnlp) + stats = MadIPM.madipm_batch(gpu_bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, + cudss_algorithm=MadNLP.LDL, + check_residual=true, tol_linear_solve=0.0) + CUDA.@allowscalar for i in 1:3 + @test stats[i].status == MadNLP.INTERNAL_ERROR + end +end + +end diff --git a/test/batch/solver.jl b/test/batch/solver.jl new file mode 100644 index 00000000..3f49f09e --- /dev/null +++ b/test/batch/solver.jl @@ -0,0 +1,518 @@ +using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel +using LinearAlgebra +using SparseArrays + +_lp() = QuadraticModel(ones(2), Int[], Int[], Float64[]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=ones(2)) + +_qp() = QuadraticModel([1.0,-2.0,0.5,1.0], [1,2,3,4], [1,2,3,4], [2.0,1.0,3.0,1.5]; + Arows=[1,1,2,2], Acols=[1,2,3,4], Avals=[1.0,1.0,1.0,1.0], + lcon=[1.0,0.5], ucon=[2.0,1.5], lvar=zeros(4), uvar=fill(Inf,4), x0=ones(4)) + +_qp_ub() = QuadraticModel([1.0,-1.0], [1,2], [1,2], [1.0,1.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[-Inf,-Inf], uvar=[5.0,5.0], x0=[2.5,2.5]) + +_qp_db() = QuadraticModel([1.0,-1.0], [1,2], [1,2], [1.0,1.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[5.0,5.0], x0=[2.5,2.5]) + +_qp_free() = QuadraticModel([1.0,-1.0], [1,2], [1,2], [1.0,1.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[-Inf,-Inf], uvar=[Inf,Inf], x0=[0.5,0.5]) + +_qp_scaled() = QuadraticModel([1e3,-2e3,5e2,1e3], [1,2,3,4], [1,2,3,4], [2e3,1e3,3e3,1.5e3]; + Arows=[1,1,2,2], Acols=[1,2,3,4], Avals=[1e3,1e3,1e3,1e3], + lcon=[1e3,5e2], ucon=[2e3,1.5e3], lvar=zeros(4), uvar=fill(Inf,4), x0=ones(4)) + +_qp_fixed() = QuadraticModel([1.0,1.0,1.0], [1,2,3], [1,2,3], [2.0,1.0,3.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[Inf], lvar=[0.0,0.0,2.0], uvar=[Inf,Inf,2.0], x0=[1.0,1.0,1.0]) + +_qp_mixed() = QuadraticModel([1.0,-1.0,0.5,1.0], [1,2,3,4], [1,2,3,4], [2.0,1.0,3.0,1.5]; + Arows=[1,1,2,2], Acols=[1,2,3,4], Avals=[1.0,1.0,1.0,1.0], + lcon=[1.0,0.5], ucon=[2.0,1.5], lvar=[0.0,-Inf,0.0,-Inf], uvar=[Inf,5.0,5.0,Inf], x0=ones(4)) + +_qp_dense() = QuadraticModel([1.0,-1.0], [1,2,2], [1,1,2], [4.0,2.0,3.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=[0.5,0.5]) + +_qp_A() = QuadraticModel([1.0,-2.0,0.5,1.0], [1,2,3,4], [1,2,3,4], [2.0,1.0,3.0,1.5]; + Arows=[1,1,2,2], Acols=[1,2,3,4], Avals=[1.0,1.0,1.0,1.0], + lcon=[1.0,0.5], ucon=[2.0,1.5], lvar=zeros(4), uvar=fill(Inf,4), x0=ones(4)) +_qp_B() = QuadraticModel([2.0,1.0,-1.0,0.5], [1,2,3,4], [1,2,3,4], [5.0,2.0,1.0,4.0]; + Arows=[1,1,2,2], Acols=[1,2,3,4], Avals=[2.0,1.0,0.5,1.0], + lcon=[1.0,0.5], ucon=[2.0,1.5], lvar=zeros(4), uvar=fill(Inf,4), x0=ones(4)) + +const PROBLEMS = [ + "LP" => _lp, + "QP mixed" => _qp_mixed, + "QP dense" => _qp_dense, + "QP" => _qp, + "QP ub-only" => _qp_ub, + "QP doubly" => _qp_db, + "QP free" => _qp_free, + "QP scaled" => _qp_scaled, + "QP fixed" => _qp_fixed, +] + +const PROBLEMS_3 = PROBLEMS[1:3] + +_solve(qp; kw...) = MadIPM.solve!(MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, rethrow_error=true, kw...)) +_batch(qps; kw...) = MadIPM.madipm_batch(ObjRHSBatchQuadraticModel(qps); print_level=MadNLP.ERROR, rethrow_error=true, kw...) +_fbatch(qps; kw...) = MadIPM.madipm_batch(BatchQuadraticModel(qps); print_level=MadNLP.ERROR, rethrow_error=true, kw...) + +function _init_solver(solver) + ws, bcb, opt = solver.workspace, solver.bcb, solver.opt + MadNLP.initialize!(bcb, solver.x, solver.xl, solver.xu, MadNLP.full(solver.y), + MadNLP.full(solver.rhs), bcb.ind_ineq, ws.bx; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(MadNLP.full(solver.jacl), 0.0) + if opt.scaling + MadNLP.set_scaling!(bcb, solver.x, solver.xl, solver.xu, MadNLP.full(solver.y), + MadNLP.full(solver.rhs), bcb.ind_ineq, Float64(opt.nlp_scaling_max_gradient), ws.bx) + end + MadNLP.initialize!(solver.kkt) + MadIPM.init_regularization!(solver, opt.regularization) + MadNLP.unpack_x!(ws.bx, bcb, solver.x) + MadNLP.eval_f_wrapper(solver, ws.bx) + MadNLP.eval_jac_wrapper!(solver, solver.kkt) + MadNLP.eval_grad_f_wrapper!(solver, ws.bx) + MadNLP.eval_cons_wrapper!(solver, ws.bx) + MadNLP.eval_lag_hess_wrapper!(solver, solver.kkt) + ws.norm_b .= maximum(abs, MadNLP.full(solver.rhs); dims=1) + ws.norm_c .= maximum(abs, MadNLP.full(solver.f); dims=1) + MadIPM.init_starting_point!(solver) + fill!(ws.mu_batch, opt.mu_init) + fill!(ws.best_complementarity, typemax(Float64)) + fill!(ws.status, MadNLP.REGULAR) + fill!(ws.inf_pr, 0.0) + fill!(ws.inf_du, 0.0) + fill!(ws.inf_compl, 0.0) + fill!(ws.dual_obj, 0.0) + fill!(ws.alpha_p, 0.0) + fill!(ws.alpha_d, 0.0) + solver.batch_cnt.start_time[] = time() + fill!(solver.batch_cnt.k, 0) + MadNLP.jtprod!(solver.jacl, solver.kkt, solver.y) + return solver +end + +function _build_bat(qp; kwargs...) + _init_solver(MadIPM.UniformBatchMPCSolver(ObjRHSBatchQuadraticModel([qp]); print_level=MadNLP.ERROR, kwargs...)) +end + +function _build_bat_n(qp, n::Int; kwargs...) + _init_solver(MadIPM.UniformBatchMPCSolver(ObjRHSBatchQuadraticModel([qp for _ in 1:n]); print_level=MadNLP.ERROR, kwargs...)) +end + +function _build_seq(qp; kwargs...) + solver = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, kwargs...) + opt = solver.opt + MadNLP.initialize!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, solver.ind_ineq; + tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) + fill!(solver.jacl, 0.0) + if opt.scaling + MadNLP.set_scaling!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, + solver.ind_ineq, Float64(opt.nlp_scaling_max_gradient)) + end + MadNLP.initialize!(solver.kkt) + MadIPM.init_regularization!(solver, opt.regularization) + MadNLP.eval_f_wrapper(solver, solver.x) + MadNLP.eval_jac_wrapper!(solver, solver.kkt, solver.x) + MadNLP.eval_grad_f_wrapper!(solver, solver.f, solver.x) + MadNLP.eval_cons_wrapper!(solver, solver.c, solver.x) + MadNLP.eval_lag_hess_wrapper!(solver, solver.kkt, solver.x, solver.y) + solver.norm_b = norm(solver.rhs, Inf) + solver.norm_c = norm(MadNLP.primal(solver.f), Inf) + MadIPM.init_starting_point!(solver) + solver.mu = opt.mu_init + solver.cnt.start_time = time() + solver.best_complementarity = typemax(Float64) + solver.status = MadNLP.REGULAR + MadNLP.jtprod!(solver.jacl, solver.kkt, solver.y) + return solver +end + +function cmp(a, b) + d = 0.0 + for (ai, bi) in zip(a, b) + if isfinite(ai) && isfinite(bi) + d = max(d, abs(ai - bi) / max(abs(ai), abs(bi), 1.0)) + elseif ai !== bi + return Inf + end + end + return d +end +col1(x::AbstractMatrix) = view(x, :, 1) +col1(x::AbstractVector) = x + +function _do_factorize!(seq, bat) + MadIPM.update_regularization!(seq, seq.opt.regularization) + MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) + MadNLP.build_kkt!(seq.kkt) + MadNLP.factorize_kkt!(seq.kkt) + MadIPM.update_regularization!(bat, bat.opt.regularization) + MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) + MadNLP.build_kkt!(bat.kkt) + MadNLP.factorize_kkt!(bat.kkt) +end + +include("fakels.jl") + +@testset "Batch solver" begin + +@testset "e2e bs=1: $name" for (name, fn) in PROBLEMS + s = _solve(fn()) + b = _batch([fn()]) + @test b.status[1] == s.status + @test isapprox(b.objective[1], s.objective; rtol=1e-6) + @test isapprox(b.solution[:,1], s.solution; atol=1e-6) +end + +@testset "e2e MehrotraAdaptiveStep: $name" for (name, fn) in [PROBLEMS[1], PROBLEMS[3], PROBLEMS[6]] + qp = fn() + s = _solve(qp; step_rule=MadIPM.MehrotraAdaptiveStep(0.99)) + b = _batch([qp]; step_rule=MadIPM.MehrotraAdaptiveStep(0.99)) + @test b.status[1] == s.status + @test isapprox(b.solution[:,1], s.solution; atol=1e-6) +end + +@testset "e2e heterogeneous FullBatch" begin + qps = [_qp_A(), _qp_A(), _qp_B()] + refs = [_solve(qp) for qp in qps] + stats = _fbatch(qps) + for i in 1:3 + @test stats[i].status == MadNLP.SOLVE_SUCCEEDED + @test stats[i].objective ≈ refs[i].objective atol=1e-6 + end +end + +@testset "e2e maximization QP" begin + qp = QuadraticModel([2.0], [1], [1], [2.0]; + Arows=[1], Acols=[1], Avals=[1.0], + lcon=[-Inf], ucon=[5.0], lvar=[0.0], uvar=[Inf], x0=[2.5], minimize=false) + s = _solve(qp) + b = _batch([qp]) + @test b.status[1] == s.status + @test isapprox(b.solution[:,1], s.solution; atol=1e-5) +end + +@testset "e2e fixed variable (MakeParameter)" begin + s = _solve(_qp_fixed(); fixed_variable_treatment=MadNLP.MakeParameter) + b = _batch([_qp_fixed()]; fixed_variable_treatment=MadNLP.MakeParameter) + @test isapprox(b.solution[:,1], s.solution; atol=1e-6) +end + +@testset "kernel parity: $name" for (name, fn) in PROBLEMS_3 + qp = fn() + seq = _build_seq(qp) + bat = _build_bat(qp) + + @test cmp(MadNLP.full(seq.x), col1(MadNLP.full(bat.x))) < 1e-10 + @test cmp(seq.y, col1(MadNLP.full(bat.y))) < 1e-10 + @test cmp(MadNLP.full(seq.zl), col1(MadNLP.full(bat.zl))) < 1e-10 + + _do_factorize!(seq, bat) + + @test cmp(seq.kkt.reg, col1(bat.kkt.reg)) < 1e-12 + @test cmp(seq.kkt.pr_diag, col1(MadIPM.pr_diag(bat.kkt))) < 1e-12 + @test cmp(seq.kkt.aug_raw.V, bat.kkt.nzVals[:,1]) < 1e-12 + @test cmp(SparseArrays.nonzeros(seq.kkt.aug_com), bat.kkt.aug_com_nzvals[:,1]) < 1e-12 + + MadIPM.set_predictive_rhs!(seq, seq.kkt) + MadIPM.set_predictive_rhs!(bat, bat.kkt) + @test cmp(MadNLP.full(seq.p), col1(MadNLP.full(bat.p))) < 1e-12 + + MadIPM.solve_system!(seq.d, seq, seq.p) + MadIPM.solve_system!(bat.d, bat, bat.p) + @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 + + w_s, w_b = seq._w1, bat._w1 + fill!(MadNLP.full(w_s), 0.0) + fill!(MadNLP.full(w_b), 0.0) + mul!(w_s, seq.kkt, seq.d) + mul!(w_b, bat.kkt, bat.d) + @test cmp(MadNLP.full(w_s), col1(MadNLP.full(w_b))) < 1e-10 + + seq_mu = MadIPM.get_complementarity_measure(seq) + MadIPM.get_complementarity_measure!(bat) + @test abs(seq_mu - bat.workspace.mu_curr[1]) < 1e-12 +end + +@testset "regularization" begin + qp = _qp() + for (reg1, reg2) in [ + (MadIPM.NoRegularization(), MadIPM.NoRegularization()), + (MadIPM.FixedRegularization(1e-8,-1e-9), MadIPM.FixedRegularization(1e-8,-1e-9)), + (MadIPM.AdaptiveRegularization(1e-8,-1e-9,1e-9), MadIPM.AdaptiveRegularization(1e-8,-1e-9,1e-9)), + ] + seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=reg1) + bat = MadIPM.UniformBatchMPCSolver(ObjRHSBatchQuadraticModel([qp]); print_level=MadNLP.ERROR, regularization=reg2) + MadIPM.init_regularization!(seq, reg1) + MadIPM.init_regularization!(bat, reg2) + @test seq.del_w == bat.del_w[1] + for _ in 1:3 + MadIPM.update_regularization!(seq, reg1) + MadIPM.update_regularization!(bat, reg2) + @test abs(seq.del_w - bat.del_w[1]) < 1e-15 + end + end +end + +@testset "bs=3 consistency: $name" for (name, fn) in PROBLEMS_3 + bat1 = _build_bat(fn()) + bat3 = _build_bat_n(fn(), 3) + _do_factorize!(bat1, bat3) + MadIPM.prediction_step!(bat1) + MadIPM.prediction_step!(bat3) + MadIPM.mehrotra_correction_direction!(bat1) + MadIPM.mehrotra_correction_direction!(bat3) + MadIPM.update_step!(bat1.opt.step_rule, bat1) + MadIPM.update_step!(bat3.opt.step_rule, bat3) + MadIPM.apply_step!(bat1) + MadIPM.apply_step!(bat3) + for i in 1:3 + @test cmp(col1(MadNLP.full(bat1.x)), view(MadNLP.full(bat3.x),:,i)) < 1e-10 + end + @test bat3.workspace.alpha_p[1] == bat3.workspace.alpha_p[2] == bat3.workspace.alpha_p[3] +end + +@testset "staggered frozen: $name" for (name, fn) in PROBLEMS_3 + bat = _build_bat_n(fn(), 3) + ws = bat.workspace + for _ in 1:5 + MadIPM.mpc_step!(bat) + end + snap_x = copy(MadNLP.full(bat.x)[:,2]) + snap_y = copy(MadNLP.full(bat.y)[:,2]) + snap_nzv = copy(bat.kkt.nzVals[:,2]) + + ws.status[2] = MadNLP.SOLVE_SUCCEEDED + MadIPM.update_active_set!(bat) + MadIPM._update_active_mask!(bat) + for _ in 1:5 + MadIPM.mpc_step!(bat) + end + @test MadNLP.full(bat.x)[:,2] ≈ snap_x atol=1e-14 + @test MadNLP.full(bat.y)[:,2] ≈ snap_y atol=1e-14 + @test bat.kkt.nzVals[:,2] ≈ snap_nzv atol=1e-14 + @test ws.alpha_p[1,1] > 0 || ws.status[1] != MadNLP.REGULAR +end + +@testset "termination" begin + bat = _build_bat(_qp()) + bat.opt.max_iter = 0 + MadIPM.update_termination_criteria!(bat) + MadIPM.update_termination_status!(bat) + @test bat.workspace.status[1] == MadNLP.MAXIMUM_ITERATIONS_EXCEEDED + + bat = _build_bat(_qp()) + bat.opt.max_wall_time = 0.0 + bat.batch_cnt.start_time[] = time() - 1.0 + MadIPM.update_termination_criteria!(bat) + MadIPM.update_termination_status!(bat) + @test bat.workspace.status[1] == MadNLP.MAXIMUM_WALLTIME_EXCEEDED +end + +@testset "compute_term_gpu!" begin + bat = _build_bat_n(_qp(), 2) + ws = bat.workspace + tol = bat.opt.tol + ds = bat.opt.divergence_scale + dt = bat.opt.divergence_tol + + function _set!(; pr=1.0, du=1.0, ic=1.0, best=1e10, obj=1.0, dobj=1.0, ls=Int32(0)) + ws.inf_pr .= pr + ws.inf_du .= du + ws.inf_compl .= ic + ws.best_complementarity .= best + ws.obj_val .= obj + ws.dual_obj .= dobj + fill!(ws._ls_error, ls) + end + _s(j) = MadNLP.Status(ws._term_gpu[1,j]) + + _set!(pr=tol/10, du=tol/10, ic=tol/10) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.SOLVE_SUCCEEDED + + _set!(ls=Int32(1)) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.INTERNAL_ERROR + + _set!(ic=1e10, best=1e-8, dobj=1e12, obj=1.0) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.INFEASIBLE_PROBLEM_DETECTED + + _set!(dobj=1.0, obj=-(dt*ds*2)) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.DIVERGING_ITERATES + + _set!(pr=tol/10, du=tol/10, ic=tol/10, ls=Int32(1)) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.INTERNAL_ERROR + + _set!() + ws.inf_pr[1,1] = tol/10 + ws.inf_du[1,1] = tol/10 + ws.inf_compl[1,1] = tol/10 + ws.obj_val[1,2] = -(dt*ds*2) + MadIPM.compute_term_gpu!(ws, bat.opt) + @test _s(1) == MadNLP.SOLVE_SUCCEEDED + @test _s(2) == MadNLP.DIVERGING_ITERATES +end + +@testset "structure mismatch" begin + qp1 = QuadraticModel([1.0,1.0], [1,2], [1,2], [2.0,2.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[1.0,1.0], x0=[0.5,0.5]) + qp2 = QuadraticModel([1.0,1.0], [1,2], [1,2], [2.0,2.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[0.0,1.0], x0=[0.0,0.5]) + @test_throws AssertionError MadIPM.UniformBatchMPCSolver(BatchQuadraticModel([qp1,qp2]); print_level=MadNLP.ERROR) + MadIPM.UniformBatchMPCSolver(BatchQuadraticModel([qp1,qp2]); print_level=MadNLP.ERROR, check_batch_structure=false) +end + +@testset "residual INTERNAL_ERROR (all fail)" begin + stats = _batch([_lp() for _ in 1:3]; check_residual=true, tol_linear_solve=0.0) + for i in 1:3 + @test stats[i].status == MadNLP.INTERNAL_ERROR + end +end + +@testset "residual INTERNAL_ERROR (partial NaN)" begin + good() = QuadraticModel([1.0,1.0], [1,2], [1,2], [2.0,2.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=[0.5,0.5]) + bad = QuadraticModel([NaN,NaN], [1,2], [1,2], [2.0,2.0]; + Arows=[1,1], Acols=[1,2], Avals=[1.0,1.0], + lcon=[1.0], ucon=[1.0], lvar=[0.0,0.0], uvar=[Inf,Inf], x0=[0.5,0.5]) + stats = _fbatch([good(), bad, good()]) + @test stats[1].status == MadNLP.SOLVE_SUCCEEDED + @test stats[3].status == MadNLP.SOLVE_SUCCEEDED + @test stats[2].status == MadNLP.INTERNAL_ERROR +end + +@testset "factorize retry" begin + function _mk(; fail_positions=Set{Int}(), fail_remaining=0) + bnlp = BatchQuadraticModel([_qp_A(), _qp_A(), _qp_B()]) + s = MadIPM.UniformBatchMPCSolver(bnlp; + print_level=MadNLP.ERROR, + uniformbatch_linear_solver=FailOnDemandLS, + looped_linear_solver=MadNLP.LDLSolver, + regularization=MadIPM.FixedRegularization(1e-8,-1e-9), + rethrow_error=true) + MadIPM.initialize!(s) + s.kkt.batch_solver.fail_positions = fail_positions + s.kkt.batch_solver.fail_remaining = fail_remaining + return s + end + + @testset "non-identity view" begin + ref = _mk() + ref.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED + MadIPM.update_active_set!(ref) + MadIPM._update_active_mask!(ref) + MadIPM.factorize_system!(ref) + MadIPM.set_predictive_rhs!(ref, ref.kkt) + copyto!(MadNLP.full(ref.d), MadNLP.full(ref.p)) + MadNLP.solve_kkt!(ref.kkt, ref) + ref_d1 = copy(MadNLP.primal(ref.d)[:,1]) + + bad = _mk(fail_positions=Set([2]), fail_remaining=1) + bad.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED + MadIPM.update_active_set!(bad) + MadIPM._update_active_mask!(bad) + MadIPM.factorize_system!(bad) + + @test MadIPM.is_factorized(bad.kkt.batch_solver.solvers[1]) + @test MadIPM.is_factorized(bad.kkt.batch_solver.solvers[2]) + + MadIPM.set_predictive_rhs!(bad, bad.kkt) + copyto!(MadNLP.full(bad.d), MadNLP.full(bad.p)) + MadNLP.solve_kkt!(bad.kkt, bad) + @test norm(ref_d1 - MadNLP.primal(bad.d)[:,1]) / norm(ref_d1) < 1e-6 + end + + @testset "identity view" begin + ref = _mk() + MadIPM.factorize_system!(ref) + MadIPM.set_predictive_rhs!(ref, ref.kkt) + copyto!(MadNLP.full(ref.d), MadNLP.full(ref.p)) + MadNLP.solve_kkt!(ref.kkt, ref) + ref_d1 = copy(MadNLP.primal(ref.d)[:,1]) + + bad = _mk(fail_positions=Set([3]), fail_remaining=1) + MadIPM.factorize_system!(bad) + @test all(MadIPM.is_factorized(bad.kkt.batch_solver.solvers[j]) for j in 1:3) + + MadIPM.set_predictive_rhs!(bad, bad.kkt) + copyto!(MadNLP.full(bad.d), MadNLP.full(bad.p)) + MadNLP.solve_kkt!(bad.kkt, bad) + @test norm(ref_d1 - MadNLP.primal(bad.d)[:,1]) / norm(ref_d1) < 1e-6 + end + + @testset "regularization bumped correctly" begin + bad = _mk(fail_positions=Set([2]), fail_remaining=1) + bad.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED + MadIPM.update_active_set!(bad) + MadIPM._update_active_mask!(bad) + MadIPM.factorize_system!(bad) + @test bad.del_w[1,1] ≈ 1e-8 + @test bad.del_w[1,2] ≈ 1.0 + @test bad.del_w[1,3] ≈ 1e-6 + end +end + +@testset "free vars: mu = mu_min" begin + bat = _build_bat(_qp_free()) + MadIPM.get_complementarity_measure!(bat) + @test bat.workspace.mu_curr[1] == 0.0 + MadIPM.update_barrier!(MadIPM.Mehrotra(), bat, bat.workspace.mu_affine) + @test bat.workspace.mu_batch[1] == bat.opt.mu_min +end + +@testset "_adjust_boundary_active!" begin + qp = _qp_mixed() + seq = _build_seq(qp) + bat = _build_bat(qp) + seq.mu = 0.1 + bat.workspace.mu_batch .= 0.1 + bat.workspace.active_mask .= 1.0 + MadNLP.adjust_boundary!(seq.x_lr, seq.xl_r, seq.x_ur, seq.xu_r, seq.mu) + MadIPM._adjust_boundary_active!( + MadIPM.lower(bat.x), MadIPM.lower(bat.xl), + MadIPM.upper(bat.x), MadIPM.upper(bat.xu), + bat.workspace.mu_batch, bat.workspace.active_mask) + @test cmp(seq.xl_r, col1(MadIPM.lower(bat.xl))) < 1e-14 + @test cmp(seq.xu_r, col1(MadIPM.upper(bat.xu))) < 1e-14 + + bat2 = _build_bat_n(_qp(), 2) + bat2.workspace.mu_batch .= 1e-12 + bat2.workspace.active_mask .= 0.0 + xl_before = copy(MadIPM.lower(bat2.xl)) + MadIPM._adjust_boundary_active!( + MadIPM.lower(bat2.x), MadIPM.lower(bat2.xl), + MadIPM.upper(bat2.x), MadIPM.upper(bat2.xu), + bat2.workspace.mu_batch, bat2.workspace.active_mask) + @test MadIPM.lower(bat2.xl) == xl_before +end + +@testset "partial active KKT solve preserves rhs" begin + solver = MadIPM.UniformBatchMPCSolver(BatchQuadraticModel([_lp() for _ in 1:3]); print_level=MadNLP.ERROR) + MadIPM.initialize!(solver) + solver.workspace.status[2] = MadNLP.INTERNAL_ERROR + MadIPM.update_active_set!(solver) + pd = MadNLP.primal_dual(solver.d) + pd .= reshape(collect(1.0:length(pd)), size(pd)) + pd_before = copy(pd) + MadNLP.build_kkt!(solver.kkt) + MadNLP.factorize_kkt!(solver.kkt) + MadNLP.solve_kkt!(solver.kkt, solver) + @test pd[:,2] == pd_before[:,2] +end + +end diff --git a/test/batch/views.jl b/test/batch/views.jl new file mode 100644 index 00000000..6f4d7058 --- /dev/null +++ b/test/batch/views.jl @@ -0,0 +1,112 @@ +struct _MockCB end # to make BatchViewState without an NLPModel +MadNLP.create_array(::_MockCB, ::Type{T}, n::Int) where T = Vector{T}(undef, n) +_bvs(bs) = MadIPM.BatchViewState(_MockCB(), bs) + +@testset "Batch views" begin + +@testset "select_local!" begin + bvs = _bvs(5) + @test MadIPM.is_identity_view(MadIPM.root_view(bvs)) + + MadIPM.select_local!(bvs, Int32[2, 4]) + @test MadIPM.active_view(bvs).local_to_root[1:2] == Int32[2, 4] + + MadIPM.select_local!(bvs, Int32[2]) + @test MadIPM.active_view(bvs).local_to_root[1] == Int32(4) +end + +@testset "select_local! by bool" begin + bvs = _bvs(5) + MadIPM.select_local!(bvs, Bool[false, true, false, true, true]) + @test MadIPM.active_view(bvs).local_to_root[1:3] == Int32[2, 4, 5] +end + +@testset "select_local! empty" begin + bvs = _bvs(3) + MadIPM.select_local!(bvs, Int32[]) + @test MadIPM.local_batch_size(MadIPM.active_view(bvs)) == 0 +end + +@testset "exclude_local!" begin + bvs = _bvs(5) + MadIPM.exclude_local!(bvs, Bool[false, true, false, true, false]) + @test MadIPM.active_view(bvs).local_to_root[1:3] == Int32[1, 3, 5] + + bvs = _bvs(3) + MadIPM.exclude_local!(bvs, Bool[true, true, true]) + @test MadIPM.local_batch_size(MadIPM.active_view(bvs)) == 0 + + bvs = _bvs(3) + MadIPM.exclude_local!(bvs, Bool[false, false, false]) + @test MadIPM.is_identity_view(MadIPM.active_view(bvs)) +end + +@testset "restore_state! and reset_active_view!" begin + bvs = _bvs(4) + saved_root = MadIPM.select_local!(bvs, Int32[1, 3]) + saved_mid = MadIPM.select_local!(bvs, Int32[2]) + @test MadIPM.active_view(bvs).local_to_root[1] == Int32(3) + + MadIPM.restore_state!(bvs, saved_mid) + @test MadIPM.active_view(bvs).local_to_root[1:2] == Int32[1, 3] + + MadIPM.restore_state!(bvs, saved_root) + @test MadIPM.is_identity_view(MadIPM.active_view(bvs)) + + MadIPM.select_local!(bvs, Int32[2, 3]) + MadIPM.reset_active_view!(bvs) + @test MadIPM.is_identity_view(MadIPM.active_view(bvs)) +end + +@testset "fill_batch_view_mask!" begin + bvs = _bvs(5) + MadIPM.select_local!(bvs, Int32[2, 4]) + mask = zeros(Float64, 1, 5) + MadIPM.fill_batch_view_mask!(mask, MadIPM.active_view(bvs)) + @test mask == [0.0 1.0 0.0 1.0 0.0] +end + +@testset "local_to_root_dev" begin + bvs = _bvs(4) + MadIPM.select_local!(bvs, Int32[2, 4]) + @test MadIPM.local_to_root_dev(MadIPM.active_view(bvs))[1:2] == Int32[2, 4] + + MadIPM.select_local!(bvs, Int32[1]) + @test MadIPM.local_to_root_dev(MadIPM.active_view(bvs))[1] == Int32(2) +end + +@testset "gather/scatter/compact" begin + bvs = _bvs(5) + MadIPM.select_local!(bvs, Int32[2, 4]) + v = MadIPM.active_view(bvs) + src = reshape(collect(1.0:15.0), 3, 5) + + gathered = zeros(3, 2) + MadIPM.gather_batch_view_columns!(gathered, src, v) + @test gathered == src[:, [2, 4]] + + dst = fill(-1.0, 3, 5) + MadIPM.scatter_batch_view_columns!(dst, gathered, v) + @test dst[:, [2, 4]] == src[:, [2, 4]] + @test all(dst[:, [1, 3, 5]] .== -1.0) + + data = copy(src) + MadIPM.compact_active_columns_inplace!(data, v) + @test data[:, 1:2] == src[:, [2, 4]] +end + +@testset "gather then scatter round-trips" begin + bvs = _bvs(5) + MadIPM.select_local!(bvs, Int32[1, 3, 5]) + v = MadIPM.active_view(bvs) + + src = reshape(collect(1.0:10.0), 2, 5) + gathered = zeros(2, 3) + MadIPM.gather_batch_view_columns!(gathered, src, v) + dst = zeros(2, 5) + MadIPM.scatter_batch_view_columns!(dst, gathered, v) + @test dst[:, [1, 3, 5]] == src[:, [1, 3, 5]] + @test dst[:, [2, 4]] == zeros(2, 2) +end + +end diff --git a/test/batch_kernels/runtests.jl b/test/batch_kernels/runtests.jl deleted file mode 100644 index 4cc433d8..00000000 --- a/test/batch_kernels/runtests.jl +++ /dev/null @@ -1,10 +0,0 @@ -@testset "Batch kernel tests" begin - include("setup.jl") - include("test_rhs.jl") - include("test_kkt.jl") - include("test_complementarity.jl") - include("test_step.jl") - include("test_regularization.jl") - include("test_batch_multi.jl") - include("test_termination.jl") -end diff --git a/test/batch_kernels/setup.jl b/test/batch_kernels/setup.jl deleted file mode 100644 index 20f5da74..00000000 --- a/test/batch_kernels/setup.jl +++ /dev/null @@ -1,340 +0,0 @@ -using Test -using LinearAlgebra -using SparseArrays -using MadNLP -using MadIPM -using QuadraticModels -using BatchQuadraticModels: ObjRHSBatchQuadraticModel - -# ────────────────────────────────────────────────────────────── -# Test problems -# ────────────────────────────────────────────────────────────── - -# LP: n=2, m=1, no upper bounds (nub=0) -function _setup_simple_lp() - c = ones(2) - Hrows = Int[]; Hcols = Int[]; Hvals = Float64[] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - QuadraticModel(c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - c0=0.0, x0=ones(2), name="simpleLP") -end - -# QP: n=4, m=2, with finite upper bounds (nub>0) -function _setup_small_qp() - n, m = 4, 2 - c = [1.0, -2.0, 0.5, 1.0] - Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 3.0, 1.5] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 3, 4]; Avals = [1.0, 1.0, 1.0, 1.0] - QuadraticModel(c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[2.0, 1.5], - lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n)) -end - -# QP with only upper bounds (nlb=0, nub>0) -function _setup_upper_only_qp() - QuadraticModel( - [1.0, -1.0], - [1, 2], [1, 2], [1.0, 1.0], - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[-Inf, -Inf], uvar=[5.0, 5.0], - x0=[2.5, 2.5], - ) -end - -# QP with doubly-bounded variables (both lvar and uvar finite) -# Regression test for has_inequalities bug where ind_llb/ind_uub are empty -# but nlb+nub > 0 (variables have bound multipliers on both sides) -function _setup_doubly_bounded_qp() - QuadraticModel( - [1.0, -1.0], - [1, 2], [1, 2], [1.0, 1.0], - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[5.0, 5.0], - x0=[2.5, 2.5], - ) -end - -# Free-variable QP (nlb=0, nub=0, no bound multipliers) -function _setup_free_qp() - QuadraticModel( - [1.0, -1.0], - [1, 2], [1, 2], [2.0, 2.0], - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[-Inf, -Inf], uvar=[Inf, Inf], - x0=[0.5, 0.5], - ) -end - -# QP with large coefficients that trigger non-unit scaling (obj_scale, con_scale < 1) -function _setup_scaled_qp() - QuadraticModel( - [500.0, -300.0, 400.0], - [1, 2, 3], [1, 2, 3], [1.0, 1.0, 1.0], - Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[200.0, 150.0, 100.0], - lcon=[1.0], ucon=[1.0], - lvar=zeros(3), uvar=fill(Inf, 3), - x0=ones(3), - ) -end - -# QP with large coefficients + inequality constraints (exercises slack scaling) -function _setup_scaled_ineq_qp() - QuadraticModel( - [500.0, -300.0, 400.0], - [1, 2, 3], [1, 2, 3], [1.0, 1.0, 1.0], - Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[200.0, 150.0, 100.0], - lcon=[0.5], ucon=[2.0], - lvar=zeros(3), uvar=fill(Inf, 3), - x0=ones(3), - ) -end - -# QP with a fixed variable (lvar[1]==uvar[1]) exercising MakeParameter -function _setup_fixed_var_qp() - QuadraticModel( - [1.0, -1.0, 0.5], - [1, 2, 3], [1, 2, 3], [2.0, 1.0, 1.0], - Arows=[1, 1, 1], Acols=[1, 2, 3], Avals=[1.0, 1.0, 1.0], - lcon=[3.0], ucon=[3.0], - lvar=[2.0, 0.0, 0.0], uvar=[2.0, Inf, Inf], - x0=[2.0, 0.5, 0.5], - ) -end - -# QP with all inequality constraints (ns == m, full slack initialization) -function _setup_all_ineq_qp() - QuadraticModel( - [1.0, -2.0, 0.5], - [1, 2, 3], [1, 2, 3], [2.0, 1.0, 1.5], - Arows=[1, 1, 2, 2], Acols=[1, 2, 2, 3], Avals=[1.0, 1.0, 1.0, 1.0], - lcon=[0.0, 0.0], ucon=[3.0, 3.0], - lvar=zeros(3), uvar=fill(Inf, 3), - x0=ones(3), - ) -end - -# QP with mixed bound types (lower-only, upper-only, doubly-bounded, free) + inequality -function _setup_mixed_bounds_qp() - QuadraticModel( - [1.0, -1.0, 0.5, -0.5], - [1, 2, 3, 4], [1, 2, 3, 4], [2.0, 1.0, 1.5, 1.0], - Arows=[1, 1, 1, 1, 2, 2], Acols=[1, 2, 3, 4, 1, 3], - Avals=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - lcon=[1.0, 0.5], ucon=[3.0, 1.5], - lvar=[0.0, -Inf, 0.0, -Inf], uvar=[Inf, 5.0, 10.0, Inf], - x0=[1.0, 2.0, 5.0, 0.0], - ) -end - -# QP with non-diagonal (dense lower-triangular) Hessian -# H = [4 .; 2 3] (lower triangular of symmetric PD matrix) -function _setup_dense_hess_qp() - QuadraticModel( - [1.0, -1.0], - [1, 2, 2], [1, 1, 2], [4.0, 2.0, 3.0], - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=[0.5, 0.5], - ) -end - -# Larger QP with non-diagonal Hessian + inequality + mixed bounds -# H = [5 . .; 1 4 .; 0 2 3] (sparse lower-triangular, off-diagonal entries) -function _setup_dense_hess_mixed_qp() - QuadraticModel( - [1.0, -2.0, 0.5], - [1, 2, 2, 3, 3], [1, 1, 2, 2, 3], [5.0, 1.0, 4.0, 2.0, 3.0], - Arows=[1, 1, 1, 2, 2], Acols=[1, 2, 3, 1, 3], - Avals=[1.0, 1.0, 1.0, 1.0, 1.0], - lcon=[1.0, 0.5], ucon=[3.0, 1.5], - lvar=[0.0, -Inf, 0.0], uvar=[Inf, 5.0, 10.0], - x0=[1.0, 2.0, 1.0], - ) -end - -# All test problems with descriptions -const ALL_TEST_PROBLEMS = [ - ("LP (nlb>0, nub=0)", _setup_simple_lp), - ("QP (nlb>0, nub>0)", _setup_small_qp), - ("QP (nlb=0, nub>0)", _setup_upper_only_qp), - ("QP doubly-bounded", _setup_doubly_bounded_qp), - ("QP free vars (nlb=0, nub=0)", _setup_free_qp), - ("QP scaled", _setup_scaled_qp), - ("QP scaled+ineq", _setup_scaled_ineq_qp), - ("QP fixed var", _setup_fixed_var_qp), - ("QP all-ineq", _setup_all_ineq_qp), - ("QP mixed bounds", _setup_mixed_bounds_qp), - ("QP dense Hessian", _setup_dense_hess_qp), - ("QP dense Hessian+mixed", _setup_dense_hess_mixed_qp), -] - -# ────────────────────────────────────────────────────────────── -# Build initialized sequential solver -# ────────────────────────────────────────────────────────────── -function build_seq(qp; kwargs...) - solver = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, kwargs...) - opt = solver.opt - - MadNLP.initialize!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, solver.ind_ineq; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(solver.jacl, 0.0) - if opt.scaling - MadNLP.set_scaling!(solver.cb, solver.x, solver.xl, solver.xu, solver.y, solver.rhs, - solver.ind_ineq, Float64(opt.nlp_scaling_max_gradient)) - end - MadNLP.initialize!(solver.kkt) - MadIPM.init_regularization!(solver, opt.regularization) - - solver.obj_val = MadNLP.eval_f_wrapper(solver, solver.x) - MadNLP.eval_jac_wrapper!(solver, solver.kkt, solver.x) - MadNLP.eval_grad_f_wrapper!(solver, solver.f, solver.x) - MadNLP.eval_cons_wrapper!(solver, solver.c, solver.x) - MadNLP.eval_lag_hess_wrapper!(solver, solver.kkt, solver.x, solver.y) - solver.norm_b = norm(solver.rhs, Inf) - solver.norm_c = norm(MadNLP.primal(solver.f), Inf) - - MadIPM.init_starting_point!(solver) - solver.mu = opt.mu_init - solver.cnt.start_time = time() - solver.best_complementarity = typemax(Float64) - solver.status = MadNLP.REGULAR - MadNLP.jtprod!(solver.jacl, solver.kkt, solver.y) - return solver -end - -# ────────────────────────────────────────────────────────────── -# Build initialized batch solver (bs=1) -# ────────────────────────────────────────────────────────────── -function build_batch(qp; kwargs...) - bnlp = ObjRHSBatchQuadraticModel([qp]) - batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, kwargs...) - ws = batch_solver.workspace - bcb = batch_solver.bcb - opt = batch_solver.opt - - MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - batch_solver.workspace.bx; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(MadNLP.full(batch_solver.jacl), 0.0) - if opt.scaling - MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - Float64(opt.nlp_scaling_max_gradient), - batch_solver.workspace.bx) - end - MadNLP.initialize!(batch_solver.kkt) - MadIPM.init_regularization!(batch_solver, opt.regularization) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP.eval_f_wrapper(batch_solver, ws.bx) - MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) - MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) - MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) - MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) - - MadIPM.init_starting_point!(batch_solver) - fill!(ws.mu_batch, opt.mu_init) - fill!(ws.best_complementarity, typemax(Float64)) - fill!(ws.status, MadNLP.REGULAR) - fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) - fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) - fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) - batch_solver.batch_cnt.start_time[] = time() - fill!(batch_solver.batch_cnt.k, 0) - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return batch_solver -end - -# ────────────────────────────────────────────────────────────── -# Build initialized batch solver with batch_size > 1 -# ────────────────────────────────────────────────────────────── -function build_batch_n(qp, n::Int) - bnlp = ObjRHSBatchQuadraticModel([qp for _ in 1:n]) - batch_solver = MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR) - ws = batch_solver.workspace - bcb = batch_solver.bcb - opt = batch_solver.opt - - MadNLP.initialize!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - batch_solver.workspace.bx; - tol=opt.bound_relax_factor, bound_push=opt.bound_push, bound_fac=opt.bound_fac) - fill!(MadNLP.full(batch_solver.jacl), 0.0) - if opt.scaling - MadNLP.set_scaling!(bcb, batch_solver.x, batch_solver.xl, batch_solver.xu, - MadNLP.full(batch_solver.y), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, - Float64(opt.nlp_scaling_max_gradient), - batch_solver.workspace.bx) - end - MadNLP.initialize!(batch_solver.kkt) - MadIPM.init_regularization!(batch_solver, opt.regularization) - - MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - MadNLP.eval_f_wrapper(batch_solver, ws.bx) - MadNLP.eval_jac_wrapper!(batch_solver, batch_solver.kkt) - MadNLP.eval_grad_f_wrapper!(batch_solver, ws.bx) - MadNLP.eval_cons_wrapper!(batch_solver, ws.bx) - MadNLP.eval_lag_hess_wrapper!(batch_solver, batch_solver.kkt) - ws.norm_b .= maximum(abs, MadNLP.full(batch_solver.rhs); dims=1) - ws.norm_c .= maximum(abs, MadNLP.full(batch_solver.f); dims=1) - - MadIPM.init_starting_point!(batch_solver) - fill!(ws.mu_batch, opt.mu_init) - fill!(ws.best_complementarity, typemax(Float64)) - fill!(ws.status, MadNLP.REGULAR) - fill!(ws.inf_pr, 0.0); fill!(ws.inf_du, 0.0) - fill!(ws.inf_compl, 0.0); fill!(ws.dual_obj, 0.0) - fill!(ws.alpha_p, 0.0); fill!(ws.alpha_d, 0.0) - batch_solver.batch_cnt.start_time[] = time() - fill!(batch_solver.batch_cnt.k, 0) - MadNLP.jtprod!(batch_solver.jacl, batch_solver.kkt, batch_solver.y) - return batch_solver -end - -# ────────────────────────────────────────────────────────────── -# Comparison helper: max absolute difference -# ────────────────────────────────────────────────────────────── -# Safe comparison that handles Inf values (e.g. xu with uvar=Inf) -function cmp(a, b) - d = 0.0 - for (ai, bi) in zip(a, b) - if isfinite(ai) && isfinite(bi) - d = max(d, abs(ai - bi) / max(abs(ai), abs(bi), 1.0)) - elseif ai !== bi - return Inf - end - end - return d -end - -# Extract column 1 from a batch matrix/vector -col1(x::AbstractMatrix) = view(x, :, 1) -col1(x::AbstractVector) = x # already a vector (scalar workspace) - -# ────────────────────────────────────────────────────────────── -# Run first factorize_system to get post-factorization state -# ────────────────────────────────────────────────────────────── -function do_first_factorize!(seq, bat) - # Sequential - MadIPM.update_regularization!(seq, seq.opt.regularization) - MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) - MadNLP.build_kkt!(seq.kkt) - MadNLP.factorize_kkt!(seq.kkt) - # Batch - MadIPM.update_regularization!(bat, bat.opt.regularization) - MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) - MadNLP.build_kkt!(bat.kkt) - MadNLP.factorize_kkt!(bat.kkt) -end - diff --git a/test/batch_kernels/test_batch_multi.jl b/test/batch_kernels/test_batch_multi.jl deleted file mode 100644 index c69f3b27..00000000 --- a/test/batch_kernels/test_batch_multi.jl +++ /dev/null @@ -1,55 +0,0 @@ -@testset "Multi-batch (batch_size > 1)" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - @testset "batch_size=3 consistency" begin - # All instances identical: every column should match the bs=1 result - qp = make_qp() - bat1 = build_batch(qp) - bat3 = build_batch_n(qp, 3) - do_first_factorize!(bat1, bat3) - - # Run one full iteration on both - MadIPM.prediction_step!(bat1) - MadIPM.prediction_step!(bat3) - MadIPM.mehrotra_correction_direction!(bat1) - MadIPM.mehrotra_correction_direction!(bat3) - MadIPM.update_step!(bat1.opt.step_rule, bat1) - MadIPM.update_step!(bat3.opt.step_rule, bat3) - MadIPM.apply_step!(bat1) - MadIPM.apply_step!(bat3) - - # Every column of batch_size=3 should match the single-instance result - for i in 1:3 - @test cmp(col1(MadNLP.full(bat1.x)), view(MadNLP.full(bat3.x), :, i)) < 1e-10 - @test cmp(col1(MadNLP.full(bat1.y)), view(MadNLP.full(bat3.y), :, i)) < 1e-10 - @test cmp(col1(MadNLP.full(bat1.zl)), view(MadNLP.full(bat3.zl), :, i)) < 1e-10 - @test cmp(col1(MadNLP.full(bat1.zu)), view(MadNLP.full(bat3.zu), :, i)) < 1e-10 - end - @test bat3.workspace.alpha_p[1] == bat3.workspace.alpha_p[2] == bat3.workspace.alpha_p[3] - @test bat3.workspace.alpha_d[1] == bat3.workspace.alpha_d[2] == bat3.workspace.alpha_d[3] - end - - @testset "active-set deactivation" begin - qp = make_qp() - bat = build_batch_n(qp, 3) - - # Mark instance 2 as converged - bat.workspace.status[2] = MadNLP.SOLVE_SUCCEEDED - MadIPM.update_active_set!(bat) - active = MadIPM.active_view(bat.batch_views) - - @test MadIPM.local_batch_size(active) == 2 - @test Int[active.local_to_root[i] for i in 1:active.n] == [1, 3] - - # zero_inactive_step! should zero out the deactivated instance - fill!(bat.workspace.active_mask, 1.0) - MadIPM._update_active_mask!(bat) - fill!(bat.workspace.alpha_p, 0.5) - fill!(bat.workspace.alpha_d, 0.5) - MadIPM.zero_inactive_step!(bat) - @test bat.workspace.alpha_p[2] == 0.0 - @test bat.workspace.alpha_d[2] == 0.0 - @test bat.workspace.alpha_p[1] == 0.5 - @test bat.workspace.alpha_d[1] == 0.5 - end - end -end diff --git a/test/batch_kernels/test_complementarity.jl b/test/batch_kernels/test_complementarity.jl deleted file mode 100644 index 20cbf671..00000000 --- a/test/batch_kernels/test_complementarity.jl +++ /dev/null @@ -1,53 +0,0 @@ -@testset "Barrier / complementarity" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - - @testset "get_complementarity_measure!" begin - seq_mu = MadIPM.get_complementarity_measure(seq) - MadIPM.get_complementarity_measure!(bat) - bat_mu = bat.workspace.mu_curr[1] - @test abs(seq_mu - bat_mu) < 1e-12 - end - - @testset "get_affine_complementarity_measure!" begin - do_first_factorize!(seq, bat) - # Compute affine direction - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - # Compute alpha with tau=1 - seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) - fill!(bat.workspace.tau, 1.0) - MadIPM.get_fraction_to_boundary_step!(bat) - - seq_mu_aff = MadIPM.get_affine_complementarity_measure(seq, seq_ap, seq_ad) - MadIPM.get_affine_complementarity_measure!(bat, bat.workspace.alpha_p, bat.workspace.alpha_d) - bat_mu_aff = bat.workspace.mu_affine[1] - @test abs(seq_mu_aff - bat_mu_aff) < 1e-10 - end - - @testset "update_barrier! (Mehrotra)" begin - do_first_factorize!(seq, bat) - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) - fill!(bat.workspace.tau, 1.0) - MadIPM.get_fraction_to_boundary_step!(bat) - - seq_mu_aff = MadIPM.get_affine_complementarity_measure(seq, seq_ap, seq_ad) - MadIPM.get_affine_complementarity_measure!(bat, bat.workspace.alpha_p, bat.workspace.alpha_d) - - seq.mu_curr = MadIPM.update_barrier!(seq.opt.barrier_update, seq, seq_mu_aff) - MadIPM.update_barrier!(bat.opt.barrier_update, bat, bat.workspace.mu_affine) - - @test abs(seq.mu - bat.workspace.mu_batch[1]) < 1e-10 - end - end -end diff --git a/test/batch_kernels/test_kkt.jl b/test/batch_kernels/test_kkt.jl deleted file mode 100644 index a3f9ff71..00000000 --- a/test/batch_kernels/test_kkt.jl +++ /dev/null @@ -1,86 +0,0 @@ -@testset "KKT / augmented system" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - - @testset "set_aug_diagonal_reg!" begin - MadIPM.update_regularization!(seq, seq.opt.regularization) - MadIPM.update_regularization!(bat, bat.opt.regularization) - MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) - MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) - - skkt = seq.kkt - bkkt = bat.kkt - @test cmp(skkt.reg, col1(bkkt.reg)) < 1e-12 - @test cmp(skkt.l_diag, col1(bkkt.l_diag)) < 1e-12 - @test cmp(skkt.u_diag, col1(bkkt.u_diag)) < 1e-12 - @test cmp(skkt.l_lower, col1(bkkt.l_lower)) < 1e-12 - @test cmp(skkt.u_lower, col1(bkkt.u_lower)) < 1e-12 - @test cmp(skkt.pr_diag, col1(MadIPM.pr_diag(bkkt))) < 1e-12 - end - - @testset "build_kkt!" begin - MadIPM.update_regularization!(seq, seq.opt.regularization) - MadIPM.update_regularization!(bat, bat.opt.regularization) - MadIPM.set_aug_diagonal_reg!(seq.kkt, seq) - MadIPM.set_aug_diagonal_reg!(bat.kkt, bat) - MadNLP.build_kkt!(seq.kkt) - MadNLP.build_kkt!(bat.kkt) - - # Compare COO values - seq_V = seq.kkt.aug_raw.V - bat_V = bat.kkt.nzVals[:, 1] - @test cmp(seq_V, bat_V) < 1e-12 - - # Compare CSC nzvals - seq_csc = SparseArrays.nonzeros(seq.kkt.aug_com) - bat_csc = bat.kkt.aug_com_nzvals[:, 1] - @test cmp(seq_csc, bat_csc) < 1e-12 - end - - @testset "factorize + solve" begin - do_first_factorize!(seq, bat) - - # Set same RHS - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 - end - - @testset "mul! (KKT)" begin - do_first_factorize!(seq, bat) - - # Set up input vector x from predictive rhs - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - # w = K * d - w_seq = seq._w1 - w_bat = bat._w1 - fill!(MadNLP.full(w_seq), 0.0) - fill!(MadNLP.full(w_bat), 0.0) - mul!(w_seq, seq.kkt, seq.d) - mul!(w_bat, bat.kkt, bat.d) - @test cmp(MadNLP.full(w_seq), col1(MadNLP.full(w_bat))) < 1e-10 - end - - @testset "jtprod!" begin - # Explicit jtprod! call with initialized solvers - seq_jacl = similar(seq.jacl) - fill!(seq_jacl, 0.0) - MadNLP.jtprod!(seq_jacl, seq.kkt, seq.y) - - bat_jacl = similar(MadNLP.full(bat.jacl)) - fill!(bat_jacl, 0.0) - MadNLP.jtprod!(bat_jacl, bat.kkt, bat.y) - - @test cmp(seq_jacl, col1(bat_jacl)) < 1e-12 - end - end -end diff --git a/test/batch_kernels/test_regularization.jl b/test/batch_kernels/test_regularization.jl deleted file mode 100644 index 49cf1a8e..00000000 --- a/test/batch_kernels/test_regularization.jl +++ /dev/null @@ -1,61 +0,0 @@ -@testset "Regularization" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - @testset "NoRegularization" begin - qp = make_qp() - seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=MadIPM.NoRegularization()) - bat_nlp = ObjRHSBatchQuadraticModel([qp]) - bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=MadIPM.NoRegularization()) - - MadIPM.init_regularization!(seq, MadIPM.NoRegularization()) - MadIPM.init_regularization!(bat, MadIPM.NoRegularization()) - @test seq.del_w == bat.del_w[1] - @test seq.del_c == bat.del_c[1] - - MadIPM.update_regularization!(seq, MadIPM.NoRegularization()) - MadIPM.update_regularization!(bat, MadIPM.NoRegularization()) - @test seq.del_w == bat.del_w[1] - @test seq.del_c == bat.del_c[1] - end - - @testset "FixedRegularization" begin - reg = MadIPM.FixedRegularization(1e-8, -1e-9) - qp = make_qp() - seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=reg) - reg2 = MadIPM.FixedRegularization(1e-8, -1e-9) - bat_nlp = ObjRHSBatchQuadraticModel([qp]) - bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=reg2) - - MadIPM.init_regularization!(seq, reg) - MadIPM.init_regularization!(bat, reg2) - @test seq.del_w == bat.del_w[1] - @test seq.del_c == bat.del_c[1] - - MadIPM.update_regularization!(seq, reg) - MadIPM.update_regularization!(bat, reg2) - @test seq.del_w == bat.del_w[1] - @test seq.del_c == bat.del_c[1] - end - - @testset "AdaptiveRegularization" begin - reg = MadIPM.AdaptiveRegularization(1e-8, -1e-9, 1e-9) - qp = make_qp() - seq = MadIPM.MPCSolver(qp; print_level=MadNLP.ERROR, regularization=reg) - reg2 = MadIPM.AdaptiveRegularization(1e-8, -1e-9, 1e-9) - bat_nlp = ObjRHSBatchQuadraticModel([qp]) - bat = MadIPM.UniformBatchMPCSolver(bat_nlp; print_level=MadNLP.ERROR, regularization=reg2) - - MadIPM.init_regularization!(seq, reg) - MadIPM.init_regularization!(bat, reg2) - @test seq.del_w == bat.del_w[1] - @test seq.del_c == bat.del_c[1] - - # Multiple updates to test the adaptive decay - for _ in 1:3 - MadIPM.update_regularization!(seq, reg) - MadIPM.update_regularization!(bat, reg2) - @test abs(seq.del_w - bat.del_w[1]) < 1e-15 - @test abs(seq.del_c - bat.del_c[1]) < 1e-15 - end - end - end -end diff --git a/test/batch_kernels/test_rhs.jl b/test/batch_kernels/test_rhs.jl deleted file mode 100644 index 57c0b7ef..00000000 --- a/test/batch_kernels/test_rhs.jl +++ /dev/null @@ -1,53 +0,0 @@ -@testset "RHS setup functions" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - do_first_factorize!(seq, bat) - - @testset "set_initial_primal_rhs!" begin - MadIPM.set_initial_primal_rhs!(seq) - MadIPM.set_initial_primal_rhs!(bat) - @test cmp(MadNLP.full(seq.p), col1(MadNLP.full(bat.p))) < 1e-12 - end - - @testset "set_initial_dual_rhs!" begin - MadIPM.set_initial_dual_rhs!(seq) - MadIPM.set_initial_dual_rhs!(bat) - @test cmp(MadNLP.full(seq.p), col1(MadNLP.full(bat.p))) < 1e-12 - end - - @testset "set_predictive_rhs!" begin - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - @test cmp(MadNLP.primal(seq.p), col1(MadNLP.primal(bat.p))) < 1e-12 - @test cmp(MadNLP.dual(seq.p), col1(MadNLP.dual(bat.p))) < 1e-12 - @test cmp(MadNLP.dual_lb(seq.p), col1(MadNLP.dual_lb(bat.p))) < 1e-12 - @test cmp(MadNLP.dual_ub(seq.p), col1(MadNLP.dual_ub(bat.p))) < 1e-12 - end - - # Need an affine direction to test correction/correction_rhs - @testset "get_correction!" begin - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - MadIPM.get_correction!(seq, seq.correction_lb, seq.correction_ub) - MadIPM.get_correction!(bat, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub)) - @test cmp(seq.correction_lb, col1(MadNLP.full(bat.correction_lb))) < 1e-10 - @test cmp(seq.correction_ub, col1(MadNLP.full(bat.correction_ub))) < 1e-10 - end - - @testset "set_correction_rhs!" begin - # Use the state from get_correction! above - mu_val = seq.mu - MadIPM.set_correction_rhs!(seq, seq.kkt, mu_val, seq.correction_lb, seq.correction_ub, seq.ind_lb, seq.ind_ub) - MadIPM.set_correction_rhs!(bat, bat.kkt, bat.workspace.mu_batch, MadNLP.full(bat.correction_lb), MadNLP.full(bat.correction_ub), nothing, nothing) - @test cmp(MadNLP.primal(seq.p), col1(MadNLP.primal(bat.p))) < 1e-10 - @test cmp(MadNLP.dual(seq.p), col1(MadNLP.dual(bat.p))) < 1e-10 - @test cmp(MadNLP.dual_lb(seq.p), col1(MadNLP.dual_lb(bat.p))) < 1e-10 - @test cmp(MadNLP.dual_ub(seq.p), col1(MadNLP.dual_ub(bat.p))) < 1e-10 - end - end -end diff --git a/test/batch_kernels/test_solver_steps.jl b/test/batch_kernels/test_solver_steps.jl deleted file mode 100644 index ae6e4c0d..00000000 --- a/test/batch_kernels/test_solver_steps.jl +++ /dev/null @@ -1,100 +0,0 @@ -@testset "High-level solver steps" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - @testset "init_starting_point!" begin - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - - @test cmp(MadNLP.full(seq.x), col1(MadNLP.full(bat.x))) < 1e-10 - @test cmp(MadNLP.full(seq.xl), col1(MadNLP.full(bat.xl))) < 1e-10 - @test cmp(MadNLP.full(seq.xu), col1(MadNLP.full(bat.xu))) < 1e-10 - @test cmp(seq.y, col1(MadNLP.full(bat.y))) < 1e-10 - @test cmp(MadNLP.full(seq.zl), col1(MadNLP.full(bat.zl))) < 1e-10 - @test cmp(MadNLP.full(seq.zu), col1(MadNLP.full(bat.zu))) < 1e-10 - end - - @testset "Full first IPM iteration" begin - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - do_first_factorize!(seq, bat) - - # --- affine_direction! --- - MadIPM.affine_direction!(seq) - MadIPM.affine_direction!(bat) - @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 - - # --- prediction_step! (from post-factorize state) --- - # Rebuild fresh solvers since affine_direction! mutated state - seq = build_seq(qp) - bat = build_batch(qp) - do_first_factorize!(seq, bat) - - MadIPM.prediction_step!(seq) - MadIPM.prediction_step!(bat) - - # Check alpha - seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, 1.0) - @test abs(seq_ap - bat.workspace.alpha_p[1]) < 1e-10 - @test abs(seq_ad - bat.workspace.alpha_d[1]) < 1e-10 - # Check mu - @test abs(seq.mu - bat.workspace.mu_batch[1]) < 1e-10 - # Check corrections - @test cmp(seq.correction_lb, col1(MadNLP.full(bat.correction_lb))) < 1e-10 - @test cmp(seq.correction_ub, col1(MadNLP.full(bat.correction_ub))) < 1e-10 - - # --- mehrotra_correction_direction! --- - MadIPM.mehrotra_correction_direction!(seq) - MadIPM.mehrotra_correction_direction!(bat) - @test cmp(MadNLP.full(seq.d), col1(MadNLP.full(bat.d))) < 1e-10 - - # --- update_step! --- - MadIPM.update_step!(seq.opt.step_rule, seq) - MadIPM.update_step!(bat.opt.step_rule, bat) - @test abs(seq.alpha_p - bat.workspace.alpha_p[1]) < 1e-10 - @test abs(seq.alpha_d - bat.workspace.alpha_d[1]) < 1e-10 - - # --- apply_step! --- - MadIPM.apply_step!(seq) - MadIPM.apply_step!(bat) - @test cmp(MadNLP.full(seq.x), col1(MadNLP.full(bat.x))) < 1e-10 - @test cmp(seq.y, col1(MadNLP.full(bat.y))) < 1e-10 - @test cmp(MadNLP.full(seq.zl), col1(MadNLP.full(bat.zl))) < 1e-10 - @test cmp(MadNLP.full(seq.zu), col1(MadNLP.full(bat.zu))) < 1e-10 - - # --- evaluate_model! --- - MadIPM.evaluate_model!(seq) - MadIPM.evaluate_model!(bat) - @test cmp(MadNLP.primal(seq.f), col1(MadNLP.primal(bat.f))) < 1e-10 - @test cmp(seq.c, col1(MadNLP.full(bat.c))) < 1e-10 - @test cmp(seq.jacl, col1(MadNLP.full(bat.jacl))) < 1e-10 - end - - @testset "update_termination_criteria!" begin - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - # Run one full iteration first - do_first_factorize!(seq, bat) - MadIPM.prediction_step!(seq) - MadIPM.prediction_step!(bat) - MadIPM.mehrotra_correction_direction!(seq) - MadIPM.mehrotra_correction_direction!(bat) - MadIPM.update_step!(seq.opt.step_rule, seq) - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.apply_step!(seq) - MadIPM.apply_step!(bat) - MadIPM.evaluate_model!(seq) - MadIPM.evaluate_model!(bat) - - # Now check termination criteria - MadIPM.update_termination_criteria!(seq) - MadIPM.update_termination_criteria!(bat) - - @test abs(seq.inf_pr - bat.workspace.inf_pr[1]) < 1e-10 - @test abs(seq.inf_du - bat.workspace.inf_du[1]) < 1e-10 - @test abs(seq.inf_compl - bat.workspace.inf_compl[1]) < 1e-10 - @test seq.status == bat.workspace.status[1] - end - end -end diff --git a/test/batch_kernels/test_step.jl b/test/batch_kernels/test_step.jl deleted file mode 100644 index b01cdc6a..00000000 --- a/test/batch_kernels/test_step.jl +++ /dev/null @@ -1,64 +0,0 @@ -@testset "Step size computation" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - @testset "get_fraction_to_boundary_step!" begin - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - do_first_factorize!(seq, bat) - MadIPM.set_predictive_rhs!(seq, seq.kkt) - MadIPM.set_predictive_rhs!(bat, bat.kkt) - MadIPM.solve_system!(seq.d, seq, seq.p) - MadIPM.solve_system!(bat.d, bat, bat.p) - - for tau_val in [1.0, 0.99, 0.995] - seq_ap, seq_ad = MadIPM.get_fraction_to_boundary_step(seq, tau_val) - fill!(bat.workspace.tau, tau_val) - MadIPM.get_fraction_to_boundary_step!(bat) - @test abs(seq_ap - bat.workspace.alpha_p[1]) < 1e-10 - @test abs(seq_ad - bat.workspace.alpha_d[1]) < 1e-10 - end - end - - @testset "set_tau! (ConservativeStep)" begin - qp = make_qp() - bat = build_batch(qp) - rule = MadIPM.ConservativeStep(0.99) - MadIPM.set_tau!(rule, bat) - @test bat.workspace.tau[1] == 0.99 - end - - @testset "set_tau! (AdaptiveStep)" begin - qp = make_qp() - bat = build_batch(qp) - rule = MadIPM.AdaptiveStep(0.99) - MadIPM.set_tau!(rule, bat) - expected_tau = max(1.0 - bat.workspace.mu_batch[1], 0.99) - @test bat.workspace.tau[1] ≈ expected_tau atol=1e-12 - end - - @testset "update_step! ($rule_name)" for (rule_name, make_rule) in [ - ("ConservativeStep", () -> MadIPM.ConservativeStep(0.99)), - ("AdaptiveStep", () -> MadIPM.AdaptiveStep(0.99)), - ("MehrotraAdaptiveStep", () -> MadIPM.MehrotraAdaptiveStep(0.99)), - ] - qp = make_qp() - seq = build_seq(qp) - bat = build_batch(qp) - do_first_factorize!(seq, bat) - - # Standard flow: prediction → correction → update_step - MadIPM.prediction_step!(seq) - MadIPM.prediction_step!(bat) - MadIPM.mehrotra_correction_direction!(seq) - MadIPM.mehrotra_correction_direction!(bat) - - rule = make_rule() - seq.opt.step_rule = rule - bat.opt.step_rule = rule - MadIPM.update_step!(rule, seq) - MadIPM.update_step!(rule, bat) - @test abs(seq.alpha_p - bat.workspace.alpha_p[1]) < 1e-10 - @test abs(seq.alpha_d - bat.workspace.alpha_d[1]) < 1e-10 - end - end -end diff --git a/test/batch_kernels/test_termination.jl b/test/batch_kernels/test_termination.jl deleted file mode 100644 index b37bb30c..00000000 --- a/test/batch_kernels/test_termination.jl +++ /dev/null @@ -1,48 +0,0 @@ -@testset "Termination status branches" begin - @testset "$label" for (label, make_qp) in ALL_TEST_PROBLEMS - @testset "SOLVE_SUCCEEDED" begin - # Run enough iterations to converge - qp = make_qp() - bat = build_batch(qp) - - # Run several iterations until convergence - for _ in 1:50 - bat.workspace.status[1] != MadNLP.REGULAR && break - MadIPM.update_termination_criteria!(bat) - changed = MadIPM.update_termination_status!(bat) - if changed - MadIPM.update_active_set!(bat) - MadIPM.local_batch_size(MadIPM.active_view(bat.batch_views)) == 0 && break - MadIPM._update_active_mask!(bat) - end - MadIPM.factorize_system!(bat) - MadIPM.prediction_step!(bat) - MadIPM.mehrotra_correction_direction!(bat) - MadIPM.update_step!(bat.opt.step_rule, bat) - MadIPM.zero_inactive_step!(bat) - MadIPM.apply_step!(bat) - MadIPM.evaluate_model!(bat) - end - @test bat.workspace.status[1] == MadNLP.SOLVE_SUCCEEDED - end - - @testset "MAXIMUM_ITERATIONS_EXCEEDED" begin - qp = make_qp() - bat = build_batch(qp) - bat.opt.max_iter = 0 # no iterations allowed - MadIPM.update_termination_criteria!(bat) - MadIPM.update_termination_status!(bat) - @test bat.workspace.status[1] == MadNLP.MAXIMUM_ITERATIONS_EXCEEDED - end - - @testset "MAXIMUM_WALLTIME_EXCEEDED" begin - qp = make_qp() - bat = build_batch(qp) - bat.opt.max_wall_time = 0.0 # zero walltime - bat.batch_cnt.start_time[] = time() - 1.0 # started 1s ago - MadIPM.update_termination_criteria!(bat) - MadIPM.update_termination_status!(bat) - @test bat.workspace.status[1] == MadNLP.MAXIMUM_WALLTIME_EXCEEDED - end - end -end diff --git a/test/runtests.jl b/test/runtests.jl index 471f5d47..3213ddb6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -214,10 +214,10 @@ end # include("MOI_wrapper.jl") # end -include("test_batch.jl") -include("batch_kernels/runtests.jl") +include("batch/views.jl") +include("batch/solver.jl") if CUDA.functional() include("test_gpu.jl") - include("test_batch_gpu.jl") + include("batch/gpu.jl") end diff --git a/test/test_batch.jl b/test/test_batch.jl deleted file mode 100644 index 742e4c7e..00000000 --- a/test/test_batch.jl +++ /dev/null @@ -1,367 +0,0 @@ -using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel - -_test_local_to_root(view) = Int[view.local_to_root[i] for i in 1:view.n] - -struct RecordingBatchLinearSolver{T, MT} <: MadNLP.AbstractLinearSolver{T} - nzvals_mat::MT - call_counts::Vector{Int} - factorized::Vector{Bool} - fail_marker::T -end - -@kwdef mutable struct RecordingBatchLinearSolverOptions <: MadNLP.AbstractOptions - fail_marker::Float64 = -999.0 -end - -MadNLP.default_options(::Type{RecordingBatchLinearSolver}) = RecordingBatchLinearSolverOptions() - -function RecordingBatchLinearSolver( - aug_com, - nzvals_mat::AbstractMatrix{T}, - n::Int; - opt::RecordingBatchLinearSolverOptions = RecordingBatchLinearSolverOptions(), -) where T - batch_size = size(nzvals_mat, 2) - return RecordingBatchLinearSolver( - nzvals_mat, - zeros(Int, batch_size), - fill(true, batch_size), - T(opt.fail_marker), - ) -end - -function MadIPM.factorize_active!( - s::RecordingBatchLinearSolver{T, MT}, - factor_view::MadIPM.BatchView, -) where {T, MT} - @inbounds for j in 1:MadIPM.local_batch_size(factor_view) - s.call_counts[j] += 1 - col = view(s.nzvals_mat, :, j) - corrupted = any(==(s.fail_marker), col) - s.factorized[j] = !(corrupted && s.call_counts[j] == 1) - end - return -end - -function MadIPM.failed_factorization_local_count!( - failed_local_buffer::Vector{Int32}, - s::RecordingBatchLinearSolver, - factor_view::MadIPM.BatchView, -) - nfailed = 0 - @inbounds for j in 1:MadIPM.local_batch_size(factor_view) - if !s.factorized[j] - nfailed += 1 - failed_local_buffer[nfailed] = j - end - end - return nfailed -end - -MadIPM.is_factorized(s::RecordingBatchLinearSolver) = all(s.factorized) - -MadIPM.solve_active!(s::RecordingBatchLinearSolver, rhs::AbstractMatrix, active::MadIPM.BatchView) = rhs - -function _make_batch_solver(qps; batch_kwargs...) - bnlp = BatchQuadraticModel(qps) - return MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) -end - -function _make_small_qp() - # Small QP: min 0.5 xᵀHx + cᵀx s.t. lcon ≤ Ax ≤ ucon, lvar ≤ x ≤ uvar - n, m = 4, 2 - c = [1.0, -2.0, 0.5, 1.0] - Hrows = [1, 2, 3, 4] - Hcols = [1, 2, 3, 4] - Hvals = [2.0, 1.0, 3.0, 1.5] - Arows = [1, 1, 2, 2] - Acols = [1, 2, 3, 4] - Avals = [1.0, 1.0, 1.0, 1.0] - return QuadraticModel( - c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[2.0, 1.5], - lvar=zeros(n), uvar=fill(Inf, n), - x0=ones(n), - ) -end - -function _test_batch_lp(; batch_kwargs...) - qp = simple_lp() - ref = MadIPM.madipm(qp; print_level=MadNLP.ERROR) - @test ref.status == MadNLP.SOLVE_SUCCEEDED - - bs = 4 - qps = [simple_lp() for _ in 1:bs] - - @testset "ObjRHSBatch" begin - bnlp = ObjRHSBatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) - for i in 1:bs - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ ref.objective atol=1e-6 - @test si.solution ≈ ref.solution atol=1e-6 - end - end - - @testset "FullBatch" begin - bnlp = BatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) - for i in 1:bs - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ ref.objective atol=1e-6 - @test si.solution ≈ ref.solution atol=1e-6 - end - end -end - -function _test_batch_qp(; batch_kwargs...) - qp = _make_small_qp() - ref = MadIPM.madipm(qp; print_level=MadNLP.ERROR) - @test ref.status == MadNLP.SOLVE_SUCCEEDED - - bs = 3 - qps = [_make_small_qp() for _ in 1:bs] - - @testset "ObjRHSBatch" begin - bnlp = ObjRHSBatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) - for i in 1:bs - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ ref.objective atol=1e-6 - @test si.solution ≈ ref.solution atol=1e-6 - end - end - - @testset "FullBatch" begin - bnlp = BatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) - for i in 1:bs - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ ref.objective atol=1e-6 - @test si.solution ≈ ref.solution atol=1e-6 - end - end -end - -function _test_fullbatch_different_data(; batch_kwargs...) - # QP with different H and A values per instance (same sparsity) - Hrows = [1, 2, 2]; Hcols = [1, 1, 2] - Arows = [1, 1]; Acols = [1, 2] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - - qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; - Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - qp3 = QuadraticModel([0.5, 0.5], Hrows, Hcols, [3.0, 1.5, 4.0]; - Arows=Arows, Acols=Acols, Avals=[0.5, 2.0], - lcon=[1.5], ucon=[1.5], lvar=lvar, uvar=uvar, x0=[0.75, 0.75]) - - qps = [qp1, qp2, qp3] - refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] - for r in refs - @test r.status == MadNLP.SOLVE_SUCCEEDED - end - - bnlp = BatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR, batch_kwargs...) - for i in 1:3 - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ refs[i].objective atol=1e-6 - @test si.solution ≈ refs[i].solution atol=1e-6 - end -end - -@testset "Batch solver (CPU)" begin - @testset "Batch views" begin - solver = _make_batch_solver([simple_lp() for _ in 1:4]) - root = MadIPM.root_view(solver.batch_views) - - @test MadIPM.local_batch_size(root) == 4 - @test MadIPM.batch_size_root(root) == 4 - @test MadIPM.is_identity_view(root) - @test _test_local_to_root(root) == [1, 2, 3, 4] - - saved_root = MadIPM.select_local!(solver.batch_views, [2, 4]) - child = MadIPM.active_view(solver.batch_views) - @test MadIPM.local_batch_size(child) == 2 - @test _test_local_to_root(child) == [2, 4] - @test !MadIPM.is_identity_view(child) - mask = zeros(Float64, 1, 4) - MadIPM.fill_batch_view_mask!(mask, child) - @test mask == [0.0 1.0 0.0 1.0] - - saved_child = MadIPM.select_local!(solver.batch_views, [2]) - grandchild = MadIPM.active_view(solver.batch_views) - @test _test_local_to_root(grandchild) == [4] - MadIPM.restore_state!(solver.batch_views, saved_child) - - MadIPM.select_local!(solver.batch_views, Int[]) - empty_child = MadIPM.active_view(solver.batch_views) - @test MadIPM.local_batch_size(empty_child) == 0 - MadIPM.restore_state!(solver.batch_views, saved_root) - - MadIPM.select_local!(solver.batch_views, [1, 2, 3, 4]) - full_child = MadIPM.active_view(solver.batch_views) - @test MadIPM.is_identity_view(full_child) - MadIPM.restore_state!(solver.batch_views, root) - end - - @testset "Batch structure mismatch throws" begin - qp1 = QuadraticModel( - [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[1.0, 1.0], - x0=[0.5, 0.5], - ) - qp2 = QuadraticModel( - [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[0.0, 1.0], - x0=[0.0, 0.5], - ) - bnlp = BatchQuadraticModel([qp1, qp2]) - @test_throws AssertionError MadIPM.UniformBatchMPCSolver(bnlp; print_level=MadNLP.ERROR) - solver = MadIPM.UniformBatchMPCSolver( - bnlp; - print_level=MadNLP.ERROR, - check_batch_structure=false, - ) - @test solver.batch_size == 2 - end - - @testset "Partial active KKT solve preserves rhs" begin - solver = _make_batch_solver([simple_lp() for _ in 1:3]) - MadIPM.initialize!(solver) - status = fill(MadNLP.REGULAR, 3) - status[2] = MadNLP.INTERNAL_ERROR - solver.workspace.status .= status - MadIPM.update_active_set!(solver) - - pd_view = MadNLP.primal_dual(solver.d) - pd_view .= reshape(collect(1.0:length(pd_view)), size(pd_view)) - pd_before = copy(pd_view) - - MadNLP.build_kkt!(solver.kkt) - MadNLP.factorize_kkt!(solver.kkt) - MadNLP.solve_kkt!(solver.kkt, solver) - - @test pd_view[:, 2] == pd_before[:, 2] - end - - @testset "Batch LP" begin - _test_batch_lp() - end - @testset "Batch QP" begin - _test_batch_qp() - end - @testset "FullBatch different H/A data" begin - _test_fullbatch_different_data() - end - @testset "Residual check marks INTERNAL_ERROR (all fail)" begin - # Force residual check failure with tol_linear_solve=0 — should mark - # instances as INTERNAL_ERROR instead of throwing a SolveException. - qps = [simple_lp() for _ in 1:3] - bnlp = ObjRHSBatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; - print_level=MadNLP.ERROR, - check_residual=true, - tol_linear_solve=0.0, - ) - for i in 1:3 - @test stats[i].status == MadNLP.INTERNAL_ERROR - end - end - - @testset "Residual check marks INTERNAL_ERROR (partial)" begin - # Use NaN objective coefficients to produce NaN residuals for one instance. - # The per-instance residual check should mark only that instance as INTERNAL_ERROR. - Hrows = [1, 2]; Hcols = [1, 2] - Arows = [1, 1]; Acols = [1, 2] - - good_qp() = QuadraticModel( - [1.0, 1.0], Hrows, Hcols, [2.0, 2.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], x0=[0.5, 0.5], - ) - bad_qp = QuadraticModel( - [NaN, NaN], Hrows, Hcols, [2.0, 2.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], x0=[0.5, 0.5], - ) - qps = [good_qp(), bad_qp, good_qp()] - bnlp = BatchQuadraticModel(qps) - stats = MadIPM.madipm_batch(bnlp; print_level=MadNLP.ERROR) - # Good instances should solve; bad instance should fail gracefully - @test stats[1].status == MadNLP.SOLVE_SUCCEEDED - @test stats[3].status == MadNLP.SOLVE_SUCCEEDED - @test stats[2].status != MadNLP.SOLVE_SUCCEEDED - end - - @testset "Factorization retry only refactorizes failed instance" begin - solver = _make_batch_solver( - [simple_lp() for _ in 1:3]; - uniformbatch_linear_solver=RecordingBatchLinearSolver, - regularization=MadIPM.FixedRegularization(1.0, -1.0), - ) - MadIPM.initialize!(solver) - fill!(solver.kkt.batch_solver.call_counts, 0) - fill!(solver.kkt.batch_solver.factorized, true) - - corrupt_k = solver.kkt.n_tot + solver.kkt.nnzh + 1 - solver.kkt.nzVals[corrupt_k, 2] = -999.0 - solver.kkt.nzVals[corrupt_k, 3] = -999.0 - # two corruptions -- should see an extra solve in slots 1 and 2 - # and bumped regularizations in instances 2 and 3 - - MadIPM.factorize_system!(solver) - - ls = solver.kkt.batch_solver - @test ls.call_counts == [2, 2, 2] - @test solver.del_w == [1.0 100.0 100.0] - @test solver.del_c == [-1.0 -100.0 -100.0] - end - - @testset "Factorization retry with one inactive instance" begin - solver = _make_batch_solver( - [simple_lp() for _ in 1:3]; - uniformbatch_linear_solver=RecordingBatchLinearSolver, - regularization=MadIPM.FixedRegularization(1.0, -1.0), - ) - MadIPM.initialize!(solver) - fill!(solver.kkt.batch_solver.call_counts, 0) - fill!(solver.kkt.batch_solver.factorized, true) - - # terminate an instance, so the last slot should be never used - status = fill(MadNLP.REGULAR, 3) - status[2] = MadNLP.SOLVE_SUCCEEDED - solver.workspace.status .= status - MadIPM.update_active_set!(solver) - MadIPM._update_active_mask!(solver) - - # corrupt instance 3, which is at position 2 after compacting [1,3] - corrupt_k = solver.kkt.n_tot + solver.kkt.nnzh + 1 - solver.kkt.nzVals[corrupt_k, 3] = -999.0 - - MadIPM.factorize_system!(solver) - ls = solver.kkt.batch_solver - @test ls.call_counts == [2, 2, 0] - - # instance 1 and 2's regularization is untouched, - # instance 3's regularization should be bumped - @test solver.del_w == [1.0 1.0 100.0] - @test solver.del_c == [-1.0 -1.0 -100.0] - end -end diff --git a/test/test_batch_gpu.jl b/test/test_batch_gpu.jl deleted file mode 100644 index a7bcfb2c..00000000 --- a/test/test_batch_gpu.jl +++ /dev/null @@ -1,348 +0,0 @@ -using KernelAbstractions -using MadNLPGPU -using BatchQuadraticModels: ObjRHSBatchQuadraticModel, BatchQuadraticModel - -# ============================================================ -# Test problem constructors for GPU batch tests -# ============================================================ - -function _gpu_small_qp() - n, m = 4, 2 - c = [1.0, -2.0, 0.5, 1.0] - Hrows = [1, 2, 3, 4] - Hcols = [1, 2, 3, 4] - Hvals = [2.0, 1.0, 3.0, 1.5] - Arows = [1, 1, 2, 2] - Acols = [1, 2, 3, 4] - Avals = [1.0, 1.0, 1.0, 1.0] - return QuadraticModel( - c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[2.0, 1.5], - lvar=zeros(n), uvar=fill(Inf, n), - x0=ones(n), - ) -end - -function _gpu_doubly_bounded_qp() - n, m = 3, 2 - c = [1.0, -1.0, 0.5] - Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 3.0] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] - return QuadraticModel( - c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[1.0, 0.5], - lvar=[0.0, 0.0, 0.0], uvar=[5.0, 5.0, 5.0], - x0=ones(n), - ) -end - -function _gpu_dense_hess_qp() - Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - return QuadraticModel( - [1.0, -1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=[0.5, 0.5], - ) -end - -# ============================================================ -# Helper: solve batch on GPU and compare with CPU reference -# ============================================================ - -function _test_gpu_batch(qps; atol=1e-6, batch_kwargs...) - bs = length(qps) - - # CPU reference: solve each QP independently - refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] - for r in refs - @test r.status == MadNLP.SOLVE_SUCCEEDED - end - - # Build CPU batch model, convert to GPU, solve - cpu_bnlp = ObjRHSBatchQuadraticModel(qps) - gpu_bnlp = convert(ObjRHSBatchQuadraticModel{Float64, CuVector{Float64}}, cpu_bnlp) - stats = try - MadIPM.madipm_batch( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - cudss_algorithm=MadNLP.LDL, - batch_kwargs..., - ) - catch e - @error "madipm_batch failed" exception=(e, catch_backtrace()) - rethrow(e) - end - - CUDA.@allowscalar for i in 1:bs - si = stats[i] - if si.status != MadNLP.SOLVE_SUCCEEDED - @error "Instance $i failed" status=si.status objective=si.objective - end - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ refs[i].objective atol=atol - @test Array(si.solution) ≈ refs[i].solution atol=atol - end -end - -# ============================================================ -# Tests -# ============================================================ - -@testset "Batch solver (CUDA)" begin - @testset "Batch views gather/scatter" begin - cpu_bnlp = BatchQuadraticModel([simple_lp() for _ in 1:4]) - gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) - solver = MadIPM.UniformBatchMPCSolver( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - ) - root = MadIPM.root_view(solver.batch_views) - saved_root = MadIPM.select_local!(solver.batch_views, [2, 4]) - child = MadIPM.active_view(solver.batch_views) - MadIPM.select_local!(solver.batch_views, [2]) - grandchild = MadIPM.active_view(solver.batch_views) - - src = cu(reshape(collect(1.0:12.0), 3, 4)) - gathered = similar(src, 3, MadIPM.local_batch_size(child)) - MadIPM.gather_batch_view_columns!(gathered, src, child) - @test Array(gathered) == Array(src[:, [2, 4]]) - - gathered_nested = similar(src, 3, MadIPM.local_batch_size(grandchild)) - MadIPM.gather_batch_view_columns!(gathered_nested, src, grandchild) - @test Array(gathered_nested) == Array(src[:, [4]]) - - scattered = CUDA.fill(-1.0, 3, 4) - MadIPM.scatter_batch_view_columns!(scattered, gathered, child) - @test Array(scattered[:, 2]) == Array(src[:, 2]) - @test Array(scattered[:, 4]) == Array(src[:, 4]) - @test Array(scattered[:, 1]) == fill(-1.0, 3) - @test Array(scattered[:, 3]) == fill(-1.0, 3) - MadIPM.restore_state!(solver.batch_views, saved_root) - end - - @testset "Batch structure mismatch throws" begin - qp1 = QuadraticModel( - [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[1.0, 1.0], - x0=[0.5, 0.5], - ) - qp2 = QuadraticModel( - [1.0, 1.0], [1, 2], [1, 2], [2.0, 2.0]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[0.0, 1.0], - x0=[0.0, 0.5], - ) - cpu_bnlp = BatchQuadraticModel([qp1, qp2]) - gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) - @test_throws AssertionError MadIPM.UniformBatchMPCSolver( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - ) - solver = MadIPM.UniformBatchMPCSolver( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - check_batch_structure=false, - ) - @test solver.batch_size == 2 - end - - # ---------------------------------------------------------- - # Identical instances (sanity check) - # ---------------------------------------------------------- - @testset "Identical LP (bs=4)" begin - _test_gpu_batch([simple_lp() for _ in 1:4]; atol=1e-5) - end - - @testset "Identical QP (bs=3)" begin - _test_gpu_batch([_gpu_small_qp() for _ in 1:3]) - end - - @testset "Identical doubly-bounded QP (bs=2)" begin - _test_gpu_batch([_gpu_doubly_bounded_qp() for _ in 1:2]) - end - - @testset "Identical dense-Hessian QP (bs=2)" begin - _test_gpu_batch([_gpu_dense_hess_qp() for _ in 1:2]) - end - - # ---------------------------------------------------------- - # Different-data instances via ObjRHSBatch - # ---------------------------------------------------------- - @testset "Different LP data (bs=3)" begin - qp1 = QuadraticModel( - [1.0, 1.0], Int[], Int[], Float64[]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=ones(2), - ) - qp2 = QuadraticModel( - [2.0, 0.5], Int[], Int[], Float64[]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[2.0], ucon=[2.0], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=ones(2), - ) - qp3 = QuadraticModel( - [0.5, 3.0], Int[], Int[], Float64[]; - Arows=[1, 1], Acols=[1, 2], Avals=[1.0, 1.0], - lcon=[0.5], ucon=[0.5], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=ones(2), - ) - _test_gpu_batch([qp1, qp2, qp3]; atol=1e-5) - end - - @testset "Different QP data (bs=2)" begin - n, m = 4, 2 - Hrows = [1, 2, 3, 4]; Hcols = [1, 2, 3, 4]; Hvals = [2.0, 1.0, 3.0, 1.5] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 3, 4]; Avals = [1.0, 1.0, 1.0, 1.0] - qp1 = QuadraticModel( - [1.0, -2.0, 0.5, 1.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[2.0, 1.5], - lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n), - ) - qp2 = QuadraticModel( - [-1.0, 1.0, -0.5, 2.0], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5, 1.0], ucon=[1.5, 2.0], - lvar=zeros(n), uvar=fill(Inf, n), x0=ones(n), - ) - _test_gpu_batch([qp1, qp2]) - end - - @testset "Different doubly-bounded QP data (bs=2)" begin - n, m = 3, 2 - Hrows = [1, 2, 3]; Hcols = [1, 2, 3]; Hvals = [2.0, 1.0, 3.0] - Arows = [1, 1, 2, 2]; Acols = [1, 2, 2, 3]; Avals = [1.0, 1.0, 1.0, 1.0] - qp1 = QuadraticModel( - [1.0, -1.0, 0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[1.0, 0.5], ucon=[1.0, 0.5], - lvar=[0.0, 0.0, 0.0], uvar=[5.0, 5.0, 5.0], - x0=ones(n), - ) - qp2 = QuadraticModel( - [-1.0, 2.0, -0.5], Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[0.5, 1.0], ucon=[0.5, 1.0], - lvar=[0.0, 0.0, 0.0], uvar=[3.0, 3.0, 3.0], - x0=ones(n), - ) - _test_gpu_batch([qp1, qp2]) - end - - @testset "Different dense-Hessian QP data (bs=4)" begin - Hrows = [1, 2, 2]; Hcols = [1, 1, 2]; Hvals = [4.0, 2.0, 3.0] - Arows = [1, 1]; Acols = [1, 2]; Avals = [1.0, 1.0] - make_qp(c, rhs) = QuadraticModel( - c, Hrows, Hcols, Hvals; - Arows=Arows, Acols=Acols, Avals=Avals, - lcon=[rhs], ucon=[rhs], - lvar=[0.0, 0.0], uvar=[Inf, Inf], - x0=[0.5, 0.5], - ) - _test_gpu_batch([ - make_qp([1.0, -1.0], 1.0), - make_qp([-1.0, 2.0], 2.0), - make_qp([0.5, 0.5], 0.5), - make_qp([2.0, -2.0], 1.5), - ]) - end - - # ---------------------------------------------------------- - # Batch size variations - # ---------------------------------------------------------- - @testset "batch_size=1" begin - _test_gpu_batch([simple_lp()]; atol=1e-5) - end - - @testset "batch_size=8" begin - _test_gpu_batch([_gpu_small_qp() for _ in 1:8]) - end - - # ---------------------------------------------------------- - # BatchQuadraticModel (different H/A values per instance) - # ---------------------------------------------------------- - @testset "FullBatch different H/A (bs=2)" begin - Hrows = [1, 2, 2]; Hcols = [1, 1, 2] - Arows = [1, 1]; Acols = [1, 2] - lvar = [0.0, 0.0]; uvar = [Inf, Inf] - - qp1 = QuadraticModel([1.0, -1.0], Hrows, Hcols, [4.0, 2.0, 3.0]; - Arows=Arows, Acols=Acols, Avals=[1.0, 1.0], - lcon=[1.0], ucon=[1.0], lvar=lvar, uvar=uvar, x0=[0.5, 0.5]) - qp2 = QuadraticModel([-1.0, 2.0], Hrows, Hcols, [6.0, 1.0, 5.0]; - Arows=Arows, Acols=Acols, Avals=[1.5, 0.5], - lcon=[2.0], ucon=[2.0], lvar=lvar, uvar=uvar, x0=[1.0, 1.0]) - - qps = [qp1, qp2] - refs = [MadIPM.madipm(qp; print_level=MadNLP.ERROR) for qp in qps] - for r in refs; @test r.status == MadNLP.SOLVE_SUCCEEDED; end - - cpu_bnlp = BatchQuadraticModel(qps) - gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) - stats = MadIPM.madipm_batch( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - cudss_algorithm=MadNLP.LDL, - ) - CUDA.@allowscalar for i in 1:2 - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ refs[i].objective atol=1e-6 - @test Array(si.solution) ≈ refs[i].solution atol=1e-6 - end - end - - @testset "Residual check marks INTERNAL_ERROR (GPU)" begin - qps = [simple_lp() for _ in 1:3] - cpu_bnlp = ObjRHSBatchQuadraticModel(qps) - gpu_bnlp = convert(ObjRHSBatchQuadraticModel{Float64, CuVector{Float64}}, cpu_bnlp) - stats = MadIPM.madipm_batch(gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - cudss_algorithm=MadNLP.LDL, - check_residual=true, - tol_linear_solve=0.0, - ) - CUDA.@allowscalar for i in 1:3 - @test stats[i].status == MadNLP.INTERNAL_ERROR - end - end - - @testset "FullBatch identical QP (bs=3)" begin - qps = [_gpu_small_qp() for _ in 1:3] - ref = MadIPM.madipm(qps[1]; print_level=MadNLP.ERROR) - @test ref.status == MadNLP.SOLVE_SUCCEEDED - - cpu_bnlp = BatchQuadraticModel(qps) - gpu_bnlp = convert(BatchQuadraticModel{Float64, CuMatrix{Float64}}, cpu_bnlp) - stats = MadIPM.madipm_batch( - gpu_bnlp; - print_level=MadNLP.ERROR, - uniformbatch_linear_solver=MadNLPGPU.CUDSSSolver, - cudss_algorithm=MadNLP.LDL, - ) - CUDA.@allowscalar for i in 1:3 - si = stats[i] - @test si.status == MadNLP.SOLVE_SUCCEEDED - @test si.objective ≈ ref.objective atol=1e-6 - @test Array(si.solution) ≈ ref.solution atol=1e-6 - end - end -end From f315c26f25fb96a440a6d329f8fd3bc0cbcad8a3 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 31 Mar 2026 01:26:28 -0400 Subject: [PATCH 48/51] fix allocations --- src/batch/KKT/Sparse/augmented.jl | 5 ++-- src/batch/madipm/solver.jl | 2 +- src/batch/madnlp/nlpmodels.jl | 7 +++-- src/batch/madnlp/rhs.jl | 46 ++++++++++++++++++------------- src/batch/structure.jl | 8 +++--- 5 files changed, 40 insertions(+), 28 deletions(-) diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 0435b27e..99a78d39 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -253,8 +253,9 @@ function MadNLP.eval_lag_hess_wrapper!( if nnzh > 0 hess = view(nzVals, n_tot+1:n_tot+nnzh, :) MadNLP.unpack_x!(ws.bx, bcb, batch_solver.x) - scale = vec(bcb.obj_sign .* bcb.obj_scale) - MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess; obj_weight=scale) + bf_mat = reshape(ws.bf, 1, batch_solver.batch_size) + @. bf_mat = bcb.obj_sign * bcb.obj_scale + MadNLP._eval_lag_hess_wrapper!(bcb, ws.bx, MadNLP.full(batch_solver.y), ws.bv, hess; obj_weight=ws.bf) end return end diff --git a/src/batch/madipm/solver.jl b/src/batch/madipm/solver.jl index 0ca4e145..840c047d 100644 --- a/src/batch/madipm/solver.jl +++ b/src/batch/madipm/solver.jl @@ -399,7 +399,7 @@ function update_solution!(stats::BatchExecutionStats, batch_solver::AbstractBatc MadNLP.unpack_y!(stats.multipliers, bcb, MadNLP.full(batch_solver.y)) MadNLP.unpack_z!(stats.multipliers_L, bcb, MadNLP.variable(zl)) MadNLP.unpack_z!(stats.multipliers_U, bcb, MadNLP.variable(zu)) - stats.objective .= MadNLP.unpack_obj(bcb, ws.obj_val) + unpack_obj!(stats.objective, bcb, ws.obj_val) MadNLP.unpack_cons!(stats.constraints, bcb, MadNLP.full(batch_solver.c), MadNLP.full(batch_solver.rhs), bcb.ind_ineq, MadNLP.slack(x)) stats.dual_feas .= vec(ws.inf_du) diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl index 1d6cfae3..2ac4d30b 100644 --- a/src/batch/madnlp/nlpmodels.jl +++ b/src/batch/madnlp/nlpmodels.jl @@ -62,10 +62,13 @@ MadNLP.get_lcon(bcb::UniformBatchCallback) = bcb.nlp.meta.lcon MadNLP.get_ucon(bcb::UniformBatchCallback) = bcb.nlp.meta.ucon function MadNLP.unpack_y!(y_full, bcb::UniformBatchCallback, y) - y_full .= y .* bcb.con_scale .* (bcb.obj_sign ./ bcb.obj_scale) + @. y_full = y * bcb.con_scale * bcb.obj_sign / bcb.obj_scale end -MadNLP.unpack_obj(bcb::UniformBatchCallback, obj_val) = vec(bcb.obj_sign .* obj_val ./ bcb.obj_scale) +function unpack_obj!(dst, bcb::UniformBatchCallback, obj_val) + dst_mat = reshape(dst, 1, length(dst)) + @. dst_mat = bcb.obj_sign * obj_val / bcb.obj_scale +end function MadNLP.unpack_cons!(c_full, bcb::UniformBatchCallback, c, rhs, ind_ineq, slack) c_full .= c ./ bcb.con_scale .+ rhs diff --git a/src/batch/madnlp/rhs.jl b/src/batch/madnlp/rhs.jl index 3bcdd301..4ccdd9f7 100644 --- a/src/batch/madnlp/rhs.jl +++ b/src/batch/madnlp/rhs.jl @@ -1,4 +1,4 @@ -struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VI} +struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VI, SV, IV} values::MT n::Int m::Int @@ -6,13 +6,13 @@ struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VI} nub::Int ind_lb::VI ind_ub::VI - _primal::SubArray - _dual::SubArray - _primal_dual::SubArray - _dual_lb::SubArray - _dual_ub::SubArray - _xp_lr::SubArray - _xp_ur::SubArray + _primal::SV + _dual::SV + _primal_dual::SV + _dual_lb::SV + _dual_ub::SV + _xp_lr::IV + _xp_ur::IV end function BatchUnreducedKKTVector( @@ -24,14 +24,18 @@ function BatchUnreducedKKTVector( values = MT(undef, total, batch_size) fill!(values, zero(T)) - return BatchUnreducedKKTVector{T, MT, typeof(ind_lb)}( + primal = view(values, 1:n, :) + xp_lr = view(values, ind_lb, :) + SV = typeof(primal) + IV = typeof(xp_lr) + return BatchUnreducedKKTVector{T, MT, typeof(ind_lb), SV, IV}( values, n, m, nlb, nub, ind_lb, ind_ub, - view(values, 1:n, :), + primal, view(values, n+1:n+m, :), view(values, 1:n+m, :), view(values, n+m+1:n+m+nlb, :), view(values, n+m+nlb+1:n+m+nlb+nub, :), - view(values, ind_lb, :), + xp_lr, view(values, ind_ub, :), ) end @@ -45,16 +49,16 @@ MadNLP.dual_ub(bv::BatchUnreducedKKTVector) = bv._dual_ub xp_lr(bv::BatchUnreducedKKTVector) = bv._xp_lr xp_ur(bv::BatchUnreducedKKTVector) = bv._xp_ur -struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VI} +struct BatchPrimalVector{T, MT<:AbstractMatrix{T}, VI, SV, IV} values::MT nx::Int ns::Int ind_lb::VI ind_ub::VI - _variable::SubArray - _slack::SubArray - _lower::SubArray - _upper::SubArray + _variable::SV + _slack::SV + _lower::IV + _upper::IV end function BatchPrimalVector( @@ -66,11 +70,15 @@ function BatchPrimalVector( values = MT(undef, total, batch_size) fill!(values, zero(T)) - return BatchPrimalVector{T, MT, typeof(ind_lb)}( + variable = view(values, 1:nx, :) + lower = view(values, ind_lb, :) + SV = typeof(variable) + IV = typeof(lower) + return BatchPrimalVector{T, MT, typeof(ind_lb), SV, IV}( values, nx, ns, ind_lb, ind_ub, - view(values, 1:nx, :), + variable, view(values, nx+1:nx+ns, :), - view(values, ind_lb, :), + lower, view(values, ind_ub, :), ) end diff --git a/src/batch/structure.jl b/src/batch/structure.jl index e8302bd3..b26e60d2 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -65,7 +65,7 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, similar(_proto, Int64, 1, batch_size), # _term_gpu zeros(Int64, 1, batch_size), # _term_cpu fill!(similar(_proto, Int64, 1, 1), Int64(MadNLP.REGULAR)), # _any_nonregular_gpu - fill!(zeros(Int64, 1, 1), Int64(MadNLP.REGULAR)), # _any_nonregular_cpu + fill(Int64(MadNLP.REGULAR), 1, 1), # _any_nonregular_cpu MT(undef, 1, batch_size), # _norm_gpu_w MT(undef, 1, batch_size), # _norm_gpu_p fill!(similar(_proto, Int32), zero(Int32)), # _ls_error @@ -78,7 +78,7 @@ function UniformBatchWorkspace(::Type{MT}, ::Type{VT}, n::Int, m::Int, nlb::Int, ) end -mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, BVS} <: AbstractBatchMPCSolver{T, MT, VT} +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, BVS, KKT<:AbstractBatchKKTSystem{T}} <: AbstractBatchMPCSolver{T, MT, VT} batch_size::Int d::BatchUnreducedKKTVector{T, MT} @@ -105,7 +105,7 @@ mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, BVS} <: AbstractBat batch_cnt::BatchCounters logger::MadNLP.MadNLPLogger batch_views::BVS - kkt::AbstractBatchKKTSystem{T} + kkt::KKT del_w::MT del_c::MT @@ -217,7 +217,7 @@ function UniformBatchMPCSolver( batch_del_w = fill!(MT(undef, 1, batch_size), zero(T)) batch_del_c = fill!(MT(undef, 1, batch_size), zero(T)) - return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb), typeof(batch_views)}( + return UniformBatchMPCSolver{T, MT, VT, VI, typeof(bnlp), typeof(bcb), typeof(batch_views), typeof(batch_kkts)}( batch_size, batch_d, batch_p, batch_w1, batch_x, batch_xl, batch_xu, batch_zl, batch_zu, batch_f, From e22a8acdace1b65056e77d97f1807e90b413c334 Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 31 Mar 2026 01:49:16 -0400 Subject: [PATCH 49/51] remove dead code --- ext/MadIPMCUDAExt/cuda_wrapper.jl | 74 ------------------------------- src/batch/madipm/kernels.jl | 4 -- src/batch/madnlp/callback.jl | 5 --- src/batch/madnlp/nlpmodels.jl | 1 + src/batch/structure.jl | 2 - 5 files changed, 1 insertion(+), 85 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index d1f247a1..44ae68fc 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -103,81 +103,7 @@ function MadIPM.assemble_normal_system!( kernel!(n_rows, n_cols, Jtp, Jtj, Jtx, Cp, Cj, Cx, Dx, Tv; ndrange = n_rows) end -@kernel function count_normal_nnz!(Cp, @Const(Jtp), @Const(Jtj), @Const(n_rows), @Const(n_cols)) - i = @index(Global, Linear) - - # thread-local binary buffer - xb = @localmem UInt8 n_cols - for k = 1:n_cols - xb[k] = 0 - end - - for c = Jtp[i]:Jtp[i+1]-1 - j = Jtj[c] - xb[j] = 1 - end - - count = 0 - for j = i:n_rows - for c = Jtp[j]:Jtp[j+1]-1 - k = Jtj[c] - if xb[k] == 1 - count += 1 - break - end - end - end - - Cp[i+1] = count - nothing -end - -@kernel function fill_normal_indices!(Cj, @Const(Cp), @Const(Jtp), @Const(Jtj), @Const(n_rows), @Const(n_cols)) - i = @index(Global, Linear) - - xb = @localmem UInt8 n_cols - for k = 1:n_cols - xb[k] = 0 - end - - for c = Jtp[i]:Jtp[i+1]-1 - j = Jtj[c] - xb[j] = 1 - end - pos = Cp[i] - for j = i:n_rows - for c = Jtp[j]:Jtp[j+1]-1 - k = Jtj[c] - if xb[k] == 1 - Cj[pos] = j - pos += 1 - break - end - end - end - nothing -end - -function MadIPM.build_normal_system( - n_rows, - n_cols, - Jtp::CuVector{Ti}, - Jtj::CuVector{Ti}, -) where {Ti} - backend = CUDABackend() - Cp = CUDA.ones(Ti, n_rows + 1) - kernel1! = count_normal_nnz!(backend) - kernel1!(Cp, Jtp, Jtj, n_rows, n_cols; ndrange = n_rows) - - Cp = cumsum(Cp) - nnz_JtJ = CUDA.@allowscalar (Cp[end] - 1) - Cj = CUDA.zeros(Ti, nnz_JtJ) - - kernel2! = fill_normal_indices!(backend) - kernel2!(Cj, Cp, Jtp, Jtj, n_rows, n_cols; ndrange = n_rows) - return (Cp, Cj) -end MadIPM.sparse_csc_format(::Type{<:CuArray}) = CuSparseMatrixCSC MadIPM._colptr(A::CuSparseMatrixCSC) = A.colPtr diff --git a/src/batch/madipm/kernels.jl b/src/batch/madipm/kernels.jl index 077c420a..185163a9 100644 --- a/src/batch/madipm/kernels.jl +++ b/src/batch/madipm/kernels.jl @@ -403,11 +403,9 @@ end corrected_p = one(T) @inbounds if max_ap < one(T) if best_xl <= best_xu && i_xl > 0 - idx = ind_lb[i_xl] zl_stepped = zl_r[i_xl, j] + max_ad * d_vals[dlb_off + i_xl, j] corrected_p = (x_lr[i_xl, j] - xl_r[i_xl, j] - mu_j / zl_stepped) / (-dx_lr[i_xl, j]) elseif i_xu > 0 - idx = ind_ub[i_xu] zu_stepped = zu_r[i_xu, j] + max_ad * d_vals[dub_off + i_xu, j] corrected_p = (xu_r[i_xu, j] - x_ur[i_xu, j] - mu_j / zu_stepped) / dx_ur[i_xu, j] end @@ -418,11 +416,9 @@ end corrected_d = one(T) @inbounds if max_ad < one(T) if best_zl <= best_zu && i_zl > 0 - idx = ind_lb[i_zl] x_gap = x_lr[i_zl, j] + max_ap * dx_lr[i_zl, j] - xl_r[i_zl, j] corrected_d = -(zl_r[i_zl, j] - mu_j / x_gap) / d_vals[dlb_off + i_zl, j] elseif i_zu > 0 - idx = ind_ub[i_zu] x_gap = xu_r[i_zu, j] - x_ur[i_zu, j] - max_ap * dx_ur[i_zu, j] corrected_d = -(zu_r[i_zu, j] - mu_j / x_gap) / d_vals[dub_off + i_zu, j] end diff --git a/src/batch/madnlp/callback.jl b/src/batch/madnlp/callback.jl index 90281710..f3669381 100644 --- a/src/batch/madnlp/callback.jl +++ b/src/batch/madnlp/callback.jl @@ -220,11 +220,6 @@ end -function MadNLP.n_variables(bcb::UniformBatchCallback{T, VT, MT, VI, BM, FH}) where {T, VT, MT, VI, BM, FH<:MadNLP.MakeParameter} - return length(bcb.fixed_handler.free) -end -MadNLP.n_variables(bcb::UniformBatchCallback) = bcb.nvar - function MadNLP._jac_sparsity_wrapper!(bcb::UniformBatchCallback, I::AbstractVector, J::AbstractVector) copyto!(I, bcb.jac_I) copyto!(J, bcb.jac_J) diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl index 2ac4d30b..2c9a7bc7 100644 --- a/src/batch/madnlp/nlpmodels.jl +++ b/src/batch/madnlp/nlpmodels.jl @@ -39,6 +39,7 @@ end function MadNLP.unpack_z!( Z_full::AbstractMatrix, bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, z_free, ) where {T,VT,MT,VI,BM,FH<:MadNLP.MakeParameter,EH} + fill!(Z_full, zero(T)) Z_full[bcb.fixed_handler.free, :] .= z_free ./ bcb.obj_scale end diff --git a/src/batch/structure.jl b/src/batch/structure.jl index b26e60d2..44db0b73 100644 --- a/src/batch/structure.jl +++ b/src/batch/structure.jl @@ -116,8 +116,6 @@ end _get_ind_lb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_lb _get_ind_ub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_ub -_get_ind_llb(bs::AbstractBatchMPCSolver) = bs.bcb.ind_llb -_get_ind_uub(bs::AbstractBatchMPCSolver) = bs.bcb.ind_uub active_batch_size(bs::AbstractBatchMPCSolver) = local_batch_size(active_view(bs.batch_views)) function update_active_set!(state::BatchViewState, status::Vector{MadNLP.Status}) From 748f37201e03a3aaa41bda9c96a5306c96c0602a Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 31 Mar 2026 01:49:30 -0400 Subject: [PATCH 50/51] fix race in gpu compact --- ext/MadIPMCUDAExt/cuda_batch_kernels.jl | 13 +++++++------ src/batch/KKT/Sparse/augmented.jl | 9 ++++++--- src/batch/views.jl | 2 +- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl index b18e71fa..e71ae2b7 100644 --- a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -26,12 +26,9 @@ end @inbounds dst[i, local_to_root[j]] = src[i, j] end -@kernel function _compact_active_columns_inplace_kernel!(dst, @Const(local_to_root)) +@kernel function _compact_active_columns_inplace_kernel!(dst, @Const(src), @Const(local_to_root)) i, j = @index(Global, NTuple) - src_j = local_to_root[j] - @inbounds if src_j != j - dst[i, j] = dst[i, src_j] - end + @inbounds dst[i, j] = src[i, local_to_root[j]] end @inline function _atomic_colreduce!(::typeof(+), out, j, value) @@ -76,10 +73,14 @@ end function MadIPM.compact_active_columns_inplace!( dst::CuMatrix{T}, batch_view::MadIPM.BatchView, + scratch::CuMatrix{T}, ) where T na = MadIPM.local_batch_size(batch_view) + na == 0 && return dst + nrows = size(dst, 1) + copyto!(view(scratch, 1:nrows, :), view(dst, 1:nrows, :)) backend = CUDABackend() - _compact_active_columns_inplace_kernel!(backend)(dst, MadIPM.local_to_root_dev(batch_view); ndrange=(size(dst, 1), na)) + _compact_active_columns_inplace_kernel!(backend)(dst, scratch, MadIPM.local_to_root_dev(batch_view); ndrange=(nrows, na)) return dst end diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl index 99a78d39..a60f3a4f 100644 --- a/src/batch/KKT/Sparse/augmented.jl +++ b/src/batch/KKT/Sparse/augmented.jl @@ -5,6 +5,7 @@ struct SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT, BVS} <: AbstractBat aug_J::VI32 # shared column indices batch_solver::LS # batched linear solver rhs_buffer::MT # (n+m) × batch_size for batch solve + compact_buffer::MT # scratch for GPU compact (max(nnz_csc, n+m) × batch_size) batch_size::Int batch_views::BVS aug_com_nzvals::MT # (nnz_csc × batch_size) CSC nonzero values @@ -114,8 +115,10 @@ function MadNLP.create_kkt_system( LS = typeof(batch_ls) VI32 = typeof(I) OPT = typeof(jt_op) + compact_buffer = similar(nzVals, max(nnz_csc, aug_vec_length), batch_size) + return SparseUniformBatchKKTSystem{T, LS, MT, VI, VI32, OPT, typeof(batch_views)}( - nzVals, I, J, batch_ls, rhs_buffer, batch_size, batch_views, + nzVals, I, J, batch_ls, rhs_buffer, compact_buffer, batch_size, batch_views, aug_com_nzvals, batch_csc_map, n_tot, m, n_hess, reg, l_diag, u_diag, l_lower, u_lower, hess_op, jt_op, j_op, @@ -125,7 +128,7 @@ end function MadNLP.factorize_kkt!(bkkt::SparseUniformBatchKKTSystem) factor_view = active_view(bkkt.batch_views) if !is_identity_view(factor_view) - compact_active_columns_inplace!(bkkt.aug_com_nzvals, factor_view) + compact_active_columns_inplace!(bkkt.aug_com_nzvals, factor_view, bkkt.compact_buffer) end factorize_active!(bkkt.batch_solver, factor_view) return @@ -134,7 +137,7 @@ end function MadNLP.solve_linear_system!(bkkt::SparseUniformBatchKKTSystem{T}, rhs::AbstractMatrix) where T active = active_view(bkkt.batch_views) if !is_identity_view(active) - compact_active_columns_inplace!(rhs, active) + compact_active_columns_inplace!(rhs, active, bkkt.compact_buffer) end solve_active!(bkkt.batch_solver, rhs, active) return rhs diff --git a/src/batch/views.jl b/src/batch/views.jl index 4f27a195..1b35ef58 100644 --- a/src/batch/views.jl +++ b/src/batch/views.jl @@ -149,7 +149,7 @@ function gather_batch_view_columns!( return dst end -function compact_active_columns_inplace!(dst::AbstractMatrix{T}, batch_view::BatchView) where T +function compact_active_columns_inplace!(dst::AbstractMatrix{T}, batch_view::BatchView, scratch=nothing) where T roots = batch_view.local_to_root @inbounds for j in 1:batch_view.n src_j = roots[j] From 135f8226ad8844b758654439afbb79032bfe17ef Mon Sep 17 00:00:00 2001 From: "Klamkin, Michael" Date: Tue, 28 Apr 2026 14:52:26 -0400 Subject: [PATCH 51/51] merge with main --- ext/MadIPMCUDAExt/operators.jl | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ext/MadIPMCUDAExt/operators.jl b/ext/MadIPMCUDAExt/operators.jl index 7ba3d79c..ac1e174b 100644 --- a/ext/MadIPMCUDAExt/operators.jl +++ b/ext/MadIPMCUDAExt/operators.jl @@ -22,13 +22,15 @@ for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T}), :BlasFloat), @eval begin function MadIPMOperator(A::$SparseMatrixType; transa::Char='N', symmetric::Bool=false, spmm_ncols::Int=0) where T <: $BlasType m, n = size(A) + op_in = transa == 'N' ? n : m + op_out = transa == 'N' ? m : n alpha = Ref{T}(one(T)) beta = Ref{T}(zero(T)) bool = symmetric && (nnz(A) > 0) mat = bool ? tril(A, -1) + A' : A descA = CUSPARSE.CuSparseMatrixDescriptor(mat, 'O') - descX = CUSPARSE.CuDenseVectorDescriptor(T, n) - descY = CUSPARSE.CuDenseVectorDescriptor(T, m) + descX = CUSPARSE.CuDenseVectorDescriptor(T, op_in) + descY = CUSPARSE.CuDenseVectorDescriptor(T, op_out) algo = CUSPARSE.CUSPARSE_SPMV_ALG_DEFAULT buffer_size = Ref{Csize_t}() CUSPARSE.cusparseSpMV_bufferSize(CUSPARSE.handle(), transa, alpha, descA, descX, beta, descY, T, algo, buffer_size) @@ -54,7 +56,7 @@ for (SparseMatrixType, BlasType) in ((:(CuSparseMatrixCSR{T}), :BlasFloat), else CuVector{UInt8}(undef, 0) end - return MadIPMOperator{T,M,M2}(T, m, n, A, mat, transa, descA, buffer, spmm_buffer, alpha, beta) + return MadIPMOperator{T,M,M2}(T, op_out, op_in, A, mat, transa, descA, buffer, spmm_buffer, alpha, beta) end end end