From fc9c280584f63236a1b97da4178a41eba65b6da2 Mon Sep 17 00:00:00 2001 From: Daniel Karrasch Date: Tue, 22 Mar 2022 21:44:56 +0100 Subject: [PATCH] Fix performance bug for `*` with `AbstractQ` (#44615) --- stdlib/LinearAlgebra/src/bidiag.jl | 10 ++-- stdlib/LinearAlgebra/src/diagonal.jl | 12 ++++- stdlib/LinearAlgebra/src/special.jl | 69 ++++++++++++++++++++++++---- stdlib/LinearAlgebra/src/tridiag.jl | 9 ++-- stdlib/LinearAlgebra/test/special.jl | 32 +++++++------ 5 files changed, 96 insertions(+), 36 deletions(-) diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/stdlib/LinearAlgebra/src/bidiag.jl index 243553ebc64c6..dfcbec69c6de2 100644 --- a/stdlib/LinearAlgebra/src/bidiag.jl +++ b/stdlib/LinearAlgebra/src/bidiag.jl @@ -168,15 +168,13 @@ end function Matrix{T}(A::Bidiagonal) where T n = size(A, 1) B = zeros(T, n, n) - if n == 0 - return B - end - for i = 1:n - 1 + n == 0 && return B + @inbounds for i = 1:n - 1 B[i,i] = A.dv[i] if A.uplo == 'U' - B[i, i + 1] = A.ev[i] + B[i,i+1] = A.ev[i] else - B[i + 1, i] = A.ev[i] + B[i+1,i] = A.ev[i] end end B[n,n] = A.dv[n] diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index 4b7d9bd9d4af1..11f3fff9cb3e2 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -77,8 +77,16 @@ Diagonal{T}(D::Diagonal{T}) where {T} = D Diagonal{T}(D::Diagonal) where {T} = Diagonal{T}(D.diag) AbstractMatrix{T}(D::Diagonal) where {T} = Diagonal{T}(D) -Matrix(D::Diagonal) = diagm(0 => D.diag) -Array(D::Diagonal) = Matrix(D) +Matrix(D::Diagonal{T}) where {T} = Matrix{T}(D) +Array(D::Diagonal{T}) where {T} = Matrix{T}(D) +function Matrix{T}(D::Diagonal) where {T} + n = size(D, 1) + B = zeros(T, n, n) + @inbounds for i in 1:n + B[i,i] = D.diag[i] + end + return B +end """ Diagonal{T}(undef, n) diff --git a/stdlib/LinearAlgebra/src/special.jl b/stdlib/LinearAlgebra/src/special.jl index 39b62d5e3ca03..beac0c524f2f4 100644 --- a/stdlib/LinearAlgebra/src/special.jl +++ b/stdlib/LinearAlgebra/src/special.jl @@ -292,16 +292,65 @@ function (-)(A::UniformScaling, B::Diagonal{<:Number}) Diagonal(A.λ .- B.diag) end -rmul!(A::AbstractTriangular, adjB::Adjoint{<:Any,<:Union{QRCompactWYQ,QRPackedQ}}) = - rmul!(full!(A), adjB) -*(A::AbstractTriangular, adjB::Adjoint{<:Any,<:Union{QRCompactWYQ,QRPackedQ}}) = - *(copyto!(similar(parent(A)), A), adjB) -*(A::BiTriSym, adjB::Adjoint{<:Any,<:Union{QRCompactWYQ, QRPackedQ}}) = - rmul!(copyto!(Array{promote_type(eltype(A), eltype(adjB))}(undef, size(A)...), A), adjB) -*(adjA::Adjoint{<:Any,<:Union{QRCompactWYQ, QRPackedQ}}, B::Diagonal) = - lmul!(adjA, copyto!(Array{promote_type(eltype(adjA), eltype(B))}(undef, size(B)...), B)) -*(adjA::Adjoint{<:Any,<:Union{QRCompactWYQ, QRPackedQ}}, B::BiTriSym) = - lmul!(adjA, copyto!(Array{promote_type(eltype(adjA), eltype(B))}(undef, size(B)...), B)) +lmul!(Q::AbstractQ, B::AbstractTriangular) = lmul!(Q, full!(B)) +lmul!(Q::QRPackedQ, B::AbstractTriangular) = lmul!(Q, full!(B)) # disambiguation +lmul!(Q::Adjoint{<:Any,<:AbstractQ}, B::AbstractTriangular) = lmul!(Q, full!(B)) +lmul!(Q::Adjoint{<:Any,<:QRPackedQ}, B::AbstractTriangular) = lmul!(Q, full!(B)) # disambiguation + +function _qlmul(Q::AbstractQ, B) + TQB = promote_type(eltype(Q), eltype(B)) + if size(Q.factors, 1) == size(B, 1) + Bnew = Matrix{TQB}(B) + elseif size(Q.factors, 2) == size(B, 1) + Bnew = [Matrix{TQB}(B); zeros(TQB, size(Q.factors, 1) - size(B,1), size(B, 2))] + else + throw(DimensionMismatch("first dimension of matrix must have size either $(size(Q.factors, 1)) or $(size(Q.factors, 2))")) + end + lmul!(convert(AbstractMatrix{TQB}, Q), Bnew) +end +function _qlmul(adjQ::Adjoint{<:Any,<:AbstractQ}, B) + TQB = promote_type(eltype(adjQ), eltype(B)) + lmul!(adjoint(convert(AbstractMatrix{TQB}, parent(adjQ))), Matrix{TQB}(B)) +end + +*(Q::AbstractQ, B::AbstractTriangular) = _qlmul(Q, B) +*(Q::Adjoint{<:Any,<:AbstractQ}, B::AbstractTriangular) = _qlmul(Q, B) +*(Q::AbstractQ, B::BiTriSym) = _qlmul(Q, B) +*(Q::Adjoint{<:Any,<:AbstractQ}, B::BiTriSym) = _qlmul(Q, B) +*(Q::AbstractQ, B::Diagonal) = _qlmul(Q, B) +*(Q::Adjoint{<:Any,<:AbstractQ}, B::Diagonal) = _qlmul(Q, B) + +rmul!(A::AbstractTriangular, Q::AbstractQ) = rmul!(full!(A), Q) +rmul!(A::AbstractTriangular, Q::Adjoint{<:Any,<:AbstractQ}) = rmul!(full!(A), Q) + +function _qrmul(A, Q::AbstractQ) + TAQ = promote_type(eltype(A), eltype(Q)) + return rmul!(Matrix{TAQ}(A), convert(AbstractMatrix{TAQ}, Q)) +end +function _qrmul(A, adjQ::Adjoint{<:Any,<:AbstractQ}) + Q = adjQ.parent + TAQ = promote_type(eltype(A), eltype(Q)) + if size(A,2) == size(Q.factors, 1) + Anew = Matrix{TAQ}(A) + elseif size(A,2) == size(Q.factors,2) + Anew = [Matrix{TAQ}(A) zeros(TAQ, size(A, 1), size(Q.factors, 1) - size(Q.factors, 2))] + else + throw(DimensionMismatch("matrix A has dimensions $(size(A)) but matrix B has dimensions $(size(Q))")) + end + return rmul!(Anew, adjoint(convert(AbstractMatrix{TAQ}, Q))) +end + +*(A::AbstractTriangular, Q::AbstractQ) = _qrmul(A, Q) +*(A::AbstractTriangular, Q::Adjoint{<:Any,<:AbstractQ}) = _qrmul(A, Q) +*(A::BiTriSym, Q::AbstractQ) = _qrmul(A, Q) +*(A::BiTriSym, Q::Adjoint{<:Any,<:AbstractQ}) = _qrmul(A, Q) +*(A::Diagonal, Q::AbstractQ) = _qrmul(A, Q) +*(A::Diagonal, Q::Adjoint{<:Any,<:AbstractQ}) = _qrmul(A, Q) + +*(Q::AbstractQ, B::AbstractQ) = _qlmul(Q, B) +*(Q::Adjoint{<:Any,<:AbstractQ}, B::AbstractQ) = _qrmul(Q, B) +*(Q::AbstractQ, B::Adjoint{<:Any,<:AbstractQ}) = _qlmul(Q, B) +*(Q::Adjoint{<:Any,<:AbstractQ}, B::Adjoint{<:Any,<:AbstractQ}) = _qrmul(Q, B) # fill[stored]! methods fillstored!(A::Diagonal, x) = (fill!(A.diag, x); A) diff --git a/stdlib/LinearAlgebra/src/tridiag.jl b/stdlib/LinearAlgebra/src/tridiag.jl index 5a3c7612f6784..4b1d3add5df5b 100644 --- a/stdlib/LinearAlgebra/src/tridiag.jl +++ b/stdlib/LinearAlgebra/src/tridiag.jl @@ -571,15 +571,16 @@ function size(M::Tridiagonal, d::Integer) end end -function Matrix{T}(M::Tridiagonal{T}) where T +function Matrix{T}(M::Tridiagonal) where {T} A = zeros(T, size(M)) - for i = 1:length(M.d) + n = length(M.d) + n == 0 && return A + for i in 1:n-1 A[i,i] = M.d[i] - end - for i = 1:length(M.d)-1 A[i+1,i] = M.dl[i] A[i,i+1] = M.du[i] end + A[n,n] = M.d[n] A end Matrix(M::Tridiagonal{T}) where {T} = Matrix{T}(M) diff --git a/stdlib/LinearAlgebra/test/special.jl b/stdlib/LinearAlgebra/test/special.jl index 624868db9ba44..9b094c267d41b 100644 --- a/stdlib/LinearAlgebra/test/special.jl +++ b/stdlib/LinearAlgebra/test/special.jl @@ -188,16 +188,21 @@ end @testset "Triangular Types and QR" begin - for typ in [UpperTriangular,LowerTriangular,LinearAlgebra.UnitUpperTriangular,LinearAlgebra.UnitLowerTriangular] + for typ in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) a = rand(n,n) atri = typ(a) + matri = Matrix(atri) b = rand(n,n) qrb = qr(b, ColumnNorm()) - @test *(atri, adjoint(qrb.Q)) ≈ Matrix(atri) * qrb.Q' - @test rmul!(copy(atri), adjoint(qrb.Q)) ≈ Matrix(atri) * qrb.Q' + @test atri * qrb.Q ≈ matri * qrb.Q ≈ rmul!(copy(atri), qrb.Q) + @test atri * qrb.Q' ≈ matri * qrb.Q' ≈ rmul!(copy(atri), qrb.Q') + @test qrb.Q * atri ≈ qrb.Q * matri ≈ lmul!(qrb.Q, copy(atri)) + @test qrb.Q' * atri ≈ qrb.Q' * matri ≈ lmul!(qrb.Q', copy(atri)) qrb = qr(b, NoPivot()) - @test *(atri, adjoint(qrb.Q)) ≈ Matrix(atri) * qrb.Q' - @test rmul!(copy(atri), adjoint(qrb.Q)) ≈ Matrix(atri) * qrb.Q' + @test atri * qrb.Q ≈ matri * qrb.Q ≈ rmul!(copy(atri), qrb.Q) + @test atri * qrb.Q' ≈ matri * qrb.Q' ≈ rmul!(copy(atri), qrb.Q') + @test qrb.Q * atri ≈ qrb.Q * matri ≈ lmul!(qrb.Q, copy(atri)) + @test qrb.Q' * atri ≈ qrb.Q' * matri ≈ lmul!(qrb.Q', copy(atri)) end end @@ -421,19 +426,18 @@ end end @testset "BiTriSym*Q' and Q'*BiTriSym" begin - dl = [1, 1, 1]; - d = [1, 1, 1, 1]; - Tri = Tridiagonal(dl, d, dl) + dl = [1, 1, 1] + d = [1, 1, 1, 1] + D = Diagonal(d) Bi = Bidiagonal(d, dl, :L) + Tri = Tridiagonal(dl, d, dl) Sym = SymTridiagonal(d, dl) F = qr(ones(4, 1)) A = F.Q' - @test Tri*A ≈ Matrix(Tri)*A - @test A*Tri ≈ A*Matrix(Tri) - @test Bi*A ≈ Matrix(Bi)*A - @test A*Bi ≈ A*Matrix(Bi) - @test Sym*A ≈ Matrix(Sym)*A - @test A*Sym ≈ A*Matrix(Sym) + for A in (F.Q, F.Q'), B in (D, Bi, Tri, Sym) + @test B*A ≈ Matrix(B)*A + @test A*B ≈ A*Matrix(B) + end end @testset "Ops on SymTridiagonal ev has the same length as dv" begin