diff --git a/stdlib/LinearAlgebra/src/triangular.jl b/stdlib/LinearAlgebra/src/triangular.jl index 07eafaa8bfc3f..aba67f24a6b28 100644 --- a/stdlib/LinearAlgebra/src/triangular.jl +++ b/stdlib/LinearAlgebra/src/triangular.jl @@ -246,10 +246,11 @@ Base.isstored(A::UpperTriangular, i::Int, j::Int) = @propagate_inbounds getindex(A::UpperTriangular, i::Int, j::Int) = i <= j ? A.data[i,j] : _zero(A.data,j,i) +_zero_triangular_half_str(::Type{<:UpperOrUnitUpperTriangular}) = "lower" +_zero_triangular_half_str(::Type{<:LowerOrUnitLowerTriangular}) = "upper" + @noinline function throw_nonzeroerror(T, @nospecialize(x), i, j) - _upper_lower_str(::Type{<:UpperOrUnitUpperTriangular}) = "upper" - _upper_lower_str(::Type{<:LowerOrUnitLowerTriangular}) = "lower" - Ts = _upper_lower_str(T) + Ts = _zero_triangular_half_str(T) Tn = nameof(T) throw(ArgumentError( lazy"cannot set index in the $Ts triangular part ($i, $j) of an $Tn matrix to a nonzero value ($x)")) @@ -301,9 +302,7 @@ end end @noinline function throw_setindex_structuralzero_error(T, @nospecialize(x)) - _struct_zero_half_str(::Type{<:UpperTriangular}) = "lower" - _struct_zero_half_str(::Type{<:LowerTriangular}) = "upper" - Ts = _struct_zero_half_str(T) + Ts = _zero_triangular_half_str(T) Tn = nameof(T) throw(ArgumentError( lazy"cannot set indices in the $Ts triangular part of an $Tn matrix to a nonzero value ($x)")) @@ -510,24 +509,40 @@ tr(A::UnitLowerTriangular) = size(A, 1) * oneunit(eltype(A)) tr(A::UpperTriangular) = tr(A.data) tr(A::UnitUpperTriangular) = size(A, 1) * oneunit(eltype(A)) +@propagate_inbounds function copyto!(dest::UpperOrLowerTriangular, U::UpperOrLowerTriangular) + if axes(dest) != axes(U) + @invoke copyto!(dest::AbstractArray, U::AbstractArray) + else + _copyto!(dest, U) + end + return dest +end + # copy and scale -function copyto!(A::T, B::T) where {T<:Union{UpperTriangular,UnitUpperTriangular}} +for T in (:UpperTriangular, :LowerTriangular) + @eval @inline function _copyto!(A::$T, B::$T) + @boundscheck checkbounds(A, axes(B)...) + copytrito!(parent(A), parent(B), uplo_char(A)) + return A + end +end +@inline function _copyto!(A::UnitUpperTriangular, B::UnitUpperTriangular) @boundscheck checkbounds(A, axes(B)...) n = size(B,1) B2 = Base.unalias(A, B) for j = 1:n - for i = 1:(isa(B, UnitUpperTriangular) ? j-1 : j) + for i = 1:j-1 @inbounds A[i,j] = B2[i,j] end end return A end -function copyto!(A::T, B::T) where {T<:Union{LowerTriangular,UnitLowerTriangular}} +@inline function _copyto!(A::UnitLowerTriangular, B::UnitLowerTriangular) @boundscheck checkbounds(A, axes(B)...) n = size(B,1) B2 = Base.unalias(A, B) for j = 1:n - for i = (isa(B, UnitLowerTriangular) ? j+1 : j):n + for i = j+1:n @inbounds A[i,j] = B2[i,j] end end @@ -537,7 +552,7 @@ end _triangularize!(::UpperOrUnitUpperTriangular) = triu! _triangularize!(::LowerOrUnitLowerTriangular) = tril! -function copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular) +@propagate_inbounds function copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular) if axes(dest) != axes(U) @invoke copyto!(dest::StridedMatrix, U::AbstractArray) else @@ -545,19 +560,20 @@ function copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular) end return dest end -function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular) +@propagate_inbounds function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular) copytrito!(dest, parent(U), U isa UpperOrUnitUpperTriangular ? 'U' : 'L') copytrito!(dest, U, U isa UpperOrUnitUpperTriangular ? 'L' : 'U') return dest end -function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular{<:Any, <:StridedMatrix}) +@propagate_inbounds function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular{<:Any, <:StridedMatrix}) U2 = Base.unalias(dest, U) copyto_unaliased!(dest, U2) return dest end # for strided matrices, we explicitly loop over the arrays to improve cache locality # This fuses the copytrito! for the two halves -function copyto_unaliased!(dest::StridedMatrix, U::UpperOrUnitUpperTriangular{<:Any, <:StridedMatrix}) +@inline function copyto_unaliased!(dest::StridedMatrix, U::UpperOrUnitUpperTriangular{<:Any, <:StridedMatrix}) + @boundscheck checkbounds(dest, axes(U)...) isunit = U isa UnitUpperTriangular for col in axes(dest,2) for row in 1:col-isunit @@ -569,7 +585,8 @@ function copyto_unaliased!(dest::StridedMatrix, U::UpperOrUnitUpperTriangular{<: end return dest end -function copyto_unaliased!(dest::StridedMatrix, L::LowerOrUnitLowerTriangular{<:Any, <:StridedMatrix}) +@inline function copyto_unaliased!(dest::StridedMatrix, L::LowerOrUnitLowerTriangular{<:Any, <:StridedMatrix}) + @boundscheck checkbounds(dest, axes(L)...) isunit = L isa UnitLowerTriangular for col in axes(dest,2) for row in 1:col-!isunit diff --git a/stdlib/LinearAlgebra/test/triangular.jl b/stdlib/LinearAlgebra/test/triangular.jl index 5a4f2fd0318d0..93077c5631390 100644 --- a/stdlib/LinearAlgebra/test/triangular.jl +++ b/stdlib/LinearAlgebra/test/triangular.jl @@ -1083,7 +1083,7 @@ end @testset "copyto! with aliasing (#39460)" begin M = Matrix(reshape(1:36, 6, 6)) - @testset for T in (UpperTriangular, LowerTriangular) + @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) A = T(view(M, 1:5, 1:5)) A2 = copy(A) B = T(view(M, 2:6, 2:6)) @@ -1091,6 +1091,25 @@ end end end +@testset "copyto! with different sizes" begin + Ap = zeros(3,3) + Bp = rand(2,2) + @testset for T in (UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular) + A = T(Ap) + B = T(Bp) + @test_throws ArgumentError copyto!(A, B) + end + @testset "error message" begin + A = UpperTriangular(Ap) + B = UpperTriangular(Bp) + @test_throws "cannot set index in the lower triangular part" copyto!(A, B) + + A = LowerTriangular(Ap) + B = LowerTriangular(Bp) + @test_throws "cannot set index in the upper triangular part" copyto!(A, B) + end +end + @testset "getindex with Integers" begin M = reshape(1:4,2,2) for Ttype in (UpperTriangular, UnitUpperTriangular)