Revert to fc1ae09, keep Sch.jl changes

Author: Felipe Tome

Project.toml

@@ -13,7 +13,6 @@ Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
-MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 NextLA = "d37ed344-79c4-486d-9307-6d11355a15a3"
@@ -79,7 +78,6 @@ Graphs = "1"
 JSON3 = "1"
 KernelAbstractions = "0.9"
 MPI = "0.20.22"
-MPIPreferences = "0.1.11"
 MacroTools = "0.5"
 MemPool = "0.4.12"
 Metal = "1.1"

src/array/alloc.jl

@@ -70,11 +70,11 @@ function partition(p::AbstractBlocks, dom::ArrayDomain)
         map(_cumlength, map(length, indexes(dom)), p.blocksize))
 end
 
-function allocate_array(f, ::Type{T}, idx, sz::NTuple{N,Int})::Array{T,N} where {T,N}
+function allocate_array(f, T, idx, sz)
     new_f = allocate_array_func(task_processor(), f)
     return new_f(idx, T, sz)
 end
-function allocate_array(f, ::Type{T}, sz::NTuple{N,Int})::Array{T,N} where {T,N}
+function allocate_array(f, T, sz)
     new_f = allocate_array_func(task_processor(), f)
     return new_f(T, sz)
 end
@@ -189,15 +189,8 @@ function Base.view(A::AbstractArray{T,N}, p::Blocks{N}; space=default_memory_spa
     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.
-    # Per-chunk space for DArray: use each chunk's owner so tochunk on owner uses
-    # local_rank == space.rank and registers refs correctly (fixes MPI aliasing).
-    if A isa DArray && size(A.chunks) == size(dc)
-        chunks = [@with(MPI_UID => eager_next_id(), tochunk(view(A, x.indexes...),
-            (c = A.chunks[I]; c isa Chunk ? memory_space(c) : space))) for (I, x) in pairs(IndexCartesian(), dc)]
-    else
-        chunks = [@with(MPI_UID => eager_next_id(), tochunk(view(A, x.indexes...), space)) for x in dc]
-    end
+    # taking views should be very fast
+    chunks = [@with(MPI_UID => eager_next_id(), tochunk(view(A, x.indexes...), space)) for x in dc]
     return DArray(T, d, dc, chunks, p)
 end
 Base.view(A::AbstractArray, ::AutoBlocks) =

src/array/copy.jl

@@ -84,10 +84,6 @@ function darray_copyto!(B::DArray{TB,NB}, A::DArray{TA,NA}, Binds=parentindices(
 
     Dagger.spawn_datadeps() do
         for Bidx in Bci
-            accel = current_acceleration()
-            if accel isa MPIAcceleration
-                service_aliasing_requests(accel.comm)
-            end
             Bpart = B.chunks[Bidx]
             Bsd_global_raw = padNmax(Bsd_all[Bidx])
             Bsd_global_shifted = shift_ranges(Bsd_global_raw, Binds_offset)

src/array/darray.jl

@@ -582,29 +582,26 @@ 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)
 
-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))
+@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))
     subdomains = partition(dist, domain)
-    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)
+    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...,))
 
 function Base.:(==)(x::ArrayOp{T,N}, y::AbstractArray{S,N}) where {T,S,N}
     collect(x) == y