From 0388dd661622ebf6c438921e6cc4f7980c1a494a Mon Sep 17 00:00:00 2001 From: Julian P Samaroo Date: Tue, 4 Jun 2024 12:37:08 -0700 Subject: [PATCH] DArray: Improve 0-dim alloc and show --- src/array/darray.jl | 28 ++++++++++++++++++++++----- test/array/allocation.jl | 42 ++++++++++++++++++++++++++++++---------- 2 files changed, 55 insertions(+), 15 deletions(-) diff --git a/src/array/darray.jl b/src/array/darray.jl index b0837fa99..d4343b8ee 100644 --- a/src/array/darray.jl +++ b/src/array/darray.jl @@ -20,6 +20,7 @@ ArrayDomain(xs::NTuple{N,Base.OneTo}) where N = ArrayDomain{N,NTuple{N,UnitRange{Int}}}(ntuple(i->UnitRange(xs[i]), N)) ArrayDomain(xs::NTuple{N,Int}) where N = ArrayDomain{N,NTuple{N,UnitRange{Int}}}(ntuple(i->xs[i]:xs[i], N)) +ArrayDomain(::Tuple{}) = ArrayDomain{0,Tuple{}}(()) ArrayDomain(xs...) = ArrayDomain((xs...,)) ArrayDomain(xs::Array) = ArrayDomain((xs...,)) @@ -31,6 +32,7 @@ chunks(a::ArrayDomain{N}) where {N} = DomainBlocks( (==)(a::ArrayDomain, b::ArrayDomain) = indexes(a) == indexes(b) Base.getindex(arr::AbstractArray, d::ArrayDomain) = arr[indexes(d)...] +Base.getindex(arr::AbstractArray{T,0} where T, d::ArrayDomain{0}) = arr function intersect(a::ArrayDomain, b::ArrayDomain) if a === b @@ -145,7 +147,6 @@ const WrappedDVector{T} = WrappedDArray{T,1} const DMatrix{T} = DArray{T,2} const DVector{T} = DArray{T,1} - # mainly for backwards-compatibility DArray{T, N}(domain, subdomains, chunks, partitioning, concat=cat) where {T,N} = DArray(T, domain, subdomains, chunks, partitioning, concat) @@ -178,6 +179,10 @@ function Base.collect(d::DArray; tree=false) return Array{eltype(d)}(undef, size(d)...) end + if ndims(d) == 0 + return fetch(a.chunks[1]) + end + dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)] if tree collect(fetch(treereduce_nd(map(x -> ((args...,) -> Dagger.@spawn x(args...)) , dimcatfuncs), a.chunks))) @@ -214,6 +219,7 @@ else Base.alignment(io::IO, x::ColorElement) = Base.alignment(io, something(x.value, "...")) #end +Base.show(io::IO, x::ColorElement) = show(io, MIME("text/plain"), x) struct ColorArray{T,N} <: DenseArray{T,N} A::DArray{T,N} color_map::Vector{Symbol} @@ -261,9 +267,21 @@ function Base.getindex(A::ColorArray{T,N}, idxs::Dims{S}) where {T,N,S} end end function Base.show(io::IO, ::MIME"text/plain", A::DArray{T,N}) where {T,N} - write(io, string(DArray{T,N})) - write(io, string(size(A))) - write(io, " with $(join(size(A.chunks), 'x')) partitions of size $(join(A.partitioning.blocksize, 'x')):") + if N == 1 + write(io, "$(length(A))-element ") + write(io, string(DVector{T})) + elseif N == 2 + write(io, string(DMatrix{T})) + elseif N == 0 + write(io, "0-dimensional ") + write(io, "DArray{$T, $N}") + else + write(io, "$(join(size(A), 'x')) ") + write(io, "DArray{$T, $N}") + end + nparts = N > 0 ? size(A.chunks) : 1 + partsize = N > 0 ? A.partitioning.blocksize : 1 + write(io, " with $(join(nparts, 'x')) partitions of size $(join(partsize, 'x')):") pct_complete = 100 * (sum(c->c isa Chunk ? true : isready(c), A.chunks) / length(A.chunks)) if pct_complete < 100 println(io) @@ -472,7 +490,7 @@ struct AutoBlocks end function auto_blocks(dims::Dims{N}) where N # TODO: Allow other partitioning schemes np = num_processors() - p = cld(dims[end], np) + p = N > 0 ? cld(dims[end], np) : 1 return Blocks(ntuple(i->i == N ? p : dims[i], N)) end auto_blocks(A::AbstractArray{T,N}) where {T,N} = auto_blocks(size(A)) diff --git a/test/array/allocation.jl b/test/array/allocation.jl index 61d4c73a5..a95ef2efc 100644 --- a/test/array/allocation.jl +++ b/test/array/allocation.jl @@ -1,9 +1,16 @@ @testset "DVector/DMatrix/DArray constructor" begin for T in [Float32, Float64, Int32, Int64] + F = fill(one(T)) V = rand(T, 64) M = rand(T, 64, 64) A = rand(T, 64, 64, 64) + # DArray ctor (empty) + DF = DArray(F, Blocks(())) + @test DF isa DArray{T,0} + @test collect(DF) == F + @test size(DF) == size(F) + # DVector ctor DV = DVector(V, Blocks(8)) @test DV isa DVector{T} @@ -26,7 +33,8 @@ end @testset "random" begin for T in [Float32, Float64, Int32, Int64] - for dims in [(100,), + for dims in [(), + (100,), (100, 100), (100, 100, 100)] dist = Blocks(ntuple(i->10, length(dims))...) @@ -53,10 +61,9 @@ end @test AXn isa Array{T,length(dims)} @test AXn == collect(Xn) @test AXn != collect(randn(dist, T, dims...)) - @test !all(AXn .> 0) end - if length(dims) <= 2 + if 1 <= length(dims) <= 2 # sprand Xsp = sprand(dist, T, dims..., 0.1) @test Xsp isa DArray{T,length(dims)} @@ -76,7 +83,8 @@ end @testset "ones/zeros" begin for T in [Float32, Float64, Int32, Int64] for (fn, value) in [(ones, one(T)), (zeros, zero(T))] - for dims in [(100,), + for dims in [(), + (100,), (100, 100), (100, 100, 100)] dist = Blocks(ntuple(i->10, length(dims))...) @@ -119,11 +127,16 @@ end for i in 1:(length(dims)-1) @test part_size[i] == 100 end - @test part_size[end] == cld(100, np) + if length(dims) > 0 + @test part_size[end] == cld(100, np) + else + @test part_size == () + end @test size(DA) == ntuple(i->100, length(dims)) end - for dims in [(100,), + for dims in [(), + (100,), (100, 100), (100, 100, 100)] fn = if length(dims) == 1 @@ -133,15 +146,23 @@ end else DArray end - DA = fn(rand(dims...), AutoBlocks()) + if length(dims) > 0 + DA = fn(rand(dims...), AutoBlocks()) + else + DA = fn(fill(rand()), AutoBlocks()) + end test_auto_blocks(DA, dims) - DA = distribute(rand(dims...), AutoBlocks()) + if length(dims) > 0 + DA = distribute(rand(dims...), AutoBlocks()) + else + DA = distribute(fill(rand()), AutoBlocks()) + end test_auto_blocks(DA, dims) for fn in [rand, randn, sprand, ones, zeros] if fn === sprand - if length(dims) > 2 + if length(dims) > 2 || length(dims) == 0 continue end DA = fn(AutoBlocks(), dims..., 0.1) @@ -155,7 +176,8 @@ end @testset "Constructor variants" begin for fn in [ones, zeros, rand, randn, sprand] - for dims in [(100,), + for dims in [(), + (100,), (100, 100), (100, 100, 100)] for dist in [Blocks(ntuple(i->10, length(dims))...),