Add vectorized jprod

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"

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
 

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

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))

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