fixup! fixup! fixup! DArray: Add cholesky implementation

Author: jpsamaroo

Diff for: src/array/cholesky.jl

@@ -1,4 +1,12 @@
 LinearAlgebra.cholcopy(A::DArray{T,2}) where T = copy(A)
+function potrf_checked!(uplo, A, info_arr)
+    _A, info = LAPACK.potrf!(uplo, A)
+    if info > 0
+        info_arr[1] = info
+        throw(PosDefException(info))
+    end
+    return _A, info
+end
 function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T
     LinearAlgebra.checksquare(A)
 
@@ -12,50 +20,72 @@ function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T
     iscomplex = T <: Complex
     trans = iscomplex ? 'C' : 'T'
 
-    Dagger.spawn_datadeps() do
-        for k in range(1, mt)
-            Dagger.@spawn LAPACK.potrf!(uplo, InOut(Ac[k, k]))
-            for n in range(k+1, nt)
-                Dagger.@spawn BLAS.trsm!('L', uplo, trans, 'N', zone, In(Ac[k, k]), InOut(Ac[k, n]))
-            end
-            for m in range(k+1, mt)
-                if iscomplex
-                    Dagger.@spawn BLAS.herk!(uplo, 'C', rmzone, In(Ac[k, m]), rzone, InOut(Ac[m, m]))
-                else
-                    Dagger.@spawn BLAS.syrk!(uplo, 'T', rmzone, In(Ac[k, m]), rzone, InOut(Ac[m, m]))
+    info = [convert(LinearAlgebra.BlasInt, 0)]
+    try
+        Dagger.spawn_datadeps() do
+            for k in range(1, mt)
+                Dagger.@spawn potrf_checked!(uplo, InOut(Ac[k, k]), Out(info))
+                for n in range(k+1, nt)
+                    Dagger.@spawn BLAS.trsm!('L', uplo, trans, 'N', zone, In(Ac[k, k]), InOut(Ac[k, n]))
                 end
-                for n in range(m+1, nt)
-                    Dagger.@spawn BLAS.gemm!(trans, 'N', mzone, In(Ac[k, m]), In(Ac[k, n]), zone, InOut(Ac[m, n]))
+                for m in range(k+1, mt)
+                    if iscomplex
+                        Dagger.@spawn BLAS.herk!(uplo, 'C', rmzone, In(Ac[k, m]), rzone, InOut(Ac[m, m]))
+                    else
+                        Dagger.@spawn BLAS.syrk!(uplo, 'T', rmzone, In(Ac[k, m]), rzone, InOut(Ac[m, m]))
+                    end
+                    for n in range(m+1, nt)
+                        Dagger.@spawn BLAS.gemm!(trans, 'N', mzone, In(Ac[k, m]), In(Ac[k, n]), zone, InOut(Ac[m, n]))
+                    end
                 end
             end
         end
+    catch err
+        err isa ThunkFailedException || rethrow()
+        err = Dagger.Sch.unwrap_nested_exception(err.ex)
+        err isa PosDefException || rethrow()
     end
 
-    return UpperTriangular(A), convert(LinearAlgebra.BlasInt, 0)
+    return UpperTriangular(A), info[1]
 end
 function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{LowerTriangular}) where T
     LinearAlgebra.checksquare(A)
 
     zone = one(T)
     mzone = -one(T)
+    rzone = one(real(T))
+    rmzone = -one(real(T))
     uplo = 'L'
     Ac = A.chunks
     mt, nt = size(Ac)
+    iscomplex = T <: Complex
+    trans = iscomplex ? 'C' : 'T'
 
-    Dagger.spawn_datadeps() do
-        for k in range(1, mt)
-            Dagger.@spawn LAPACK.potrf!(uplo, InOut(Ac[k, k]))
-            for m in range(k+1, mt)
-                Dagger.@spawn BLAS.trsm!('R', uplo, 'T', 'N', zone, In(Ac[k, k]), InOut(Ac[m, k]))
-            end
-            for n in range(k+1, nt)
-                Dagger.@spawn BLAS.syrk!(uplo, 'N', mzone, In(Ac[n, k]), zone, InOut(Ac[n, n]))
-                for m in range(n+1, mt)
-                    Dagger.@spawn BLAS.gemm!('N', 'T', mzone, In(Ac[m, k]), In(Ac[n, k]), zone, InOut(Ac[m, n]))
+    info = [convert(LinearAlgebra.BlasInt, 0)]
+    try
+        Dagger.spawn_datadeps() do
+            for k in range(1, mt)
+                Dagger.@spawn potrf_checked!(uplo, InOut(Ac[k, k]), Out(info))
+                for m in range(k+1, mt)
+                    Dagger.@spawn BLAS.trsm!('R', uplo, trans, 'N', zone, In(Ac[k, k]), InOut(Ac[m, k]))
+                end
+                for n in range(k+1, nt)
+                    if iscomplex
+                        Dagger.@spawn BLAS.herk!(uplo, 'N', rmzone, In(Ac[n, k]), rzone, InOut(Ac[n, n]))
+                    else
+                        Dagger.@spawn BLAS.syrk!(uplo, 'N', rmzone, In(Ac[n, k]), rzone, InOut(Ac[n, n]))
+                    end
+                    for m in range(n+1, mt)
+                        Dagger.@spawn BLAS.gemm!('N', trans, mzone, In(Ac[m, k]), In(Ac[n, k]), zone, InOut(Ac[m, n]))
+                    end
                 end
             end
         end
+    catch err
+        err isa ThunkFailedException || rethrow()
+        err = Dagger.Sch.unwrap_nested_exception(err.ex)
+        err isa PosDefException || rethrow()
     end
 
-    return LowerTriangular(A), convert(LinearAlgebra.BlasInt, 0)
+    return LowerTriangular(A), info[1]
 end

Diff for: test/linalg.jl

@@ -8,6 +8,7 @@
 
         A = rand(T, 128, 128)
         A = A * A'
+        A[diagind(A)] .+= size(A, 1)
         DA = view(A, Blocks(32, 32))
         if !(T <: Complex)
             @test issymmetric(DA)
@@ -30,5 +31,17 @@
         @test chol_A.U ≈ chol_DA.U
         # Check that changes propagated to A
         @test UpperTriangular(collect(DA)) ≈ UpperTriangular(collect(A))
+
+        # Non-PosDef matrix
+        A = rand(T, 128, 128)
+        A = A * A'
+        A[diagind(A)] .+= size(A, 1)
+        A[1, 1] = -100
+        DA = view(A, Blocks(32, 32))
+        if !(T <: Complex)
+            @test issymmetric(DA)
+        end
+        @test ishermitian(DA)
+        @test_throws_unwrap PosDefException cholesky(DA).U
     end
 end