Faster dot product for sparse matrices and dense vectors (#39889)

(cherry picked from commit bf0364b)

stdlib/LinearAlgebra/src/diagonal.jl

@@ -697,6 +697,14 @@ end
 *(x::Transpose{<:Any,<:AbstractVector}, D::Diagonal, y::AbstractVector) = _mapreduce_prod(*, x, D, y)
 dot(x::AbstractVector, D::Diagonal, y::AbstractVector) = _mapreduce_prod(dot, x, D, y)
 
+dot(A::Diagonal, B::Diagonal) = dot(A.diag, B.diag)
+function dot(D::Diagonal, B::AbstractMatrix)
+    size(D) == size(B) || throw(DimensionMismatch("Matrix sizes $(size(D)) and $(size(B)) differ"))
+    return dot(D.diag, view(B, diagind(B)))
+end
+
+dot(A::AbstractMatrix, B::Diagonal) = conj(dot(B, A))
+
 function _mapreduce_prod(f, x, D::Diagonal, y)
     if isempty(x) && isempty(D) && isempty(y)
         return zero(Base.promote_op(f, eltype(x), eltype(D), eltype(y)))

stdlib/LinearAlgebra/test/diagonal.jl

@@ -733,4 +733,13 @@ end
     @test dot(zeros(Int32, 0), Diagonal(zeros(Int, 0)), zeros(Int16, 0)) === 0
 end
 
+@testset "Inner product" begin
+    A = Diagonal(rand(10) .+ im)
+    B = Diagonal(rand(10) .+ im)
+    @test dot(A, B) ≈ dot(Matrix(A), B)
+    @test dot(A, B) ≈ dot(A, Matrix(B))
+    @test dot(A, B) ≈ dot(Matrix(A), Matrix(B))
+    @test dot(A, B) ≈ conj(dot(B, A))
+end
+
 end # module TestDiagonal

stdlib/SparseArrays/src/linalg.jl

@@ -376,6 +376,42 @@ function dot(x::SparseVector, A::AbstractSparseMatrixCSC, y::SparseVector)
     r
 end
 
+const WrapperMatrixTypes{T,MT} = Union{
+    SubArray{T,2,MT},
+    Adjoint{T,MT},
+    Transpose{T,MT},
+    AbstractTriangular{T,MT},
+    UpperHessenberg{T,MT},
+    Symmetric{T,MT},
+    Hermitian{T,MT},
+}
+
+function dot(A::MA, B::AbstractSparseMatrixCSC{TB}) where {MA<:Union{DenseMatrixUnion,WrapperMatrixTypes{<:Any,Union{DenseMatrixUnion,AbstractSparseMatrix}}},TB}
+    T = promote_type(eltype(A), TB)
+    (m, n) = size(A)
+    if (m, n) != size(B)
+        throw(DimensionMismatch())
+    end
+    s = zero(T)
+    if m * n == 0
+        return s
+    end
+    rows = rowvals(B)
+    vals = nonzeros(B)
+    @inbounds for j in 1:n
+        for ridx in nzrange(B, j)
+            i = rows[ridx]
+            v = vals[ridx]
+            s += dot(A[i,j], v)
+        end
+    end
+    return s
+end
+
+function dot(A::AbstractSparseMatrixCSC{TA}, B::MB) where {TA,MB<:Union{DenseMatrixUnion,WrapperMatrixTypes{<:Any,Union{DenseMatrixUnion,AbstractSparseMatrix}}}}
+    return conj(dot(B, A))
+end
+
 ## triangular sparse handling
 
 possible_adjoint(adj::Bool, a::Real) = a

stdlib/SparseArrays/test/sparse.jl

@@ -469,12 +469,42 @@ end
 end
 
 @testset "sparse Frobenius dot/inner product" begin
+    full_view = M -> view(M, :, :)
     for i = 1:5
         A = sprand(ComplexF64,10,15,0.4)
         B = sprand(ComplexF64,10,15,0.5)
-        @test dot(A,B) ≈ dot(Matrix(A),Matrix(B))
+        C = rand(10,15) .> 0.3
+        @test dot(A,B) ≈ dot(Matrix(A), Matrix(B))
+        @test dot(A,B) ≈ dot(A, Matrix(B))
+        @test dot(A,B) ≈ dot(Matrix(A), B)
+        @test dot(A,C) ≈ dot(Matrix(A), C)
+        @test dot(C,A) ≈ dot(C, Matrix(A))
+        # square matrices required by most linear algebra wrappers
+        SA = A * A'
+        SB = B * B'
+        SC = C * C'
+        for W in (full_view, LowerTriangular, UpperTriangular, UpperHessenberg, Symmetric, Hermitian)
+            WA = W(Matrix(SA))
+            WB = W(Matrix(SB))
+            WC = W(Matrix(SC))
+            @test dot(WA,SB) ≈ dot(WA, Matrix(SB))
+            @test dot(SA,WB) ≈ dot(Matrix(SA), WB)
+            @test dot(SA,WC) ≈ dot(Matrix(SA), WC)
+        end
+        for W in (transpose, adjoint)
+            WA = W(Matrix(A))
+            WB = W(Matrix(B))
+            WC = W(Matrix(C))
+            TA = copy(W(A))
+            TB = copy(W(B))
+            @test dot(WA,TB) ≈ dot(WA, Matrix(TB))
+            @test dot(TA,WB) ≈ dot(Matrix(TA), WB)
+            @test dot(TA,WC) ≈ dot(Matrix(TA), WC)
+        end
     end
     @test_throws DimensionMismatch dot(sprand(5,5,0.2),sprand(5,6,0.2))
+    @test_throws DimensionMismatch dot(rand(5,5),sprand(5,6,0.2))
+    @test_throws DimensionMismatch dot(sprand(5,5,0.2),rand(5,6))
 end
 
 @testset "generalized dot product" begin