DArray: Operations return DArrays, local partition

To enable MPI support in the DArray (which is best implemented as each
MPI rank holding only a single local partition), this commit splits
AbstractBlock further by multi-partition or single-partition storage
schemes, where `Blocks <: AbstractMultiBlocks`, and a future `MPIBlocks
<: AbstractSingleBlocks`. Additionally, a DArray ctor is added for when
only a single subdomain and chunk/thunk is provided.

For easier post-hoc repartitioning of DArrays, we now store the original
user-provided partitioning scheme within the DArray, and also add it as
a type parameter. This also assists future MPI integration by allowing
for operations to dispatch on the partitioning scheme.

Finally, this commit also adjusts all DArray operations to return
DArrays, so that no lazy operators like MatMul or Map are returned to
the user. While it would be nice to be able to work with these
operators directly, the array ecosystem has generally settled on
propagating arrays of the same or similar types as the inputs to many
operations. The operators themselves are still present behind a single
materializing call, so it should be possible to expose them again in the
future if optimization opportunities become feasible.

Author: Felipe de Alcântara Tomé

src/array/alloc.jl

@@ -8,6 +8,7 @@ mutable struct AllocateArray{T,N} <: ArrayOp{T,N}
     f::Function
     domain::ArrayDomain{N}
     domainchunks
+    partitioning::AbstractBlocks
 end
 size(a::AllocateArray) = size(a.domain)
 
@@ -18,15 +19,15 @@ function _cumlength(len, step)
     cumsum(extra > 0 ? vcat(ps, extra) : ps)
 end
 
-function partition(p::Blocks, dom::ArrayDomain)
+function partition(p::AbstractBlocks, dom::ArrayDomain)
     DomainBlocks(map(first, indexes(dom)),
         map(_cumlength, map(length, indexes(dom)), p.blocksize))
 end
 
 function stage(ctx, a::AllocateArray)
     alloc(idx, sz) = a.f(idx, a.eltype, sz)
     thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(a.domainchunks)]
-    DArray(a.eltype,a.domain, a.domainchunks, thunks)
+    return DArray(a.eltype, a.domain, a.domainchunks, thunks, a.partitioning)
 end
 
 function Base.rand(p::Blocks, eltype::Type, dims)
@@ -35,7 +36,8 @@ function Base.rand(p::Blocks, eltype::Type, dims)
         rand(MersenneTwister(s+idx), x...)
     end
     d = ArrayDomain(map(x->1:x, dims))
-    AllocateArray(eltype, f, d, partition(p, d))
+    a = AllocateArray(eltype, f, d, partition(p, d), p)
+    return _to_darray(a)
 end
 
 Base.rand(p::Blocks, t::Type, dims::Integer...) = rand(p, t, dims)
@@ -48,21 +50,24 @@ function Base.randn(p::Blocks, dims)
         randn(MersenneTwister(s+idx), x...)
     end
     d = ArrayDomain(map(x->1:x, dims))
-    AllocateArray(Float64, f, d, partition(p, d))
+    a = AllocateArray(Float64, f, d, partition(p, d), p)
+    return _to_darray(a)
 end
 Base.randn(p::Blocks, dims::Integer...) = randn(p, dims)
 
 function Base.ones(p::Blocks, eltype::Type, dims)
     d = ArrayDomain(map(x->1:x, dims))
-    AllocateArray(eltype, (_, x...) -> ones(x...), d, partition(p, d))
+    a = AllocateArray(eltype, (_, x...) -> ones(x...), d, partition(p, d), p)
+    return _to_darray(a)
 end
 Base.ones(p::Blocks, t::Type, dims::Integer...) = ones(p, t, dims)
 Base.ones(p::Blocks, dims::Integer...) = ones(p, Float64, dims)
 Base.ones(p::Blocks, dims::Tuple) = ones(p, Float64, dims)
 
 function Base.zeros(p::Blocks, eltype::Type, dims)
     d = ArrayDomain(map(x->1:x, dims))
-    AllocateArray(eltype, (_, x...) -> zeros(x...), d, partition(p, d))
+    a = AllocateArray(eltype, (_, x...) -> zeros(x...), d, partition(p, d), p)
+    return _to_darray(a)
 end
 Base.zeros(p::Blocks, t::Type, dims::Integer...) = zeros(p, t, dims)
 Base.zeros(p::Blocks, dims::Integer...) = zeros(p, Float64, dims)
@@ -73,7 +78,7 @@ function Base.zero(x::DArray{T,N}) where {T,N}
     sd = first(x.subdomains)
     part_size = ntuple(i->sd.indexes[i].stop, N)
     a = zeros(Blocks(part_size...), T, dims)
-    return cached_stage(Context(global_context()), a)
+    return _to_darray(a)
 end
 
 function sprand(p::Blocks, m::Integer, n::Integer, sparsity::Real)
@@ -82,13 +87,15 @@ function sprand(p::Blocks, m::Integer, n::Integer, sparsity::Real)
         sprand(MersenneTwister(s+idx), sz...,sparsity)
     end
     d = ArrayDomain((1:m, 1:n))
-    AllocateArray(Float64, f, d, partition(p, d))
+    a = AllocateArray(Float64, f, d, partition(p, d), p)
+    return _to_darray(a)
 end
 
 function sprand(p::Blocks, n::Integer, sparsity::Real)
     s = rand(UInt)
     f = function (idx,t,sz)
         sprand(MersenneTwister(s+idx), sz...,sparsity)
     end
-    AllocateArray(Float64, f, d, partition(p, ArrayDomain((1:n,))))
+    a = AllocateArray(Float64, f, d, partition(p, ArrayDomain((1:n,))), p)
+    return _to_darray(a)
 end

src/array/darray.jl

@@ -81,7 +81,8 @@ Base.IndexStyle(::Type{<:ArrayOp}) = IndexCartesian()
 
 collect(x::ArrayOp) = collect(fetch(x))
 
-Base.fetch(x::ArrayOp) = fetch(cached_stage(Context(global_context()), x)::DArray)
+_to_darray(x::ArrayOp) = cached_stage(Context(global_context()), x)::DArray
+Base.fetch(x::ArrayOp) = fetch(_to_darray(x))
 
 collect(x::Computation) = collect(fetch(x))
 
@@ -97,6 +98,27 @@ function Base.show(io::IO, x::ArrayOp)
     show(io, m, x)
 end
 
+export BlockPartition, Blocks
+
+abstract type AbstractBlocks{N} end
+
+abstract type AbstractMultiBlocks{N}<:AbstractBlocks{N} end
+
+abstract type AbstractSingleBlocks{N}<:AbstractBlocks{N} end
+
+struct Blocks{N} <: AbstractMultiBlocks{N}
+    blocksize::NTuple{N, Int}
+end
+
+"""
+    Blocks(xs...)
+
+Indicates the size of an array operation, specified as `xs`, whose length
+indicates the number of dimensions in the resulting array.
+"""
+Blocks(xs::Int...) = Blocks(xs)
+
+
 """
     DArray{T,N,F}(domain, subdomains, chunks, concat)
     DArray(T, domain, subdomains, chunks, [concat=cat])
@@ -111,23 +133,35 @@ An N-dimensional distributed array of element type T, with a concatenation funct
 - `concat::F`: a function of type `F`. `concat(x, y; dims=d)` takes two chunks `x` and `y`
   and concatenates them along dimension `d`. `cat` is used by default.
 """
-mutable struct DArray{T,N,F} <: ArrayOp{T, N}
+mutable struct DArray{T,N,B<:AbstractBlocks{N},F} <: ArrayOp{T, N}
     domain::ArrayDomain{N}
     subdomains::AbstractArray{ArrayDomain{N}, N}
     chunks::AbstractArray{Any, N}
+    partitioning::B
     concat::F
-    function DArray{T,N,F}(domain, subdomains, chunks, concat::Function) where {T,N,F}
-        new{T,N,F}(domain, subdomains, chunks, concat)
+    function DArray{T,N,B,F}(domain, subdomains, chunks, partitioning::B, concat::Function) where {T,N,B,F}
+        new{T,N,B,F}(domain, subdomains, chunks, partitioning, concat)
     end
 end
 
 # mainly for backwards-compatibility
-DArray{T, N}(domain, subdomains, chunks) where {T,N} = DArray(T, domain, subdomains, chunks)
+DArray{T, N}(domain, subdomains, chunks, partitioning, concat=cat) where {T,N} =
+    DArray(T, domain, subdomains, chunks, partitioning, concat)
+
+function DArray(T, domain::ArrayDomain{N},
+                subdomains::AbstractArray{ArrayDomain{N}, N},
+                chunks::AbstractArray{<:Any, N}, partitioning::B, concat=cat) where {N,B<:AbstractMultiBlocks{N}}
+    DArray{T,N,B,typeof(concat)}(domain, subdomains, chunks, partitioning, concat)
+end
 
 function DArray(T, domain::ArrayDomain{N},
-             subdomains::AbstractArray{ArrayDomain{N}, N},
-             chunks::AbstractArray{<:Any, N}, concat=cat) where N
-    DArray{T, N, typeof(concat)}(domain, subdomains, chunks, concat)
+                subdomains::ArrayDomain{N},
+                chunks::Any, partitioning::B, concat=cat) where {N,B<:AbstractSingleBlocks{N}}
+    _subdomains = Array{ArrayDomain{N}, N}(undef, ntuple(i->1, N)...)
+    _subdomains[1] = subdomains
+    _chunks = Array{Any, N}(undef, ntuple(i->1, N)...)
+    _chunks[1] = chunks
+    DArray{T,N,B,typeof(concat)}(domain, _subdomains, _chunks, partitioning, concat)
 end
 
 domain(d::DArray) = d.domain
@@ -136,9 +170,8 @@ domainchunks(d::DArray) = d.subdomains
 size(x::DArray) = size(domain(x))
 stage(ctx, c::DArray) = c
 
-function collect(d::DArray; tree=false)
+function Base.collect(d::DArray; tree=false)
     a = fetch(d)
-
     if isempty(d.chunks)
         return Array{eltype(d)}(undef, size(d)...)
     end
@@ -163,18 +196,18 @@ function Base.isequal(x::ArrayOp, y::ArrayOp)
     x === y
 end
 
-function Base.similar(x::DArray{T,N,F}) where {T,N,F}
+function Base.similar(x::DArray{T,N}) where {T,N}
     alloc(idx, sz) = Array{T,N}(undef, sz)
     thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(x.subdomains)]
-    return DArray{T,N,F}(x.domain, x.subdomains, thunks, x.concat)
+    return DArray(T, x.domain, x.subdomains, thunks, x.partitioning, x.concat)
 end
 
-Base.copy(x::DArray{T,N,F}) where {T,N,F} =
-    cached_stage(Context(global_context()), map(identity, x))::DArray{T,N,F}
+Base.copy(x::DArray{T,N,B,F}) where {T,N,B,F} =
+    map(identity, x)::DArray{T,N,B,F}
 
 # Because OrdinaryDiffEq uses `Base.promote_op(/, ::DArray, ::Real)`
-Base.:(/)(x::DArray{T,N,F}, y::U) where {T<:Real,U<:Real,N,F} =
-    (x ./ y)::DArray{Base.promote_op(/, T, U),N,F}
+Base.:(/)(x::DArray{T,N,B,F}, y::U) where {T<:Real,U<:Real,N,B,F} =
+    (x ./ y)::DArray{Base.promote_op(/, T, U),N,B,F}
 
 """
     view(c::DArray, d)
@@ -184,7 +217,7 @@ A `view` of a `DArray` chunk returns a `DArray` of `Thunk`s.
 function Base.view(c::DArray, d)
     subchunks, subdomains = lookup_parts(chunks(c), domainchunks(c), d)
     d1 = alignfirst(d)
-    DArray(eltype(c), d1, subdomains, subchunks)
+    DArray(eltype(c), d1, subdomains, subchunks, c.partitioning, c.concat)
 end
 
 function group_indices(cumlength, idxs,at=1, acc=Any[])
@@ -246,12 +279,13 @@ function Base.fetch(c::DArray{T}) where T
         sz = size(thunks)
         dmn = domain(c)
         dmnchunks = domainchunks(c)
-        fetch(Dagger.spawn(Options(meta=true), thunks...) do results...
+        return fetch(Dagger.spawn(Options(meta=true), thunks...) do results...
             t = eltype(fetch(results[1]))
-            DArray(t, dmn, dmnchunks, reshape(Any[results...], sz))
+            DArray(t, dmn, dmnchunks, reshape(Any[results...], sz),
+                   c.partitioning, c.concat)
         end)
     else
-        c
+        return c
     end
 end
 
@@ -290,31 +324,31 @@ Base.@deprecate_binding ComputedArray DArray
 
 export Distribute, distribute
 
-struct Distribute{T, N} <: ArrayOp{T, N}
+struct Distribute{T,N,B<:AbstractBlocks} <: ArrayOp{T, N}
     domainchunks
+    partitioning::B
     data::AbstractArray{T,N}
 end
 
 size(x::Distribute) = size(domain(x.data))
 
-export BlockPartition, Blocks
-
-"""
-    Blocks(xs...)
-
-Indicates the size of an array operation, specified as `xs`, whose length
-indicates the number of dimensions in the resulting array.
-"""
-struct Blocks{N}
-    blocksize::NTuple{N, Int}
-end
-Blocks(xs::Int...) = Blocks(xs)
-
 Base.@deprecate BlockPartition Blocks
 
 
 Distribute(p::Blocks, data::AbstractArray) =
-    Distribute(partition(p, domain(data)), data)
+    Distribute(partition(p, domain(data)), p, data)
+
+function Distribute(domainchunks::DomainBlocks{N}, data::AbstractArray{T,N}) where {T,N}
+    p = Blocks(ntuple(i->first(domainchunks.cumlength[i]), N))
+    Distribute(domainchunks, p, data)
+end
+
+function Distribute(data::AbstractArray{T,N}) where {T,N}
+    nprocs = sum(w->length(Dagger.get_processors(OSProc(w))),
+                 Distributed.procs())
+    p = Blocks(ntuple(i->max(cld(size(data, i), nprocs), 1), N))
+    return Distribute(partition(p, domain(data)), p, data)
+end
 
 function stage(ctx::Context, d::Distribute)
     if isa(d.data, ArrayOp)
@@ -329,6 +363,8 @@ function stage(ctx::Context, d::Distribute)
         cs = map(d.domainchunks) do idx
             chunks = cached_stage(ctx, x[idx]).chunks
             shape = size(chunks)
+            # TODO: fix hashing
+            #hash = uhash(idx, Base.hash(Distribute, Base.hash(d.data)))
             Dagger.spawn(shape, chunks...) do shape, parts...
                 if prod(shape) == 0
                     return Array{T}(undef, shape)
@@ -339,18 +375,21 @@ function stage(ctx::Context, d::Distribute)
             end
         end
     else
-        cs = map(c -> (Dagger.@spawn identity(d.data[c])), d.domainchunks)
+        cs = map(d.domainchunks) do c
+            # TODO: fix hashing
+            #hash = uhash(c, Base.hash(Distribute, Base.hash(d.data)))
+            Dagger.@spawn identity(d.data[c])
+        end
     end
-    DArray(
-           eltype(d.data),
-           domain(d.data),
-           d.domainchunks,
-           cs
-    )
+    return DArray(eltype(d.data),
+                  domain(d.data),
+                  d.domainchunks,
+                  cs,
+                  d.partitioning)
 end
 
-function distribute(x::AbstractArray, dist)
-    fetch(Distribute(dist, x))
+function distribute(x::AbstractArray, dist::Blocks)
+    _to_darray(Distribute(dist, x))
 end
 
 function distribute(x::AbstractArray{T,N}, n::NTuple{N}) where {T,N}

src/array/getindex.jl

@@ -37,6 +37,6 @@ function stage(ctx::Context, gidx::GetIndexScalar)
     Dagger.@spawn identity(collect(s)[1])
 end
 
-Base.getindex(c::ArrayOp, idx::ArrayDomain) = GetIndex(c, indexes(idx))
-Base.getindex(c::ArrayOp, idx...)           = GetIndex(c, idx)
+Base.getindex(c::ArrayOp, idx::ArrayDomain) = _to_darray(GetIndex(c, indexes(idx)))
+Base.getindex(c::ArrayOp, idx...)           = _to_darray(GetIndex(c, idx))
 Base.getindex(c::ArrayOp, idx::Integer...)  = fetch(GetIndexScalar(c, idx))