Faster add_sub_jprod

src/buffer.jl

@@ -120,8 +120,8 @@ _vec(x::FillArrays.Zeros{T}) where {T} = FillArrays.Zeros{T}(length(x))
 _vec(x::AbstractArray) = UnsafeArrays.uview(x, :)
 _vec(x::Base.ReshapedArray) = _vec(parent(x))
 
-function _add_jprod!(V, Jv::AbstractArray{T}, A) where {T}
-    return LinearAlgebra.mul!(Jv, A', _vec(V), true, true)
+function _add_jprod!(V, Jv::AbstractArray, A)
+    return _add_mul!(Jv, A', _vec(V), true)
 end
 
 function add_sub_jprod!(

src/factorization.jl

@@ -457,6 +457,54 @@ function _add_mul!(
     end
 end
 
+# `SparseArrays` has a specialized (CSC-traversing) `mul!` method for the
+# adjoint of a `SparseMatrixCSC` so we can safely delegate. It requires a
+# strided destination, so we restrict `res` to `StridedVector` to get a
+# `MethodError` for a non-strided destination rather than silently hitting
+# the generic `getindex`-based fallback.
+function _add_mul!(
+    res::StridedVector,
+    Ft::LinearAlgebra.Adjoint{<:Any,<:SparseArrays.SparseMatrixCSC},
+    C::AbstractVector,
+    α,
+)
+    return LinearAlgebra.mul!(res, Ft, C, α, true)
+end
+
+# The adjoint of a column subset of a `SparseMatrixCSC` has no specialized
+# `mul!` method so it would hit the generic `getindex`-based fallback
+# that does a binary search for each entry of the full dense index space;
+# iterate the stored entries of the selected columns of the parent instead.
+function _add_mul!(
+    res::AbstractVector,
+    Ft::LinearAlgebra.Adjoint{
+        <:Any,
+        <:SubArray{
+            <:Any,
+            2,
+            <:SparseArrays.SparseMatrixCSC,
+            <:Tuple{Base.Slice,Any},
+        },
+    },
+    C::AbstractVector,
+    α,
+)
+    F = parent(parent(Ft))
+    cols = parentindices(parent(Ft))[2]
+    @assert axes(C, 1) == axes(F, 1)
+    @assert axes(res, 1) == axes(cols, 1)
+    rows = SparseArrays.rowvals(F)
+    vals = SparseArrays.nonzeros(F)
+    @inbounds for (k, col) in enumerate(cols)
+        acc = zero(eltype(res))
+        for i in SparseArrays.nzrange(F, col)
+            acc += vals[i] * C[rows[i]]
+        end
+        res[k] += acc * α
+    end
+    return res
+end
+
 function _add_mul!(
     res::AbstractMatrix,
     F::SparseArrays.SparseMatrixCSC,

test/BurerMonteiro.jl

@@ -134,7 +134,7 @@ end;
     set_attribute(model, "max_iter", 0)
     optimize!(model)
     @test termination_status(model) == MOI.ITERATION_LIMIT
-    nlp = unsafe_backend(model).model;
+    nlp = unsafe_backend(model).model
     @test nlp.C isa Vector{SparseMatrixCSC{T,Int}}
     @test nlp.C[1] == [3 1; 1 4]
     @test nlp.A isa Matrix{
@@ -210,7 +210,7 @@ end
     optimize!(model)
     @test termination_status(model) == MOI.ITERATION_LIMIT
     diff_check(model)
-    nlp = unsafe_backend(model).model;
+    nlp = unsafe_backend(model).model
     T = Float64
     MT = LRO.Factorization{T,Matrix{T},LRO.Ones{T}}
     @test nlp.C isa Vector{MT}