Triangular of Tridiagonal, use unalias (#37)

* Triangular of Tridiagonal

* Use unalias, improve array-valued mul

* use fill! for MulAdd

* Increase coverage

Author: dlfivefifty

src/ArrayLayouts.jl

@@ -27,7 +27,7 @@ import Base: AbstractArray, AbstractMatrix, AbstractVector,
       similar, @_gc_preserve_end, @_gc_preserve_begin,
       @nexprs, @ncall, @ntuple, tuple_type_tail,
       all, any, isbitsunion, issubset, replace_in_print_matrix, replace_with_centered_mark,
-      strides, unsafe_convert, first_index
+      strides, unsafe_convert, first_index, unalias
 
 import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcasted,
                         combine_eltypes, DefaultArrayStyle, instantiate, materialize,

src/ldiv.jl

@@ -138,6 +138,9 @@ macro _layoutldiv(Typ)
 
         LinearAlgebra.ldiv!(A::Factorization, x::$Typ) = ArrayLayouts.ldiv!(A,x)
 
+        LinearAlgebra.ldiv!(A::Bidiagonal, B::$Typ) = ArrayLayouts.ldiv!(A,B)
+
+
         Base.:\(A::$Typ, x::AbstractVector) = ArrayLayouts.ldiv(A,x)
         Base.:\(A::$Typ, x::AbstractMatrix) = ArrayLayouts.ldiv(A,x)
 

src/memorylayout.jl

@@ -485,7 +485,7 @@ MemoryLayout(A::Type{UpperTriangular{T,P}}) where {T,P} = triangularlayout(Upper
 MemoryLayout(A::Type{UnitUpperTriangular{T,P}}) where {T,P} = triangularlayout(UnitUpperTriangularLayout, MemoryLayout(P))
 MemoryLayout(A::Type{LowerTriangular{T,P}}) where {T,P} = triangularlayout(LowerTriangularLayout, MemoryLayout(P))
 MemoryLayout(A::Type{UnitLowerTriangular{T,P}}) where {T,P} = triangularlayout(UnitLowerTriangularLayout, MemoryLayout(P))
-triangularlayout(_::Type{Tri}, ::MemoryLayout) where Tri = Tri{UnknownLayout}()
+triangularlayout(::Type{Tri}, ::MemoryLayout) where Tri = Tri{UnknownLayout}()
 triangularlayout(::Type{Tri}, ::ML) where {Tri, ML<:AbstractColumnMajor} = Tri{ML}()
 triangularlayout(::Type{Tri}, ::ML) where {Tri, ML<:AbstractRowMajor} = Tri{ML}()
 triangularlayout(::Type{Tri}, ::ML) where {Tri, ML<:ConjLayout{<:AbstractRowMajor}} = Tri{ML}()
@@ -574,6 +574,17 @@ transposelayout(ml::SymTridiagonalLayout) = ml
 transposelayout(ml::TridiagonalLayout) = ml
 transposelayout(ml::ConjLayout{DiagonalLayout}) = ml
 
+triangularlayout(::Type{<:TriangularLayout{UPLO,'N'}}, ::TridiagonalLayout{ML}) where {UPLO,ML} = BidiagonalLayout{ML}()
+triangularlayout(::Type{<:TriangularLayout{UPLO,'N'}}, ::TridiagonalLayout{FillLayout}) where UPLO = BidiagonalLayout{FillLayout}()
+triangularlayout(::Type{<:TriangularLayout{UPLO,'U'}}, ::TridiagonalLayout{FillLayout}) where UPLO = BidiagonalLayout{FillLayout}()
+
+bidiagonaluplo(::Union{UpperTriangular,UnitUpperTriangular}) = 'U'
+bidiagonaluplo(::Union{LowerTriangular,UnitLowerTriangular}) = 'L'
+diagonaldata(U::Union{UnitUpperTriangular{T},UnitLowerTriangular{T}}) where T = Ones{T}(size(U,1))
+diagonaldata(U::Union{UpperTriangular{T},LowerTriangular{T}}) where T = diagonaldata(triangulardata(U))
+supdiagonaldata(U::Union{UnitUpperTriangular,UpperTriangular}) = supdiagonaldata(triangulardata(U))
+subdiagonaldata(U::Union{UnitLowerTriangular,LowerTriangular}) = subdiagonaldata(triangulardata(U))
+
 adjointlayout(::Type{<:Real}, ml::SymTridiagonalLayout) = ml
 adjointlayout(::Type{<:Real}, ml::TridiagonalLayout) = ml
 adjointlayout(::Type{<:Real}, ml::BidiagonalLayout) = ml

src/mul.jl

@@ -19,9 +19,12 @@ axes(M::Mul, p::Int) = axes(M)[p]
 length(M::Mul) = prod(size(M))
 size(M::Mul) = map(length,axes(M))
 
-_mul_axes(A::Tuple{<:Any,<:Any}, ::Tuple{<:Any}) = (A[1],)
-_mul_axes(A::Tuple{<:Any}, B::Tuple{<:Any,<:Any}) = (A[1],B[2])
-_mul_axes(A::Tuple{<:Any,<:Any}, B::Tuple{<:Any,<:Any}) = (A[1],B[2])
+_mul_axes(A::Tuple{<:Any,<:Any}, ::Tuple{<:Any}) = (A[1],) # matrix * vector
+_mul_axes(A::Tuple{<:Any}, B::Tuple{<:Any,<:Any}) = (A[1],B[2]) # vector * matrix
+_mul_axes(A::Tuple{<:Any,<:Any}, B::Tuple{<:Any,<:Any}) = (A[1],B[2]) # matrix * matrix
+_mul_axes(::Tuple{}, ::Tuple{}) = () # scalar * scalar
+_mul_axes(::Tuple{}, B::Tuple) = B # scalar * B
+_mul_axes(A::Tuple, B::Tuple{}) = A # A * scalar
 
 axes(M::Mul) = _mul_axes(axes(M.A), axes(M.B))
 

src/muladd.jl

@@ -19,13 +19,13 @@ end
     MulAdd{StyleA,StyleB,StyleC}(α, A, B, β, C)
 end
 
-@inline MulAdd(α, A::AA, B::BB, β, C::CC) where {AA,BB,CC} = 
+@inline MulAdd(α, A::AA, B::BB, β, C::CC) where {AA,BB,CC} =
     MulAdd{typeof(MemoryLayout(AA)), typeof(MemoryLayout(BB)), typeof(MemoryLayout(CC))}(α, A, B, β, C)
 
 MulAdd(A, B) = MulAdd(Mul(A, B))
 function MulAdd(M::Mul)
     TV = eltype(M)
-    MulAdd(scalarone(TV), M.A, M.B, scalarzero(TV), fillzeros(TV,axes(M)))
+    MulAdd(scalarone(TV), M.A, M.B, scalarzero(TV), mulzeros(TV,M))
 end
 
 @inline eltype(::MulAdd{StyleA,StyleB,StyleC,T,AA,BB,CC}) where {StyleA,StyleB,StyleC,T,AA,BB,CC} =
@@ -69,18 +69,11 @@ muladd!(α, A, B, β, C) = materialize!(MulAdd(α, A, B, β, C))
 materialize(M::MulAdd) = copy(instantiate(M))
 copy(M::MulAdd) = copyto!(similar(M), M)
 
-@inline function copyto!(dest::AbstractArray{T}, M::MulAdd) where T
-    M.C === dest  || copyto!(dest, M.C)
-    muladd!(M.α, M.A, M.B, M.β, dest)
-end
+_fill_copyto!(dest, C) = copyto!(dest, C)
+_fill_copyto!(dest, C::Zeros) = zero!(dest) # exploit special fill! overload
 
-@inline function copyto!(dest::AbstractArray{T}, M::MulAdd{<:Any,<:Any,ZerosLayout}) where T
-    α,A,B,β,C = M.α, M.A, M.B, M.β, M.C  
-    if !isbitstype(T) # instantiate
-        dest .= β .* view(A,:,1) .* Ref(B[1])  # get shape right
-    end
-    muladd!(α, A, B, β, dest)
-end
+@inline copyto!(dest::AbstractArray{T}, M::MulAdd) where T = 
+    muladd!(M.α, unalias(dest,M.A), unalias(dest,M.B), M.β, _fill_copyto!(dest, M.C))
 
 # Modified from LinearAlgebra._generic_matmatmul!
 function tile_size(T, S, R)
@@ -226,32 +219,28 @@ function _default_blasmul!(::IndexCartesian, α, A::AbstractMatrix, B::AbstractV
     C
 end
 
-default_blasmul!(α, A::AbstractMatrix, B::AbstractVector, β, C::AbstractVector) = 
+default_blasmul!(α, A::AbstractMatrix, B::AbstractVector, β, C::AbstractVector) =
     _default_blasmul!(Base.IndexStyle(typeof(A)), α, A, B, β, C)
 
 function materialize!(M::MatMulMatAdd)
     α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
-    if C ≡ B
-        B = copy(B)
-    end
-    default_blasmul!(α, A, B, iszero(β) ? false : β, C)
+    default_blasmul!(α, unalias(C,A), unalias(C,B), iszero(β) ? false : β, C)
 end
 
 function materialize!(M::MatMulMatAdd{<:AbstractStridedLayout,<:AbstractStridedLayout,<:AbstractStridedLayout})
-    α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
-    if C ≡ B
-        B = copy(B)
-    end
+    α, Ain, Bin, β, C = M.α, M.A, M.B, M.β, M.C
+    A = unalias(C, Ain)
+    B = unalias(C, Bin)
     ts = tile_size(eltype(A), eltype(B), eltype(C))
     if iszero(β) # false is a "strong" zero to wipe out NaNs
         if ts == 0 || !(axes(A) isa NTuple{2,OneTo{Int}}) || !(axes(B) isa NTuple{2,OneTo{Int}}) || !(axes(C) isa NTuple{2,OneTo{Int}})
-            default_blasmul!(α, A, B, false, C) 
-        else 
+            default_blasmul!(α, A, B, false, C)
+        else
             tiled_blasmul!(ts, α, A, B, false, C)
         end
     else
         if ts == 0 || !(axes(A) isa NTuple{2,OneTo{Int}}) || !(axes(B) isa NTuple{2,OneTo{Int}}) || !(axes(C) isa NTuple{2,OneTo{Int}})
-            default_blasmul!(α, A, B, β, C) 
+            default_blasmul!(α, A, B, β, C)
         else
             tiled_blasmul!(ts, α, A, B, β, C)
         end
@@ -260,29 +249,11 @@ end
 
 function materialize!(M::MatMulVecAdd)
     α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
-    if C ≡ B
-        B = copy(B)
-    end
-    default_blasmul!(α, A, B, iszero(β) ? false : β, C)
+    default_blasmul!(α, unalias(C,A), unalias(C,B), iszero(β) ? false : β, C)
 end
 
-# make copy to make sure always works
-@inline function _gemv!(tA, α, A, x, β, y)
-    if x ≡ y
-        BLAS.gemv!(tA, α, A, copy(x), β, y)
-    else
-        BLAS.gemv!(tA, α, A, x, β, y)
-    end
-end
-
-# make copy to make sure always works
-@inline function _gemm!(tA, tB, α, A, B, β, C)
-    if B ≡ C
-        BLAS.gemm!(tA, tB, α, A, copy(B), β, C)
-    else
-        BLAS.gemm!(tA, tB, α, A, B, β, C)
-    end
-end
+@inline _gemv!(tA, α, A, x, β, y) = BLAS.gemv!(tA, α, unalias(y,A), unalias(y,x), β, y)
+@inline _gemm!(tA, tB, α, A, B, β, C) = BLAS.gemm!(tA, tB, α, unalias(C,A), unalias(C,B), β, C)
 
 # work around pointer issues
 @inline materialize!(M::BlasMatMulVecAdd{<:AbstractColumnMajor,<:AbstractStridedLayout,<:AbstractStridedLayout}) =
@@ -350,21 +321,8 @@ end
 ###
 
 # make copy to make sure always works
-@inline function _symv!(tA, α, A, x, β, y)
-    if x ≡ y
-        BLAS.symv!(tA, α, A, copy(x), β, y)
-    else
-        BLAS.symv!(tA, α, A, x, β, y)
-    end
-end
-
-@inline function _hemv!(tA, α, A, x, β, y)
-    if x ≡ y
-        BLAS.hemv!(tA, α, A, copy(x), β, y)
-    else
-        BLAS.hemv!(tA, α, A, x, β, y)
-    end
-end
+@inline _symv!(tA, α, A, x, β, y) = BLAS.symv!(tA, α, unalias(y,A), unalias(y,x), β, y)
+@inline _hemv!(tA, α, A, x, β, y) = BLAS.hemv!(tA, α, unalias(y,A), unalias(y,x), β, y)
 
 
 materialize!(M::BlasMatMulVecAdd{<:SymmetricLayout{<:AbstractColumnMajor},<:AbstractStridedLayout,<:AbstractStridedLayout}) =
@@ -411,10 +369,28 @@ scalarone(::Type{<:AbstractArray{T}}) where T = scalarone(T)
 scalarzero(::Type{T}) where T = zero(T)
 scalarzero(::Type{<:AbstractArray{T}}) where T = scalarzero(T)
 
-fillzeros(::Type{T}, ax) where T = Zeros{T}(ax)
+fillzeros(::Type{T}, ax) where T<:Number = Zeros{T}(ax)
+mulzeros(::Type{T}, M) where T<:Number = fillzeros(T, axes(M))
+
+# initiate array-valued MulAdd
+function _mulzeros!(dest::AbstractVector{T}, A, B) where T
+    for k in axes(dest,1)
+        dest[k] = similar(Mul(A[k,1],B[1]), eltype(T))
+    end
+    dest
+end
+
+function _mulzeros!(dest::AbstractMatrix{T}, A, B) where T
+    for j in axes(dest,2), k in axes(dest,1)
+        dest[k,j] = similar(Mul(A[k,1],B[1,j]), eltype(T))
+    end
+    dest
+end
+
+mulzeros(::Type{T}, M) where T<:AbstractArray = _mulzeros!(similar(Array{T}, axes(M)), M.A, M.B)
 
 ###
-# Fill 
+# Fill
 ###
 
 copy(M::MulAdd{<:AbstractFillLayout,<:AbstractFillLayout,<:AbstractFillLayout}) = M.α*M.A*M.B + M.β*M.C

test/test_layouts.jl

@@ -126,6 +126,19 @@ struct FooNumber <: Number end
         @test MemoryLayout(Bidiagonal(view(randn(10),[1,2,3]), view(randn(10),[1,2]), :U)) isa BidiagonalLayout{UnknownLayout}
         @test MemoryLayout(SymTridiagonal(view(randn(10),[1,2,3]), view(randn(10),[1,2]))) isa SymTridiagonalLayout{UnknownLayout}
         @test MemoryLayout(Tridiagonal(view(randn(10),[1,2]), view(randn(10),[1,2,3]), view(randn(10),[1,2]))) isa TridiagonalLayout{UnknownLayout}
+
+        @testset "Triangular of Tridiagonal" begin
+            U = UpperTriangular(T)
+            L = LowerTriangular(T)
+            @test MemoryLayout(U) isa BidiagonalLayout{DenseColumnMajor}
+            @test MemoryLayout(L) isa BidiagonalLayout{DenseColumnMajor}
+            @test diagonaldata(U) == diagonaldata(L) == diagonaldata(T)
+            @test subdiagonaldata(L) == subdiagonaldata(T)
+            @test supdiagonaldata(U) == supdiagonaldata(T)
+            @test bidiagonaluplo(U) == 'U'
+            @test bidiagonaluplo(L) == 'L'
+            @test_throws MethodError subdiagonaldata(U)
+        end
     end
 
     @testset "Symmetric/Hermitian" begin

test/test_muladd.jl

@@ -74,6 +74,14 @@ Random.seed!(0)
                 @test all(c .=== BLAS.gemv!('N', 2one(T), A, b, one(T), copy(b)))
             end
         end
+
+        @testset "Matrix{Int} * Vector{Vector{Int}}" begin
+            A, x =  [1 2; 3 4] , [[1,2],[3,4]]
+            X = reshape([[1 2],[3 4], [5 6], [7 8]],2,2)
+            @test mul(A,x) == A*x
+            @test mul(A,X) == A*X
+            @test mul(X,A) == X*A
+        end
     end
 
     @testset "Matrix * Matrix" begin