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..17257433 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'] @@ -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/Project.toml b/Project.toml index 1ab66531..49d8cac7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,9 +1,10 @@ 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] +BatchQuadraticModels = "412afcf0-3ec6-4826-be17-8d792afa05a8" LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" @@ -13,17 +14,25 @@ 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" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +[sources] +BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl.git"} + [extensions] -MadIPMCUDAExt = ["CUDA", "KernelAbstractions", "MadNLPGPU"] +MadIPMCUDAExt = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPGPU"] MadIPMMathOptInterfaceExt = "MathOptInterface" [compat] +Atomix = "1" +BatchQuadraticModels = "0.1" CUDA = "5.4.0" +CUDSS = "0.6" KernelAbstractions = "0.9" LDLFactorizations = "0.10.1" LinearAlgebra = "1.10" @@ -39,8 +48,14 @@ 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" +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 = ["Atomix", "CUDA", "CUDSS", "KernelAbstractions", "MadNLPTests", "MadNLPGPU", "MathOptInterface", "Test"] diff --git a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl index aec17c5a..a0569f88 100644 --- a/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl +++ b/ext/MadIPMCUDAExt/MadIPMCUDAExt.jl @@ -3,15 +3,25 @@ module MadIPMCUDAExt using LinearAlgebra using SparseArrays using NLPModels +using BatchQuadraticModels using QuadraticModels using CUDA using CUDA.CUSPARSE +using CUDSS using KernelAbstractions +import Atomix import QuadraticModels: SparseMatrixCOO +import LinearAlgebra: BlasFloat import MadIPM +import MadNLP -include("cuda_wrapper.jl") include("operators.jl") +include("cuda_wrapper.jl") +include("cuda_batch_kernels.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) @@ -138,4 +148,3 @@ function Base.convert(::Type{QuadraticModel{T, S}}, qp::QuadraticModel{T}) where end end - diff --git a/ext/MadIPMCUDAExt/cuda_batch_kernels.jl b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl new file mode 100644 index 00000000..e71ae2b7 --- /dev/null +++ b/ext/MadIPMCUDAExt/cuda_batch_kernels.jl @@ -0,0 +1,441 @@ +@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 + +@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(src), @Const(local_to_root)) + i, j = @index(Global, NTuple) + @inbounds dst[i, j] = src[i, local_to_root[j]] +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 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, + 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, scratch, MadIPM.local_to_root_dev(batch_view); ndrange=(nrows, na)) + return dst +end + +function MadNLP._set_con_scale_sparse!( + con_scale::CuMatrix{T}, + jac_I::CuVector{<:Integer}, + jac_buffer::CuMatrix{T}, +) where T + nnzj = length(jac_I) + if nnzj > 0 + backend = CUDABackend() + _set_con_scale_kernel!(backend)(con_scale, jac_I, jac_buffer; ndrange=nnzj) + end + return con_scale +end + +_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 + +_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 + 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 + return +end + +_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 _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 + 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 + return +end + +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, +) + 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, + 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)) + i, j = @index(Global, NTuple) + @inbounds values[ind_lb[i], j] -= values[lb_off + i, j] / l_diag[i, j] +end + +@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)) + 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)) + return +end diff --git a/ext/MadIPMCUDAExt/cuda_wrapper.jl b/ext/MadIPMCUDAExt/cuda_wrapper.jl index 56e173e3..44ae68fc 100644 --- a/ext/MadIPMCUDAExt/cuda_wrapper.jl +++ b/ext/MadIPMCUDAExt/cuda_wrapper.jl @@ -4,8 +4,7 @@ import LinearAlgebra: BlasFloat @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 @@ -14,11 +13,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)) - KernelAbstractions.synchronize(backend) + _transfer_to_map!(backend)(dest, map, src; ndrange=length(map)) end return end @@ -99,88 +101,46 @@ 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)) - 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 +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 - Cp[i+1] = count - nothing +# 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 -@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 +MadIPM.is_factorized(::MadNLPGPU.CUDSSSolver) = true - 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 +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.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) - KernelAbstractions.synchronize(backend) - - 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) - KernelAbstractions.synchronize(backend) - return (Cp, Cj) +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) + CUDSS.cudss("solve", s.inner, s.x_gpu, s.b_gpu, asynchronous=s.opt.cudss_asynchronous) + return end - -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 diff --git a/ext/MadIPMCUDAExt/operators.jl b/ext/MadIPMCUDAExt/operators.jl index 6a288e90..ac1e174b 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) op_in = transa == 'N' ? n : m op_out = transa == 'N' ? m : n @@ -41,11 +42,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, op_out, op_in, 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, op_out, op_in, 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")) @@ -54,3 +81,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/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 85fb907e..28517c05 100644 --- a/src/MadIPM.jl +++ b/src/MadIPM.jl @@ -7,7 +7,15 @@ 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!, + batch_maximum! include("utils.jl") include("structure.jl") @@ -16,7 +24,20 @@ include("KKT/normalkkt.jl") include("linear_solver.jl") include("solver.jl") -export MPCSolver, madipm +include("batch/utils.jl") +include("batch/views.jl") +include("batch/madnlp/rhs.jl") +include("batch/madnlp/callback.jl") +include("batch/KKT/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 MadNLP.madsuite(::Val{:madipm}, args...; kwargs...) = madipm(args...; kwargs...) diff --git a/src/batch/KKT/KKT.jl b/src/batch/KKT/KKT.jl new file mode 100644 index 00000000..2ff64d46 --- /dev/null +++ b/src/batch/KKT/KKT.jl @@ -0,0 +1,18 @@ +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") diff --git a/src/batch/KKT/Sparse/augmented.jl b/src/batch/KKT/Sparse/augmented.jl new file mode 100644 index 00000000..a60f3a4f --- /dev/null +++ b/src/batch/KKT/Sparse/augmented.jl @@ -0,0 +1,300 @@ +"""Batched version of SparseKKTSystem""" +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 + 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 + 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 + # Operators for batch SpMV (jtprod! and KKT mul!) + hess_op::OPT + jt_op::OPT + j_op::OPT +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, 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 + + 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 = 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 + + 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) + u_diag = similar(nzVals, nub, batch_size) + l_lower = similar(nzVals, nlb, batch_size) + u_lower = similar(nzVals, nub, batch_size) + + 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, 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, + ) +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, bkkt.compact_buffer) + end + factorize_active!(bkkt.batch_solver, factor_view) + return +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, bkkt.compact_buffer) + end + solve_active!(bkkt.batch_solver, rhs, active) + 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) + 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 + 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 + +function MadNLP.solve_kkt!(bkkt::SparseUniformBatchKKTSystem, batch_solver::AbstractBatchMPCSolver) + d = batch_solver.d + + MadNLP.reduce_rhs!(bkkt, d) + + 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) + 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 +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) + batch_spmv!(res, bkkt.jt_op, MadNLP.full(y)) + 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) + 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 + +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 + xv = MadNLP.full(x) + # 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/madipm/kernels.jl b/src/batch/madipm/kernels.jl new file mode 100644 index 00000000..185163a9 --- /dev/null +++ b/src/batch/madipm/kernels.jl @@ -0,0 +1,502 @@ +function dual_objective!(dual_obj, y_vals, rhs_vals, zl_r, xl_r, zu_r, xu_r, + 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 + batch_mapreduce!(*, +, zero(T), sum_lb, zl_r, xl_r) + dual_obj .+= sum_lb + end + if nub > 0 + batch_mapreduce!(*, +, zero(T), sum_ub, zu_r, xu_r) + 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) + 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_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) + 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 = 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 + 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) + + 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 + +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) + + _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 + +function update_barrier!(::Mehrotra, solver::AbstractBatchMPCSolver, mu_affine) + ws = solver.workspace + T = eltype(ws.mu_curr) + + has_inequalities = (solver.d.nlb + solver.d.nub) > 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) + + if nlb > 0 + _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 + _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)) + 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 _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 +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 _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, +) + 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 + 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 + 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) + + # dual step + corrected_d = one(T) + @inbounds if max_ad < one(T) + if best_zl <= best_zu && i_zl > 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 > 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) + return +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 + + dlb_off = d.n + d.m + dub_off = d.n + d.m + d.nlb + bs = batch_solver.batch_size + + _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 +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, 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, 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) + 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 new file mode 100644 index 00000000..840c047d --- /dev/null +++ b/src/batch/madipm/solver.jl @@ -0,0 +1,619 @@ +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 + 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 + batch_mapreduce!(-, min, T(Inf), _s1, xl, lb) + @. _s1 = min(_s1, zero(T)) + else + fill!(_s1, zero(T)) + end + if nub_init > 0 + batch_mapreduce!(-, min, T(Inf), _s2, ub, xu) + @. _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 + batch_mapreduce!(identity, min, T(Inf), _s1, zl) + @. _s1 = min(_s1, zero(T)) + else + fill!(_s1, zero(T)) + end + if nub_init > 0 + batch_mapreduce!(identity, min, T(Inf), _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 + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_lb, xl, zl) + μ .+= ws.sum_lb + batch_mapreduce!((a, b) -> a * b, +, zero(T), ws.sum_lb, lb, zl) + μ .-= ws.sum_lb + end + if nub_init > 0 + 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 + batch_mapreduce!(identity, +, zero(T), ws.sum_lb, zl) + else + fill!(ws.sum_lb, zero(T)) + end + if nub_init > 0 + batch_mapreduce!(identity, +, zero(T), 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 + batch_mapreduce!(-, +, zero(T), ws.sum_lb, xl, lb) + else + fill!(ws.sum_lb, zero(T)) + end + if nub_init > 0 + batch_mapreduce!(-, +, zero(T), ws.sum_ub, ub, xu) + 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, + ws.bx; + 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), + ws.bx, + ) + 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) + + 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) + 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) + 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)) + 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 + 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 + 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 + + 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)) + @. ws.inf_du /= max(one(T), ws.norm_c) + + get_inf_compl!(ws.inf_compl, x, xl, zl, xu, zu, + 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), + ws.sum_lb, ws.sum_ub, nlb, nub) + + 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) + 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]) + 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 + ws.status[i] = MadNLP.MAXIMUM_WALLTIME_EXCEEDED + end + end + return true +end + +function solve_system!( + d::BatchUnreducedKKTVector{T}, + batch_solver::AbstractBatchMPCSolver{T}, + p::BatchUnreducedKKTVector{T}, +) where T + 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)) + + 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) + batch_maximum!(ws._norm_gpu_w, _fw) # (1,bs) per-instance norm_w + _fw .= abs.(MadNLP.full(p)) + 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 + +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)) + 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) + 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 _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 + # 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 + +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 + 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 = is_factorized!( + failed_locals, batch_solver.kkt.batch_solver, factor_view, + ) + nfailed == 0 && break + _bump_failed_regularization!(batch_solver, failed_locals, nfailed) + end + return +end + + +function apply_step!(batch_solver::AbstractBatchMPCSolver) + ws = batch_solver.workspace + 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 + MadNLP.full(x) .+= ws.alpha_p .* MadNLP.primal(d) + + # y += alpha_d * d_dual + MadNLP.full(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 + + _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 + +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) + fill!(batch_solver.workspace._ls_error, zero(Int32)) + 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 + buf = ws.active_mask_cpu + fill_batch_view_mask!(buf, active_view(batch_solver.batch_views)) + copyto!(ws.active_mask, buf) +end + +function mpc!(batch_solver::AbstractBatchMPCSolver) + while true + MadNLP.print_iter(batch_solver) + update_termination_criteria!(batch_solver) + changed = update_termination_status!(batch_solver) + if changed + update_active_set!(batch_solver) + active_batch_size(batch_solver) == 0 && return + _update_active_mask!(batch_solver) + end + mpc_step!(batch_solver) + end +end + +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 + 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) + status_counts = Dict{MadNLP.Status, Int}() + for i in 1:bs + 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 + + 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 = active_batch_size(batch_solver) + 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 ", 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), + )) + return +end diff --git a/src/batch/madnlp/callback.jl b/src/batch/madnlp/callback.jl new file mode 100644 index 00000000..f3669381 --- /dev/null +++ b/src/batch/madnlp/callback.jl @@ -0,0 +1,240 @@ +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, + check_batch_structure::Bool=true, +) 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 + + 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) + 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 diff --git a/src/batch/madnlp/initialization.jl b/src/batch/madnlp/initialization.jl new file mode 100644 index 00000000..c086c975 --- /dev/null +++ b/src/batch/madnlp/initialization.jl @@ -0,0 +1,108 @@ +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, + bx_buffer; + 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) + 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, :) + 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, + 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, 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, 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, :]) + 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..5164966c --- /dev/null +++ b/src/batch/madnlp/kernels.jl @@ -0,0 +1,29 @@ +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) + 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, sum_lb, sum_ub, nlb, nub) + T = eltype(inf_compl) + if nlb > 0 + 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 + 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 + @. inf_compl = max(sum_lb, sum_ub) + return inf_compl +end + diff --git a/src/batch/madnlp/linear_solver.jl b/src/batch/madnlp/linear_solver.jl new file mode 100644 index 00000000..208dac11 --- /dev/null +++ b/src/batch/madnlp/linear_solver.jl @@ -0,0 +1,65 @@ +# FIXME: threads/polyester version + +struct LoopedBatchLinearSolver{T, VT, LS<:MadNLP.AbstractLinearSolver{T}} <: MadNLP.AbstractLinearSolver{T} + solvers::Vector{LS} + batch_size::Int +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{T, VT, eltype(individual_solvers)}(individual_solvers, batch_size) +end + +function is_factorized!( + 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 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 +is_factorized!(::Vector{Int32}, ::MadNLP.AbstractLinearSolver, ::BatchView) = 0 diff --git a/src/batch/madnlp/nlpmodels.jl b/src/batch/madnlp/nlpmodels.jl new file mode 100644 index 00000000..2c9a7bc7 --- /dev/null +++ b/src/batch/madnlp/nlpmodels.jl @@ -0,0 +1,206 @@ +function MadNLP._update_x!( + bcb::UniformBatchCallback{T,VT,MT,VI,BM,FH,EH}, x, +) where {T,VT,MT,VI,BM,FH,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,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,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} + fill!(Z_full, zero(T)) + 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,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,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 + +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 + 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,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,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; + 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, obj_weight, 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; + 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, obj_weight, 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..4ccdd9f7 --- /dev/null +++ b/src/batch/madnlp/rhs.jl @@ -0,0 +1,91 @@ +struct BatchUnreducedKKTVector{T, MT<:AbstractMatrix{T}, VI, SV, IV} + values::MT + n::Int + m::Int + nlb::Int + nub::Int + ind_lb::VI + ind_ub::VI + _primal::SV + _dual::SV + _primal_dual::SV + _dual_lb::SV + _dual_ub::SV + _xp_lr::IV + _xp_ur::IV +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)) + + 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, + 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, :), + xp_lr, + view(values, ind_ub, :), + ) +end + +MadNLP.full(bv::BatchUnreducedKKTVector) = bv.values +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}, VI, SV, IV} + values::MT + nx::Int + ns::Int + ind_lb::VI + ind_ub::VI + _variable::SV + _slack::SV + _lower::IV + _upper::IV +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)) + + 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, + variable, + view(values, nx+1:nx+ns, :), + lower, + view(values, ind_ub, :), + ) +end + +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 diff --git a/src/batch/structure.jl b/src/batch/structure.jl new file mode 100644 index 00000000..44db0b73 --- /dev/null +++ b/src/batch/structure.jl @@ -0,0 +1,232 @@ +struct UniformBatchWorkspace{T, VT<:AbstractVector{T}, MT<:AbstractMatrix{T}, MI<:AbstractMatrix{Int32}, MI64<:AbstractMatrix{Int64}} + alpha_xl::MT + alpha_xu::MT + alpha_zl::MT + alpha_zu::MT + alpha_p::MT + alpha_d::MT + 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_gpu::MI64 + _term_cpu::Matrix{Int64} + _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 + bg::MT + bv::MT +end + +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)) + 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 + 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 + 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(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 + 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, m, batch_size), # bv + ) +end + +mutable struct UniformBatchMPCSolver{T, MT, VT, VI, BM, BCB, BVS, KKT<:AbstractBatchKKTSystem{T}} <: AbstractBatchMPCSolver{T, MT, VT} + batch_size::Int + + d::BatchUnreducedKKTVector{T, MT} + p::BatchUnreducedKKTVector{T, MT} + _w1::BatchUnreducedKKTVector{T, MT} + + 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} + + opt::IPMOptions + batch_cnt::BatchCounters + logger::MadNLP.MadNLPLogger + batch_views::BVS + kkt::KKT + + 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 +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) + + +""" + UniformBatchMPCSolver(bnlp::AbstractBatchNLPModel; linear_solver, kwargs...) + +Construct a batch solver from a `AbstractBatchNLPModel`. +""" +function UniformBatchMPCSolver( + 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, MT} + 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, + check_batch_structure=check_batch_structure, + ) + + 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_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) + 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_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, 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)) + + 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, + batch_y, batch_c, batch_jacl, batch_rhs, + batch_correction_lb, batch_correction_ub, + workspace, + ipm_opt, batch_cnt, logger, + batch_views, + 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..bbabdd49 --- /dev/null +++ b/src/batch/utils.jl @@ -0,0 +1,159 @@ +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) +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 _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, + ) +end + +function _build_jt_op( + aug_I, aug_J, jac_range, n_tot, + 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] + 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 _build_batch_op(nzVals, nz_map, con_map, coo_I, n_tot) +end + +function _build_j_op( + aug_I, aug_J, jac_range, n_tot, m, + 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) + 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 _build_batch_op(nzVals, nz_map, var_map, coo_I, m) +end + +function _build_hess_op( + aug_I, aug_J, n_tot, n_hess, + nzVals::AbstractMatrix{T}, aug_csc_map, +) where T + if n_hess == 0 + nz_map = similar(aug_csc_map, 0) + var_map = similar(aug_csc_map, 0) + 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 + 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) + + 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 + 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 _build_batch_op(nzVals, nz_map, var_map, coo_rows, n_tot) +end + +struct BatchVector{T, MT<:AbstractMatrix{T}} + values::MT +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)) + return BatchVector{T, MT}(values) +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)) diff --git a/src/batch/views.jl b/src/batch/views.jl new file mode 100644 index 00000000..1b35ef58 --- /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, scratch=nothing) 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/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 new file mode 100644 index 00000000..e2129e29 --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,71 @@ +[deps] +AMD = "14f7f29c-3bd6-536c-9a0b-7339e30b5a3e" +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" +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" +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" +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" +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" +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] +BatchQuadraticModels = {rev = "main", url = "https://github.com/klamike/BatchQuadraticModels.jl"} +MadIPM = {path = ".."} 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/runtests.jl b/test/runtests.jl index e68a1dfd..3213ddb6 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("batch/views.jl") +include("batch/solver.jl") if CUDA.functional() include("test_gpu.jl") + include("batch/gpu.jl") end