Skip to content

Commit

Permalink
Optimized arithmetic methods for strided triangular matrices (#52571)
Browse files Browse the repository at this point in the history
This uses broadcasting for operations like `A::UpperTriangular +
B::UpperTriangular` in case the parents are `StridedMatrix`es. Looping
only over the triangular part is usually faster for large matrices,
where presumably memory is the bottleneck.

Some performance comparisons, using
```julia
julia> U = UpperTriangular(rand(1000,1000));

julia> U1 = UnitUpperTriangular(rand(size(U)...));
```
| Operation | master | PR |
| --------------- | ---------- | ----- |
|`-U` |`1.011 ms (3 allocations: 7.63 MiB)` |`559.680 μs (3 allocations:
7.63 MiB)` |
|`U + U`/`U - U` |`971.740 μs (3 allocations: 7.63 MiB)` | `560.063 μs
(3 allocations: 7.63 MiB)` |
|`U + U1`/`U - U1` |`3.014 ms (9 allocations: 22.89 MiB)` | `944.772 μs
(3 allocations: 7.63 MiB)` |
|`U1 + U1` |`4.509 ms (12 allocations: 30.52 MiB)` | `1.687 ms (3
allocations: 7.63 MiB)` |
|`U1 - U1` |`3.357 ms (9 allocations: 22.89 MiB)` | `1.763 ms (3
allocations: 7.63 MiB)` |

I've retained the existing methods as fallback, in case there's current
code that works without broadcasting.
  • Loading branch information
jishnub authored Dec 30, 2023
1 parent fe0db7d commit 89cae45
Show file tree
Hide file tree
Showing 2 changed files with 107 additions and 14 deletions.
67 changes: 53 additions & 14 deletions stdlib/LinearAlgebra/src/triangular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ for t in (:LowerTriangular, :UnitLowerTriangular, :UpperTriangular, :UnitUpperTr

real(A::$t{<:Real}) = A
real(A::$t{<:Complex}) = (B = real(A.data); $t(B))
real(A::$t{<:Complex, <:StridedMaybeAdjOrTransMat}) = $t(real.(A))
end
end

Expand Down Expand Up @@ -156,8 +157,26 @@ const UpperOrLowerTriangular{T,S} = Union{UpperOrUnitUpperTriangular{T,S}, Lower

imag(A::UpperTriangular) = UpperTriangular(imag(A.data))
imag(A::LowerTriangular) = LowerTriangular(imag(A.data))
imag(A::UnitLowerTriangular) = LowerTriangular(tril!(imag(A.data),-1))
imag(A::UnitUpperTriangular) = UpperTriangular(triu!(imag(A.data),1))
imag(A::UpperTriangular{<:Any,<:StridedMaybeAdjOrTransMat}) = imag.(A)
imag(A::LowerTriangular{<:Any,<:StridedMaybeAdjOrTransMat}) = imag.(A)
function imag(A::UnitLowerTriangular)
L = LowerTriangular(A.data)
Lim = similar(L) # must be mutable to set diagonals to zero
Lim .= imag.(L)
for i in 1:size(Lim,1)
Lim[i,i] = zero(Lim[i,i])
end
return Lim
end
function imag(A::UnitUpperTriangular)
U = UpperTriangular(A.data)
Uim = similar(U) # must be mutable to set diagonals to zero
Uim .= imag.(U)
for i in 1:size(Uim,1)
Uim[i,i] = zero(Uim[i,i])
end
return Uim
end

Array(A::AbstractTriangular) = Matrix(A)
parent(A::UpperOrLowerTriangular) = A.data
Expand Down Expand Up @@ -481,6 +500,11 @@ function -(A::UnitUpperTriangular)
UpperTriangular(Anew)
end

# use broadcasting if the parents are strided, where we loop only over the triangular part
for TM in (:LowerTriangular, :UpperTriangular)
@eval -(A::$TM{<:Any, <:StridedMaybeAdjOrTransMat}) = broadcast(-, A)
end

tr(A::LowerTriangular) = tr(A.data)
tr(A::UnitLowerTriangular) = size(A, 1) * oneunit(eltype(A))
tr(A::UpperTriangular) = tr(A.data)
Expand Down Expand Up @@ -719,6 +743,16 @@ fillstored!(A::UnitUpperTriangular, x) = (fillband!(A.data, x, 1, size(A,2)-1);
-(A::UnitLowerTriangular, B::UnitLowerTriangular) = LowerTriangular(tril(A.data, -1) - tril(B.data, -1))
-(A::AbstractTriangular, B::AbstractTriangular) = copyto!(similar(parent(A)), A) - copyto!(similar(parent(B)), B)

# use broadcasting if the parents are strided, where we loop only over the triangular part
for op in (:+, :-)
for TM1 in (:LowerTriangular, :UnitLowerTriangular), TM2 in (:LowerTriangular, :UnitLowerTriangular)
@eval $op(A::$TM1{<:Any, <:StridedMaybeAdjOrTransMat}, B::$TM2{<:Any, <:StridedMaybeAdjOrTransMat}) = broadcast($op, A, B)
end
for TM1 in (:UpperTriangular, :UnitUpperTriangular), TM2 in (:UpperTriangular, :UnitUpperTriangular)
@eval $op(A::$TM1{<:Any, <:StridedMaybeAdjOrTransMat}, B::$TM2{<:Any, <:StridedMaybeAdjOrTransMat}) = broadcast($op, A, B)
end
end

######################
# BlasFloat routines #
######################
Expand Down Expand Up @@ -918,47 +952,52 @@ end

for (t, unitt) in ((UpperTriangular, UnitUpperTriangular),
(LowerTriangular, UnitLowerTriangular))
tstrided = t{<:Any, <:StridedMaybeAdjOrTransMat}
@eval begin
(*)(A::$t, x::Number) = $t(A.data*x)
(*)(A::$tstrided, x::Number) = A .* x

function (*)(A::$unitt, x::Number)
B = A.data*x
B = $t(A.data)*x
for i = 1:size(A, 1)
B[i,i] = x
B.data[i,i] = x
end
$t(B)
return B
end

(*)(x::Number, A::$t) = $t(x*A.data)
(*)(x::Number, A::$tstrided) = x .* A

function (*)(x::Number, A::$unitt)
B = x*A.data
B = x*$t(A.data)
for i = 1:size(A, 1)
B[i,i] = x
B.data[i,i] = x
end
$t(B)
return B
end

(/)(A::$t, x::Number) = $t(A.data/x)
(/)(A::$tstrided, x::Number) = A ./ x

function (/)(A::$unitt, x::Number)
B = A.data/x
B = $t(A.data)/x
invx = inv(x)
for i = 1:size(A, 1)
B[i,i] = invx
B.data[i,i] = invx
end
$t(B)
return B
end

(\)(x::Number, A::$t) = $t(x\A.data)
(\)(x::Number, A::$tstrided) = x .\ A

function (\)(x::Number, A::$unitt)
B = x\A.data
B = x\$t(A.data)
invx = inv(x)
for i = 1:size(A, 1)
B[i,i] = invx
B.data[i,i] = invx
end
$t(B)
return B
end
end
end
Expand Down
54 changes: 54 additions & 0 deletions stdlib/LinearAlgebra/test/triangular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -526,6 +526,23 @@ for elty1 in (Float32, Float64, BigFloat, ComplexF32, ComplexF64, Complex{BigFlo
end
end

@testset "non-strided arithmetic" begin
for (T,T1) in ((UpperTriangular, UnitUpperTriangular), (LowerTriangular, UnitLowerTriangular))
U = T(reshape(1:16, 4, 4))
M = Matrix(U)
@test -U == -M
U1 = T1(reshape(1:16, 4, 4))
M1 = Matrix(U1)
@test -U1 == -M1
for op in (+, -)
for (A, MA) in ((U, M), (U1, M1)), (B, MB) in ((U, M), (U1, M1))
@test op(A, B) == op(MA, MB)
end
end
@test imag(U) == zero(U)
end
end

# Matrix square root
Atn = UpperTriangular([-1 1 2; 0 -2 2; 0 0 -3])
Atp = UpperTriangular([1 1 2; 0 2 2; 0 0 3])
Expand Down Expand Up @@ -894,6 +911,11 @@ end
U = UT(F)
@test -U == -Array(U)
end

F = FillArrays.Fill(3im, (4,4))
for U in (UnitUpperTriangular(F), UnitLowerTriangular(F))
@test imag(F) == imag(collect(F))
end
end

@testset "error paths" begin
Expand All @@ -911,4 +933,36 @@ end
end
end

@testset "arithmetic with partly uninitialized matrices" begin
@testset "$(typeof(A))" for A in (Matrix{BigFloat}(undef,2,2), Matrix{Complex{BigFloat}}(undef,2,2)')
A[1,1] = A[2,2] = A[2,1] = 4
B = Matrix{eltype(A)}(undef, size(A))
for MT in (LowerTriangular, UnitLowerTriangular)
L = MT(A)
B .= 0
copyto!(B, L)
@test L * 2 == 2 * L == 2B
@test L/2 == B/2
@test 2\L == 2\B
@test real(L) == real(B)
@test imag(L) == imag(B)
end
end

@testset "$(typeof(A))" for A in (Matrix{BigFloat}(undef,2,2), Matrix{Complex{BigFloat}}(undef,2,2)')
A[1,1] = A[2,2] = A[1,2] = 4
B = Matrix{eltype(A)}(undef, size(A))
for MT in (UpperTriangular, UnitUpperTriangular)
U = MT(A)
B .= 0
copyto!(B, U)
@test U * 2 == 2 * U == 2B
@test U/2 == B/2
@test 2\U == 2\B
@test real(U) == real(B)
@test imag(U) == imag(B)
end
end
end

end # module TestTriangular

0 comments on commit 89cae45

Please sign in to comment.