Modernise distributed broadcast implementation

There are a number of challenges that distributed broadcast needs to
address.

1. The distribution of chunks is not guranteed to be homogeneous
2. The array type which the DArray wraps can be pretty much anthing
   and even customise Broadcasting. So instead of implementing it
   completly ourselves we need to defer to that local implementation
3. We want to minimise data copies, and only move chunks that are
   actually required.

This implementation first `distributes` all array-esque data to nodes
and then localises (through views) the data before defering to the
underlying broadcast implementation.

Author: vchuravy

src/DistributedArrays.jl

@@ -27,6 +27,7 @@ export close, d_closeall
 include("darray.jl")
 include("core.jl")
 include("serialize.jl")
+include("broadcast.jl")
 include("mapreduce.jl")
 include("linalg.jl")
 include("sort.jl")

src/broadcast.jl

@@ -0,0 +1,145 @@
+###
+# Distributed broadcast implementation
+##
+
+using Base.Broadcast
+import Base.Broadcast: BroadcastStyle, Broadcasted, _max
+
+# We define a custom ArrayStyle here since we need to keep track of
+# the fact that it is Distributed and what kind of underlying broadcast behaviour
+# we will encounter.
+struct DArrayStyle{Style <: BroadcastStyle} <: Broadcast.AbstractArrayStyle{Any} end
+DArrayStyle(::S) where {S} = DArrayStyle{S}()
+DArrayStyle(::S, ::Val{N}) where {S,N} = DArrayStyle(S(Val(N)))
+DArrayStyle(::Val{N}) where N = DArrayStyle{Broadcast.DefaultArrayStyle{N}}()
+
+BroadcastStyle(::Type{<:DArray{<:Any, N, A}}) where {N, A} = DArrayStyle(BroadcastStyle(A), Val(N))
+
+# promotion rules
+function BroadcastStyle(::DArrayStyle{AStyle}, ::DArrayStyle{BStyle}) where {AStyle, BStyle}
+    DArrayStyle{BroadcastStyle(AStyle, BStyle)}()
+end
+
+# # deal with one layer deep lazy arrays
+# BroadcastStyle(::Type{<:LinearAlgebra.Transpose{<:Any,T}}) where T <: DArray = BroadcastStyle(T)
+# BroadcastStyle(::Type{<:LinearAlgebra.Adjoint{<:Any,T}}) where T <: DArray = BroadcastStyle(T)
+# BroadcastStyle(::Type{<:SubArray{<:Any,<:Any,<:T}}) where T <: DArray = BroadcastStyle(T)
+
+# # This Union is a hack. Ideally Base would have a Transpose <: WrappedArray <: AbstractArray
+# # and we could define our methods in terms of Union{DArray, WrappedArray{<:Any, <:DArray}}
+# const DDestArray = Union{DArray,
+#                          LinearAlgebra.Transpose{<:Any,<:DArray},
+#                          LinearAlgebra.Adjoint{<:Any,<:DArray},
+#                          SubArray{<:Any, <:Any, <:DArray}}
+const DDestArray = DArray
+
+# This method is responsible for selection the output type of broadcast
+function Base.similar(bc::Broadcasted{<:DArrayStyle{Style}}, ::Type{ElType}) where {Style, ElType}
+    DArray(map(length, axes(bc))) do I 
+        # create fake Broadcasted for underlying ArrayStyle
+        bc′ = Broadcasted{Style}(identity, (), map(length, I))
+        similar(bc′, ElType)
+    end
+end
+
+##
+# We purposefully only specialise `copyto!`,
+# Broadcast implementation that defers to the underlying BroadcastStyle. We can't 
+# assume that `getindex` is fast, furthermore  we can't assume that the distribution of
+# DArray accross workers is equal or that the underlying array type is consistent.
+#
+# Implementation:
+#   - first distribute all arguments
+#     - Q: How do decide on the cuts
+#   - then localise arguments on each node
+##
+@inline function Base.copyto!(dest::DDestArray, bc::Broadcasted)
+    axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
+
+    # Distribute Broadcasted
+    # This will turn local AbstractArrays into DArrays
+    dbc = bcdistribute(bc)
+
+    asyncmap(procs(dest)) do p
+        remotecall_fetch(p) do
+            # get the indices for the localpart
+            lpidx = localpartindex(dest)
+            @assert lpidx != 0
+            # create a local version of the broadcast, by constructing views
+            # Note: creates copies of the argument
+            lbc = bclocal(dbc, dest.indices[lpidx])
+            Base.copyto!(localpart(dest), lbc)
+            return nothing
+        end
+    end
+    return dest
+end
+
+@inline function Base.copy(bc::Broadcasted{<:DArrayStyle})
+    dbc = bcdistribute(bc)
+    # TODO: teach DArray about axes since this is wrong for OffsetArrays
+    DArray(map(length, axes(bc))) do I
+        lbc = bclocal(dbc, I)
+        copy(lbc)
+    end
+end
+
+# _bcview creates takes the shapes of a view and the shape of a broadcasted argument,
+# and produces the view over that argument that constitutes part of the broadcast
+# it is in a sense the inverse of _bcs in Base.Broadcast
+_bcview(::Tuple{}, ::Tuple{}) = ()
+_bcview(::Tuple{}, view::Tuple) = ()
+_bcview(shape::Tuple, ::Tuple{}) = (shape[1], _bcview(tail(shape), ())...)
+function _bcview(shape::Tuple, view::Tuple)
+    return (_bcview1(shape[1], view[1]), _bcview(tail(shape), tail(view))...)
+end
+
+# _bcview1 handles the logic for a single dimension
+function _bcview1(a, b)
+    if a == 1 || a == 1:1
+        return 1:1
+    elseif first(a) <= first(b) <= last(a) &&
+           first(a) <= last(b)  <= last(b)
+        return b
+    else
+        throw(DimensionMismatch("broadcast view could not be constructed"))
+    end
+end
+
+# Distribute broadcast
+# TODO: How to decide on cuts
+@inline bcdistribute(bc::Broadcasted{Style}) where Style = Broadcasted{DArrayStyle{Style}}(bc.f, bcdistribute_args(bc.args), bc.axes)
+@inline bcdistribute(bc::Broadcasted{Style}) where Style<:DArrayStyle = Broadcasted{Style}(bc.f, bcdistribute_args(bc.args), bc.axes)
+
+# ask BroadcastStyle to decide if argument is in need of being distributed
+bcdistribute(x::T) where T = _bcdistribute(BroadcastStyle(T), x)
+_bcdistribute(::DArrayStyle, x) = x
+# Don't bother distributing singletons
+_bcdistribute(::Broadcast.AbstractArrayStyle{0}, x) = x
+_bcdistribute(::Broadcast.AbstractArrayStyle, x) = distribute(x)
+_bcdistribute(::Any, x) = x
+
+@inline bcdistribute_args(args::Tuple) = (bcdistribute(args[1]), bcdistribute_args(tail(args))...)
+bcdistribute_args(args::Tuple{Any}) = (bcdistribute(args[1]),)
+bcdistribute_args(args::Tuple{}) = ()
+
+# dropping axes here since recomputing is easier
+@inline bclocal(bc::Broadcasted{DArrayStyle{Style}}, idxs) where Style = Broadcasted{Style}(bc.f, bclocal_args(_bcview(axes(bc), idxs), bc.args))
+
+# bclocal will do a view of the data and the copy it over, except
+# when the shard match precisly (TODO: make sure that the invariant holds more often)
+function bclocal(x::DArray{T, N, AT}, idxs) where {T, N, AT}
+    bcidxs = _bcview(axes(x), idxs)
+    lpidx = localpartindex(x)
+    if lpidx != 0 && all(x.indices[lpidx] .== bcidxs)
+        return localpart(x)
+    end
+    return convert(__type(AT), view(x, bcidxs...))
+end
+bclocal(x, idxs) = x
+
+@inline bclocal_args(idxs, args::Tuple) = (bclocal(args[1], idxs), bclocal_args(idxs, tail(args))...)
+bclocal_args(idxs, args::Tuple{Any}) = (bclocal(args[1], idxs),)
+bclocal_args(idxs, args::Tuple{}) = ()
+
+__type(T) = Base.typename(T).wrapper

src/mapreduce.jl

@@ -15,30 +15,6 @@ function Base.map!(f::F, dest::DArray, src::DArray{<:Any,<:Any,A}) where {F,A}
     return dest
 end
 
-# new broadcasting implementation for julia-0.7
-
-Base.BroadcastStyle(::Type{<:DArray}) = Broadcast.ArrayStyle{DArray}()
-Base.BroadcastStyle(::Type{<:DArray}, ::Any) = Broadcast.ArrayStyle{DArray}()
-
-function Base.copy(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{DArray}})
-    T     = Base.Broadcast.combine_eltypes(bc.f, bc.args)
-    shape = Base.Broadcast.combine_axes(bc.args...)
-    iter  = Base.CartesianIndices(shape)
-    D     = DArray(map(length, shape)) do I
-        A = map(bc.args) do a
-            if isa(a, Union{Number,Ref})
-                return a
-            else
-                return localtype(a)(
-                    a[ntuple(i -> i > ndims(a) ? 1 : (size(a, i) == 1 ? (1:1) : I[i]), length(shape))...]
-                    )
-            end
-        end
-        broadcast(bc.f, A...)
-    end
-    return D
-end
-
 function Base.reduce(f, d::DArray)
     results = asyncmap(procs(d)) do p
         remotecall_fetch(p) do

test/darray.jl

@@ -846,6 +846,16 @@ check_leaks()
     f = 2 .* e
     @test Array(f) == 2 .* Array(e)
     @test Array(map(x -> sum(x) .+ 2, e)) == map(x -> sum(x) .+ 2, e)
+
+    @testset "test nested broadcast" begin
+       g = a .- m .* sin.(c)
+       @test Array(g) == Array(a) .- Array(m) .* sin.(Array(c))
+    end
+
+    # @testset "lazy wrapped broadcast" begin
+    #    l = similar(a)
+    #    l[1:10, :] .= view(a, 1:10, : )
+    # end
     d_closeall()
 end