Skip to content

Commit e279303

Browse files
jpsamaroojpsamaroo
authored
Merge pull request #505 from JuliaParallel/jps/darray-show-repl
DArray: Add automatic partitioning and show array elements
2 parents 4cae113 + 5ad2868 commit e279303

File tree

9 files changed

+732
-250
lines changed

9 files changed

+732
-250
lines changed

‎docs/src/darray.md‎

Lines changed: 268 additions & 54 deletions
Large diffs are not rendered by default.

‎src/array/alloc.jl‎

Lines changed: 46 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -29,42 +29,71 @@ function stage(ctx, a::AllocateArray)
2929
return DArray(a.eltype, a.domain, a.domainchunks, thunks, a.partitioning)
3030
end
3131

32-
function Base.rand(p::Blocks, eltype::Type, dims)
32+
const BlocksOrAuto = Union{Blocks{N} where N, AutoBlocks}
33+
34+
function Base.rand(p::Blocks, eltype::Type, dims::Dims)
3335
d = ArrayDomain(map(x->1:x, dims))
3436
a = AllocateArray(eltype, (_, x...) -> rand(x...), d, partition(p, d), p)
3537
return _to_darray(a)
3638
end
39+
Base.rand(p::BlocksOrAuto, T::Type, dims::Integer...) = rand(p, T, dims)
40+
Base.rand(p::BlocksOrAuto, T::Type, dims::Dims) = rand(p, T, dims)
41+
Base.rand(p::BlocksOrAuto, dims::Integer...) = rand(p, Float64, dims)
42+
Base.rand(p::BlocksOrAuto, dims::Dims) = rand(p, Float64, dims)
43+
Base.rand(::AutoBlocks, eltype::Type, dims::Dims) =
44+
rand(auto_blocks(dims), eltype, dims)
3745

38-
Base.rand(p::Blocks, t::Type, dims::Integer...) = rand(p, t, dims)
39-
Base.rand(p::Blocks, dims::Integer...) = rand(p, Float64, dims)
40-
Base.rand(p::Blocks, dims::Tuple) = rand(p, Float64, dims)
46+
function Base.randn(p::Blocks, eltype::Type, dims::Dims)
47+
d = ArrayDomain(map(x->1:x, dims))
48+
a = AllocateArray(eltype, (_, x...) -> randn(x...), d, partition(p, d), p)
49+
return _to_darray(a)
50+
end
51+
Base.randn(p::BlocksOrAuto, T::Type, dims::Integer...) = randn(p, T, dims)
52+
Base.randn(p::BlocksOrAuto, T::Type, dims::Dims) = randn(p, T, dims)
53+
Base.randn(p::BlocksOrAuto, dims::Integer...) = randn(p, Float64, dims)
54+
Base.randn(p::BlocksOrAuto, dims::Dims) = randn(p, Float64, dims)
55+
Base.randn(::AutoBlocks, eltype::Type, dims::Dims) =
56+
randn(auto_blocks(dims), eltype, dims)
4157

42-
function Base.randn(p::Blocks, eltype::Type, dims)
58+
function sprand(p::Blocks, eltype::Type, dims::Dims, sparsity::AbstractFloat)
4359
d = ArrayDomain(map(x->1:x, dims))
44-
a = AllocateArray(Float64, (_, x...) -> randn(x...), d, partition(p, d), p)
60+
a = AllocateArray(eltype, (_, T, _dims) -> sprand(T, _dims..., sparsity), d, partition(p, d), p)
4561
return _to_darray(a)
4662
end
47-
Base.randn(p::Blocks, t::Type, dims::Integer...) = randn(p, t, dims)
48-
Base.randn(p::Blocks, dims::Integer...) = randn(p, dims)
49-
Base.randn(p::Blocks, dims::Tuple) = randn(p, Float64, dims)
63+
sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...) =
64+
sprand(p, T, dims_and_sparsity[1:end-1], dims_and_sparsity[end])
65+
sprand(p::BlocksOrAuto, T::Type, dims::Dims, sparsity::AbstractFloat) =
66+
sprand(p, T, dims, sparsity)
67+
sprand(p::BlocksOrAuto, dims_and_sparsity::Real...) =
68+
sprand(p, Float64, dims_and_sparsity[1:end-1], dims_and_sparsity[end])
69+
sprand(p::BlocksOrAuto, dims::Dims, sparsity::AbstractFloat) =
70+
sprand(p, Float64, dims, sparsity)
71+
sprand(::AutoBlocks, eltype::Type, dims::Dims, sparsity::AbstractFloat) =
72+
sprand(auto_blocks(dims), eltype, dims, sparsity)
5073

51-
function Base.ones(p::Blocks, eltype::Type, dims)
74+
function Base.ones(p::Blocks, eltype::Type, dims::Dims)
5275
d = ArrayDomain(map(x->1:x, dims))
5376
a = AllocateArray(eltype, (_, x...) -> ones(x...), d, partition(p, d), p)
5477
return _to_darray(a)
5578
end
56-
Base.ones(p::Blocks, t::Type, dims::Integer...) = ones(p, t, dims)
57-
Base.ones(p::Blocks, dims::Integer...) = ones(p, Float64, dims)
58-
Base.ones(p::Blocks, dims::Tuple) = ones(p, Float64, dims)
79+
Base.ones(p::BlocksOrAuto, T::Type, dims::Integer...) = ones(p, T, dims)
80+
Base.ones(p::BlocksOrAuto, T::Type, dims::Dims) = ones(p, T, dims)
81+
Base.ones(p::BlocksOrAuto, dims::Integer...) = ones(p, Float64, dims)
82+
Base.ones(p::BlocksOrAuto, dims::Dims) = ones(p, Float64, dims)
83+
Base.ones(::AutoBlocks, eltype::Type, dims::Dims) =
84+
ones(auto_blocks(dims), eltype, dims)
5985

60-
function Base.zeros(p::Blocks, eltype::Type, dims)
86+
function Base.zeros(p::Blocks, eltype::Type, dims::Dims)
6187
d = ArrayDomain(map(x->1:x, dims))
6288
a = AllocateArray(eltype, (_, x...) -> zeros(x...), d, partition(p, d), p)
6389
return _to_darray(a)
6490
end
65-
Base.zeros(p::Blocks, t::Type, dims::Integer...) = zeros(p, t, dims)
66-
Base.zeros(p::Blocks, dims::Integer...) = zeros(p, Float64, dims)
67-
Base.zeros(p::Blocks, dims::Tuple) = zeros(p, Float64, dims)
91+
Base.zeros(p::BlocksOrAuto, T::Type, dims::Integer...) = zeros(p, T, dims)
92+
Base.zeros(p::BlocksOrAuto, T::Type, dims::Dims) = zeros(p, T, dims)
93+
Base.zeros(p::BlocksOrAuto, dims::Integer...) = zeros(p, Float64, dims)
94+
Base.zeros(p::BlocksOrAuto, dims::Dims) = zeros(p, Float64, dims)
95+
Base.zeros(::AutoBlocks, eltype::Type, dims::Dims) =
96+
zeros(auto_blocks(dims), eltype, dims)
6897

6998
function Base.zero(x::DArray{T,N}) where {T,N}
7099
dims = ntuple(i->x.domain.indexes[i].stop, N)
@@ -82,22 +111,3 @@ function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N}
82111
chunks = [tochunk(view(A, x.indexes...)) for x in dc]
83112
return DArray(T, d, dc, chunks, p)
84113
end
85-
86-
function sprand(p::Blocks, m::Integer, n::Integer, sparsity::Real)
87-
s = rand(UInt)
88-
f = function (idx, t,sz)
89-
sprand(MersenneTwister(s+idx), sz...,sparsity)
90-
end
91-
d = ArrayDomain((1:m, 1:n))
92-
a = AllocateArray(Float64, f, d, partition(p, d), p)
93-
return _to_darray(a)
94-
end
95-
96-
function sprand(p::Blocks, n::Integer, sparsity::Real)
97-
s = rand(UInt)
98-
f = function (idx,t,sz)
99-
sprand(MersenneTwister(s+idx), sz...,sparsity)
100-
end
101-
a = AllocateArray(Float64, f, d, partition(p, ArrayDomain((1:n,))), p)
102-
return _to_darray(a)
103-
end

‎src/array/darray.jl‎

Lines changed: 129 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import Base: ==, fetch
22

3+
export DArray, DVector, DMatrix, Blocks, AutoBlocks
4+
export distribute
5+
6+
37
###### Array Domains ######
48

59
"""
@@ -90,18 +94,6 @@ collect(x::Computation) = collect(fetch(x))
9094

9195
Base.fetch(x::Computation) = fetch(stage(Context(global_context()), x))
9296

93-
function Base.show(io::IO, ::MIME"text/plain", x::ArrayOp)
94-
write(io, string(typeof(x)))
95-
write(io, string(size(x)))
96-
end
97-
98-
function Base.show(io::IO, x::ArrayOp)
99-
m = MIME"text/plain"()
100-
show(io, m, x)
101-
end
102-
103-
export BlockPartition, Blocks
104-
10597
abstract type AbstractBlocks{N} end
10698

10799
abstract type AbstractMultiBlocks{N}<:AbstractBlocks{N} end
@@ -147,11 +139,11 @@ mutable struct DArray{T,N,B<:AbstractBlocks{N},F} <: ArrayOp{T, N}
147139
end
148140
end
149141

150-
WrappedDArray{T,N} = Union{<:DArray{T,N}, Transpose{<:DArray{T,N}}, Adjoint{<:DArray{T,N}}}
151-
WrappedDMatrix{T} = WrappedDArray{T,2}
152-
WrappedDVector{T} = WrappedDArray{T,1}
153-
DMatrix{T} = DArray{T,2}
154-
DVector{T} = DArray{T,1}
142+
const WrappedDArray{T,N} = Union{<:DArray{T,N}, Transpose{<:DArray{T,N}}, Adjoint{<:DArray{T,N}}}
143+
const WrappedDMatrix{T} = WrappedDArray{T,2}
144+
const WrappedDVector{T} = WrappedDArray{T,1}
145+
const DMatrix{T} = DArray{T,2}
146+
const DVector{T} = DArray{T,1}
155147

156148

157149
# mainly for backwards-compatibility
@@ -194,6 +186,95 @@ function Base.collect(d::DArray; tree=false)
194186
end
195187
end
196188

189+
### show
190+
191+
#= FIXME
192+
@static if isdefined(Base, :AnnotatedString)
193+
# FIXME: Import StyledStrings
194+
struct ColorElement{T}
195+
color::Symbol
196+
value::T
197+
end
198+
function Base.show(io::IO, ::MIME"text/plain", x::ColorElement)
199+
print(io, styled"{(foreground=$(x.color)):$(x.value)}")
200+
end
201+
else
202+
=#
203+
struct ColorElement{T}
204+
color::Symbol
205+
value::Union{Some{T},Nothing}
206+
end
207+
function Base.show(io::IO, ::MIME"text/plain", x::ColorElement)
208+
if x.value !== nothing
209+
printstyled(io, something(x.value); color=x.color)
210+
else
211+
printstyled(io, "..."; color=x.color)
212+
end
213+
end
214+
Base.alignment(io::IO, x::ColorElement) =
215+
Base.alignment(io, something(x.value, "..."))
216+
#end
217+
struct ColorArray{T,N} <: DenseArray{T,N}
218+
A::DArray{T,N}
219+
color_map::Vector{Symbol}
220+
seen_values::Dict{NTuple{N,Int},Union{Some{T},Nothing}}
221+
function ColorArray(A::DArray{T,N}) where {T,N}
222+
colors = [:red, :green, :yellow, :blue, :magenta, :cyan]
223+
color_map = [colors[mod1(idx, length(colors))] for idx in 1:length(A.chunks)]
224+
return new{T,N}(A, color_map, Dict{NTuple{N,Int},Union{Some{T},Nothing}}())
225+
end
226+
end
227+
Base.size(A::ColorArray) = size(A.A)
228+
Base.getindex(A::ColorArray, idx::Integer) = getindex(A, (idx,))
229+
Base.getindex(A::ColorArray, idxs::Integer...) = getindex(A, (idxs...,))
230+
function Base.getindex(A::ColorArray{T,N}, idxs::NTuple{N,Int}) where {T,N}
231+
sd_idx_tuple, _ = partition_for(A.A, idxs)
232+
sd_idx = CartesianIndex(sd_idx_tuple)
233+
sd_idx_linear = LinearIndices(A.A.chunks)[sd_idx]
234+
if !haskey(A.seen_values, idxs)
235+
chunk = A.A.chunks[sd_idx]
236+
if chunk isa Chunk || isready(chunk)
237+
value = A.seen_values[idxs] = Some(getindex(A.A, idxs))
238+
else
239+
# Show a placeholder instead
240+
value = A.seen_values[idxs] = nothing
241+
end
242+
else
243+
value = A.seen_values[idxs]
244+
end
245+
if value !== nothing
246+
color = A.color_map[sd_idx_linear]
247+
else
248+
color = :light_black
249+
end
250+
return ColorElement{T}(color, value)
251+
end
252+
function Base.getindex(A::ColorArray{T,N}, idxs::Dims{S}) where {T,N,S}
253+
if S > N
254+
if all(idxs[(N+1):end] .== 1)
255+
return getindex(A, idxs[1:N])
256+
else
257+
throw(BoundsError(A, idxs))
258+
end
259+
elseif S < N
260+
throw(BoundsError(A, idxs))
261+
end
262+
end
263+
function Base.show(io::IO, ::MIME"text/plain", A::DArray{T,N}) where {T,N}
264+
write(io, string(DArray{T,N}))
265+
write(io, string(size(A)))
266+
write(io, " with $(join(size(A.chunks), 'x')) partitions of size $(join(A.partitioning.blocksize, 'x')):")
267+
pct_complete = 100 * (sum(c->c isa Chunk ? true : isready(c), A.chunks) / length(A.chunks))
268+
if pct_complete < 100
269+
println(io)
270+
printstyled(io, "~$(round(Int, pct_complete))% completed"; color=:yellow)
271+
end
272+
println(io)
273+
with_index_caching(1) do
274+
Base.print_array(IOContext(io, :compact=>true), ColorArray(A))
275+
end
276+
end
277+
197278
function (==)(x::ArrayOp, y::ArrayOp)
198279
x === y || reduce((a,b)->a&&b, map(==, x, y))
199280
end
@@ -317,8 +398,6 @@ end
317398
Base.@deprecate_binding Cat DArray
318399
Base.@deprecate_binding ComputedArray DArray
319400

320-
export Distribute, distribute
321-
322401
struct Distribute{T,N,B<:AbstractBlocks} <: ArrayOp{T, N}
323402
domainchunks
324403
partitioning::B
@@ -383,22 +462,44 @@ function stage(ctx::Context, d::Distribute)
383462
d.partitioning)
384463
end
385464

386-
function distribute(x::AbstractArray, dist::Blocks)
387-
_to_darray(Distribute(dist, x))
388-
end
465+
"""
466+
AutoBlocks
389467
468+
Automatically determines the size and number of blocks for a distributed array.
469+
This may construct any kind of `Dagger.AbstractBlocks` partitioning.
470+
"""
471+
struct AutoBlocks end
472+
function auto_blocks(dims::Dims{N}) where N
473+
# TODO: Allow other partitioning schemes
474+
np = num_processors()
475+
p = cld(dims[end], np)
476+
return Blocks(ntuple(i->i == N ? p : dims[i], N))
477+
end
478+
auto_blocks(A::AbstractArray{T,N}) where {T,N} = auto_blocks(size(A))
479+
480+
distribute(A::AbstractArray) = distribute(A, AutoBlocks())
481+
distribute(A::AbstractArray{T,N}, dist::Blocks{N}) where {T,N} =
482+
_to_darray(Distribute(dist, A))
483+
distribute(A::AbstractArray, ::AutoBlocks) = distribute(A, auto_blocks(A))
390484
function distribute(x::AbstractArray{T,N}, n::NTuple{N}) where {T,N}
391485
p = map((d, dn)->ceil(Int, d / dn), size(x), n)
392486
distribute(x, Blocks(p))
393487
end
488+
distribute(x::AbstractVector, n::Int) = distribute(x, (n,))
489+
distribute(x::AbstractVector, n::Vector{<:Integer}) =
490+
distribute(x, DomainBlocks((1,), (cumsum(n),)))
394491

395-
function distribute(x::AbstractVector, n::Int)
396-
distribute(x, (n,))
397-
end
492+
DVector(A::AbstractVector{T}, part::Blocks{1}) where T = distribute(A, part)
493+
DMatrix(A::AbstractMatrix{T}, part::Blocks{2}) where T = distribute(A, part)
494+
DArray(A::AbstractArray{T,N}, part::Blocks{N}) where {T,N} = distribute(A, part)
398495

399-
function distribute(x::AbstractVector, n::Vector{<:Integer})
400-
distribute(x, DomainBlocks((1,), (cumsum(n),)))
401-
end
496+
DVector(A::AbstractVector{T}) where T = DVector(A, AutoBlocks())
497+
DMatrix(A::AbstractMatrix{T}) where T = DMatrix(A, AutoBlocks())
498+
DArray(A::AbstractArray) = DArray(A, AutoBlocks())
499+
500+
DVector(A::AbstractVector{T}, ::AutoBlocks) where T = DVector(A, auto_blocks(A))
501+
DMatrix(A::AbstractMatrix{T}, ::AutoBlocks) where T = DMatrix(A, auto_blocks(A))
502+
DArray(A::AbstractArray, ::AutoBlocks) = DArray(A, auto_blocks(A))
402503

403504
function Base.:(==)(x::ArrayOp{T,N}, y::AbstractArray{S,N}) where {T,S,N}
404505
collect(x) == y

‎src/array/indexing.jl‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,18 @@ Base.getindex(A::DArray, idx::Integer) =
9595
getindex(A, Base._ind2sub(A, idx))
9696
Base.getindex(A::DArray, idx::CartesianIndex) =
9797
getindex(A, Tuple(idx))
98+
function Base.getindex(A::DArray{T,N}, idxs::Dims{S}) where {T,N,S}
99+
if S > N
100+
if all(idxs[(N+1):end] .== 1)
101+
return getindex(A, idxs[1:N])
102+
else
103+
throw(BoundsError(A, idxs))
104+
end
105+
elseif S < N
106+
throw(BoundsError(A, idxs))
107+
end
108+
error()
109+
end
98110

99111
### setindex!
100112

@@ -126,6 +138,18 @@ Base.setindex!(A::DArray, value, idx::Integer) =
126138
setindex!(A, value, Base._ind2sub(A, idx))
127139
Base.setindex!(A::DArray, value, idx::CartesianIndex) =
128140
setindex!(A, value, Tuple(idx))
141+
function Base.setindex!(A::DArray{T,N}, value, idxs::Dims{S}) where {T,N,S}
142+
if S > N
143+
if all(idxs[(N+1):end] .== 1)
144+
return setindex!(A, value, idxs[1:N])
145+
else
146+
throw(BoundsError(A, idxs))
147+
end
148+
elseif S < N
149+
throw(BoundsError(A, idxs))
150+
end
151+
error()
152+
end
129153

130154
### Allow/disallow scalar indexing
131155

0 commit comments

Comments
 (0)