Merge pull request #469 from JuliaParallel/jps/darray-datadeps

DArray: Add Cholesky and view implementations

Author: jpsamaroo

Diff for: src/Dagger.jl

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

Diff for: src/array/alloc.jl

@@ -81,6 +81,15 @@ function Base.zero(x::DArray{T,N}) where {T,N}
     return _to_darray(a)
 end
 
+function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N}
+    d = ArrayDomain(Base.index_shape(A))
+    dc = partition(p, d)
+    # N.B. We use `tochunk` because we only want to take the view locally, and
+    # taking views should be very fast
+    chunks = [tochunk(view(A, x.indexes...)) for x in dc]
+    return DArray(T, d, dc, chunks, p)
+end
+
 function sprand(p::Blocks, m::Integer, n::Integer, sparsity::Real)
     s = rand(UInt)
     f = function (idx, t,sz)

Diff for: src/array/cholesky.jl

@@ -0,0 +1,91 @@
+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)
+
+    zone = one(T)
+    mzone = -one(T)
+    rzone = one(real(T))
+    rmzone = -one(real(T))
+    uplo = 'U'
+    Ac = A.chunks
+    mt, nt = size(Ac)
+    iscomplex = T <: Complex
+    trans = iscomplex ? 'C' : 'T'
+
+    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 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), 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'
+
+    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), info[1]
+end

Diff for: src/array/linalg.jl

@@ -0,0 +1,61 @@
+function LinearAlgebra.norm2(A::DArray{T,2}) where T
+    Ac = A.chunks
+    norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]
+    return sqrt(sum(map(fetch, norms)))
+end
+function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T
+    Ac = parent(A).chunks
+    Ac_upper = []
+    for i in 1:size(Ac, 1)
+        append!(Ac_upper, Ac[i, (i+1):end])
+    end
+    upper_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_upper]
+    Ac_diag = [Dagger.spawn(UpperTriangular, Ac[i,i]) for i in 1:size(Ac, 1)]
+    diag_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_diag]
+    return sqrt(sum(map(fetch, upper_norms)) + sum(map(fetch, diag_norms)))
+end
+function LinearAlgebra.norm2(A::LowerTriangular{T,<:DArray{T,2}}) where T
+    Ac = parent(A).chunks
+    Ac_lower = []
+    for i in 1:size(Ac, 1)
+        append!(Ac_lower, Ac[(i+1):end, i])
+    end
+    lower_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_lower]
+    Ac_diag = [Dagger.spawn(LowerTriangular, Ac[i,i]) for i in 1:size(Ac, 1)]
+    diag_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_diag]
+    return sqrt(sum(map(fetch, lower_norms)) + sum(map(fetch, diag_norms)))
+end
+
+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

Diff for: test/array.jl

@@ -38,6 +38,16 @@ end
     @test r[1:10] != r[11:20]
 end
 
+@testset "view" begin
+    A = rand(64, 64)
+    DA = view(A, Blocks(8, 8))
+    @test collect(DA) == A
+    @test size(DA) == (64, 64)
+    A_v = fetch(first(DA.chunks))
+    @test A_v isa SubArray
+    @test A_v == A[1:8, 1:8]
+end
+
 @testset "map" begin
     X1 = ones(Blocks(10, 10), 100, 100)
     X2 = map(x->x+1, X1)

Diff for: test/linalg.jl

@@ -0,0 +1,47 @@
+@testset "Linear Algebra" begin
+    @testset "Cholesky: $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+        D = rand(Blocks(4, 4), T, 32, 32)
+        if !(T <: Complex)
+            @test !issymmetric(D)
+        end
+        @test !ishermitian(D)
+
+        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)
+        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))
+
+        # 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

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")