Skip to content

Commit

Permalink
Add @inline for Diagonal's 2-arg l/rdiv! to enable auto vectori…
Browse files Browse the repository at this point in the history
…zation (#43171)
  • Loading branch information
N5N3 authored Nov 24, 2021
1 parent 01f0965 commit a40d8c4
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions stdlib/LinearAlgebra/src/diagonal.jl
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,7 @@ end

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

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

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

ldiv!(D::Diagonal, B::AbstractVecOrMat) = ldiv!(B, D, B)
ldiv!(D::Diagonal, B::AbstractVecOrMat) = @inline ldiv!(B, D, B)
function ldiv!(B::AbstractVecOrMat, D::Diagonal, A::AbstractVecOrMat)
require_one_based_indexing(A, B)
d = length(D.diag)
dd = D.diag
d = length(dd)
m, n = size(A, 1), size(A, 2)
m′, n′ = size(B, 1), size(B, 2)
m == d || throw(DimensionMismatch("right hand side has $m rows but D is $d by $d"))
(m, n) == (m′, n′) || throw(DimensionMismatch("expect output to be $m by $n, but got $m′ by $n′"))
j = findfirst(iszero, D.diag)
isnothing(j) || throw(SingularException(j))
B .= D.diag .\ A
@inbounds for j = 1:n, i = 1:m
B[i, j] = dd[i] \ A[i, j]
end
B
end

# Optimizations for \, / between Diagonals
Expand Down

0 comments on commit a40d8c4

Please sign in to comment.