Speicalize copy! for triangular, and use copy! in ldiv (#1263)

Currently, we specialize `copyto!` for a triangular source. However,
this has two branches, depending on whether the axes match. In the
branch where the axes do match, we may use `copy!` instead, and
specialize this in terms of the internal `_copyto!` which is identical
in implementation.

We also use `copy!` in `ldiv` instead of `copyto!`. These should be
equivalent in most cases, barring one extra axes check in `copy!`, as
the fallback method for `copy!` calls `copyto!` internally. However, the
advantage comes for triangular matrices, where `copy!` doesn't have
branches, as the axes necessarily match.

As a consequence, this reduces the TTFX in operations like
```julia
julia> using Random, LinearAlgebra

julia> A = rand(4,4);

julia> @time A \ UpperTriangular(A);
  0.598575 seconds (1.28 M allocations: 61.788 MiB, 98.57% compilation time: 3% of which was recompilation) # master
  0.487267 seconds (1.01 M allocations: 49.726 MiB, 3.54% gc time, 98.95% compilation time) # this PR
```
```julia
julia> @time UpperTriangular(A) / A;
  0.826212 seconds (1.45 M allocations: 71.427 MiB, 16.96% gc time, 84.47% compilation time) # master
  0.616258 seconds (1.28 M allocations: 63.135 MiB, 2.65% gc time, 99.19% compilation time) # this PR
```

I've renamed the internal functions as `_copyto!` -> `_copy!` and
`copyto_unaliased!` -> `copy_unaliased!`, as these are closer to the
meaning of `copy!` than to `copyto!`.

Author: jishnub

src/LinearAlgebra.jl

@@ -10,7 +10,7 @@ module LinearAlgebra
 import Base: \, /, //, *, ^, +, -, ==
 import Base: USE_BLAS64, abs, acos, acosh, acot, acoth, acsc, acsch, adjoint, asec, asech,
     asin, asinh, atan, atanh, axes, big, broadcast, cbrt, ceil, cis, collect, conj, convert,
-    copy, copyto!, copymutable, cos, cosh, cot, coth, csc, csch, eltype, exp, fill!, floor,
+    copy, copy!, copyto!, copymutable, cos, cosh, cot, coth, csc, csch, eltype, exp, fill!, floor,
     getindex, hcat, getproperty, imag, inv, invpermuterows!, isapprox, isequal, isone, iszero,
     IndexStyle, kron, kron!, length, log, map, ndims, one, oneunit, parent, permutecols!,
     permutedims, permuterows!, power_by_squaring, promote_rule, real, isreal, sec, sech, setindex!,
@@ -706,9 +706,9 @@ function ldiv(F::Factorization, B::AbstractVecOrMat)
 
     if n > size(B, 1)
         # Underdetermined
-        copyto!(view(BB, 1:m, :), B)
+        copy!(view(BB, axes(B,1), ntuple(_->:, ndims(B)-1)...), B)
     else
-        copyto!(BB, B)
+        copy!(BB, B)
     end
 
     ldiv!(FF, BB)

src/triangular.jl

@@ -178,7 +178,7 @@ end
 parent(A::UpperOrLowerTriangular) = A.data
 
 # For strided matrices, we may only loop over the filled triangle
-copy(A::UpperOrLowerTriangular{<:Any, <:StridedMaybeAdjOrTransMat}) = copyto!(similar(A), A)
+copy(A::UpperOrLowerTriangular{<:Any, <:StridedMaybeAdjOrTransMat}) = copy!(similar(A), A)
 
 # then handle all methods that requires specific handling of upper/lower and unit diagonal
 
@@ -532,30 +532,34 @@ for T in (:UpperOrUnitUpperTriangular, :LowerOrUnitLowerTriangular)
         if axes(dest) != axes(U)
             @invoke copyto!(dest::AbstractArray, U::AbstractArray)
         else
-            _copyto!(dest, U)
+            copy!(dest, U)
         end
         return dest
     end
+    @eval function copy!(dest::$T, U::$T)
+        axes(dest) == axes(U) || throw(ArgumentError(
+            "arrays must have the same axes for copy! (consider using `copyto!`)"))
+        _copy!(dest, U)
+    end
 end
 
 # copy and scale
 for (T, UT) in ((:UpperTriangular, :UnitUpperTriangular), (:LowerTriangular, :UnitLowerTriangular))
-    @eval @inline function _copyto!(A::$T, B::$T)
-        @boundscheck checkbounds(A, axes(B)...)
+    @eval @inline function _copy!(A::$T, B::$T)
         copytrito!(parent(A), parent(B), uplo_char(A))
         return A
     end
-    @eval @inline function _copyto!(A::$UT, B::$T)
+    @eval @inline function _copy!(A::$UT, B::$T)
         for dind in diagind(A, IndexStyle(A))
             if A[dind] != B[dind]
                 throw_nononeerror(typeof(A), B[dind], Tuple(dind)...)
             end
         end
-        _copyto!($T(parent(A)), B)
+        _copy!($T(parent(A)), B)
         return A
     end
 end
-@inline function _copyto!(A::UpperOrUnitUpperTriangular, B::UnitUpperTriangular)
+@inline function _copy!(A::UpperOrUnitUpperTriangular, B::UnitUpperTriangular)
     @boundscheck checkbounds(A, axes(B)...)
     B2 = Base.unalias(A, B)
     Ap = parent(A)
@@ -570,7 +574,7 @@ end
     end
     return A
 end
-@inline function _copyto!(A::LowerOrUnitLowerTriangular, B::UnitLowerTriangular)
+@inline function _copy!(A::LowerOrUnitLowerTriangular, B::UnitLowerTriangular)
     @boundscheck checkbounds(A, axes(B)...)
     B2 = Base.unalias(A, B)
     Ap = parent(A)
@@ -595,23 +599,30 @@ _triangularize!(::LowerOrUnitLowerTriangular) = tril!
     if axes(dest) != axes(U)
         @invoke copyto!(dest::StridedMatrix, U::AbstractArray)
     else
-        _copyto!(dest, U)
+        copy!(dest, U)
     end
     return dest
 end
-@propagate_inbounds function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular)
+
+function copy!(dest::StridedMatrix, U::UpperOrLowerTriangular)
+    axes(dest) == axes(U) || throw(ArgumentError(
+            "arrays must have the same axes for copy! (consider using `copyto!`)"))
+    _copy!(dest, U)
+end
+
+@propagate_inbounds function _copy!(dest::StridedMatrix, U::UpperOrLowerTriangular)
     copytrito!(dest, parent(U), U isa UpperOrUnitUpperTriangular ? 'U' : 'L')
     copytrito!(dest, U, U isa UpperOrUnitUpperTriangular ? 'L' : 'U')
     return dest
 end
-@propagate_inbounds function _copyto!(dest::StridedMatrix, U::UpperOrLowerTriangular{<:Any, <:StridedMatrix})
+@propagate_inbounds function _copy!(dest::StridedMatrix, U::UpperOrLowerTriangular{<:Any, <:StridedMatrix})
     U2 = Base.unalias(dest, U)
-    copyto_unaliased!(dest, U2)
+    copy_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
-@inline function copyto_unaliased!(dest::StridedMatrix, U::UpperOrUnitUpperTriangular{<:Any, <:StridedMatrix})
+@inline function copy_unaliased!(dest::StridedMatrix, U::UpperOrUnitUpperTriangular{<:Any, <:StridedMatrix})
     @boundscheck checkbounds(dest, axes(U)...)
     isunit = U isa UnitUpperTriangular
     for col in axes(dest,2)
@@ -624,7 +635,7 @@ end
     end
     return dest
 end
-@inline function copyto_unaliased!(dest::StridedMatrix, L::LowerOrUnitLowerTriangular{<:Any, <:StridedMatrix})
+@inline function copy_unaliased!(dest::StridedMatrix, L::LowerOrUnitLowerTriangular{<:Any, <:StridedMatrix})
     @boundscheck checkbounds(dest, axes(L)...)
     isunit = L isa UnitLowerTriangular
     for col in axes(dest,2)
@@ -645,7 +656,7 @@ Base.@constprop :aggressive function copytrito_triangular!(Bdata, Adata, uplo, u
         BLAS.chkuplo(uplo)
         LAPACK.lacpy_size_check(size(Bdata), sz)
         # only the diagonal is copied in this case
-        copyto!(diagview(Bdata), diagview(Adata))
+        copy!(diagview(Bdata), diagview(Adata))
     end
     return Bdata
 end
@@ -1055,15 +1066,17 @@ isunit_char(::UnitUpperTriangular) = 'U'
 isunit_char(::LowerTriangular) = 'N'
 isunit_char(::UnitLowerTriangular) = 'U'
 
+_copy_or_copyto!(dest, src) = ndims(dest) == ndims(src) ? copy!(dest, src) : copyto!(dest, src)
+
 # generic fallback for AbstractTriangular matrices outside of the four subtypes provided here
 _trimul!(C::AbstractVecOrMat, A::AbstractTriangular, B::AbstractVector) =
-    lmul!(A, copyto!(C, B))
+    lmul!(A, _copy_or_copyto!(C, B))
 _trimul!(C::AbstractMatrix, A::AbstractTriangular, B::AbstractMatrix) =
-    lmul!(A, copyto!(C, B))
+    lmul!(A, copy!(C, B))
 _trimul!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractTriangular) =
-    rmul!(copyto!(C, A), B)
+    rmul!(copy!(C, A), B)
 _trimul!(C::AbstractMatrix, A::AbstractTriangular, B::AbstractTriangular) =
-    lmul!(A, copyto!(C, B))
+    lmul!(A, copy!(C, B))
 # redirect for UpperOrLowerTriangular
 _trimul!(C::AbstractVecOrMat, A::UpperOrLowerTriangular, B::AbstractVector) =
     generic_trimatmul!(C, uplo_char(A), isunit_char(A), wrapperop(parent(A)), _unwrap_at(parent(A)), B)
@@ -1124,9 +1137,9 @@ end
 ldiv!(C::AbstractVecOrMat, A::AbstractTriangular, B::AbstractVecOrMat) = _ldiv!(C, A, B)
 # generic fallback for AbstractTriangular, directs to 2-arg [l/r]div!
 _ldiv!(C::AbstractVecOrMat, A::AbstractTriangular, B::AbstractVecOrMat) =
-    ldiv!(A, copyto!(C, B))
+    ldiv!(A, _copy_or_copyto!(C, B))
 _rdiv!(C::AbstractMatrix, A::AbstractMatrix, B::AbstractTriangular) =
-    rdiv!(copyto!(C, A), B)
+    rdiv!(copy!(C, A), B)
 # redirect for UpperOrLowerTriangular to generic_*div!
 _ldiv!(C::AbstractVecOrMat, A::UpperOrLowerTriangular, B::AbstractVecOrMat) =
     generic_trimatdiv!(C, uplo_char(A), isunit_char(A), wrapperop(parent(A)), _unwrap_at(parent(A)), B)
@@ -1204,40 +1217,40 @@ for (t, uploc, isunitc) in ((:LowerTriangular, 'L', 'N'),
             elseif p == Inf
                 return inv(LAPACK.trcon!('I', $uploc, $isunitc, A.data))
             else # use fallback
-                return cond(copyto!(similar(parent(A)), A), p)
+                return cond(copy!(similar(parent(A)), A), p)
             end
         end
     end
 end
 
 # multiplication
 generic_trimatmul!(c::StridedVector{T}, uploc, isunitc, tfun::Function, A::StridedMatrix{T}, b::AbstractVector{T}) where {T<:BlasFloat} =
-    BLAS.trmv!(uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, A, c === b ? c : copyto!(c, b))
+    BLAS.trmv!(uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, A, c === b ? c : copy!(c, b))
 function generic_trimatmul!(C::StridedMatrix{T}, uploc, isunitc, tfun::Function, A::StridedMatrix{T}, B::AbstractMatrix{T}) where {T<:BlasFloat}
     if stride(C,1) == stride(A,1) == 1
-        BLAS.trmm!('L', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), A, C === B ? C : copyto!(C, B))
+        BLAS.trmm!('L', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), A, C === B ? C : copy!(C, B))
     else # incompatible with BLAS
         @invoke generic_trimatmul!(C::AbstractMatrix, uploc, isunitc, tfun::Function, A::AbstractMatrix, B::AbstractMatrix)
     end
 end
 function generic_mattrimul!(C::StridedMatrix{T}, uploc, isunitc, tfun::Function, A::AbstractMatrix{T}, B::StridedMatrix{T}) where {T<:BlasFloat}
     if stride(C,1) == stride(B,1) == 1
-        BLAS.trmm!('R', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), B, C === A ? C : copyto!(C, A))
+        BLAS.trmm!('R', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), B, C === A ? C : copy!(C, A))
     else # incompatible with BLAS
         @invoke generic_mattrimul!(C::AbstractMatrix, uploc, isunitc, tfun::Function, A::AbstractMatrix, B::AbstractMatrix)
     end
 end
 # division
 function generic_trimatdiv!(C::StridedVecOrMat{T}, uploc, isunitc, tfun::Function, A::StridedMatrix{T}, B::AbstractVecOrMat{T}) where {T<:BlasFloat}
     if stride(C,1) == stride(A,1) == 1
-        LAPACK.trtrs!(uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, A, C === B ? C : copyto!(C, B))
+        LAPACK.trtrs!(uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, A, C === B ? C : _copy_or_copyto!(C, B))
     else # incompatible with LAPACK
         @invoke generic_trimatdiv!(C::AbstractVecOrMat, uploc, isunitc, tfun::Function, A::AbstractMatrix, B::AbstractVecOrMat)
     end
 end
 function generic_mattridiv!(C::StridedMatrix{T}, uploc, isunitc, tfun::Function, A::AbstractMatrix{T}, B::StridedMatrix{T}) where {T<:BlasFloat}
     if stride(C,1) == stride(B,1) == 1
-        BLAS.trsm!('R', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), B, C === A ? C : copyto!(C, A))
+        BLAS.trsm!('R', uploc, tfun === identity ? 'N' : tfun === transpose ? 'T' : 'C', isunitc, one(T), B, C === A ? C : copy!(C, A))
     else # incompatible with BLAS
         @invoke generic_mattridiv!(C::AbstractMatrix, uploc, isunitc, tfun::Function, A::AbstractMatrix, B::AbstractMatrix)
     end
@@ -1962,22 +1975,22 @@ function powm!(A0::UpperTriangular, p::Real)
         for i in axes(S,1)
             @inbounds S[i, i] = S[i, i] + 1
         end
-        copyto!(Stmp, S)
+        copy!(Stmp, S)
         mul!(S, A, c)
         ldiv!(Stmp, S)
 
         c = (p - j) / (j4 - 2)
         for i in axes(S,1)
             @inbounds S[i, i] = S[i, i] + 1
         end
-        copyto!(Stmp, S)
+        copy!(Stmp, S)
         mul!(S, A, c)
         ldiv!(Stmp, S)
     end
     for i in axes(S,1)
         S[i, i] = S[i, i] + 1
     end
-    copyto!(Stmp, S)
+    copy!(Stmp, S)
     mul!(S, A, -p)
     ldiv!(Stmp, S)
     for i in axes(S,1)
@@ -1987,7 +2000,7 @@ function powm!(A0::UpperTriangular, p::Real)
     blockpower!(A0, S, p/(2^s))
     for m = 1:s
         mul!(Stmp.data, S, S)
-        copyto!(S, Stmp)
+        copy!(S, Stmp)
         blockpower!(A0, S, p/(2^(s-m)))
     end
     rmul!(S, normA0^p)
@@ -2174,7 +2187,7 @@ function _find_params_log_quasitriu!(A)
             break
         end
         _sqrt_quasitriu!(A isa UpperTriangular ? parent(A) : A, A)
-        copyto!(AmI, A)
+        copy!(AmI, A)
         for i in axes(AmI,1)
             @inbounds AmI[i,i] -= 1
         end