Added check option to cholesky

Author: chriselrod

src/cholesky.jl

@@ -1,21 +1,33 @@
 # Generic Cholesky decomposition for fixed-size matrices, mostly unrolled
 non_hermitian_error() = throw(LinearAlgebra.PosDefException(-1))
-@inline function LinearAlgebra.cholesky(A::StaticMatrix)
+@inline function LinearAlgebra.cholesky(A::StaticMatrix; check::Bool = true)
     ishermitian(A) || non_hermitian_error()
-    C = _cholesky(Size(A), A)
-    return Cholesky(C, 'U', 0)
+    _cholesky(Size(A), A, check)
+    # (check && (info ≠ 0)) && throw(LinearAlgebra.PosDefException(info))
+    # return Cholesky(C, 'U', info)
 end
 
-@inline function LinearAlgebra.cholesky(A::LinearAlgebra.RealHermSymComplexHerm{<:Real, <:StaticMatrix})
-    C = _cholesky(Size(A), A.data)
-    return Cholesky(C, 'U', 0)
+@inline function LinearAlgebra.cholesky(A::LinearAlgebra.RealHermSymComplexHerm{<:Real, <:StaticMatrix}; check::Bool = true)
+    C = _cholesky(Size(A), A.data, check)
+    # (check && (info ≠ 0)) && throw(LinearAlgebra.PosDefException(info))
+    # return Cholesky(C, 'U', 0)
 end
 @inline LinearAlgebra._chol!(A::StaticMatrix, ::Type{UpperTriangular}) = (cholesky(A).U, 0)
 
-@generated function _cholesky(::Size{S}, A::StaticMatrix{M,M}) where {S,M}
+@inline function _check_chol(A, info, check)
+    if check
+        throw(LinearAlgebra.PosDefException(info))
+    else
+        return Cholesky(A, 'U', info)
+    end
+end
+@inline _nonpdcheck(x::Real) = x < zero(x)
+@inline _nonpdcheck(x) = false
+
+@generated function _cholesky(::Size{S}, A::StaticMatrix{M,M}, check::Bool) where {S,M}
     @assert (M,M) == S
     M > 24 && return :(_cholesky_large(Size{$S}(), A))
-    q = Expr(:block)
+    q = Expr(:block, :(info = 0), :(failure = false))
     for n ∈ 1:M
         for m ∈ n:M
             L_m_n = Symbol(:L_,m,:_,n)
@@ -28,7 +40,13 @@ end
             push!(q.args, :($L_m_n = muladd(-$L_m_k, $L_n_k', $L_m_n)))
         end
         L_n_n = Symbol(:L_,n,:_,n)
-        push!(q.args, :($L_n_n = sqrt($L_n_n)))
+        L_n_n_ltz = Symbol(:L_,n,:_,n,:_,:ltz)
+        # x < 0.0 is check used in `sqrt`, letting LLVM eliminate that check and remove error code.
+        # push!(q.args, :($L_n_n_ltz = )
+        push!(q.args, :($L_n_n = _nonpdcheck($L_n_n) ? (return _check_chol(A, $n, check)) : sqrt($L_n_n)))
+        # push!(q.args, :(info = ($L_n_n_ltz & (!failure)) ? $n : info))
+        # push!(q.args, :(failure |= $L_n_n_ltz))
+        # push!(q.args, :($L_n_n = $L_n_n_ltz ? float(typeof($L_n_n))(NaN) : sqrt($L_n_n)))
         Linv_n_n = Symbol(:Linv_,n,:_,n)
         push!(q.args, :($Linv_n_n = inv($L_n_n)))
         for m ∈ n+1:M
@@ -46,13 +64,13 @@ end
             push!(ret.args, :(zero(T)))
         end
     end
-    push!(q.args, :(similar_type(A, T)($ret)))
+    push!(q.args, :(Cholesky(similar_type(A, T)($ret), 'U', 0)))
     return Expr(:block, Expr(:meta, :inline), q)
 end
 
 # Otherwise default algorithm returning wrapped SizedArray
 @inline _cholesky_large(::Size{S}, A::StaticArray) where {S} =
-    similar_type(A)(cholesky(Hermitian(Matrix(A))).U)
+    Cholesky(similar_type(A)(cholesky(Hermitian(Matrix(A))).U), 'U', 0)
 
 LinearAlgebra.hermitian_type(::Type{SA}) where {T, S, SA<:SArray{S,T}} = Hermitian{T,SA}
 

test/chol.jl

@@ -34,6 +34,16 @@ using LinearAlgebra: PosDefException
             m_a = randn(elty, 4,4)
             #non hermitian
             @test_throws PosDefException cholesky(SMatrix{4,4}(m_a))
+            nonpd = SVector{4}(@view(m_a[:,1])) |> x -> x * x'
+            if elty <: Real
+                @test_throws PosDefException cholesky(nonpd)
+                @test !issuccess(cholesky(nonpd,check=false))
+                @test_throws PosDefException cholesky(Hermitian(nonpd))
+                @test !issuccess(cholesky(Symmetric(nonpd),check=false))
+            else
+                @test issuccess(cholesky(Hermitian(nonpd),check=false))
+            end
+            
             m_a = m_a*m_a'
             m = SMatrix{4,4}(m_a)
             @test cholesky(m).L ≈ cholesky(m_a).L