LinearAlgebra: copyto! between banded matrix types (#54041)

This specialized `copyto!` for combinations of banded structured matrix
types so that the copy may be O(N) instead of the fallback O(N^2)
implementation.

E.g.:
```julia
julia> T = Tridiagonal(zeros(999), zeros(1000), zeros(999));

julia> B = Bidiagonal(ones(1000), fill(2.0, 999), :U);

julia> @Btime copyto!($T, $B);
  1.927 ms (0 allocations: 0 bytes) # master
  229.870 ns (0 allocations: 0 bytes) # PR
```

This also changes the `copyto!` implementation for mismatched matrix
sizes, bringing it closer to the docstring. So, the following works on
master:
```julia
julia> Ddest = Diagonal(zeros(4));

julia> Dsrc = Diagonal(ones(2));

julia> copyto!(Ddest, Dsrc)
4×4 Diagonal{Float64, Vector{Float64}}:
 1.0   ⋅    ⋅    ⋅ 
  ⋅   1.0   ⋅    ⋅ 
  ⋅    ⋅   0.0   ⋅ 
  ⋅    ⋅    ⋅   0.0
```
but this won't work anymore with this PR. This was inconsistent anyway,
as materializing the matrices produces a different result, which
shouldn't be the case:
```julia
julia> copyto!(Matrix(Ddest), Dsrc)
4×4 Matrix{Float64}:
 1.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0
 0.0  0.0  0.0  0.0
 1.0  0.0  0.0  0.0
```
After this PR, the way to carry out the copy would be
```julia
julia> copyto!(Ddest, CartesianIndices(Dsrc), Dsrc, CartesianIndices(Dsrc))
4×4 Diagonal{Float64, Vector{Float64}}:
 1.0   ⋅    ⋅    ⋅ 
  ⋅   1.0   ⋅    ⋅ 
  ⋅    ⋅   0.0   ⋅ 
  ⋅    ⋅    ⋅   0.0
```
This change fixes https://github.com/JuliaLang/julia/issues/46005.

Also fixes https://github.com/JuliaLang/julia/issues/53997
After this,
```julia
julia> @Btime copyto!(C, B) setup=(n = 1_000; B = Bidiagonal(randn(n), randn(n-1), :L); C = Bidiagonal(randn(n), randn(n-1), :L));
  158.405 ns (0 allocations: 0 bytes)

julia> @Btime copyto!(C, B) setup=(n = 10_000; B = Bidiagonal(randn(n), randn(n-1), :L); C = Bidiagonal(randn(n), randn(n-1), :L));
  4.706 μs (0 allocations: 0 bytes)

julia> @Btime copyto!(C, B) setup=(n = 100_000; B = Bidiagonal(randn(n), randn(n-1), :L); C = Bidiagonal(randn(n), randn(n-1), :L));
  120.880 μs (0 allocations: 0 bytes)
```
which is roughly linear scaling.

Taken along with #54027, the
speed-ups would also apply to the adjoints of banded matrices.

Author: jishnub

stdlib/LinearAlgebra/src/bidiag.jl

@@ -291,6 +291,22 @@ function Base.copy(tB::Transpose{<:Any,<:Bidiagonal})
     return Bidiagonal(map(x -> copy.(transpose.(x)), (B.dv, B.ev))..., B.uplo == 'U' ? :L : :U)
 end
 
+# copyto! for matching axes
+function _copyto_banded!(A::Bidiagonal, B::Bidiagonal)
+    A.dv .= B.dv
+    if A.uplo == B.uplo
+        A.ev .= B.ev
+    elseif iszero(B.ev) # diagonal source
+        A.ev .= zero.(A.ev)
+    else
+        zeroband = istriu(A) ? "lower" : "upper"
+        uplo = A.uplo
+        throw(ArgumentError(string("cannot set the ",
+            zeroband, " bidiagonal band to a nonzero value for uplo=:", uplo)))
+    end
+    return A
+end
+
 iszero(M::Bidiagonal) = iszero(M.dv) && iszero(M.ev)
 isone(M::Bidiagonal) = all(isone, M.dv) && iszero(M.ev)
 function istriu(M::Bidiagonal, k::Integer=0)
@@ -332,6 +348,8 @@ function istril(M::Bidiagonal, k::Integer=0)
     end
 end
 isdiag(M::Bidiagonal) = iszero(M.ev)
+issymmetric(M::Bidiagonal) = isdiag(M) && all(issymmetric, M.dv)
+ishermitian(M::Bidiagonal) = isdiag(M) && all(ishermitian, M.dv)
 
 function tril!(M::Bidiagonal{T}, k::Integer=0) where T
     n = length(M.dv)

stdlib/LinearAlgebra/src/diagonal.jl

@@ -136,7 +136,8 @@ Diagonal{T}(::UndefInitializer, n::Integer) where T = Diagonal(Vector{T}(undef,
 similar(D::Diagonal, ::Type{T}) where {T} = Diagonal(similar(D.diag, T))
 similar(D::Diagonal, ::Type{T}, dims::Union{Dims{1},Dims{2}}) where {T} = similar(D.diag, T, dims)
 
-copyto!(D1::Diagonal, D2::Diagonal) = (copyto!(D1.diag, D2.diag); D1)
+# copyto! for matching axes
+_copyto_banded!(D1::Diagonal, D2::Diagonal) = (copyto!(D1.diag, D2.diag); D1)
 
 size(D::Diagonal) = (n = length(D.diag); (n,n))
 

stdlib/LinearAlgebra/src/special.jl

@@ -307,6 +307,92 @@ isdiag(A::HermOrSym{<:Any,<:Diagonal}) = isdiag(parent(A))
 dot(x::AbstractVector, A::RealHermSymComplexSym{<:Real,<:Diagonal}, y::AbstractVector) =
     dot(x, A.data, y)
 
+# O(N) implementations using the banded structure
+function copyto!(dest::BandedMatrix, src::BandedMatrix)
+    if axes(dest) == axes(src)
+        _copyto_banded!(dest, src)
+    else
+        @invoke copyto!(dest::AbstractMatrix, src::AbstractMatrix)
+    end
+    return dest
+end
+function _copyto_banded!(T::Tridiagonal, D::Diagonal)
+    T.d .= D.diag
+    T.dl .= zero.(T.dl)
+    T.du .= zero.(T.du)
+    return T
+end
+function _copyto_banded!(SymT::SymTridiagonal, D::Diagonal)
+    issymmetric(D) || throw(ArgumentError("cannot copy a non-symmetric Diagonal matrix to a SymTridiagonal"))
+    SymT.dv .= D.diag
+    _ev = _evview(SymT)
+    _ev .= zero.(_ev)
+    return SymT
+end
+function _copyto_banded!(B::Bidiagonal, D::Diagonal)
+    B.dv .= D.diag
+    B.ev .= zero.(B.ev)
+    return B
+end
+function _copyto_banded!(D::Diagonal, B::Bidiagonal)
+    isdiag(B) ||
+        throw(ArgumentError("cannot copy a Bidiagonal with a non-zero off-diagonal band to a Diagonal"))
+    D.diag .= B.dv
+    return D
+end
+function _copyto_banded!(D::Diagonal, T::Tridiagonal)
+    isdiag(T) ||
+        throw(ArgumentError("cannot copy a Tridiagonal with a non-zero off-diagonal band to a Diagonal"))
+    D.diag .= T.d
+    return D
+end
+function _copyto_banded!(D::Diagonal, SymT::SymTridiagonal)
+    isdiag(SymT) ||
+        throw(ArgumentError("cannot copy a SymTridiagonal with a non-zero off-diagonal band to a Diagonal"))
+    # we broadcast identity for numbers using the fact that symmetric(x::Number) = x
+    # this potentially allows us to access faster copyto! paths
+    _symmetric = eltype(SymT) <: Number ? identity : symmetric
+    D.diag .= _symmetric.(SymT.dv)
+    return D
+end
+function _copyto_banded!(T::Tridiagonal, B::Bidiagonal)
+    T.d .= B.dv
+    if B.uplo == 'U'
+        T.du .= B.ev
+        T.dl .= zero.(T.dl)
+    else
+        T.dl .= B.ev
+        T.du .= zero.(T.du)
+    end
+    return T
+end
+function _copyto_banded!(SymT::SymTridiagonal, B::Bidiagonal)
+    issymmetric(B) || throw(ArgumentError("cannot copy a non-symmetric Bidiagonal matrix to a SymTridiagonal"))
+    SymT.dv .= B.dv
+    _ev = _evview(SymT)
+    _ev .= zero.(_ev)
+    return SymT
+end
+function _copyto_banded!(B::Bidiagonal, T::Tridiagonal)
+    if B.uplo == 'U' && !iszero(T.dl)
+        throw(ArgumentError("cannot copy a Tridiagonal with a non-zero subdiagonal to a Bidiagonal with uplo=:U"))
+    elseif B.uplo == 'L' && !iszero(T.du)
+        throw(ArgumentError("cannot copy a Tridiagonal with a non-zero superdiagonal to a Bidiagonal with uplo=:L"))
+    end
+    B.dv .= T.d
+    B.ev .= B.uplo == 'U' ? T.du : T.dl
+    return B
+end
+function _copyto_banded!(B::Bidiagonal, SymT::SymTridiagonal)
+    isdiag(SymT) ||
+        throw(ArgumentError("cannot copy a SymTridiagonal with a non-zero off-diagonal band to a Bidiagonal"))
+    # we broadcast identity for numbers using the fact that symmetric(x::Number) = x
+    # this potentially allows us to access faster copyto! paths
+    _symmetric = eltype(SymT) <: Number ? identity : symmetric
+    B.dv .= _symmetric.(SymT.dv)
+    return B
+end
+
 # equals and approx equals methods for structured matrices
 # SymTridiagonal == Tridiagonal is already defined in tridiag.jl
 

stdlib/LinearAlgebra/src/tridiag.jl

@@ -157,7 +157,8 @@ axes(M::SymTridiagonal) = (ax = axes(M.dv, 1); (ax, ax))
 similar(S::SymTridiagonal, ::Type{T}) where {T} = SymTridiagonal(similar(S.dv, T), similar(S.ev, T))
 similar(S::SymTridiagonal, ::Type{T}, dims::Union{Dims{1},Dims{2}}) where {T} = similar(S.dv, T, dims)
 
-copyto!(dest::SymTridiagonal, src::SymTridiagonal) =
+# copyto! for matching axes
+_copyto_banded!(dest::SymTridiagonal, src::SymTridiagonal) =
     (copyto!(dest.dv, src.dv); copyto!(dest.ev, _evview(src)); dest)
 
 #Elementary operations
@@ -606,7 +607,13 @@ similar(M::Tridiagonal, ::Type{T}) where {T} = Tridiagonal(similar(M.dl, T), sim
 similar(M::Tridiagonal, ::Type{T}, dims::Union{Dims{1},Dims{2}}) where {T} = similar(M.d, T, dims)
 
 # Operations on Tridiagonal matrices
-copyto!(dest::Tridiagonal, src::Tridiagonal) = (copyto!(dest.dl, src.dl); copyto!(dest.d, src.d); copyto!(dest.du, src.du); dest)
+# copyto! for matching axes
+function _copyto_banded!(dest::Tridiagonal, src::Tridiagonal)
+    copyto!(dest.dl, src.dl)
+    copyto!(dest.d, src.d)
+    copyto!(dest.du, src.du)
+    dest
+end
 
 #Elementary operations
 for func in (:conj, :copy, :real, :imag)
@@ -981,3 +988,21 @@ function ldiv!(A::Tridiagonal, B::AbstractVecOrMat)
     end
     return B
 end
+
+# combinations of Tridiagonal and Symtridiagonal
+# copyto! for matching axes
+function _copyto_banded!(A::Tridiagonal, B::SymTridiagonal)
+    Bev = _evview(B)
+    A.du .= Bev
+    # Broadcast identity for numbers to access the faster copyto! path
+    # This uses the fact that transpose(x::Number) = x and symmetric(x::Number) = x
+    A.dl .= (eltype(B) <: Number ? identity : transpose).(Bev)
+    A.d .= (eltype(B) <: Number ? identity : symmetric).(B.dv)
+    return A
+end
+function _copyto_banded!(A::SymTridiagonal, B::Tridiagonal)
+    issymmetric(B) || throw(ArgumentError("cannot copy a non-symmetric Tridiagonal matrix to a SymTridiagonal"))
+    A.dv .= B.d
+    _evview(A) .= B.du
+    return A
+end

stdlib/LinearAlgebra/test/bidiag.jl

@@ -834,6 +834,26 @@ end
     end
 end
 
+@testset "copyto!" begin
+    ev, dv = [1:4;], [1:5;]
+    B = Bidiagonal(dv, ev, :U)
+    B2 = copyto!(zero(B), B)
+    @test B2 == B
+    for (ul1, ul2) in ((:U, :L), (:L, :U))
+        B3 = Bidiagonal(dv, zero(ev), ul1)
+        B2 = Bidiagonal(zero(dv), zero(ev), ul2)
+        @test copyto!(B2, B3) == B3
+    end
+
+    @testset "mismatched sizes" begin
+        dv2 = [4; @view dv[2:end]]
+        @test copyto!(B, Bidiagonal([4], Int[], :U)) == Bidiagonal(dv2, ev, :U)
+        @test copyto!(B, Bidiagonal([4], Int[], :L)) == Bidiagonal(dv2, ev, :U)
+        @test copyto!(B, Bidiagonal(Int[], Int[], :U)) == Bidiagonal(dv, ev, :U)
+        @test copyto!(B, Bidiagonal(Int[], Int[], :L)) == Bidiagonal(dv, ev, :U)
+    end
+end
+
 @testset "copyto! with UniformScaling" begin
     @testset "Fill" begin
         for len in (4, InfiniteArrays.Infinity())