Skip to content

Commit

Permalink
copytrito! for triangular matrices (#56620)
Browse files Browse the repository at this point in the history
This does two things:

1. Forward `copytrito!` for triangular matrices to the parent in case
the specified `uplo` corresponds to the stored part. This works because
these matrices share their elements with the parents for the stored
part.
2. Make `copytrito!` only copy the diagonal if the `uplo` corresponds to
the non-stored part.

This makes `copytrito!` involving a triangular matrix equivalent to that
involving its parent if the filled part is copied, and O(N) otherwise.

Examples of improvements in performance:
```julia
julia> using LinearAlgebra

julia> A1 = UpperTriangular(rand(400,400));

julia> A2 = similar(A1);

julia> @Btime copytrito!($A2, $A1, 'U');
  70.753 μs (0 allocations: 0 bytes) # nightly v"1.12.0-DEV.1657"
  26.143 μs (0 allocations: 0 bytes) # this PR

julia> @Btime copytrito!(parent($A2), $A1, 'U');
  56.025 μs (0 allocations: 0 bytes) # nightly
  26.633 μs (0 allocations: 0 bytes) # this PR
```
  • Loading branch information
jishnub authored Nov 21, 2024
1 parent 034e609 commit 4709b6c
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 5 deletions.
13 changes: 8 additions & 5 deletions stdlib/LinearAlgebra/src/generic.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2069,17 +2069,20 @@ function copytrito!(B::AbstractMatrix, A::AbstractMatrix, uplo::AbstractChar)
require_one_based_indexing(A, B)
BLAS.chkuplo(uplo)
m,n = size(A)
m1,n1 = size(B)
A = Base.unalias(B, A)
if uplo == 'U'
LAPACK.lacpy_size_check((m1, n1), (n < m ? n : m, n))
LAPACK.lacpy_size_check(size(B), (n < m ? n : m, n))
for j in axes(A,2), i in axes(A,1)[begin : min(j,end)]
@inbounds B[i,j] = A[i,j]
# extract the parents for UpperTriangular matrices
Bv, Av = uppertridata(B), uppertridata(A)
@inbounds Bv[i,j] = Av[i,j]
end
else # uplo == 'L'
LAPACK.lacpy_size_check((m1, n1), (m, m < n ? m : n))
LAPACK.lacpy_size_check(size(B), (m, m < n ? m : n))
for j in axes(A,2), i in axes(A,1)[j:end]
@inbounds B[i,j] = A[i,j]
# extract the parents for LowerTriangular matrices
Bv, Av = lowertridata(B), lowertridata(A)
@inbounds Bv[i,j] = Av[i,j]
end
end
return B
Expand Down
30 changes: 30 additions & 0 deletions stdlib/LinearAlgebra/src/triangular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -633,6 +633,36 @@ end
return dest
end

Base.@constprop :aggressive function copytrito_triangular!(Bdata, Adata, uplo, uplomatch, sz)
if uplomatch
copytrito!(Bdata, Adata, uplo)
else
BLAS.chkuplo(uplo)
LAPACK.lacpy_size_check(size(Bdata), sz)
# only the diagonal is copied in this case
copyto!(diagview(Bdata), diagview(Adata))
end
return Bdata
end

function copytrito!(B::UpperTriangular, A::UpperTriangular, uplo::AbstractChar)
m,n = size(A)
copytrito_triangular!(B.data, A.data, uplo, uplo == 'U', (m, m < n ? m : n))
return B
end
function copytrito!(B::LowerTriangular, A::LowerTriangular, uplo::AbstractChar)
m,n = size(A)
copytrito_triangular!(B.data, A.data, uplo, uplo == 'L', (n < m ? n : m, n))
return B
end

uppertridata(A) = A
lowertridata(A) = A
# we restrict these specializations only to strided matrices to avoid cases where an UpperTriangular type
# doesn't share its indexing with the parent
uppertridata(A::UpperTriangular{<:Any, <:StridedMatrix}) = parent(A)
lowertridata(A::LowerTriangular{<:Any, <:StridedMatrix}) = parent(A)

@inline _rscale_add!(A::AbstractTriangular, B::AbstractTriangular, C::Number, alpha::Number, beta::Number) =
@stable_muladdmul _triscale!(A, B, C, MulAddMul(alpha, beta))
@inline _lscale_add!(A::AbstractTriangular, B::Number, C::AbstractTriangular, alpha::Number, beta::Number) =
Expand Down
20 changes: 20 additions & 0 deletions stdlib/LinearAlgebra/test/triangular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1396,4 +1396,24 @@ end
@test exp(log(M)) M
end

@testset "copytrito!" begin
for T in (UpperTriangular, LowerTriangular)
M = Matrix{BigFloat}(undef, 2, 2)
M[1,1] = M[2,2] = 3
U = T(M)
isupper = U isa UpperTriangular
M[1+!isupper, 1+isupper] = 4
uplo, loup = U isa UpperTriangular ? ('U', 'L') : ('L', 'U' )
@test copytrito!(similar(U), U, uplo) == U
@test copytrito!(zero(M), U, uplo) == U
@test copytrito!(similar(U), Array(U), uplo) == U
@test copytrito!(zero(U), U, loup) == Diagonal(U)
@test copytrito!(similar(U), MyTriangular(U), uplo) == U
@test copytrito!(zero(M), MyTriangular(U), uplo) == U
Ubig = T(similar(M, (3,3)))
copytrito!(Ubig, U, uplo)
@test Ubig[axes(U)...] == U
end
end

end # module TestTriangular

2 comments on commit 4709b6c

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executing the daily package evaluation, I will reply here when finished:

@nanosoldier runtests(isdaily = true)

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The package evaluation job you requested has completed - possible new issues were detected.
The full report is available.

Please sign in to comment.