Skip to content

Commit c234bed

Browse files
authored
Loop over diagind in diag for banded matrices (#1211)
This improves performance in obtaining the zero bands for banded matrices. This is mainly because of `@inbounds` annotations. Using ```julia D = Diagonal(1:6000); B = Bidiagonal(1:6000, 1:5999, :U); ``` | Operation | master | PR | | -------- | ------- | ------- | |`@btime diag($D, 2);` |2.832 μs |1.925 μs | |`@btime diag($B, 2);` |2.283 μs |1.937 μs | The performance difference in `Tridiagonal` and `SymTridiagonal` seems pretty negligible, for reasons that I don't fully understand. Perhaps the indexing is a more complicated operation for these, so that the bounds-checking isn't the dominant contributor.
1 parent 57785c7 commit c234bed

File tree

3 files changed

+14
-10
lines changed

3 files changed

+14
-10
lines changed

src/bidiag.jl

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -405,14 +405,15 @@ end
405405
function diag(M::Bidiagonal, n::Integer=0)
406406
# every branch call similar(..., ::Int) to make sure the
407407
# same vector type is returned independent of n
408-
v = similar(M.dv, max(0, length(M.dv)-abs(n)))
408+
dinds = diagind(M, n, IndexStyle(M))
409+
v = similar(M.dv, length(dinds))
409410
if n == 0
410411
copyto!(v, M.dv)
411412
elseif (n == 1 && M.uplo == 'U') || (n == -1 && M.uplo == 'L')
412413
copyto!(v, M.ev)
413414
elseif -size(M,1) <= n <= size(M,1)
414-
for i in eachindex(v)
415-
v[i] = M[BandIndex(n,i)]
415+
for i in eachindex(v, dinds)
416+
@inbounds v[i] = M[BandIndex(n,i)]
416417
end
417418
end
418419
return v

src/diagonal.jl

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -836,12 +836,13 @@ permutedims(D::Diagonal, perm) = (Base.checkdims_perm(axes(D), axes(D), perm); D
836836
function diag(D::Diagonal, k::Integer=0)
837837
# every branch call similar(..., ::Int) to make sure the
838838
# same vector type is returned independent of k
839-
v = similar(D.diag, max(0, length(D.diag)-abs(k)))
839+
dinds = diagind(D, k, IndexStyle(D))
840+
v = similar(D.diag, length(dinds))
840841
if k == 0
841842
copyto!(v, D.diag)
842843
else
843-
for i in eachindex(v)
844-
v[i] = D[BandIndex(k, i)]
844+
for i in eachindex(v, dinds)
845+
@inbounds v[i] = D[dinds[i]]
845846
end
846847
end
847848
return v

src/tridiag.jl

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -191,15 +191,16 @@ _eviter_transposed(M::SymTridiagonal) = (transpose(x) for x in _evview(M))
191191
function diag(M::SymTridiagonal, n::Integer=0)
192192
# every branch call similar(..., ::Int) to make sure the
193193
# same vector type is returned independent of n
194-
v = similar(M.dv, max(0, length(M.dv)-abs(n)))
194+
dinds = diagind(M, n, IndexStyle(M))
195+
v = similar(M.dv, length(dinds))
195196
if n == 0
196197
return copyto!(v, _diagiter(M))
197198
elseif n == 1
198199
return copyto!(v, _evview(M))
199200
elseif n == -1
200201
return copyto!(v, _eviter_transposed(M))
201202
else
202-
for i in eachindex(v)
203+
for i in eachindex(v, dinds)
203204
v[i] = M[BandIndex(n,i)]
204205
end
205206
end
@@ -662,15 +663,16 @@ issymmetric(S::Tridiagonal) = all(issymmetric, S.d) && all(Iterators.map((x, y)
662663
function diag(M::Tridiagonal, n::Integer=0)
663664
# every branch call similar(..., ::Int) to make sure the
664665
# same vector type is returned independent of n
665-
v = similar(M.d, max(0, length(M.d)-abs(n)))
666+
dinds = diagind(M, n, IndexStyle(M))
667+
v = similar(M.d, length(dinds))
666668
if n == 0
667669
copyto!(v, M.d)
668670
elseif n == -1
669671
copyto!(v, M.dl)
670672
elseif n == 1
671673
copyto!(v, M.du)
672674
elseif abs(n) <= size(M,1)
673-
for i in eachindex(v)
675+
for i in eachindex(v, dinds)
674676
v[i] = M[BandIndex(n,i)]
675677
end
676678
end

0 commit comments

Comments
 (0)