Skip to content

Commit a40d8c4

Browse files
authored
Add @inline for Diagonal's 2-arg l/rdiv! to enable auto vectorization (#43171)
1 parent 01f0965 commit a40d8c4

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

stdlib/LinearAlgebra/src/diagonal.jl

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -357,7 +357,7 @@ end
357357

358358
/(A::AbstractVecOrMat, D::Diagonal) = _rdiv!(similar(A, promote_op(/, eltype(A), eltype(D)), size(A)), A, D)
359359

360-
rdiv!(A::AbstractVecOrMat, D::Diagonal) = _rdiv!(A, A, D)
360+
rdiv!(A::AbstractVecOrMat, D::Diagonal) = @inline _rdiv!(A, A, D)
361361
# avoid copy when possible via internal 3-arg backend
362362
function _rdiv!(B::AbstractVecOrMat, A::AbstractVecOrMat, D::Diagonal)
363363
require_one_based_indexing(A)
@@ -378,17 +378,21 @@ end
378378

379379
\(D::Diagonal, B::AbstractVecOrMat) = ldiv!(similar(B, promote_op(\, eltype(D), eltype(B)), size(B)), D, B)
380380

381-
ldiv!(D::Diagonal, B::AbstractVecOrMat) = ldiv!(B, D, B)
381+
ldiv!(D::Diagonal, B::AbstractVecOrMat) = @inline ldiv!(B, D, B)
382382
function ldiv!(B::AbstractVecOrMat, D::Diagonal, A::AbstractVecOrMat)
383383
require_one_based_indexing(A, B)
384-
d = length(D.diag)
384+
dd = D.diag
385+
d = length(dd)
385386
m, n = size(A, 1), size(A, 2)
386387
m′, n′ = size(B, 1), size(B, 2)
387388
m == d || throw(DimensionMismatch("right hand side has $m rows but D is $d by $d"))
388389
(m, n) == (m′, n′) || throw(DimensionMismatch("expect output to be $m by $n, but got $m′ by $n′"))
389390
j = findfirst(iszero, D.diag)
390391
isnothing(j) || throw(SingularException(j))
391-
B .= D.diag .\ A
392+
@inbounds for j = 1:n, i = 1:m
393+
B[i, j] = dd[i] \ A[i, j]
394+
end
395+
B
392396
end
393397

394398
# Optimizations for \, / between Diagonals

0 commit comments

Comments
 (0)