copytrito! for triangular matrices (#56620)

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
```

Author: jishnub

src/generic.jl

@@ -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

src/triangular.jl

@@ -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) =

test/triangular.jl

@@ -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