Define necessary methods to make tests pass without hitting scalar

indexing and some other cleanup.

src/DistributedArrays.jl

@@ -8,7 +8,7 @@ using LinearAlgebra
 
 import Base: +, -, *, div, mod, rem, &, |, xor
 import Base.Callable
-import LinearAlgebra: axpy!, dot, norm,
+import LinearAlgebra: axpy!, dot, norm
 
 import Primes
 import Primes: factor

src/core.jl

@@ -55,7 +55,8 @@ end
 
 Get the vector of processes storing pieces of DArray `d`.
 """
-Distributed.procs(d::DArray) = d.pids
+Distributed.procs(d::DArray) = d.pids
+Distributed.procs(d::SubDArray) = procs(parent(d))
 
 """
  localpart(A)

src/darray.jl

@@ -364,9 +364,41 @@ function localindices(d::DArray)
  return d.indices[lpidx]
 end
 
-# find which piece holds index (I...)
-locate(d::DArray, I::Int...) =
- ntuple(i -> searchsortedlast(d.cuts[i], I[i]), ndims(d))
+# Equality
+function Base.:(==)(d::DArray{<:Any,<:Any,A}, a::AbstractArray) where A
+ if size(d) != size(a)
+ return false
+ else
+ b = asyncmap(procs(d)) do p
+ remotecall_fetch(p) do
+ localpart(d) == A(a[localindices(d)...])
+ end
+ end
+ return all(b)
+ end
+end
+Base.:(==)(d::SubDArray, a::AbstractArray) = copy(d) == a
+Base.:(==)(a::AbstractArray, d::DArray) = d == a
+Base.:(==)(a::AbstractArray, d::SubDArray) = d == a
+Base.:(==)(d1::DArray, d2::DArray) = invoke(==, Tuple{DArray, AbstractArray}, d1, d2)
+Base.:(==)(d1::SubDArray, d2::DArray) = copy(d1) == d2
+Base.:(==)(d1::DArray, d2::SubDArray) = d1 == copy(d2)
+Base.:(==)(d1::SubDArray, d2::SubDArray) = copy(d1) == copy(d2)
+
+"""
+ locate(d::DArray, I::Int...)
+
+Determine the index of `procs(d)` that hold element `I`.
+"""
+function locate(d::DArray, I::Int...)
+ ntuple(ndims(d)) do i
+ fi = searchsortedlast(d.cuts[i], I[i])
+ if fi >= length(d.cuts[i])
+ throw(ArgumentError("element not contained in array"))
+ end
+ return fi
+ end
+end
 
 chunk(d::DArray{T,N,A}, i...) where {T,N,A} = remotecall_fetch(localpart, d.pids[i...], d)::A
 
@@ -479,15 +511,15 @@ end
 function (::Type{Array{S,N}})(s::SubDArray{T,N}) where {S,T,N}
  I = s.indices
  d = s.parent
- if isa(I,Tuple{Vararg{UnitRange{Int}}}) && S<:T && T<:S
+ if isa(I,Tuple{Vararg{UnitRange{Int}}}) && S<:T && T<:S && !isempty(s)
  l = locate(d, map(first, I)...)
  if isequal(d.indices[l...], I)
  # SubDArray corresponds to a chunk
  return chunk(d, l...)
  end
  end
  a = Array{S}(undef, size(s))
- a[[1:size(a,i) for i=1:N]...] .= s
+ a[[1:size(a,i) for i=1:N]...] = s
  return a
 end
 
@@ -540,15 +572,15 @@ end
 
 function Base.getindex(d::DArray, i::Int)
  _scalarindexingallowed()
- return getindex_tuple(d, CartesianIndices(d)[i])
+ return getindex_tuple(d, Tuple(CartesianIndices(d)[i]))
 end
 function Base.getindex(d::DArray, i::Int...)
  _scalarindexingallowed()
  return getindex_tuple(d, i)
 end
 
 Base.getindex(d::DArray) = d[1]
-Base.getindex(d::DArray, I::Union{Int,UnitRange{Int},Colon,Vector{Int},StepRange{Int,Int}}...) = view(d, I...)
+Base.getindex(d::SubOrDArray, I::Union{Int,UnitRange{Int},Colon,Vector{Int},StepRange{Int,Int}}...) = view(d, I...)
 
 function Base.isassigned(D::DArray, i::Integer...)
  try
@@ -564,15 +596,15 @@ function Base.isassigned(D::DArray, i::Integer...)
 end
 
 
-Base.copyto!(dest::SubOrDArray, src::SubOrDArray) = begin
+function Base.copyto!(dest::SubOrDArray, src::AbstractArray)
  asyncmap(procs(dest)) do p
  remotecall_fetch(p) do
- localpart(dest)[:] = src[localindices(dest)...]
+ ldest = localpart(dest)
+ ldest[:] = Array(view(src, localindices(dest)...))
  end
  end
  return dest
 end
-Base.copy!(dest::SubOrDArray, src::SubOrDArray) = copyto!(dest, src)
 
 function Base.deepcopy(src::DArray)
  dest = similar(src)

src/linalg.jl

@@ -1,12 +1,14 @@
-function Base.copy(D::Adjoint{T,<:DArray{T,2}}) where T
+function Base.copy(Dadj::Adjoint{T,<:DArray{T,2}}) where T
+ D = parent(Dadj)
  DArray(reverse(size(D)), procs(D)) do I
  lp = Array{T}(undef, map(length, I))
  rp = convert(Array, D[reverse(I)...])
  adjoint!(lp, rp)
  end
 end
 
-function Base.copy(D::Transpose{T,<:DArray{T,2}}) where T
+function Base.copy(Dtr::Transpose{T,<:DArray{T,2}}) where T
+ D = parent(Dtr)
  DArray(reverse(size(D)), procs(D)) do I
  lp = Array{T}(undef, map(length, I))
  rp = convert(Array, D[reverse(I)...])
@@ -49,7 +51,7 @@ function dot(x::DVector, y::DVector)
  return reduce(+, results)
 end
 
-function norm(x::DVector, p::Real = 2)
+function norm(x::DArray, p::Real = 2)
  results = []
  @sync begin
  for pp in procs(x)
@@ -83,7 +85,7 @@ function add!(dest, src, scale = one(dest[1]))
  return dest
 end
 
-function A_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVector)
+function mul!(y::DVector, A::DMatrix, x::AbstractVector, α::Number = 1, β::Number = 0)
 
  # error checks
  if size(A, 2) != length(x)
@@ -106,11 +108,14 @@ function A_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVec
 
  # Scale y if necessary
  if β != one(β)
- @sync for p in y.pids
- if β != zero(β)
- @async remotecall_fetch(y -> (rmul!(localpart(y), β); nothing), p, y)
- else
- @async remotecall_fetch(y -> (fill!(localpart(y), 0); nothing), p, y)
+ asyncmap(procs(y)) do p
+ remotecall_fetch(p) do
+ if !iszero(β)
+ rmul!(localpart(y), β)
+ else
+ fill!(localpart(y), 0)
+ end
+ return nothing
  end
  end
  end
@@ -127,7 +132,9 @@ function A_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVec
  return y
 end
 
-function Ac_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVector)
+function mul!(y::DVector, adjA::Adjoint{<:Number,<:DMatrix}, x::AbstractVector, α::Number = 1, β::Number = 0)
+
+ A = parent(adjA)
 
  # error checks
  if size(A, 1) != length(x)
@@ -148,11 +155,14 @@ function Ac_mul_B!(α::Number, A::DMatrix, x::AbstractVector, β::Number, y::DVe
 
  # Scale y if necessary
  if β != one(β)
- @sync for p in y.pids
- if β != zero(β)
- @async remotecall_fetch(() -> (rmul!(localpart(y), β); nothing), p)
- else
- @async remotecall_fetch(() -> (fill!(localpart(y), 0); nothing), p)
+ asyncmap(procs(y)) do p
+ remotecall_fetch(p) do
+ if !iszero(β)
+ rmul!(localpart(y), β)
+ else
+ fill!(localpart(y), 0)
+ end
+ return nothing
  end
  end
  end
@@ -189,7 +199,7 @@ function LinearAlgebra.rmul!(DA::DMatrix, D::Diagonal)
 end
 
 # Level 3
-function _matmatmul!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::DMatrix, tA)
+function _matmatmul!(C::DMatrix, A::DMatrix, B::AbstractMatrix, α::Number, β::Number, tA)
  # error checks
  Ad1, Ad2 = (tA == 'N') ? (1,2) : (2,1)
  mA, nA = (size(A, Ad1), size(A, Ad2))
@@ -254,40 +264,60 @@ function _matmatmul!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::D
  return C
 end
 
-A_mul_B!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::DMatrix) = _matmatmul!(α, A, B, β, C, 'N')
-Ac_mul_B!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::DMatrix) = _matmatmul!(α, A, B, β, C, 'C')
-At_mul_B!(α::Number, A::DMatrix, B::AbstractMatrix, β::Number, C::DMatrix) = _matmatmul!(α, A, B, β, C, 'T')
-At_mul_B!(C::DMatrix, A::DMatrix, B::AbstractMatrix) = At_mul_B!(one(eltype(C)), A, B, zero(eltype(C)), C)
+mul!(C::DMatrix, A::DMatrix, B::AbstractMatrix, α::Number = 1, β::Number = 0) = _matmatmul!(C, A, B, α, β, 'N')
+mul!(C::DMatrix, A::Adjoint{<:Number,<:DMatrix}, B::AbstractMatrix, α::Number = 1, β::Number = 0) = _matmatmul!(C, parent(A), B, α, β, 'C')
+mul!(C::DMatrix, A::Transpose{<:Number,<:DMatrix}, B::AbstractMatrix, α::Number = 1, β::Number = 0) = _matmatmul!(C, parent(A), B, α, β, 'T')
 
 _matmul_op = (t,s) -> t*s + t*s
 
 function Base.:*(A::DMatrix, x::AbstractVector)
  T = Base.promote_op(_matmul_op, eltype(A), eltype(x))
  y = DArray(I -> Array{T}(undef, map(length, I)), (size(A, 1),), procs(A)[:,1], (size(procs(A), 1),))
- return A_mul_B!(one(T), A, x, zero(T), y)
+ return mul!(y, A, x)
 end
 function Base.:*(A::DMatrix, B::AbstractMatrix)
  T = Base.promote_op(_matmul_op, eltype(A), eltype(B))
  C = DArray(I -> Array{T}(undef, map(length, I)),
  (size(A, 1), size(B, 2)),
  procs(A)[:,1:min(size(procs(A), 2), size(procs(B), 2))],
  (size(procs(A), 1), min(size(procs(A), 2), size(procs(B), 2))))
- return A_mul_B!(one(T), A, B, zero(T), C)
+ return mul!(C, A, B)
+end
+
+function Base.:*(adjA::Adjoint{<:Any,<:DMatrix}, x::AbstractVector)
+ A = parent(adjA)
+ T = Base.promote_op(_matmul_op, eltype(A), eltype(x))
+ y = DArray(I -> Array{T}(undef, map(length, I)),
+ (size(A, 2),),
+ procs(A)[1,:],
+ (size(procs(A), 2),))
+ return mul!(y, adjA, x)
+end
+function Base.:*(adjA::Adjoint{<:Any,<:DMatrix}, B::AbstractMatrix)
+ A = parent(adjA)
+ T = Base.promote_op(_matmul_op, eltype(A), eltype(B))
+ C = DArray(I -> Array{T}(undef, map(length, I)), (size(A, 2),
+ size(B, 2)),
+ procs(A)[1:min(size(procs(A), 1), size(procs(B), 2)),:],
+ (size(procs(A), 2), min(size(procs(A), 1), size(procs(B), 2))))
+ return mul!(C, adjA, B)
 end
 
-function Ac_mul_B(A::DMatrix, x::AbstractVector)
+function Base.:*(trA::Transpose{<:Any,<:DMatrix}, x::AbstractVector)
+ A = parent(trA)
  T = Base.promote_op(_matmul_op, eltype(A), eltype(x))
  y = DArray(I -> Array{T}(undef, map(length, I)),
  (size(A, 2),),
  procs(A)[1,:],
  (size(procs(A), 2),))
- return Ac_mul_B!(one(T), A, x, zero(T), y)
+ return mul!(y, trA, x)
 end
-function Ac_mul_B(A::DMatrix, B::AbstractMatrix)
+function Base.:*(trA::Transpose{<:Any,<:DMatrix}, B::AbstractMatrix)
+ A = parent(trA)
  T = Base.promote_op(_matmul_op, eltype(A), eltype(B))
  C = DArray(I -> Array{T}(undef, map(length, I)), (size(A, 2),
  size(B, 2)),
  procs(A)[1:min(size(procs(A), 1), size(procs(B), 2)),:],
  (size(procs(A), 2), min(size(procs(A), 1), size(procs(B), 2))))
- return Ac_mul_B!(one(T), A, B, zero(T), C)
+ return mul!(C, trA, B)
 end

src/mapreduce.jl

@@ -5,10 +5,10 @@ import SparseArrays: nnz
 
 Base.map(f, d0::DArray, ds::AbstractArray...) = broadcast(f, d0, ds...)
 
-function Base.map!(f::F, dest::DArray, src::DArray) where {F}
+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), src[localindices(dest)...])
+ map!(f, localpart(dest), A(view(src, localindices(dest)...)))
  return nothing
  end
  end
@@ -53,7 +53,7 @@ rewrite_local(x) = x
 
 function Base.reduce(f, d::DArray)
  results = asyncmap(procs(d)) do p
- remotecall_fetch(p, f, d) do (f, d)
+ remotecall_fetch(p) do
  return reduce(f, localpart(d))
  end
  end
@@ -122,12 +122,39 @@ function Base.mapreducedim!(f, op, R::DArray, A::DArray)
  end
  region = tuple(collect(1:ndims(A))[[size(R)...] .!= [size(A)...]]...)
  if isempty(region)
- return copy!(R, A)
+ return copyto!(R, A)
  end
  B = mapreducedim_within(f, op, A, region)
  return mapreducedim_between!(identity, op, R, B, region)
 end
 
+function Base._all(f, A::DArray, ::Colon)
+ B = asyncmap(procs(A)) do p
+ remotecall_fetch(p) do
+ all(f, localpart(A))
+ end
+ end
+ return all(B)
+end
+
+function Base._any(f, A::DArray, ::Colon)
+ B = asyncmap(procs(A)) do p
+ remotecall_fetch(p) do
+ any(f, localpart(A))
+ end
+ end
+ return any(B)
+end
+
+function Base.count(f, A::DArray)
+ B = asyncmap(procs(A)) do p
+ remotecall_fetch(p) do
+ count(f, localpart(A))
+ end
+ end
+ return sum(B)
+end
+
 function nnz(A::DArray)
  B = asyncmap(A.pids) do p
  remotecall_fetch(nnz∘localpart, p, A)

src/serialize.jl

@@ -2,7 +2,7 @@ function Serialization.serialize(S::AbstractSerializer, d::DArray{T,N,A}) where
  # Only send the ident for participating workers - we expect the DArray to exist in the
  # remote registry. DO NOT send the localpart.
  destpid = worker_id_from_socket(S.io)
- serialize_type(S, typeof(d))
+ Serialization.serialize_type(S, typeof(d))
  if (destpid in d.pids) || (destpid == d.id[1])
  serialize(S, (true, d.id)) # (id_only, id)
  else
@@ -64,7 +64,7 @@ function Serialization.serialize(S::AbstractSerializer, s::DestinationSerializer
  pid = worker_id_from_socket(S.io)
  pididx = findfirst(isequal(pid), s.pids)
  @assert pididx !== nothing
- serialize_type(S, typeof(s))
+ Serialization.serialize_type(S, typeof(s))
  serialize(S, s.generate(pididx))
 end