DArray: Add cholesky implementation

Co-authored-by: rabab53 <[email protected]>

Author: jpsamaroo

Diff for: src/Dagger.jl

@@ -60,6 +60,7 @@ include("array/setindex.jl")
 include("array/matrix.jl")
 include("array/sparse_partition.jl")
 include("array/sort.jl")
+include("array/cholesky.jl")
 
 # Other
 include("ui/graph-core.jl")

Diff for: src/array/cholesky.jl

@@ -0,0 +1,82 @@
+is_cross_symmetric(A1, A2) = A1 == A2'
+function LinearAlgebra.issymmetric(A::DArray{T,2}) where T
+    Ac = A.chunks
+    if size(Ac, 1) != size(Ac, 2)
+        return false
+    end
+
+    to_check = [Dagger.@spawn issymmetric(Ac[i, i]) for i in 1:size(Ac, 1)]
+    for i in 2:(size(Ac, 1)-1)
+        j_pre_diag = i - 1
+        for j in 1:j_pre_diag
+            push!(to_check, Dagger.@spawn is_cross_symmetric(Ac[i, j], Ac[j, i]))
+        end
+    end
+
+    return all(fetch, to_check)
+end
+function LinearAlgebra.ishermitian(A::DArray{T,2}) where T
+    Ac = A.chunks
+    if size(Ac, 1) != size(Ac, 2)
+        return false
+    end
+
+    to_check = [Dagger.@spawn ishermitian(Ac[i, i]) for i in 1:size(Ac, 1)]
+    for i in 2:(size(Ac, 1)-1)
+        j_pre_diag = i - 1
+        for j in 1:j_pre_diag
+            push!(to_check, Dagger.@spawn is_cross_symmetric(Ac[i, j], Ac[j, i]))
+        end
+    end
+
+    return all(fetch, to_check)
+end
+
+LinearAlgebra.cholcopy(A::DArray{T,2}) where T = copy(A)
+function LinearAlgebra.cholesky!(A::DArray{T,2}, ::NoPivot; check::Bool=true) where T
+    if check
+        ishermitian(A) || throw(PosDefException(-1))
+    end
+
+    zone = one(T)
+    mzone = -one(T)
+    uplo = 'U'
+    Ac = A.chunks
+    mt, nt = size(Ac)
+
+    if uplo == 'L'
+        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', 'L', '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]))
+                    end
+                end
+            end
+        end
+
+        return Cholesky(A, 'L', 0)
+    elseif uplo == 'U'
+        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, 'T', 'N', zone, In(Ac[k, k]), InOut(Ac[k, n]))
+                end
+                for m in range(k+1, mt)
+                    Dagger.@spawn BLAS.syrk!(uplo, 'T', mzone, In(Ac[k, m]), zone, InOut(Ac[m, m]))
+                    for n in range(m+1, nt)
+                        Dagger.@spawn BLAS.gemm!('T', 'N', mzone, In(Ac[k, m]), In(Ac[k, n]), zone, InOut(Ac[m, n]))
+                    end
+                end
+            end
+        end
+
+        return Cholesky(A, 'U', 0)
+    end
+end

Diff for: test/linalg.jl

@@ -0,0 +1,35 @@
+@testset "Linear Algebra" begin
+    @testset "Cholesky: $T" for T in (Float32, Float64, ComplexF64)
+        A = rand(T, 128, 128)
+        A = A * A'
+
+        if !(T <: AbstractComplex)
+            @test issymmetric(A)
+        end
+        @test !issymmetric(rand(Blocks(4, 4), T, 32, 32))
+        @test !ishermitian(rand(Blocks(4, 4), T, 32, 32))
+
+        DA = view(A, Blocks(32, 32))
+        if !(T <: AbstractComplex)
+            @test issymmetric(DA)
+        end
+        @test ishermitian(DA)
+
+        # Out-of-place
+        chol_A = cholesky(A)
+        chol_DA = cholesky(DA)
+        @test chol_DA isa Cholesky
+        @test chol_A.L ≈ chol_DA.L
+        @test chol_A.U ≈ chol_DA.U
+
+        # In-place
+        A_copy = copy(A)
+        chol_A = cholesky!(A_copy)
+        chol_DA = cholesky!(DA)
+        @test chol_DA isa Cholesky
+        @test chol_A.L ≈ chol_DA.L
+        @test chol_A.U ≈ chol_DA.U
+        # Check that changes propagated to A
+        @test UpperTriangular(collect(DA)) ≈ UpperTriangular(collect(A))
+    end
+end

Diff for: test/runtests.jl

@@ -38,6 +38,7 @@ include("task-queues.jl")
 include("datadeps.jl")
 include("domain.jl")
 include("array.jl")
+include("linalg.jl")
 include("cache.jl")
 include("diskcaching.jl")
 include("file-io.jl")