From 88e78afda8683d043bf1bd6b78ce6b137a95857a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Fri, 6 Jun 2025 22:13:53 +0200
Subject: [PATCH 01/13] Reduce allocation of Schur

---
 src/schur.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/schur.jl b/src/schur.jl
index 16303d4..defb6d4 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -94,9 +94,11 @@ function add_schur_complement!(
                 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]
+                for jj in (ii+1):n
+                    j = σ[jj,ilmi]
+                    H[j, i] += tmp2[j]
+                    H[i, j] += tmp2[j]
+                end
             else
                 if SparseArrays.nnz(Ai) > 1
                     @inbounds for jj in ii:n

From 09401eab8a6ec1396940e5aba8d2136b11bcdeab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Sat, 7 Jun 2025 07:30:11 +0200
Subject: [PATCH 02/13] Fix format

---
 src/schur.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/schur.jl b/src/schur.jl
index defb6d4..35a98da 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -95,7 +95,7 @@ function add_schur_complement!(
                 add_jprod!(model, mat_idx, tmp, tmp2)
                 H[i, i] += tmp2[i]
                 for jj in (ii+1):n
-                    j = σ[jj,ilmi]
+                    j = σ[jj, ilmi]
                     H[j, i] += tmp2[j]
                     H[i, j] += tmp2[j]
                 end

From fc510ab9639996085a9a0c3a23d602ee182a3265 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Sat, 7 Jun 2025 11:31:38 +0200
Subject: [PATCH 03/13] Add vectorized jprod

---
 src/model.jl | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/src/model.jl b/src/model.jl
index 64c709e..ed98fec 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -304,6 +304,42 @@ function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector)
     return cx
 end
 
+function buffer_for_add_jprod(
+    model::Model{T},
+    i::MatrixIndex,
+) where {T}
+    nnz = sum(1:model.meta.ncon; init = 0) do j
+        SparseArrays.nnz(model.A[i.value, j])
+    end
+    I = zeros(Int64, nnz)
+    J = zeros(Int64, nnz)
+    V = zeros(Float64, nnz)
+    offset = 0
+    for j = 1:model.meta.ncon
+        Ai, Av = findnz(Ai[i.value, j][:])
+        K = offset .+ eachindex(Ai)
+        iii[K] = Ai
+        jjj[K] .= j
+        vvv[K] = Av
+        offset += length(Ai)
+    end
+    A = sparse(I, J, K, side_dimension(model, i)^2, model.meta.ncon)
+    return (A, zeros(model.meta.ncon))
+end
+
+function add_jprod!(
+    model::Model,
+    i::MatrixIndex,
+    V::AbstractMatrix,
+    Jv::AbstractVector,
+    buffer,
+)
+    A, cache = buffer
+    LinearAlgebra.mul!(cache, A', reshape(V, length(v)))
+    Jv .+= cache
+    return Jv
+end
+
 function add_jprod!(
     model::Model,
     i::MatrixIndex,

From 1e347dc466660c47d119a1a92c2889c76813dc8a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Sat, 7 Jun 2025 18:40:53 +0200
Subject: [PATCH 04/13] Fixes

---
 src/model.jl | 21 +++++++++++++--------
 src/schur.jl | 16 ++++++++--------
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index ed98fec..dc92eb5 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -304,7 +304,7 @@ function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector)
     return cx
 end
 
-function buffer_for_add_jprod(
+function buffer_for_jprod(
     model::Model{T},
     i::MatrixIndex,
 ) where {T}
@@ -316,17 +316,21 @@ function buffer_for_add_jprod(
     V = zeros(Float64, nnz)
     offset = 0
     for j = 1:model.meta.ncon
-        Ai, Av = findnz(Ai[i.value, j][:])
+        Ai, Av = SparseArrays.findnz(model.A[i.value, j][:])
         K = offset .+ eachindex(Ai)
-        iii[K] = Ai
-        jjj[K] .= j
-        vvv[K] = Av
+        I[K] = Ai
+        J[K] .= j
+        V[K] = Av
         offset += length(Ai)
     end
-    A = sparse(I, J, K, side_dimension(model, i)^2, model.meta.ncon)
+    A = SparseArrays.sparse(I, J, V, side_dimension(model, i)^2, model.meta.ncon)
     return (A, zeros(model.meta.ncon))
 end
 
+function buffer_for_jprod(model::Model)
+    return [buffer_for_jprod(model, i) for i in matrix_indices(model)]
+end
+
 function add_jprod!(
     model::Model,
     i::MatrixIndex,
@@ -356,10 +360,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, args...)
     end
     return Jv
 end
diff --git a/src/schur.jl b/src/schur.jl
index 35a98da..abad220 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -57,25 +57,25 @@ function buffer_for_schur_complement(model::Model, κ)
         last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0)
     end
 
-    return σ, last_dense
+    return σ, last_dense, buffer_for_jprod(model)
 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
+    σ, last_dense, jprod_buffer = buffer
     ilmi = mat_idx.value
     n = model.meta.ncon
     dim = side_dimension(model, mat_idx)
@@ -92,7 +92,7 @@ function add_schur_complement!(
                 LinearAlgebra.mul!(tmp1, W, Ai)
                 LinearAlgebra.mul!(tmp, tmp1, W)
                 fill!(tmp2, zero(T))
-                add_jprod!(model, mat_idx, tmp, tmp2)
+                add_jprod!(model, mat_idx, tmp, tmp2, jprod_buffer)
                 H[i, i] += tmp2[i]
                 for jj in (ii+1):n
                     j = σ[jj, ilmi]
@@ -151,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)

From 4c09c4b6aa8f320192f9ca9d34e2909dc926f359 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 11 Jun 2025 21:35:50 +0200
Subject: [PATCH 05/13] Improve performance of jprod

---
 Project.toml |  2 ++
 src/model.jl | 45 ++++++++++++++++++++++++++++++++++++---------
 src/schur.jl | 10 ++++------
 3 files changed, 42 insertions(+), 15 deletions(-)

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/model.jl b/src/model.jl
index dc92eb5..aa4eae2 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
 
@@ -304,6 +306,13 @@ function NLPModels.cons!(model::Model, x::AbstractVector, cx::AbstractVector)
     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,
@@ -324,24 +333,40 @@ function buffer_for_jprod(
         offset += length(Ai)
     end
     A = SparseArrays.sparse(I, J, V, side_dimension(model, i)^2, model.meta.ncon)
-    return (A, zeros(model.meta.ncon))
+    return A
 end
 
 function buffer_for_jprod(model::Model)
-    return [buffer_for_jprod(model, i) for i in matrix_indices(model)]
+    return ([buffer_for_jprod(model, i) for i in matrix_indices(model)], zeros(model.meta.ncon))
 end
 
-function add_jprod!(
-    model::Model,
-    i::MatrixIndex,
+function _add_jprod!(V, Jv, A, cache)
+    LinearAlgebra.mul!(cache, A', UnsafeArrays.uview(V, :))
+    Jv .+= cache
+    return Jv
+end
+
+function add_sub_jprod!(
+    _::Model,
+    _::MatrixIndex,
     V::AbstractMatrix,
     Jv::AbstractVector,
+    I,
     buffer,
 )
     A, cache = buffer
-    LinearAlgebra.mul!(cache, A', reshape(V, length(v)))
-    Jv .+= cache
-    return Jv
+    # `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!(
@@ -355,6 +380,8 @@ function add_jprod!(
     end
 end
 
+_buffer_getindex(A_cache, i) = (A_cache[1][i.value], A_cache[2])
+
 function NLPModels.jprod!(
     model::Model,
     _::AbstractVector,
@@ -364,7 +391,7 @@ function NLPModels.jprod!(
 ) where {N}
     LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex])
     for i in matrix_indices(model)
-        add_jprod!(model, i, v[i], Jv, args...)
+        add_jprod!(model, i, v[i], Jv, _buffer_getindex.(args, i)...)
     end
     return Jv
 end
diff --git a/src/schur.jl b/src/schur.jl
index abad220..d4e6e3c 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -81,7 +81,6 @@ function add_schur_complement!(
     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)
@@ -91,13 +90,12 @@ function add_schur_complement!(
             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, jprod_buffer)
-                H[i, i] += tmp2[i]
+                buf = _buffer_getindex(jprod_buffer, mat_idx)
+                I = view(σ, ii:n, ilmi)
+                add_sub_jprod!(model, mat_idx, tmp, view(H, I, i), I, buf)
                 for jj in (ii+1):n
                     j = σ[jj, ilmi]
-                    H[j, i] += tmp2[j]
-                    H[i, j] += tmp2[j]
+                    H[i, j] = H[j, i]
                 end
             else
                 if SparseArrays.nnz(Ai) > 1

From be6d8fa94fa3cd3df464d3bfcd3f966f43f14505 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 12 Jun 2025 08:32:20 +0200
Subject: [PATCH 06/13] Fixes

---
 src/BurerMonteiro.jl |  1 +
 src/model.jl         | 24 +++++++++++++++++-------
 src/schur.jl         |  8 +++++---
 3 files changed, 23 insertions(+), 10 deletions(-)

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 aa4eae2..97c451f 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -300,8 +300,8 @@ 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
@@ -322,7 +322,7 @@ function buffer_for_jprod(
     end
     I = zeros(Int64, nnz)
     J = zeros(Int64, nnz)
-    V = zeros(Float64, nnz)
+    V = zeros(T, nnz)
     offset = 0
     for j = 1:model.meta.ncon
         Ai, Av = SparseArrays.findnz(model.A[i.value, j][:])
@@ -336,12 +336,24 @@ function buffer_for_jprod(
     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 ([buffer_for_jprod(model, i) for i in matrix_indices(model)], zeros(model.meta.ncon))
+    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', UnsafeArrays.uview(V, :))
+    LinearAlgebra.mul!(cache, A', _vec(V))
     Jv .+= cache
     return Jv
 end
@@ -380,8 +392,6 @@ function add_jprod!(
     end
 end
 
-_buffer_getindex(A_cache, i) = (A_cache[1][i.value], A_cache[2])
-
 function NLPModels.jprod!(
     model::Model,
     _::AbstractVector,
diff --git a/src/schur.jl b/src/schur.jl
index d4e6e3c..492d9cc 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -42,6 +42,8 @@ function _dot(
     return result
 end
 
+# 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, κ)
     n = model.meta.ncon
     σ = zeros(Int64, n, num_matrices(model))
@@ -57,7 +59,7 @@ function buffer_for_schur_complement(model::Model, κ)
         last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0)
     end
 
-    return σ, last_dense, buffer_for_jprod(model)
+    return buffer_for_jprod(model), σ, last_dense
 end
 
 function add_schur_complement!(model::Model, W, ::Type{MatrixIndex}, H, buffer)
@@ -75,7 +77,8 @@ function add_schur_complement!(
     H,
     buffer,
 ) where {T}
-    σ, last_dense, jprod_buffer = buffer
+    jprod_buffer, σ, last_dense = buffer
+    buf = jprod_buffer[mat_idx]
     ilmi = mat_idx.value
     n = model.meta.ncon
     dim = side_dimension(model, mat_idx)
@@ -90,7 +93,6 @@ function add_schur_complement!(
             if ii <= last_dense[ilmi]
                 LinearAlgebra.mul!(tmp1, W, Ai)
                 LinearAlgebra.mul!(tmp, tmp1, W)
-                buf = _buffer_getindex(jprod_buffer, mat_idx)
                 I = view(σ, ii:n, ilmi)
                 add_sub_jprod!(model, mat_idx, tmp, view(H, I, i), I, buf)
                 for jj in (ii+1):n

From 1ac2d0381a67e3f2013e3478fa767fb50ffca12f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 13:58:14 +0200
Subject: [PATCH 07/13] Fix format

---
 src/model.jl | 29 ++++++++++++++++++++---------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 97c451f..7ba0ad7 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -300,7 +300,12 @@ 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, args::Vararg{Any,N}) where {N}
+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
@@ -313,18 +318,15 @@ end
 # 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}
+function buffer_for_jprod(model::Model{T}, i::MatrixIndex) where {T}
     nnz = sum(1:model.meta.ncon; init = 0) do j
-        SparseArrays.nnz(model.A[i.value, 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 = 1:model.meta.ncon
+    for j in 1:model.meta.ncon
         Ai, Av = SparseArrays.findnz(model.A[i.value, j][:])
         K = offset .+ eachindex(Ai)
         I[K] = Ai
@@ -332,7 +334,13 @@ function buffer_for_jprod(
         V[K] = Av
         offset += length(Ai)
     end
-    A = SparseArrays.sparse(I, J, V, side_dimension(model, i)^2, model.meta.ncon)
+    A = SparseArrays.sparse(
+        I,
+        J,
+        V,
+        side_dimension(model, i)^2,
+        model.meta.ncon,
+    )
     return A
 end
 
@@ -343,7 +351,10 @@ struct JProdBuffer{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))
+    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)

From 79d5cf9d857ec9a45101be49439cbfe006106a49 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 14:19:02 +0200
Subject: [PATCH 08/13] Fix tests

---
 src/model.jl          | 2 +-
 src/schur.jl          | 2 +-
 test/BurerMonteiro.jl | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index 7ba0ad7..c8a46e0 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -412,7 +412,7 @@ function NLPModels.jprod!(
 ) where {N}
     LinearAlgebra.mul!(Jv, model.C_lin, v[ScalarIndex])
     for i in matrix_indices(model)
-        add_jprod!(model, i, v[i], Jv, _buffer_getindex.(args, i)...)
+        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 492d9cc..0d4350b 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -165,7 +165,7 @@ 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, buffer)
     fill!(result, zero(eltype(result)))
     for i in matrix_indices(model)
         add_jprod!(
diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl
index 844c3f7..0c0def1 100644
--- a/test/BurerMonteiro.jl
+++ b/test/BurerMonteiro.jl
@@ -265,9 +265,9 @@ function schur_test(model, w, κ)
     n = model.meta.ncon
     y = rand(n)
     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, jtprod_buffer)
     @test Hy ≈ H * y
     for i in LRO.matrix_indices(model)
         ret = LRO.dual_cons!(jtprod_buffer, model, i, y)

From 72bd952757a160e1c395811acbb21d052a72a703 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 15:29:52 +0200
Subject: [PATCH 09/13] Add alloc test

---
 src/schur.jl          | 22 ++++++++++++----------
 test/BurerMonteiro.jl |  9 +++++++++
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/src/schur.jl b/src/schur.jl
index 0d4350b..e524022 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -44,7 +44,7 @@ end
 
 # 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, κ)
+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))
@@ -59,7 +59,13 @@ function buffer_for_schur_complement(model::Model, κ)
         last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0)
     end
 
-    return buffer_for_jprod(model), σ, 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!(model::Model, W, ::Type{MatrixIndex}, H, buffer)
@@ -77,24 +83,20 @@ function add_schur_complement!(
     H,
     buffer,
 ) where {T}
-    jprod_buffer, σ, 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)
-    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)
+                LinearAlgebra.mul!(AW[ilmi], W, Ai)
+                LinearAlgebra.mul!(WAW[ilmi], AW[ilmi], W)
                 I = view(σ, ii:n, ilmi)
-                add_sub_jprod!(model, mat_idx, tmp, view(H, I, i), I, buf)
+                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]
diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl
index 0c0def1..0a31faf 100644
--- a/test/BurerMonteiro.jl
+++ b/test/BurerMonteiro.jl
@@ -259,6 +259,11 @@ function SolverCore.solve!(::ConvexSolver, ::LRO.Model)
     return
 end
 
+function _alloc_schur_complement(model, i, Wi, H, schur_buffer)
+    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)
@@ -269,6 +274,10 @@ function schur_test(model, w, κ)
     Hy = similar(y)
     LRO.eval_schur_complement!(Hy, model, w, y, 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

From a26cd6c41f3b46c77cf8200ac292f9623c0445fa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 15:29:58 +0200
Subject: [PATCH 10/13] Fix format

---
 src/schur.jl          | 6 ++----
 test/BurerMonteiro.jl | 3 ++-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/schur.jl b/src/schur.jl
index e524022..d39a0db 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -59,10 +59,8 @@ function buffer_for_schur_complement(model::Model{T}, κ) where {T}
         last_dense[i] = something(findlast(Base.Fix1(isless, κ), sorted), 0)
     end
 
-    AW = [
-        zeros(T, dim, dim) # /!\ it's the same zero everywhere, might be an issue with BigFloat
-        for dim in model.msizes
-    ]
+    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
diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl
index 0a31faf..1e85777 100644
--- a/test/BurerMonteiro.jl
+++ b/test/BurerMonteiro.jl
@@ -261,7 +261,8 @@ end
 
 function _alloc_schur_complement(model, i, Wi, H, schur_buffer)
     LRO.add_schur_complement!(model, i, Wi, H, schur_buffer)
-    @test 0 == @allocated 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, κ)

From 0bcfcb9b4ebdc7d4e5d705ab672051b94074d7fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 15:44:00 +0200
Subject: [PATCH 11/13] Disable alloc tests on v1.10

---
 test/BurerMonteiro.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl
index 1e85777..163feb5 100644
--- a/test/BurerMonteiro.jl
+++ b/test/BurerMonteiro.jl
@@ -260,6 +260,9 @@ function SolverCore.solve!(::ConvexSolver, ::LRO.Model)
 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)

From e1619212766d7a4b5235256b5c4819c245053473 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 17:41:17 +0200
Subject: [PATCH 12/13] Improve code coverage

---
 src/model.jl          | 14 ++++++++++++++
 src/schur.jl          |  5 +++--
 test/BurerMonteiro.jl |  9 ++++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/model.jl b/src/model.jl
index c8a46e0..32c20c1 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -295,6 +295,20 @@ 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]
 
diff --git a/src/schur.jl b/src/schur.jl
index d39a0db..583ac85 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -165,14 +165,15 @@ end
 # [HKS24, (5b)]
 # Returns the matrix equal to the sum, for each equation, of
 # ⟨A_i, WA(y)W⟩
-function eval_schur_complement!(result, model::Model, W, y, buffer)
+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 163feb5..575e045 100644
--- a/test/BurerMonteiro.jl
+++ b/test/BurerMonteiro.jl
@@ -273,10 +273,17 @@ function schur_test(model, w, κ)
     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!(model, w, H, schur_buffer)
     Hy = similar(y)
-    LRO.eval_schur_complement!(Hy, model, w, y, jtprod_buffer)
+    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]

From e74680bee45265abda117a4d2c391619e3880dd8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 25 Jun 2025 17:41:27 +0200
Subject: [PATCH 13/13] Fix format

---
 src/schur.jl | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/schur.jl b/src/schur.jl
index 583ac85..8760ca6 100644
--- a/src/schur.jl
+++ b/src/schur.jl
@@ -165,7 +165,14 @@ end
 # [HKS24, (5b)]
 # Returns the matrix equal to the sum, for each equation, of
 # ⟨A_i, WA(y)W⟩
-function eval_schur_complement!(result, model::Model, W, y, jprod_buffer, jtprod_buffer)
+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!(