WIP: MPI works

Author: Felipe Tome

Project.toml

@@ -12,8 +12,8 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 NextLA = "d37ed344-79c4-486d-9307-6d11355a15a3"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
@@ -77,8 +77,8 @@ GraphViz = "0.2"
 Graphs = "1"
 JSON3 = "1"
 KernelAbstractions = "0.9"
-MPI = "0.20.22"
 MacroTools = "0.5"
+MPI = "0.20.22"
 MemPool = "0.4.12"
 Metal = "1.1"
 NextLA = "0.2.2"

src/Dagger.jl

@@ -10,7 +10,7 @@ import MemPool: DRef, FileRef, poolget, poolset
 import Base: collect, reduce, view
 import NextLA
 import LinearAlgebra
-import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LU, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, diagind, ishermitian, issymmetric, I, norm, dot
+import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LU, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, Cholesky, diagind, ishermitian, issymmetric, I
 import Random
 import Random: AbstractRNG
 
@@ -53,7 +53,7 @@ import Adapt
 include("lib/util.jl")
 include("utils/dagdebug.jl")
 
-# Type definitions
+# Type definitions (for MPI/acceleration)
 include("types/processor.jl")
 include("types/scope.jl")
 include("types/memory-space.jl")
@@ -71,6 +71,7 @@ include("context.jl")
 include("utils/processors.jl")
 include("scopes.jl")
 include("utils/scopes.jl")
+include("chunks.jl")
 include("utils/signature.jl")
 include("thunkid.jl")
 include("utils/lfucache.jl")
@@ -82,11 +83,7 @@ include("argument.jl")
 include("queue.jl")
 include("thunk.jl")
 include("utils/fetch.jl")
-include("chunks.jl")
-include("affinity.jl")
-include("tochunk.jl")
-include("mutable.jl")
-include("shard.jl")
+include("utils/chunks.jl")
 include("weakchunk.jl")
 include("utils/logging.jl")
 include("submission.jl")
@@ -101,6 +98,7 @@ include("utils/clock.jl")
 include("utils/system_uuid.jl")
 include("utils/caching.jl")
 include("sch/Sch.jl"); using .Sch
+include("tochunk.jl")
 
 # Data dependency task queue
 include("datadeps/aliasing.jl")
@@ -138,7 +136,7 @@ include("array/mul.jl")
 include("array/cholesky.jl")
 include("array/trsm.jl")
 include("array/lu.jl")
-include("array/gmres.jl")
+include("array/qr.jl")
 
 # GPU
 include("gpu.jl")
@@ -167,8 +165,9 @@ function set_distributed_package!(value)
     @info "Dagger.jl preference has been set, restart your Julia session for this change to take effect!"
 end
 
-# MPI
+# MPI (mpi.jl loads MPI; mpi_mempool uses it)
 include("mpi.jl")
+include("mpi_mempool.jl")
 
 # Precompilation
 import PrecompileTools: @compile_workload

src/array/alloc.jl

@@ -184,13 +184,24 @@ function Base.zero(x::DArray{T,N}) where {T,N}
     return _to_darray(a)
 end
 
-@warn "Consider a better way to provide a unique ID for each chunk" maxlog=1
-function Base.view(A::AbstractArray{T,N}, p::Blocks{N}; space=default_memory_space(current_acceleration(), A)) where {T,N}
+# Weird LinearAlgebra dispatch in `\` needs this
+function LinearAlgebra._zeros(::Type{T}, B::DVector, n::Integer) where T
+    m = max(size(B, 1), n)
+    sz = (m,)
+    return zeros(auto_blocks(sz), T, sz)
+end
+function LinearAlgebra._zeros(::Type{T}, B::DMatrix, n::Integer) where T
+    m = max(size(B, 1), n)
+    sz = (m, size(B, 2))
+    return zeros(auto_blocks(sz), T, sz)
+end
+
+function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N}
     d = ArrayDomain(Base.index_shape(A))
     dc = partition(p, d)
     # N.B. We use `tochunk` because we only want to take the view locally, and
     # taking views should be very fast
-    chunks = [@with(MPI_UID => eager_next_id(), tochunk(view(A, x.indexes...), space)) for x in dc]
+    chunks = [tochunk(view(A, x.indexes...)) for x in dc]
     return DArray(T, d, dc, chunks, p)
 end
 Base.view(A::AbstractArray, ::AutoBlocks) =

src/array/copy.jl

@@ -119,14 +119,7 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(
                 Arange_local = Arange_global_clamped .- CartesianIndex(Arange_start) .+ CartesianIndex{Nmax}(1)
 
                 # Perform local view copy
-                space = (Bpart isa DTask ? fetch(Bpart; move_value=false, unwrap=false) : Bpart).space
-                procs = processors(space)
-                scope = UnionScope([ExactScope(proc) for proc in procs])
-                check_uniform(space)
-                for proc in procs
-                    check_uniform(proc)
-                end
-                Dagger.@spawn scope = scope copyto_view!(Out(Bpart), Brange_local, In(Apart), Arange_local)
+                Dagger.@spawn copyto_view!(Out(Bpart), Brange_local, In(Apart), Arange_local)
             end
         end
     end

src/array/darray.jl

@@ -1,4 +1,4 @@
-import Base: ==, fetch
+import Base: ==, fetch, length, isempty, size
 
 export DArray, DVector, DMatrix, DVecOrMat, Blocks, AutoBlocks
 export distribute
@@ -83,7 +83,8 @@ isempty(a::ArrayDomain) = length(a) == 0
 The domain of an array is an ArrayDomain.
 """
 domain(x::AbstractArray) = ArrayDomain([1:l for l in size(x)])
-
+# Scalar / non-array values (e.g. for Chunk of immediate data)
+domain(x::Any) = ArrayDomain(())
 
 abstract type ArrayOp{T, N} <: AbstractArray{T, N} end
 Base.IndexStyle(::Type{<:ArrayOp}) = IndexCartesian()
@@ -176,46 +177,28 @@ domainchunks(d::DArray) = d.subdomains
 size(x::DArray) = size(domain(x))
 stage(ctx, c::DArray) = c
 
-@warn "Dispatch uniform on acceleration" maxlog=1
-@warn "Take D.concat into account" maxlog=1
-function Base.collect(D::DArray{T,N}; tree=false, copyto=false, uniform::Bool=true) where {T,N}
-    if isempty(D.chunks)
-        return Array{eltype(D)}(undef, size(D)...)
+function Base.collect(d::DArray{T,N}; tree=false, copyto=false) where {T,N}
+    a = fetch(d)
+    if isempty(d.chunks)
+        return Array{eltype(d)}(undef, size(d)...)
     end
 
-    # Return a scalar, as required by Julia's array interface
-    if ndims(D) == 0
-        return fetch(D.chunks[1]; unwrap=true)
+    if ndims(d) == 0
+        return fetch(a.chunks[1])
     end
 
-    if uniform
-        @assert D.concat === cat "FIXME: Handle non-cat"
-        A = Array{eltype(D)}(undef, size(D)...)
-        DA = view(A, D.partitioning; space=CPURAMMemorySpace())
-
-        # Perform the equivalent of `copyto!(DA, D)`, but force local updates
-        # FIXME: Be more parallel?
-        for idx in eachindex(DA.chunks)
-            dest = fetch(DA.chunks[idx]; move_value=false, unwrap=true, uniform=true)::AbstractArray
-            src = fetch(D.chunks[idx]; move_value=true, unwrap=true, uniform=true)::AbstractArray
-            copyto!(dest, src)
-        end
+    if copyto
+        C = Array{T,N}(undef, size(a))
+        DC = view(C, Blocks(size(a)...))
+        copyto!(DC, a)
+        return C
+    end
 
-        return A
+    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)))
     else
-        if copyto
-            C = Array{T,N}(undef, size(D))
-            DC = view(C, Blocks(size(D)...))
-            copyto!(DC, D)
-            return C
-        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), D.chunks)))
-        else
-            treereduce_nd(dimcatfuncs, asyncmap(fetch, D.chunks))
-        end
+        collect(treereduce_nd(dimcatfuncs, asyncmap(fetch, a.chunks)))
     end
 end
 Array{T,N}(A::DArray{S,N}) where {T,N,S} = convert(Array{T,N}, collect(A))
@@ -339,8 +322,8 @@ function Base.isequal(x::ArrayOp, y::ArrayOp)
     x === y
 end
 
-Base.similar(::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} =
-    DArray{S,N}(undef, dims)
+Base.similar(D::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N} =
+    DArray{S,N}(undef, D.partitioning, dims)
 
 Base.copy(x::DArray{T,N,B,F}) where {T,N,B,F} =
     map(identity, x)::DArray{T,N,B,F}
@@ -406,18 +389,23 @@ function lookup_parts(A::DArray, ps::AbstractArray, subdmns::DomainBlocks{N}, d:
 end
 
 """
-    Base.fetch(A::DArray; unwrap::Bool=false, kwargs...) -> DArray
+    Base.fetch(c::DArray)
 
-Returns a new `DArray` with the same data as `A`, but where all values are
-fully computed.
+If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk.
 """
-function Base.fetch(A::DArray{T}; unwrap::Bool=false, kwargs...) where T
-    if any(unwrappable, chunks(A))
-        tasks = map(t->unwrappable(t) ? fetch(t; unwrap, kwargs...) : t, chunks(A))
-        B = DArray(T, A.domain, A.subdomains, tasks, A.partitioning, A.concat)
-        return B
+function Base.fetch(c::DArray{T}) where T
+    if any(istask, chunks(c))
+        thunks = chunks(c)
+        sz = size(thunks)
+        dmn = domain(c)
+        dmnchunks = domainchunks(c)
+        return fetch(Dagger.spawn(Options(meta=true), thunks...) do results...
+            t = eltype(fetch(results[1]))
+            DArray(t, dmn, dmnchunks, reshape(Any[results...], sz),
+                   c.partitioning, c.concat)
+        end)
     else
-        return A
+        return c
     end
 end
 
@@ -518,6 +506,7 @@ auto_blocks(A::AbstractArray{T,N}) where {T,N} = auto_blocks(size(A))
 
 const AssignmentType{N} = Union{Symbol, AbstractArray{<:Int, N}, AbstractArray{<:Processor, N}}
 
+distribute(A::AbstractArray, assignment::AssignmentType = :arbitrary) = distribute(A, AutoBlocks(), assignment)
 function distribute(A::AbstractArray{T,N}, dist::Blocks{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N}
     procgrid = nothing
     availprocs = collect(Dagger.compatible_processors())
@@ -558,10 +547,8 @@ function distribute(A::AbstractArray{T,N}, dist::Blocks{N}, assignment::Assignme
         procgrid = assignment
     end
 
-    return _distribute(current_acceleration(), A, dist, procgrid)
+    return _to_darray(Distribute(dist, A, procgrid))
 end
-_distribute(::DistributedAcceleration, A::AbstractArray{T,N}, dist::Blocks{N}, procgrid) where {T,N} =
-    _to_darray(Distribute(dist, A, procgrid))
 
 distribute(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = distribute(A, auto_blocks(A), assignment)
 function distribute(x::AbstractArray{T,N}, n::NTuple{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N}
@@ -570,6 +557,7 @@ function distribute(x::AbstractArray{T,N}, n::NTuple{N}, assignment::AssignmentT
 end
 distribute(x::AbstractVector, n::Int, assignment::AssignmentType{1} = :arbitrary) = distribute(x, (n,), assignment)
 
+
 DVector(A::AbstractVector{T}, part::Blocks{1}, assignment::AssignmentType{1} = :arbitrary) where T = distribute(A, part, assignment)
 DMatrix(A::AbstractMatrix{T}, part::Blocks{2}, assignment::AssignmentType{2} = :arbitrary) where T = distribute(A, part, assignment)
 DArray(A::AbstractArray{T,N}, part::Blocks{N}, assignment::AssignmentType{N} = :arbitrary) where {T,N} = distribute(A, part, assignment)
@@ -582,26 +570,29 @@ DVector(A::AbstractVector{T}, ::AutoBlocks, assignment::AssignmentType{1} = :arb
 DMatrix(A::AbstractMatrix{T}, ::AutoBlocks, assignment::AssignmentType{2} = :arbitrary) where T = DMatrix(A, auto_blocks(A), assignment)
 DArray(A::AbstractArray, ::AutoBlocks, assignment::AssignmentType = :arbitrary) = DArray(A, auto_blocks(A), assignment)
 
-@warn "Add assignment to undef initializer" maxlog=1
-function DArray{T,N}(::UndefInitializer, dims::NTuple{N,Int}) where {T,N}
-    dist = auto_blocks(dims)
-    return DArray{T,N}(undef, dist, dims...)
-end
-function DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}) where {T,N}
-    domain = ArrayDomain(ntuple(i->1:dims[i], N))
+struct AllocateUndef{S} end
+(::AllocateUndef{S})(T, dims::Dims{N}) where {S,N} = Array{S,N}(undef, dims)
+function DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N}
+    domain = ArrayDomain(map(x->1:x, dims))
     subdomains = partition(dist, domain)
-    tasks = Array{DTask,N}(undef, size(subdomains)...)
-    Dagger.spawn_datadeps() do
-        for (i, x) in enumerate(subdomains)
-            tasks[i] = Dagger.@spawn allocate_array_undef(T, size(x))
-        end
-    end
-    return DArray(T, domain, subdomains, tasks, dist)
-end
-DArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}) where {T,N} =
-    DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,))
-DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}) where {T,N} =
-    DArray{T,N}(undef, dist, (dims...,))
+    a = AllocateArray(T, AllocateUndef{T}(), false, domain, subdomains, dist, assignment)
+    return _to_darray(a)
+end
+DArray{T,N}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} =
+    DArray{T,N}(undef, dist, (dims...,); assignment)
+DArray{T,N}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N}  =
+    DArray{T,N}(undef, auto_blocks(dims), dims; assignment)
+DArray{T,N}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} =
+    DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment)
+
+DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N} =
+    DArray{T,N}(undef, dist, dims; assignment)
+DArray{T}(::UndefInitializer, dist::Blocks{N}, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} =
+    DArray{T,N}(undef, dist, (dims...,); assignment)
+DArray{T}(::UndefInitializer, dims::NTuple{N,Int}; assignment::AssignmentType{N} = :arbitrary) where {T,N}  =
+    DArray{T,N}(undef, auto_blocks(dims), dims; assignment)
+DArray{T}(::UndefInitializer, dims::Vararg{Int,N}; assignment::AssignmentType{N} = :arbitrary) where {T,N} =
+    DArray{T,N}(undef, auto_blocks((dims...,)), (dims...,); assignment)
 
 function Base.:(==)(x::ArrayOp{T,N}, y::AbstractArray{S,N}) where {T,S,N}
     collect(x) == y
@@ -622,7 +613,7 @@ end
 mapchunk(f, chunk) = tochunk(f(poolget(chunk.handle)))
 function mapchunks(f, d::DArray{T,N,F}) where {T,N,F}
     chunks = map(d.chunks) do chunk
-        owner = get_parent(chunk.processor).pid
+        owner = root_worker_id(chunk.processor)
         remotecall_fetch(mapchunk, owner, f, chunk)
     end
     DArray{T,N,F}(d.domain, d.subdomains, chunks, d.concat)