Skip to content

Commit f1fd50b

Browse files
committed
Merge branch 'fastchol' of https://github.com/chriselrod/StaticArrays.jl into mbaran/matmul-symmetric
2 parents 5e999c7 + 87bb647 commit f1fd50b

File tree

1 file changed

+34
-36
lines changed

1 file changed

+34
-36
lines changed

src/cholesky.jl

Lines changed: 34 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -12,48 +12,46 @@ end
1212
end
1313
@inline LinearAlgebra._chol!(A::StaticMatrix, ::Type{UpperTriangular}) = (cholesky(A).U, 0)
1414

15-
16-
@generated function _cholesky(::Size{(1,1)}, A::StaticMatrix)
17-
@assert size(A) == (1,1)
18-
19-
quote
20-
$(Expr(:meta, :inline))
21-
T = promote_type(typeof(sqrt(one(eltype(A)))), Float32)
22-
similar_type(A,T)((sqrt(A[1]), ))
15+
@generated function _cholesky(::Size{S}, A::StaticMatrix{M,M}) where {S,M}
16+
@assert (M,M) == S
17+
M > 24 && return :(_cholesky_large(Size{$S}(), :A))
18+
q = Expr(:block)
19+
for n in 1:M
20+
for m n:M
21+
r = Expr(:ref, :A, n, m)
22+
ln = LineNumberNode(@__LINE__,Symbol(@__FILE__))
23+
r = Expr(:macrocall, Symbol("@inbounds"), ln, r)
24+
push!(q.args, Expr(:(=), Symbol(:L_,m,:_,n), r))
25+
end
26+
for k 1:n-1, m in n:M
27+
L_m_n = Symbol(:L_,m,:_,n)
28+
push!(q.args, Expr(:(=), L_m_n, Expr(:call, :muladd, Expr(:call, :(-), Symbol(:L_,m,:_,k)), Symbol(:L_,n,:_,k), L_m_n)))
29+
end
30+
L_n_n = Symbol(:L_,n,:_,n)
31+
push!(q.args, Expr(:(=), L_n_n, Expr(:call, :sqrt, L_n_n)))
32+
Linv_n_n = Symbol(:Linv_,n,:_,n)
33+
push!(q.args, Expr(:(=), Linv_n_n, Expr(:call, :inv, L_n_n)))
34+
for m n+1:M
35+
push!(q.args, Expr(:(*=), Symbol(:L_,m,:_,n), Linv_n_n))
36+
end
2337
end
24-
end
25-
26-
@generated function _cholesky(::Size{(2,2)}, A::StaticMatrix)
27-
@assert size(A) == (2,2)
28-
29-
quote
30-
$(Expr(:meta, :inline))
31-
@inbounds a = sqrt(A[1])
32-
@inbounds b = A[3] / a
33-
@inbounds c = sqrt(A[4] - b'*b)
34-
T = promote_type(typeof(sqrt(one(eltype(A)))), Float32)
35-
similar_type(A,T)((a, zero(T), b, c))
38+
push!(q.args, :(T = promote_type(typeof(sqrt(one(eltype(A)))), Float32)))
39+
ret = Expr(:tuple)
40+
for n 1:M
41+
for m 1:n
42+
push!(ret.args, Symbol(:L_,n,:_,m))
43+
end
44+
for m n+1:M
45+
push!(ret.args, Expr(:call, :zero, :T))
46+
end
3647
end
48+
push!(q.args, Expr(:call, Expr(:call, :similar_type, :A, :T), ret))
49+
Expr(:block, Expr(:meta, :inline), q)
3750
end
3851

39-
@generated function _cholesky(::Size{(3,3)}, A::StaticMatrix)
40-
@assert size(A) == (3,3)
41-
42-
quote
43-
$(Expr(:meta, :inline))
44-
@inbounds a11 = sqrt(A[1])
45-
@inbounds a12 = A[4] / a11
46-
@inbounds a22 = sqrt(A[5] - a12'*a12)
47-
@inbounds a13 = A[7] / a11
48-
@inbounds a23 = (A[8] - a12'*a13) / a22
49-
@inbounds a33 = sqrt(A[9] - a13'*a13 - a23'*a23)
50-
T = promote_type(typeof(sqrt(one(eltype(A)))), Float32)
51-
similar_type(A,T)((a11, zero(T), zero(T), a12, a22, zero(T), a13, a23, a33))
52-
end
53-
end
5452

5553
# Otherwise default algorithm returning wrapped SizedArray
56-
@inline _cholesky(::Size{S}, A::StaticArray) where {S} =
54+
@inline _cholesky_large(::Size{S}, A::StaticArray) where {S} =
5755
SizedArray{Tuple{S...}}(Matrix(cholesky(Hermitian(Matrix(A))).U))
5856

5957
LinearAlgebra.hermitian_type(::Type{SA}) where {T, S, SA<:SArray{S,T}} = Hermitian{T,SA}

0 commit comments

Comments
 (0)