diff --git a/Project.toml b/Project.toml index d30aa1e..7752aea 100644 --- a/Project.toml +++ b/Project.toml @@ -14,6 +14,7 @@ NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" SolverCore = "ff4d7338-4cf1-434d-91df-b86cb86fb843" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6" [compat] FillArrays = "1.13" @@ -26,4 +27,5 @@ NLPModelsJuMP = "0.13.2" SolverCore = "0.3.8" SparseArrays = "1.10" Test = "1.10" +UnsafeArrays = "1.0.8" julia = "1.10" diff --git a/src/BurerMonteiro.jl b/src/BurerMonteiro.jl index 0dc6b36..fe90fcd 100644 --- a/src/BurerMonteiro.jl +++ b/src/BurerMonteiro.jl @@ -127,6 +127,7 @@ function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector) X = Solution(x, model.dim) # We don't call `cons!` as we don't want to include `-b` since the constraint # is encoded as `b <= c(x) <= b` and we just need to specify `c(x)` here. + # We don't use the version with buffers because that destroys the low-rank structure of `x` return NLPModels.jprod!(model.model, X, X, cx) end diff --git a/src/model.jl b/src/model.jl index 64c709e..32c20c1 100644 --- a/src/model.jl +++ b/src/model.jl @@ -5,6 +5,7 @@ import LinearAlgebra import MutableArithmetics as MA import MathOptInterface as MOI import NLPModels +import UnsafeArrays struct Dimensions num_scalars::Int64 @@ -22,6 +23,7 @@ end struct MatrixIndex value::Int64 end +Base.broadcastable(i::MatrixIndex) = Ref(i) abstract type AbstractSolution{T} <: AbstractVector{T} end @@ -293,17 +295,117 @@ function dual_cons!( return model.C[i] - jtprod!(buffer[i], model, mat_idx, y) end +function NLPModels.jtprod!( + model::Model, + _::AbstractVector, + y::AbstractVector, + vJ::AbstractVector, + buffer, +) + vJ[ScalarIndex] .= jtprod(model, ScalarIndex, y) + for mat_idx in matrix_indices(model) + i = mat_idx.value + vJ[mat_idx] .= jtprod!(buffer[i], model, mat_idx, y) + end +end + NLPModels.grad(model::Model, ::Type{ScalarIndex}) = model.d_lin NLPModels.grad(model::Model, i::MatrixIndex) = model.C[i.value] cons_constant(model::Model) = model.b -function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector) - NLPModels.jprod!(model, x, x, cx) +function NLPModels.cons!( + model::Model, + x::AbstractVector, + cx::AbstractVector, + args::Vararg{Any,N}, +) where {N} + NLPModels.jprod!(model, x, x, cx, args...) cx .-= model.b return cx end +# `SparseMatrixCSC` is stored with an offset by column. +# This means that getting view `view(A, :, I)` can be handles efficently, +# these give `SparseMatrixCSCView` (if `I` is a `UnitRange`) and +# `SparseMatrixCSCColumnSubset` otherwise. +# In `schur.jl`, we therefore get a `SparseMatrixCSCColumnSubset`. +# Since we want to use subsets of constraint indices, we use the columns +# of `A` for constraint indices and the rows of `A` for matrix indices. +function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T} + nnz = sum(1:model.meta.ncon; init = 0) do j + return SparseArrays.nnz(model.A[i.value, j]) + end + I = zeros(Int64, nnz) + J = zeros(Int64, nnz) + V = zeros(T, nnz) + offset = 0 + for j in 1:model.meta.ncon + Ai, Av = SparseArrays.findnz(model.A[i.value, j][:]) + K = offset .+ eachindex(Ai) + I[K] = Ai + J[K] .= j + V[K] = Av + offset += length(Ai) + end + A = SparseArrays.sparse( + I, + J, + V, + side_dimension(model, i)^2, + model.meta.ncon, + ) + return A +end + +# We define a new type so that we can define a custom `getindex` +struct JProdBuffer{T} + A::Vector{SparseArrays.SparseMatrixCSC{T,Int64}} + cache::Vector{T} +end + +function buffer_for_jprod(model::Model) + return JProdBuffer( + [buffer_for_jprod(model, i) for i in matrix_indices(model)], + zeros(model.meta.ncon), + ) +end + +Base.getindex(buf::JProdBuffer, i::MatrixIndex) = (buf.A[i.value], buf.cache) + +_vec(x::AbstractVector) = x +_vec(x::AbstractArray) = UnsafeArrays.uview(x, :) +_vec(x::Base.ReshapedArray) = _vec(parent(x)) + +function _add_jprod!(V, Jv, A, cache) + LinearAlgebra.mul!(cache, A', _vec(V)) + Jv .+= cache + return Jv +end + +function add_sub_jprod!( + _::Model, + _::MatrixIndex, + V::AbstractMatrix, + Jv::AbstractVector, + I, + buffer, +) + A, cache = buffer + # `view(cache, I)` would be terribly slow, only the number of elements of `I` matter here + return _add_jprod!(V, Jv, view(A, :, I), view(cache, eachindex(I))) +end + +function add_jprod!( + ::Model, + ::MatrixIndex, + V::AbstractMatrix, + Jv::AbstractVector, + buffer, +) + return _add_jprod!(V, Jv, buffer...) +end + function add_jprod!( model::Model, i::MatrixIndex, @@ -320,10 +422,11 @@ function NLPModels.jprod!( _::AbstractVector, v::AbstractVector, Jv::AbstractVector, -) + args::Vararg{Any,N}, # Optional buffer +) where {N} LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex]) for i in matrix_indices(model) - add_jprod!(model, i, v[i], Jv) + add_jprod!(model, i, v[i], Jv, getindex.(args, i)...) end return Jv end diff --git a/src/schur.jl b/src/schur.jl index 16303d4..8760ca6 100644 --- a/src/schur.jl +++ b/src/schur.jl @@ -42,7 +42,9 @@ function _dot( return result end -function buffer_for_schur_complement(model::Model, κ) +# The `jprod!` buffer is guaranteed to be the first argument of the tuple. +# This assumption is used by Loraine. +function buffer_for_schur_complement(model::Model{T}, κ) where {T} n = model.meta.ncon σ = zeros(Int64, n, num_matrices(model)) last_dense = zeros(Int64, num_matrices(model)) @@ -57,46 +59,46 @@ function buffer_for_schur_complement(model::Model, κ) last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0) end - return σ, last_dense + AW = [zeros(T, dim, dim) # /!\ it's the same zero everywhere, might be an issue with BigFloat + for dim in model.msizes] + WAW = copy.(AW) + + return buffer_for_jprod(model), AW, WAW, σ, last_dense end -function add_schur_complement!(buffer, model::Model, W, ::Type{MatrixIndex}, H) +function add_schur_complement!(model::Model, W, ::Type{MatrixIndex}, H, buffer) for i in matrix_indices(model) - add_schur_complement!(buffer, model, i, W[i], H) + add_schur_complement!(model, i, W[i], H, buffer) end return H end # /!\ W needs to be symmetric function add_schur_complement!( - buffer, model::Model, mat_idx::MatrixIndex, W::AbstractMatrix{T}, H, + buffer, ) where {T} - σ, last_dense = buffer + jprod_buffer, AW, WAW, σ, last_dense = buffer + buf = jprod_buffer[mat_idx] ilmi = mat_idx.value n = model.meta.ncon - dim = side_dimension(model, mat_idx) - @assert dim == LinearAlgebra.checksquare(W) - tmp1 = Matrix{T}(undef, dim, dim) - tmp2 = Vector{T}(undef, n) - tmp = zeros(T, dim, dim) for ii in axes(H, 1) i = σ[ii, ilmi] Ai = model.A[ilmi, i] if SparseArrays.nnz(Ai) > 0 if ii <= last_dense[ilmi] - LinearAlgebra.mul!(tmp1, W, Ai) - LinearAlgebra.mul!(tmp, tmp1, W) - fill!(tmp2, zero(T)) - add_jprod!(model, mat_idx, tmp, tmp2) - H[i, i] += tmp2[i] - indi = σ[(ii+1):end, ilmi] - H[indi, i] .+= tmp2[indi] - H[i, indi] .+= tmp2[indi] + LinearAlgebra.mul!(AW[ilmi], W, Ai) + LinearAlgebra.mul!(WAW[ilmi], AW[ilmi], W) + I = view(σ, ii:n, ilmi) + add_sub_jprod!(model, mat_idx, WAW[ilmi], view(H, I, i), I, buf) + for jj in (ii+1):n + j = σ[jj, ilmi] + H[i, j] = H[j, i] + end else if SparseArrays.nnz(Ai) > 1 @inbounds for jj in ii:n @@ -149,10 +151,10 @@ end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA_jW⟩ -function schur_complement!(buffer, model::Model, W::AbstractVector, H) +function schur_complement!(model::Model, W::AbstractVector, H, buffer) fill!(H, zero(eltype(H))) if num_matrices(model) > 0 - add_schur_complement!(buffer, model, W, MatrixIndex, H) + add_schur_complement!(model, W, MatrixIndex, H, buffer) end if num_scalars(model) > 0 add_schur_complement!(model, W[ScalarIndex], ScalarIndex, H) @@ -163,14 +165,22 @@ end # [HKS24, (5b)] # Returns the matrix equal to the sum, for each equation, of # ⟨A_i, WA(y)W⟩ -function eval_schur_complement!(buffer, result, model::Model, W, y) +function eval_schur_complement!( + result, + model::Model, + W, + y, + jprod_buffer, + jtprod_buffer, +) fill!(result, zero(eltype(result))) for i in matrix_indices(model) add_jprod!( model, i, - W[i] * jtprod!(buffer[i.value], model, i, y) * W[i], + W[i] * jtprod!(jtprod_buffer[i.value], model, i, y) * W[i], result, + jprod_buffer[i], ) end result .+= model.C_lin * (W[ScalarIndex] .* (model.C_lin' * y)) diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index 844c3f7..575e045 100644 --- a/test/BurerMonteiro.jl +++ b/test/BurerMonteiro.jl @@ -259,16 +259,36 @@ function SolverCore.solve!(::ConvexSolver, ::LRO.Model) return end +function _alloc_schur_complement(model, i, Wi, H, schur_buffer) + if VERSION < v"1.11" + return + end + LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) + @test 0 == + @allocated LRO.add_schur_complement!(model, i, Wi, H, schur_buffer) +end + function schur_test(model, w, κ) schur_buffer = LRO.buffer_for_schur_complement(model, κ) jtprod_buffer = LRO.buffer_for_jtprod(model) n = model.meta.ncon y = rand(n) + + Jv = similar(y) + vJ = similar(w) + NLPModels.jprod!(model, w, w, Jv, schur_buffer[1]) + NLPModels.jtprod!(model, w, y, vJ, jtprod_buffer) + @test dot(Jv, y) ≈ dot(vJ, w) + H = zeros(n, n) - H = LRO.schur_complement!(schur_buffer, model, w, H) + H = LRO.schur_complement!(model, w, H, schur_buffer) Hy = similar(y) - LRO.eval_schur_complement!(jtprod_buffer, Hy, model, w, y) + LRO.eval_schur_complement!(Hy, model, w, y, schur_buffer[1], jtprod_buffer) @test Hy ≈ H * y + for i in LRO.matrix_indices(model) + Wi = @inferred w[i] + _alloc_schur_complement(model, i, Wi, H, schur_buffer) + end for i in LRO.matrix_indices(model) ret = LRO.dual_cons!(jtprod_buffer, model, i, y) @test ret isa SparseMatrixCSC