diff --git a/src/buffer.jl b/src/buffer.jl index 0338724..4d04632 100644 --- a/src/buffer.jl +++ b/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!( diff --git a/src/factorization.jl b/src/factorization.jl index a0884b0..d587ff0 100644 --- a/src/factorization.jl +++ b/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, diff --git a/test/BurerMonteiro.jl b/test/BurerMonteiro.jl index 9fb325a..67f5c01 100644 --- a/test/BurerMonteiro.jl +++ b/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}