DArray: Use eager API (#396)

Co-authored-by: Julian P Samaroo <[email protected]>

Author: fda-tome

src/array/alloc.jl

@@ -25,7 +25,7 @@ end
 
 function stage(ctx, a::AllocateArray)
     alloc(idx, sz) = a.f(idx, a.eltype, sz)
-    thunks = [delayed(alloc)(i, size(x)) for (i, x) in enumerate(a.domainchunks)]
+    thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(a.domainchunks)]
     DArray(a.eltype,a.domain, a.domainchunks, thunks)
 end
 

src/array/darray.jl

@@ -1,4 +1,4 @@
-import Base: ==
+import Base: ==, fetch
 using Serialization
 import Serialization: serialize, deserialize
 
@@ -78,13 +78,14 @@ domain(x::AbstractArray) = ArrayDomain([1:l for l in size(x)])
 abstract type ArrayOp{T, N} <: AbstractArray{T, N} end
 Base.IndexStyle(::Type{<:ArrayOp}) = IndexCartesian()
 
-compute(ctx, x::ArrayOp; options=nothing) =
-    compute(ctx, cached_stage(ctx, x)::DArray; options=options)
 
-collect(ctx::Context, x::ArrayOp; options=nothing) =
-    collect(ctx, compute(ctx, x; options=options); options=options)
+collect(x::ArrayOp) = collect(fetch(x))
 
-collect(x::ArrayOp; options=nothing) = collect(Context(global_context()), x; options=options)
+Base.fetch(x::ArrayOp) = fetch(cached_stage(Context(global_context()), x)::DArray)
+
+collect(x::Computation) = collect(fetch(x))
+
+Base.fetch(x::Computation) = fetch(cached_stage(Context(global_context()), x))
 
 function Base.show(io::IO, ::MIME"text/plain", x::ArrayOp)
     write(io, string(typeof(x)))
@@ -113,7 +114,7 @@ An N-dimensional distributed array of element type T, with a concatenation funct
 mutable struct DArray{T,N,F} <: ArrayOp{T, N}
     domain::ArrayDomain{N}
     subdomains::AbstractArray{ArrayDomain{N}, N}
-    chunks::AbstractArray{Union{Chunk,Thunk}, N}
+    chunks::AbstractArray{Any, N}
     concat::F
     function DArray{T,N,F}(domain, subdomains, chunks, concat::Function) where {T, N,F}
         new(domain, subdomains, chunks, concat)
@@ -135,18 +136,18 @@ domainchunks(d::DArray) = d.subdomains
 size(x::DArray) = size(domain(x))
 stage(ctx, c::DArray) = c
 
-function collect(ctx::Context, d::DArray; tree=false, options=nothing)
-    a = compute(ctx, d; options=options)
+function collect(d::DArray; tree=false)
+    a = fetch(d)
 
     if isempty(d.chunks)
         return Array{eltype(d)}(undef, size(d)...)
     end
 
     dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)]
     if tree
-        collect(treereduce_nd(delayed.(dimcatfuncs), a.chunks))
+        collect(fetch(treereduce_nd(map(x -> ((args...,) -> Dagger.@spawn x(args...)) , dimcatfuncs), a.chunks)))
     else
-        treereduce_nd(dimcatfuncs, asyncmap(collect, a.chunks))
+        treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks))
     end
 end
 
@@ -209,53 +210,33 @@ _cumsum(x::AbstractArray) = length(x) == 0 ? Int[] : cumsum(x)
 function lookup_parts(ps::AbstractArray, subdmns::DomainBlocks{N}, d::ArrayDomain{N}) where N
     groups = map(group_indices, subdmns.cumlength, indexes(d))
     sz = map(length, groups)
-    pieces = Array{Union{Chunk,Thunk}}(undef, sz)
+    pieces = Array{Any}(undef, sz)
     for i = CartesianIndices(sz)
         idx_and_dmn = map(getindex, groups, i.I)
         idx = map(x->x[1], idx_and_dmn)
         dmn = ArrayDomain(map(x->x[2], idx_and_dmn))
-        pieces[i] = delayed(getindex)(ps[idx...], project(subdmns[idx...], dmn))
+        pieces[i] = Dagger.@spawn getindex(ps[idx...], project(subdmns[idx...], dmn))
     end
     out_cumlength = map(g->_cumsum(map(x->length(x[2]), g)), groups)
     out_dmn = DomainBlocks(ntuple(x->1,Val(N)), out_cumlength)
     pieces, out_dmn
 end
 
-
 """
-    compute(ctx::Context, x::DArray; persist=true, options=nothing)
-
-A `DArray` object may contain a thunk in it, in which case
-we first turn it into a `Thunk` and then compute it.
-"""
-function compute(ctx::Context, x::DArray; persist=true, options=nothing)
-    thunk = thunkize(ctx, x, persist=persist)
-    if isa(thunk, Thunk)
-        compute(ctx, thunk; options=options)
-    else
-        x
-    end
-end
-
-"""
-    thunkize(ctx::Context, c::DArray; persist=true)
+    Base.fetch(c::DArray)
 
 If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk.
 """
-function thunkize(ctx::Context, c::DArray; persist=true)
+function Base.fetch(c::DArray)
     if any(istask, chunks(c))
         thunks = chunks(c)
         sz = size(thunks)
         dmn = domain(c)
         dmnchunks = domainchunks(c)
-        if persist
-            foreach(persist!, thunks)
-        end
-        Thunk(map(thunk->nothing=>thunk, thunks)...; meta=true) do results...
+        fetch(Dagger.spawn(Options(meta=true), thunks...) do results...
             t = eltype(results[1])
-            DArray(t, dmn, dmnchunks,
-                                  reshape(Union{Chunk,Thunk}[results...], sz))
-        end
+            DArray(t, dmn, dmnchunks, reshape(Any[results...], sz))
+        end)
     else
         c
     end
@@ -335,19 +316,18 @@ function stage(ctx::Context, d::Distribute)
         cs = map(d.domainchunks) do idx
             chunks = cached_stage(ctx, x[idx]).chunks
             shape = size(chunks)
-            (delayed() do shape, parts...
+            Dagger.spawn(shape, chunks...) do shape, parts...
                 if prod(shape) == 0
                     return Array{T}(undef, shape)
                 end
                 dimcatfuncs = [(x...) -> concat(x..., dims=i) for i in 1:length(shape)]
                 ps = reshape(Any[parts...], shape)
                 collect(treereduce_nd(dimcatfuncs, ps))
-            end)(shape, chunks...)
+            end
         end
     else
-        cs = map(c -> delayed(identity)(d.data[c]), d.domainchunks)
+        cs = map(c -> (Dagger.@spawn identity(d.data[c])), d.domainchunks)
     end
-
     DArray(
            eltype(d.data),
            domain(d.data),
@@ -357,7 +337,7 @@ function stage(ctx::Context, d::Distribute)
 end
 
 function distribute(x::AbstractArray, dist)
-    compute(Distribute(dist, x))
+    fetch(Distribute(dist, x))
 end
 
 function distribute(x::AbstractArray{T,N}, n::NTuple{N}) where {T,N}

src/array/getindex.jl

@@ -34,9 +34,9 @@ end
 function stage(ctx::Context, gidx::GetIndexScalar)
     inp = cached_stage(ctx, gidx.input)
     s = view(inp, ArrayDomain(gidx.idx))
-    delayed(identity)(collect(s)[1])
+    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::Integer...)  = compute(GetIndexScalar(c, idx))
+Base.getindex(c::ArrayOp, idx::Integer...)  = fetch(GetIndexScalar(c, idx))

src/array/map-reduce.jl

@@ -19,7 +19,7 @@ function stage(ctx::Context, node::Map)
     f = node.f
     for i=eachindex(domains)
         inps = map(x->chunks(x)[i], inputs)
-        thunks[i] = Thunk((args...) -> map(f, args...), map(inp->nothing=>inp, inps)...)
+        thunks[i] = Dagger.@spawn map(f, inps...) 
     end
     DArray(Any, domain(primary), domainchunks(primary), thunks)
 end
@@ -40,16 +40,16 @@ end
 
 function stage(ctx::Context, r::ReduceBlock)
     inp = stage(ctx, r.input)
-    reduced_parts = map(x -> Thunk(r.op, nothing=>x; get_result=r.get_result), chunks(inp))
-    Thunk((xs...) -> r.op_master(xs), map(part->nothing=>part, reduced_parts)...; meta=true)
+    reduced_parts = map(x -> (Dagger.@spawn get_result=r.get_result r.op(x)), chunks(inp))
+    r_op_master(args...,) = r.op_master(args)
+    Dagger.@spawn meta=true r_op_master(reduced_parts...) 
 end
 
 reduceblock_async(f, x::ArrayOp; get_result=true) = ReduceBlock(f, f, x, get_result)
 reduceblock_async(f, g::Function, x::ArrayOp; get_result=true) = ReduceBlock(f, g, x, get_result)
 
-reduceblock(f, x::ArrayOp) = compute(reduceblock_async(f, x))
-reduceblock(f, g::Function, x::ArrayOp) =
-    compute(reduceblock_async(f, g, x))
+reduceblock(f, x::ArrayOp) = fetch(reduceblock_async(f, x))
+reduceblock(f, g::Function, x::ArrayOp) = fetch(reduceblock_async(f, g, x))
 
 reduce_async(f::Function, x::ArrayOp) = reduceblock_async(xs->reduce(f,xs), xs->reduce(f,xs), x)
 
@@ -62,7 +62,7 @@ prod(f::Function, x::ArrayOp) = reduceblock(a->prod(f, a), prod, x)
 
 mean(x::ArrayOp) = reduceblock(mean, mean, x)
 
-mapreduce(f::Function, g::Function, x::ArrayOp) = reduce(g, map(f, x))
+mapreduce(f::Function, g::Function, x::ArrayOp) = reduce(g, map(f, x)) #think about fetching
 
 function mapreducebykey_seq(f, op,  itr, dict=Dict())
     for x in itr
@@ -115,7 +115,7 @@ end
 
 function reduce(f::Function, x::ArrayOp; dims = nothing)
     if dims === nothing
-        return compute(reduce_async(f,x))
+        return fetch(reduce_async(f,x))
     elseif dims isa Int
         dims = (dims,)
     end
@@ -126,10 +126,10 @@ function stage(ctx::Context, r::Reducedim)
     inp = cached_stage(ctx, r.input)
     thunks = let op = r.op, dims=r.dims
         # do reducedim on each block
-        tmp = map(p->Thunk(b->reduce(op,b,dims=dims), nothing=>p), chunks(inp))
+        tmp = map(p->Dagger.spawn(b->reduce(op,b,dims=dims), p), chunks(inp))
         # combine the results in tree fashion
         treereducedim(tmp, r.dims) do x,y
-            Thunk(op, nothing=>x, nothing=>y)
+            Dagger.@spawn op(x,y)
         end
     end
     c = domainchunks(inp)

src/array/matrix.jl

@@ -17,10 +17,10 @@ function size(x::Transpose)
 end
 
 transpose(x::ArrayOp) = Transpose(transpose, x)
-transpose(x::Union{Chunk, Thunk}) = Thunk(transpose, nothing=>x)
+transpose(x::Union{Chunk, EagerThunk}) = @spawn transpose(x)
 
 adjoint(x::ArrayOp) = Transpose(adjoint, x)
-adjoint(x::Union{Chunk, Thunk}) = Thunk(adjoint, nothing=>x)
+adjoint(x::Union{Chunk, EagerThunk}) = @spawn adjoint(x)
 
 function adjoint(x::ArrayDomain{2})
     d = indexes(x)
@@ -32,7 +32,7 @@ function adjoint(x::ArrayDomain{1})
 end
 
 function _ctranspose(x::AbstractArray)
-    Any[delayed(adjoint)(x[j,i]) for i=1:size(x,2), j=1:size(x,1)]
+    Any[Dagger.@spawn adjoint(x[j,i]) for i=1:size(x,2), j=1:size(x,1)]
 end
 
 function stage(ctx::Context, node::Transpose)
@@ -91,8 +91,12 @@ function (+)(a::ArrayDomain, b::ArrayDomain)
     a
 end
 
-(*)(a::Union{Chunk, Thunk}, b::Union{Chunk, Thunk}) = Thunk(*, nothing=>a, nothing=>b)
-(+)(a::Union{Chunk, Thunk}, b::Union{Chunk, Thunk}) = Thunk(+, nothing=>a, nothing=>b)
+struct BinaryComputeOp{F} end
+BinaryComputeOp{F}(x::Union{Chunk,EagerThunk}, y::Union{Chunk,EagerThunk}) where F = @spawn F(x, y)
+BinaryComputeOp{F}(x, y) where F = F(x, y)
+
+const AddComputeOp = BinaryComputeOp{+}
+const MulComputeOp = BinaryComputeOp{*}
 
 # we define our own matmat and matvec multiply
 # for computing the new domains and thunks.
@@ -101,7 +105,7 @@ function _mul(a::Matrix, b::Matrix; T=eltype(a))
     n = size(a, 2)
     for i=1:size(a,1)
         for j=1:size(b, 2)
-            c[i,j] = treereduce(+, map(*, reshape(a[i,:], (n,)), b[:, j]))
+            c[i,j] = treereduce(AddComputeOp, map(MulComputeOp, reshape(a[i,:], (n,)), b[:, j]))
         end
     end
     c
@@ -111,14 +115,14 @@ function _mul(a::Matrix, b::Vector; T=eltype(b))
     c = Array{T}(undef, size(a,1))
     n = size(a,2)
     for i=1:size(a,1)
-        c[i] = treereduce(+, map(*, reshape(a[i, :], (n,)), b))
+        c[i] = treereduce(AddComputeOp, map(MulComputeOp, reshape(a[i, :], (n,)), b))
     end
     c
 end
 
 function _mul(a::Vector, b::Vector; T=eltype(b))
     @assert length(b) == 1
-    [x * b[1] for x in a]
+    [MulComputeOp(x, b[1]) for x in a]
 end
 
 function promote_distribution(ctx::Context, m::MatMul, a,b)
@@ -176,7 +180,7 @@ function stage(ctx::Context, mul::MatMul)
     a, b = stage_operands(ctx, mul, mul.a, mul.b)
     d = domain(a)*domain(b)
     DArray(Any, d, domainchunks(a)*domainchunks(b),
-                          _mul(chunks(a), chunks(b); T=Thunk))
+                          _mul(chunks(a), chunks(b); T=Any))
 end
 
 Base.power_by_squaring(x::DArray, i::Int) = foldl(*, ntuple(idx->x, i))
@@ -211,7 +215,7 @@ end
 function _scale(l, r)
     res = similar(r, Any)
     for i=1:length(l)
-        res[i,:] = map(x->Thunk((a,b) -> Diagonal(a)*b, nothing=>l[i], nothing=>x), r[i,:])
+        res[i,:] = map(x->Dagger.spawn((a,b) -> Diagonal(a)*b, l[i], x), r[i,:])
     end
     res
 end

src/array/operators.jl

@@ -90,8 +90,7 @@ function stage(ctx::Context, node::BCast)
     end
     blcks = DomainBlocks(map(_->1, size(node)), cumlengths)
 
-    thunks = broadcast(delayed((args...)->broadcast(bc.f, args...); ),
-                       args2...)
+    thunks = broadcast((args3...)->Dagger.spawn((args...)->broadcast(bc.f, args...), args3...), args2...)
     DArray(eltype(node), domain(node), blcks, thunks)
 end
 
@@ -107,7 +106,7 @@ Base.@deprecate mappart(args...) mapchunk(args...)
 function stage(ctx::Context, node::MapChunk)
     inputs = map(x->cached_stage(ctx, x), node.input)
     thunks = map(map(chunks, inputs)...) do ps...
-        Thunk(node.f, map(p->nothing=>p, ps)...)
+        Dagger.spawn(node.f, map(p->nothing=>p, ps)...)
     end
 
     DArray(Any, domain(inputs[1]), domainchunks(inputs[1]), thunks)

src/array/setindex.jl

@@ -28,14 +28,14 @@ function stage(ctx::Context, sidx::SetIndex)
 
     groups = map(group_indices, subdmns.cumlength, indexes(d))
     sz = map(length, groups)
-    pieces = Array{Union{Chunk, Thunk}}(undef, sz)
+    pieces = Array{Any}(undef, sz)
     for i = CartesianIndices(sz)
         idx_and_dmn = map(getindex, groups, i.I)
         idx = map(x->x[1], idx_and_dmn)
         local_dmn = ArrayDomain(map(x->x[2], idx_and_dmn))
         s = subdmns[idx...]
         part_to_set = sidx.val
-        ps[idx...] = Thunk(nothing=>ps[idx...]) do p
+        ps[idx...] = Dagger.spawn(ps[idx...]) do p
             q = copy(p)
             q[indexes(project(s, local_dmn))...] .= part_to_set
             q