IndexCartesian implementation of generic muladd (#24)

* Cartesian blasmul!, symmetric col/rowsupport, copyto! in Ldiv * v0.3.5 * Triangular getindex * Add Bidiagonal support * Add tests
JuliaLinearAlgebra · Jun 23, 2020 · 5e0f275 · 5e0f275 · dlfivefifty · Jun 23, 2020
1 parent d8b862b
commit 5e0f275
Show file tree

Hide file tree

Showing 8 changed files with 628 additions and 503 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 authors = ["Sheehan Olver <solver@mac.com>"]
-version = "0.3.4"
+version = "0.3.5"
 
 [deps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

diff --git a/src/ArrayLayouts.jl b/src/ArrayLayouts.jl
@@ -104,15 +104,8 @@ include("factorizations.jl")
 
 @inline layout_getindex(A, I...) = sub_materialize(view(A, I...))
 
-
-
-macro layoutmatrix(Typ)
+macro _layoutgetindex(Typ)
     esc(quote
-        ArrayLayouts.@layoutldiv $Typ
-        ArrayLayouts.@layoutmul $Typ
-        ArrayLayouts.@layoutlmul $Typ
-        ArrayLayouts.@layoutfactorizations $Typ
-
         @inline Base.getindex(A::$Typ, kr::Colon, jr::Colon) = ArrayLayouts.layout_getindex(A, kr, jr)
         @inline Base.getindex(A::$Typ, kr::Colon, jr::AbstractUnitRange) = ArrayLayouts.layout_getindex(A, kr, jr)
         @inline Base.getindex(A::$Typ, kr::AbstractUnitRange, jr::Colon) = ArrayLayouts.layout_getindex(A, kr, jr)
@@ -123,6 +116,24 @@ macro layoutmatrix(Typ)
     end)
 end
 
+macro layoutgetindex(Typ)
+    esc(quote
+        ArrayLayouts.@_layoutgetindex $Typ
+        ArrayLayouts.@_layoutgetindex LinearAlgebra.AbstractTriangular{<:Any,<:$Typ}
+    end)
+end
+
+
+macro layoutmatrix(Typ)
+    esc(quote
+        ArrayLayouts.@layoutldiv $Typ
+        ArrayLayouts.@layoutmul $Typ
+        ArrayLayouts.@layoutlmul $Typ
+        ArrayLayouts.@layoutfactorizations $Typ
+        ArrayLayouts.@layoutgetindex $Typ
+    end)
+end
+
 @layoutmatrix LayoutMatrix
 
 getindex(A::LayoutVector, kr::AbstractVector) = layout_getindex(A, kr)

diff --git a/src/diagonal.jl b/src/diagonal.jl
@@ -1,7 +1,4 @@
 
-rowsupport(::DiagonalLayout, _, k) = isempty(k) ? (1:0) : minimum(k):maximum(k)
-colsupport(::DiagonalLayout, _, j) = isempty(j) ? (1:0) : minimum(j):maximum(j)
-
 ###
 # Lmul
 ####

diff --git a/src/memorylayout.jl b/src/memorylayout.jl
@@ -500,32 +500,47 @@ abstract type AbstractBandedLayout <: MemoryLayout end
 abstract type AbstractTridiagonalLayout <: AbstractBandedLayout end
 
 struct DiagonalLayout{ML} <: AbstractBandedLayout end
+struct BidiagonalLayout{ML} <: AbstractBandedLayout end
 struct SymTridiagonalLayout{ML} <: AbstractTridiagonalLayout end
 struct TridiagonalLayout{ML} <: AbstractTridiagonalLayout end
 
 diagonallayout(_) = DiagonalLayout{UnknownLayout}()
 diagonallayout(::ML) where ML<:AbstractStridedLayout = DiagonalLayout{ML}()
 MemoryLayout(D::Type{Diagonal{T,P}}) where {T,P} = diagonallayout(MemoryLayout(P))
-diagonaldata(D::Diagonal) = parent(D)
-
+MemoryLayout(::Type{Bidiagonal{T,V}}) where {T,V} = BidiagonalLayout{typeof(MemoryLayout(V))}()
 MemoryLayout(::Type{SymTridiagonal{T,P}}) where {T,P} = SymTridiagonalLayout{typeof(MemoryLayout(P))}()
+MemoryLayout(::Type{Tridiagonal{T,P}}) where {T,P} = TridiagonalLayout{typeof(MemoryLayout(P))}()
+
+bidiagonaluplo(A::Bidiagonal) = A.uplo
+bidiagonaluplo(A::AdjOrTrans) = bidiagonaluplo(parent(A)) == 'L' ? 'U' : 'L'
+
+diagonaldata(D::Diagonal) = parent(D)
+diagonaldata(D::Bidiagonal) = D.dv
 diagonaldata(D::SymTridiagonal) = D.dv
+diagonaldata(D::Tridiagonal) = D.d
+
+supdiagonaldata(D::Bidiagonal) = D.uplo == 'U' ? D.ev : throw(ArgumentError("$D is lower-bidiagonal"))
+subdiagonaldata(D::Bidiagonal) = D.uplo == 'L' ? D.ev : throw(ArgumentError("$D is upper-bidiagonal"))
+
 supdiagonaldata(D::SymTridiagonal) = D.ev
 subdiagonaldata(D::SymTridiagonal) = D.ev
 
-MemoryLayout(::Type{Tridiagonal{T,P}}) where {T,P} = TridiagonalLayout{typeof(MemoryLayout(P))}()
-diagonaldata(D::Tridiagonal) = D.d
 subdiagonaldata(D::Tridiagonal) = D.dl
 supdiagonaldata(D::Tridiagonal) = D.du
 
-
 transposelayout(ml::DiagonalLayout) = ml
+transposelayout(ml::BidiagonalLayout) = ml
 transposelayout(ml::SymTridiagonalLayout) = ml
 transposelayout(ml::TridiagonalLayout) = ml
 transposelayout(ml::ConjLayout{DiagonalLayout}) = ml
 
 adjointlayout(::Type{<:Real}, ml::SymTridiagonalLayout) = ml
 adjointlayout(::Type{<:Real}, ml::TridiagonalLayout) = ml
+adjointlayout(::Type{<:Real}, ml::BidiagonalLayout) = ml
+
+symmetriclayout(B::BidiagonalLayout{ML}) where ML = SymTridiagonalLayout{ML}()
+hermitianlayout(::Type{<:Real}, B::BidiagonalLayout{ML}) where ML = SymTridiagonalLayout{ML}()
+hermitianlayout(_, B::BidiagonalLayout) = HermitianLayout{typeof(B)}()
 
 subdiagonaldata(D::Transpose) = supdiagonaldata(parent(D))
 supdiagonaldata(D::Transpose) = subdiagonaldata(parent(D))
@@ -535,6 +550,10 @@ subdiagonaldata(D::Adjoint{<:Real}) = supdiagonaldata(parent(D))
 supdiagonaldata(D::Adjoint{<:Real}) = subdiagonaldata(parent(D))
 diagonaldata(D::Adjoint{<:Real}) = diagonaldata(parent(D))
 
+diagonaldata(S::HermOrSym) = diagonaldata(parent(S))
+subdiagonaldata(S::HermOrSym) = symmetricuplo(S) == 'L' ? subdiagonaldata(parent(S)) : supdiagonaldata(parent(S))
+supdiagonaldata(S::HermOrSym) = symmetricuplo(S) == 'L' ? subdiagonaldata(parent(S)) : supdiagonaldata(parent(S))
+
 ###
 # Fill
 ####
@@ -575,3 +594,20 @@ colsupport(A) = colsupport(A, axes(A,2))
 
 rowsupport(::ZerosLayout, A, _) = 1:0
 colsupport(::ZerosLayout, A, _) = 1:0
+
+rowsupport(::DiagonalLayout, _, k) = isempty(k) ? (1:0) : minimum(k):maximum(k)
+colsupport(::DiagonalLayout, _, j) = isempty(j) ? (1:0) : minimum(j):maximum(j)
+
+colsupport(::BidiagonalLayout, A, j) = 
+    bidiagonaluplo(A) == 'L' ? (minimum(j):min(size(A,1),maximum(j)+1)) : (max(minimum(j)-1,1):maximum(j))
+rowsupport(::BidiagonalLayout, A, j) = 
+    bidiagonaluplo(A) == 'U' ? (minimum(j):min(size(A,2),maximum(j)+1)) : (max(minimum(j)-1,1):maximum(j))
+
+colsupport(::AbstractTridiagonalLayout, A, j) = max(minimum(j)-1,1):min(size(A,1),maximum(j)+1)
+rowsupport(::AbstractTridiagonalLayout, A, j) = max(minimum(j)-1,1):min(size(A,2),maximum(j)+1)
+
+colsupport(::SymmetricLayout, A, j) = first(colsupport(symmetricdata(A),j)):last(rowsupport(symmetricdata(A),j))
+rowsupport(::SymmetricLayout, A, j) = colsupport(A, j)
+
+colsupport(::HermitianLayout, A, j) = first(colsupport(hermitiandata(A),j)):last(rowsupport(hermitiandata(A),j))
+rowsupport(::HermitianLayout, A, j) = colsupport(A, j)
diff --git a/src/muladd.jl b/src/muladd.jl
@@ -194,7 +194,7 @@ function default_blasmul!(α, A::AbstractMatrix, B::AbstractMatrix, β, C::Abstr
     C
 end
 
-function default_blasmul!(α, A::AbstractMatrix, B::AbstractVector, β, C::AbstractVector)
+function _default_blasmul!(::IndexLinear, α, A::AbstractMatrix, B::AbstractVector, β, C::AbstractVector)
     mA, nA = size(A)
     mB = length(B)
     nA == mB || throw(DimensionMismatch("Dimensions must match"))
@@ -217,6 +217,30 @@ function default_blasmul!(α, A::AbstractMatrix, B::AbstractVector, β, C::Abstr
     C
 end
 
+function _default_blasmul!(::IndexCartesian, α, A::AbstractMatrix, B::AbstractVector, β, C::AbstractVector)
+    mA, nA = size(A)
+    mB = length(B)
+    nA == mB || throw(DimensionMismatch("Dimensions must match"))
+    length(C) == mA || throw(DimensionMismatch("Dimensions must match"))
+
+    lmul!(β, C)
+    (nA == 0 || mB == 0)  && return C
+
+    z = zero(A[1,1]*B[1] + A[1,1]*B[1])
+
+    @inbounds for k in colsupport(B,1)
+        b = B[k]
+        for i = colsupport(A,k)
+            C[i] += α * A[i,k] * b
+        end
+    end
+
+    C
+end
+
+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

diff --git a/src/triangular.jl b/src/triangular.jl
@@ -118,7 +118,7 @@ materialize!(M::MatRmulMat{<:AbstractStridedLayout,<:TriangularLayout}) = Linear
 
 @inline function copyto!(dest::AbstractArray, M::Ldiv{<:TriangularLayout})
     A, B = M.A, M.B
-    dest ≡ B || (dest .= B)
+    dest ≡ B || copyto!(dest, B)
     ldiv!(A, dest)
 end
 

diff --git a/test/test_layouts.jl b/test/test_layouts.jl
@@ -6,7 +6,7 @@ import ArrayLayouts: MemoryLayout, DenseRowMajor, DenseColumnMajor, StridedLayou
                         UnitLowerTriangularLayout, ScalarLayout, UnknownLayout,
                         hermitiandata, symmetricdata, FillLayout, ZerosLayout,
                         DiagonalLayout, TridiagonalLayout, SymTridiagonalLayout, colsupport, rowsupport,
-                        diagonaldata, subdiagonaldata, supdiagonaldata
+                        diagonaldata, subdiagonaldata, supdiagonaldata, BidiagonalLayout, bidiagonaluplo
 
 struct FooBar end
 struct FooNumber <: Number end
@@ -80,7 +80,45 @@ struct FooNumber <: Number end
         @test MemoryLayout(view(randn(5)',[1,3])) == UnknownLayout()
     end
 
-    @testset "Symmetric/Hermitian MemoryLayout" begin
+    @testset "Bi/Tridiagonal" begin
+        T = Tridiagonal(randn(5),randn(6),randn(5))
+        S = SymTridiagonal(T.d, T.du)
+        Bl = Bidiagonal(T.d, T.dl, :L)
+        Bu = Bidiagonal(T.d, T.du, :U)
+
+        @test MemoryLayout(T) isa TridiagonalLayout
+        @test MemoryLayout(Adjoint(T)) isa TridiagonalLayout
+        @test MemoryLayout(Transpose(T)) isa TridiagonalLayout
+        @test MemoryLayout(S) isa SymTridiagonalLayout
+        @test MemoryLayout(Adjoint(S)) isa SymTridiagonalLayout
+        @test MemoryLayout(Transpose(S)) isa SymTridiagonalLayout
+        @test MemoryLayout(Bl) isa BidiagonalLayout
+        @test MemoryLayout(Adjoint(Bl)) isa BidiagonalLayout
+        @test MemoryLayout(Transpose(Bl)) isa BidiagonalLayout
+        @test MemoryLayout(Bu) isa BidiagonalLayout
+        @test MemoryLayout(Adjoint(Bu)) isa BidiagonalLayout
+        @test MemoryLayout(Transpose(Bu)) isa BidiagonalLayout
+
+        @test bidiagonaluplo(Bl) == bidiagonaluplo(Adjoint(Bu)) == 'L'
+        @test bidiagonaluplo(Bu) == bidiagonaluplo(Adjoint(Bl)) == 'U'
+
+        @test diagonaldata(T) == diagonaldata(T') == diagonaldata(S) == diagonaldata(Bl) == diagonaldata(Bu)
+        @test supdiagonaldata(T) == subdiagonaldata(Adjoint(T)) == subdiagonaldata(Transpose(T)) == 
+                    supdiagonaldata(S) == subdiagonaldata(S) == 
+                    supdiagonaldata(Bu) == subdiagonaldata(Adjoint(Bu)) == subdiagonaldata(Transpose(Bu))
+        @test subdiagonaldata(T) == supdiagonaldata(Adjoint(T)) == supdiagonaldata(Transpose(T)) == 
+                subdiagonaldata(Bl) == supdiagonaldata(Adjoint(Bl)) == supdiagonaldata(Transpose(Bl)) ==
+                T.dl
+
+        @test colsupport(T,3) == rowsupport(T,3) == colsupport(S,3) == rowsupport(S,3) == 2:4
+        @test colsupport(T,3:6) == rowsupport(T,3:6) == colsupport(S,3:6) == rowsupport(S,3:6) == 2:6
+        @test colsupport(Bl,3) == rowsupport(Bu,3) == rowsupport(Adjoint(Bl),3) == 3:4
+        @test rowsupport(Bl,3) == colsupport(Bu,3) == colsupport(Adjoint(Bl),3) == 2:3
+        @test colsupport(Bl,3:6) == rowsupport(Bu,3:6) == 3:6
+        @test colsupport(Bu,3:6) == rowsupport(Bl,3:6) == 2:6
+    end
+
+    @testset "Symmetric/Hermitian" begin
         A = [1.0 2; 3 4]
         @test MemoryLayout(Symmetric(A)) == SymmetricLayout{DenseColumnMajor}()
         @test MemoryLayout(Hermitian(A)) == SymmetricLayout{DenseColumnMajor}()
@@ -108,6 +146,11 @@ struct FooNumber <: Number end
         @test symmetricdata(Symmetric(transpose(A))) ≡ transpose(A)
         @test symmetricdata(Hermitian(transpose(A))) ≡ transpose(A)
 
+        @test colsupport(Symmetric(A),2) ≡ colsupport(Symmetric(A),1:2) ≡ 
+                rowsupport(Symmetric(A),2) ≡ rowsupport(Symmetric(A),1:2) ≡ 1:2
+                @test colsupport(Hermitian(A),2) ≡ colsupport(Hermitian(A),1:2) ≡ 
+                rowsupport(Hermitian(A),2) ≡ rowsupport(Hermitian(A),1:2) ≡ 1:2
+
         B = [1.0+im 2; 3 4]
         @test MemoryLayout(Symmetric(B)) == SymmetricLayout{DenseColumnMajor}()
         @test MemoryLayout(Hermitian(B)) == HermitianLayout{DenseColumnMajor}()
@@ -132,6 +175,26 @@ struct FooNumber <: Number end
         @test hermitiandata(Hermitian(B')) ≡ B'
         @test symmetricdata(Symmetric(transpose(B))) ≡ transpose(B)
         @test hermitiandata(Hermitian(transpose(B))) ≡ transpose(B)
+
+        @testset "Bidiagonal" begin
+            B = Bidiagonal(randn(6),randn(5),:U)
+            Bc = Bidiagonal(randn(6) .+ 0im,randn(5) .+ 1im,:U)
+            S = Symmetric(B)
+            H = Hermitian(B)
+            Sc = Symmetric(Bc)
+            Hc = Hermitian(Bc)
+
+            @test MemoryLayout(S) isa SymTridiagonalLayout
+            @test MemoryLayout(H) isa SymTridiagonalLayout
+            @test MemoryLayout(Sc) isa SymTridiagonalLayout
+            @test MemoryLayout(Hc) isa HermitianLayout
+
+            @test diagonaldata(S) == diagonaldata(B)
+            @test subdiagonaldata(S) == supdiagonaldata(S) == supdiagonaldata(B)
+
+            @test colsupport(S,3) == colsupport(H,3) == colsupport(Sc,3) == colsupport(Hc,3) == 2:4
+            @test rowsupport(S,3) == rowsupport(H,3) == rowsupport(Sc,3) == rowsupport(Hc,3) == 2:4
+        end
      end
 
     @testset "triangular MemoryLayout" begin
@@ -233,19 +296,4 @@ struct FooNumber <: Number end
         MemoryLayout(revD)
         @test 0 == @allocated MemoryLayout(revD)
     end
-
-    @testset "Tridiagonal" begin
-        T = Tridiagonal(randn(5),randn(6),randn(5))
-        S = SymTridiagonal(T.d, T.du)
-        @test MemoryLayout(T) isa TridiagonalLayout
-        @test MemoryLayout(Adjoint(T)) isa TridiagonalLayout
-        @test MemoryLayout(Transpose(T)) isa TridiagonalLayout
-        @test MemoryLayout(S) isa SymTridiagonalLayout
-        @test MemoryLayout(Adjoint(S)) isa SymTridiagonalLayout
-        @test MemoryLayout(Transpose(S)) isa SymTridiagonalLayout
-
-        @test diagonaldata(T) == diagonaldata(T') == diagonaldata(S)
-        @test supdiagonaldata(T) == subdiagonaldata(Adjoint(T)) == subdiagonaldata(Transpose(T)) == supdiagonaldata(S) == subdiagonaldata(S)
-        @test subdiagonaldata(T) == supdiagonaldata(Adjoint(T)) == supdiagonaldata(Transpose(T)) == T.dl
-    end
 end