Merge pull request #173 from JuliaParallel/vc/bcview

Introduce makelocal

Author: vchuravy

src/broadcast.jl

@@ -126,20 +126,14 @@ 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)
+# bclocal will do a view of the data and the copy it over
+# except when the data already is local
 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...))
+    makelocal(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/darray.jl

@@ -334,6 +334,36 @@ end
 
 localpart(d::DArray, localidx...) = localpart(d)[localidx...]
 
+_localindex(i::Integer, offset) = i - offset
+_localindex(i::AbstractRange, offset) = (first(i)-offset):step(i):(last(i)-offset)
+_localindex(i::AbstractUnitRange, offset) = (first(i)-offset):(last(i)-offset)
+
+"""
+    makelocal(A::DArray, I...)
+
+Equivalent to `Array(view(A, I...))` but optimised for the case that the data is local.
+Can return a view into `localpart(A)`
+"""
+function makelocal(A::DArray{<:Any, <:Any, AT}, I::Vararg{Any, N}) where {N, AT}
+    Base.@_inline_meta
+    J = map(i->Base.unalias(A, i), to_indices(A, I))
+    J = map(j-> isa(j, Base.Slice) ? j.indices : j, J)
+    @boundscheck checkbounds(A, J...)
+
+    lidcs = localindices(A)
+    if Base.checkbounds_indices(Bool, lidcs, J)
+        # data we want is local
+        viewidcs = ntuple(i -> _localindex(J[i], first(lidcs[i]) - 1), ndims(A))
+        view(localpart(A), viewidcs...)
+    else
+        # Make more efficient (?maybe) by allocating new memory
+        # only for the remote part
+        viewidcs = ntuple(i -> _localindex(J[i], 0), ndims(A))
+        arr = similar(AT, map(length, viewidcs)...)
+        copyto!(arr, view(A, viewidcs...))
+    end
+end
+
 # shortcut to set/get localparts of a distributed object
 function Base.getindex(d::DArray, s::Symbol)
     @assert s in [:L, :l, :LP, :lp]
@@ -539,6 +569,11 @@ function Array{S,N}(s::SubDArray{T,N}) where {S,T,N}
         end
     end
     a = Array{S}(undef, size(s))
+    copyto!(a, s)
+end
+
+function Base.copyto!(a::Array, s::SubDArray)
+    N = ndims(a)
     a[[1:size(a,i) for i=1:N]...] = s
     return a
 end
@@ -615,12 +650,11 @@ function Base.isassigned(D::DArray, i::Integer...)
     end
 end
 
-
 function Base.copyto!(dest::SubOrDArray, src::AbstractArray)
     asyncmap(procs(dest)) do p
         remotecall_fetch(p) do
             ldest = localpart(dest)
-            ldest[:] = Array(view(src, localindices(dest)...))
+            copyto!(ldest, view(src, localindices(dest)...))
         end
     end
     return dest

src/linalg.jl

@@ -38,15 +38,10 @@ function dot(x::DVector, y::DVector)
     if length(x) != length(y)
         throw(DimensionMismatch(""))
     end
-    if (procs(x) != procs(y)) || (x.cuts != y.cuts)
-        throw(ArgumentError("vectors don't have the same distribution. Not handled for efficiency reasons."))
-    end
 
     results=Any[]
-    @sync begin
-        for i = eachindex(x.pids)
-            @async push!(results, remotecall_fetch((x, y, i) -> dot(localpart(x), fetch(y, i)), x.pids[i], x, y, i))
-        end
+    asyncmap(procs(x)) do p
+            push!(results, remotecall_fetch((x, y) -> dot(localpart(x), makelocal(y, localindices(x)...)), x, y))
     end
     return reduce(+, results)
 end

src/mapreduce.jl

@@ -8,7 +8,7 @@ Base.map(f, d0::DArray, ds::AbstractArray...) = broadcast(f, d0, ds...)
 function Base.map!(f::F, dest::DArray, src::DArray{<:Any,<:Any,A}) where {F,A}
     asyncmap(procs(dest)) do p
         remotecall_fetch(p) do
-            map!(f, localpart(dest), A(view(src, localindices(dest)...)))
+            map!(f, localpart(dest), makelocal(src, localindices(dest)...))
             return nothing
         end
     end

test/darray.jl

@@ -712,6 +712,25 @@ end
 
 check_leaks()
 
+@testset "makelocal" begin
+    A = randn(5*nprocs(), 5*nprocs())
+    dA = distribute(A, procs=procs())
+    for i in 1:size(dA, 2)
+        a = DistributedArrays.makelocal(dA, :, i)
+        @test all(Array(view(dA, :, i)) .== a)
+        @test all(      view( A, :, i) .== a)
+    end
+    for i in 1:size(dA, 1)
+        a = DistributedArrays.makelocal(dA, i, :)
+        @test all(Array(view(dA, i:i, :)) .== a)
+        @test all(      view( A, i:i, :) .== a)
+    end
+    a = DistributedArrays.makelocal(dA, 1:5, 1:5)
+    @test all(Array(view(dA, 1:5, 1:5)) .== a)
+    @test all(      view( A, 1:5, 1:5) .== a)
+    close(dA)
+end
+
 @testset "test convert from subdarray" begin
     a = drand(20, 20);