Skip to content

Commit eb7daee

Browse files
authored
Specialize lmul!/rmul! for strided triangular matrices (#1228)
This improves performance, as we only need to loop over one half of the matrix. On master, `rmul!` for an `UpperTriangular` forwards the multiplication to the parent, so the entire array is looped over. We therefore obtain identical performance for `rmul!(::Matrix, ::Diagonal)` and `rmul!(::UpperTriangular{<:Any, <:Matrix}, ::Diagonal)`. ```julia julia> using LinearAlgebra, Chairmarks julia> D = Diagonal(rand(4000)); julia> A = rand(size(D)...); julia> U = UpperTriangular(A); julia> @b (A, D) rmul!(_[1], _[2]) 10.309 ms julia> @b (U, D) rmul!(_[1], _[2]) 10.370 ms ``` On this PR, the latter is faster, as the loop is only over half the indices: ```julia julia> @b (U, D) rmul!(_[1], _[2]) 7.216 ms ```
1 parent 508e77f commit eb7daee

File tree

2 files changed

+43
-8
lines changed

2 files changed

+43
-8
lines changed

src/diagonal.jl

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,26 @@ function rmul!(T::Tridiagonal, D::Diagonal)
379379
end
380380
return T
381381
end
382+
for T in [:UpperTriangular, :UnitUpperTriangular,
383+
:LowerTriangular, :UnitLowerTriangular]
384+
@eval rmul!(A::$T{<:Any, <:StridedMatrix}, D::Diagonal) = _rmul!(A, D)
385+
@eval lmul!(D::Diagonal, A::$T{<:Any, <:StridedMatrix}) = _lmul!(D, A)
386+
end
387+
function _rmul!(A::UpperOrLowerTriangular, D::Diagonal)
388+
P = parent(A)
389+
isunit = A isa UnitUpperOrUnitLowerTriangular
390+
isupper = A isa UpperOrUnitUpperTriangular
391+
for col in axes(A,2)
392+
rowstart = isupper ? firstindex(A,1) : col+isunit
393+
rowstop = isupper ? col-isunit : lastindex(A,1)
394+
for row in rowstart:rowstop
395+
P[row, col] *= D.diag[col]
396+
end
397+
end
398+
isunit && _setdiag!(P, identity, D.diag)
399+
TriWrapper = isupper ? UpperTriangular : LowerTriangular
400+
return TriWrapper(P)
401+
end
382402

383403
function lmul!(D::Diagonal, B::AbstractVecOrMat)
384404
matmul_size_check(size(D), size(B))
@@ -388,6 +408,13 @@ function lmul!(D::Diagonal, B::AbstractVecOrMat)
388408
end
389409
return B
390410
end
411+
# A' = D * A' => A = A * D'
412+
# This uses the fact that D' is a Diagonal
413+
function lmul!(D::Diagonal, A::AdjOrTransAbsMat)
414+
f = wrapperop(A)
415+
rmul!(f(A), f(D))
416+
A
417+
end
391418

392419
# in-place multiplication with a diagonal
393420
# T .= D * T
@@ -402,12 +429,20 @@ function lmul!(D::Diagonal, T::Tridiagonal)
402429
end
403430
return T
404431
end
405-
# A' = D * A' => A = A * D'
406-
# This uses the fact that D' is a Diagonal
407-
function lmul!(D::Diagonal, A::AdjOrTransAbsMat)
408-
f = wrapperop(A)
409-
rmul!(f(A), f(D))
410-
A
432+
function _lmul!(D::Diagonal, A::UpperOrLowerTriangular)
433+
P = parent(A)
434+
isunit = A isa UnitUpperOrUnitLowerTriangular
435+
isupper = A isa UpperOrUnitUpperTriangular
436+
for col in axes(A,2)
437+
rowstart = isupper ? firstindex(A,1) : col+isunit
438+
rowstop = isupper ? col-isunit : lastindex(A,1)
439+
for row in rowstart:rowstop
440+
P[row, col] = D.diag[row] * P[row, col]
441+
end
442+
end
443+
isunit && _setdiag!(P, identity, D.diag)
444+
TriWrapper = isupper ? UpperTriangular : LowerTriangular
445+
return TriWrapper(P)
411446
end
412447

413448
@inline function __muldiag_nonzeroalpha!(out, D::Diagonal, B, alpha::Number, beta::Number)

test/diagonal.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1197,11 +1197,11 @@ end
11971197
outTri = similar(TriA)
11981198
out = similar(A)
11991199
# 2 args
1200-
for fun in (*, rmul!, rdiv!, /)
1200+
@testset for fun in (*, rmul!, rdiv!, /)
12011201
@test fun(copy(TriA), D)::Tri == fun(Matrix(TriA), D)
12021202
@test fun(copy(UTriA), D)::Tri == fun(Matrix(UTriA), D)
12031203
end
1204-
for fun in (*, lmul!, ldiv!, \)
1204+
@testset for fun in (*, lmul!, ldiv!, \)
12051205
@test fun(D, copy(TriA))::Tri == fun(D, Matrix(TriA))
12061206
@test fun(D, copy(UTriA))::Tri == fun(D, Matrix(UTriA))
12071207
end

0 commit comments

Comments
 (0)