Orthogonalize materialize for tribanded (#393)

Author: jishnub

src/tribanded.jl

@@ -12,19 +12,17 @@ triangularlayout(::Type{Tri}, ::ML) where {Tri,ML<:ConjLayout{<:BandedRows}} = T
 
 sublayout(::TriangularLayout{UPLO,UNIT,ML}, inds) where {UPLO,UNIT,ML} = sublayout(ML(), inds)
 
+getuplo(::Type{<:TriangularLayout{UPLO}}) where {UPLO} = UPLO
+getunit(::Type{<:TriangularLayout{<:Any,UNIT}}) where {UNIT} = UNIT
+getbwdim(uplo::Char) = uplo == 'U' ? 2 : 1
 
-function bandeddata(::TriangularLayout{'U'}, A)
-    B = triangulardata(A)
-    u = bandwidth(B,2)
-    D = bandeddata(B)
-    view(D, 1:u+1, :)
-end
-
-function bandeddata(::TriangularLayout{'L'}, A)
+function bandeddata(::TriLayout, A) where {TriLayout<:TriangularLayout}
     B = triangulardata(A)
     l,u = bandwidths(B)
     D = bandeddata(B)
-    view(D, u+1:l+u+1, :)
+    uplo = getuplo(TriLayout)
+    rowinds = uplo == 'U' ? (1:u+1) : (u+1:l+u+1)
+    view(D, rowinds, :)
 end
 
 # function bandedrowsdata(::TriangularLayout{'U'}, A)
@@ -46,16 +44,13 @@ end
 copy(M::Lmul{<:TriangularLayout{uplo,trans,<:AbstractBandedLayout},<:AbstractBandedLayout}) where {uplo,trans} =
     ArrayLayouts.lmul!(M.A, BandedMatrix(M.B, bandwidths(M)))
 
-@inline function materialize!(M::BlasMatLmulVec{<:TriangularLayout{'U',UNIT,<:BandedColumnMajor},
-                                                <:AbstractStridedLayout}) where UNIT
-    A,x = M.A,M.B
-    tbmv!('U', 'N', UNIT, size(A,1), bandwidth(A,2), bandeddata(A), x)
-end
+@inline function materialize!(M::BlasMatLmulVec{<:TriLayout,
+        <:AbstractStridedLayout}) where {UNIT, TriLayout<:TriangularLayout{<:Any,UNIT,<:BandedColumnMajor}}
 
-@inline function materialize!(M::BlasMatLmulVec{<:TriangularLayout{'L',UNIT,<:BandedColumnMajor},
-                                                <:AbstractStridedLayout}) where UNIT
     A,x = M.A,M.B
-    tbmv!('L', 'N', UNIT, size(A,1), bandwidth(A,1), bandeddata(A), x)
+    uplo = getuplo(TriLayout)
+    bwdim = getbwdim(uplo)
+    tbmv!(uplo, 'N', UNIT, size(A,1), bandwidth(A,bwdim), bandeddata(A), x)
 end
 
 # @inline function materialize!(M::BlasMatLmulVec{<:TriangularLayout{'U',UNIT,<:BandedRowMajor},
@@ -77,44 +72,25 @@ end
 # end
 
 # Ldiv
-for UNIT in ('N', 'U')
-    @eval begin
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'U',$UNIT,<:BandedColumnMajor}, <:AbstractStridedLayout})
-            A,x = M.A,M.B
-            tbsv!('U', 'N', $UNIT, size(A,1), bandwidth(A,2), bandeddata(A), x)
-        end
-
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'L',$UNIT,<:BandedColumnMajor}, <:AbstractStridedLayout})
-            A,x = M.A,M.B
-            tbsv!('L', 'N', $UNIT, size(A,1), bandwidth(A,1), bandeddata(A), x)
-        end
-
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'U',$UNIT,<:BandedRowMajor}, <:AbstractStridedLayout})
-            U,x = M.A,M.B
-            A = triangulardata(U)
-            tbsv!('L', 'T', $UNIT, size(A,1), bandwidth(A,2), bandeddata(LowerTriangular(A')), x)
-        end
-
-        @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{'L',$UNIT,<:BandedRowMajor}, <:AbstractStridedLayout})
-            L,x = M.A,M.B
-            A = triangulardata(L)
-            tbsv!('U', 'T', $UNIT, size(A,1), bandwidth(A,1), bandeddata(UpperTriangular(A')), x)
-        end
-    end
-    # for UPLO in ('U', 'L')
-    #     @eval begin
-    #         @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{$UPLO,$UNIT,BandedRowMajor},
-    #                                                     <:AbstractStridedLayout})
-    #             A,x = M.A,M.B
-    #             tbsv!($UPLO, 'T', $UNIT, transpose(bandeddata(A)), x)
-    #         end
-
-    #         @inline function materialize!(M::BlasMatLdivVec{<:TriangularLayout{$UPLO,$UNIT,ConjLayout{BandedRowMajor}},
-    #                                                     <:AbstractStridedLayout})
-    #             A,x = M.A,M.B
-    #             tbsv!($UPLO, 'C', $UNIT, bandeddata(A)', x)
-    #         end
-    #     end
-    # end
+@inline function materialize!(M::BlasMatLdivVec{<:TriLayout,
+        <:AbstractStridedLayout}) where {TriLayout<:TriangularLayout{<:Any,<:Any,<:BandedColumnMajor}}
+
+    uplo = getuplo(TriLayout)
+    unit = getunit(TriLayout)
+    bwdim = getbwdim(uplo)
+    A,x = M.A,M.B
+    tbsv!(uplo, 'N', unit, size(A,1), bandwidth(A,bwdim), bandeddata(A), x)
 end
 
+@inline function materialize!(M::BlasMatLdivVec{<:TriLayout,
+        <:AbstractStridedLayout}) where {TriLayout<:TriangularLayout{<:Any,<:Any,<:BandedRowMajor}}
+
+    U,x = M.A,M.B
+    A = triangulardata(U)
+    uplo = getuplo(TriLayout)
+    uploT = uplo == 'U' ? 'L' : 'U'
+    unit = getunit(TriLayout)
+    bwdim = getbwdim(uplo)
+    LUT = uplo == 'U' ? LowerTriangular : UpperTriangular
+    tbsv!(uploT, 'T', unit, size(A,1), bandwidth(A,bwdim), bandeddata(LUT(A')), x)
+end