Merge pull request #535 from JuliaParallel/jps/datadeps-gpu

DArray: Make allocations configurable via Processor
JuliaParallel · Jun 22, 2024 · 3ee4091 · 3ee4091
2 parents ee6b0ce + c822206
commit 3ee4091
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 20 deletions.
diff --git a/src/array/alloc.jl b/src/array/alloc.jl
@@ -4,7 +4,8 @@ export partition
 
 mutable struct AllocateArray{T,N} <: ArrayOp{T,N}
     eltype::Type{T}
-    f::Function
+    f
+    want_index::Bool
     domain::ArrayDomain{N}
     domainchunks
     partitioning::AbstractBlocks
@@ -23,17 +24,29 @@ function partition(p::AbstractBlocks, dom::ArrayDomain)
         map(_cumlength, map(length, indexes(dom)), p.blocksize))
 end
 
+function allocate_array(f, T, idx, sz)
+    new_f = allocate_array_func(thunk_processor(), f)
+    return new_f(idx, T, sz)
+end
+function allocate_array(f, T, sz)
+    new_f = allocate_array_func(thunk_processor(), f)
+    return new_f(T, sz)
+end
+allocate_array_func(::Processor, f) = f
 function stage(ctx, a::AllocateArray)
-    alloc(idx, sz) = a.f(idx, a.eltype, sz)
-    thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(a.domainchunks)]
+    if a.want_index
+        thunks = [Dagger.@spawn allocate_array(a.f, a.eltype, i, size(x)) for (i, x) in enumerate(a.domainchunks)]
+    else
+        thunks = [Dagger.@spawn allocate_array(a.f, a.eltype, size(x)) for (i, x) in enumerate(a.domainchunks)]
+    end
     return DArray(a.eltype, a.domain, a.domainchunks, thunks, a.partitioning)
 end
 
 const BlocksOrAuto = Union{Blocks{N} where N, AutoBlocks}
 
 function Base.rand(p::Blocks, eltype::Type, dims::Dims)
     d = ArrayDomain(map(x->1:x, dims))
-    a = AllocateArray(eltype, (_, x...) -> rand(x...), d, partition(p, d), p)
+    a = AllocateArray(eltype, rand, false, d, partition(p, d), p)
     return _to_darray(a)
 end
 Base.rand(p::BlocksOrAuto, T::Type, dims::Integer...) = rand(p, T, dims)
@@ -45,7 +58,7 @@ Base.rand(::AutoBlocks, eltype::Type, dims::Dims) =
 
 function Base.randn(p::Blocks, eltype::Type, dims::Dims)
     d = ArrayDomain(map(x->1:x, dims))
-    a = AllocateArray(eltype, (_, x...) -> randn(x...), d, partition(p, d), p)
+    a = AllocateArray(eltype, randn, false, d, partition(p, d), p)
     return _to_darray(a)
 end
 Base.randn(p::BlocksOrAuto, T::Type, dims::Integer...) = randn(p, T, dims)
@@ -57,7 +70,7 @@ Base.randn(::AutoBlocks, eltype::Type, dims::Dims) =
 
 function sprand(p::Blocks, eltype::Type, dims::Dims, sparsity::AbstractFloat)
     d = ArrayDomain(map(x->1:x, dims))
-    a = AllocateArray(eltype, (_, T, _dims) -> sprand(T, _dims..., sparsity), d, partition(p, d), p)
+    a = AllocateArray(eltype, (T, _dims) -> sprand(T, _dims..., sparsity), false, d, partition(p, d), p)
     return _to_darray(a)
 end
 sprand(p::BlocksOrAuto, T::Type, dims_and_sparsity::Real...) =
@@ -73,7 +86,7 @@ sprand(::AutoBlocks, eltype::Type, dims::Dims, sparsity::AbstractFloat) =
 
 function Base.ones(p::Blocks, eltype::Type, dims::Dims)
     d = ArrayDomain(map(x->1:x, dims))
-    a = AllocateArray(eltype, (_, x...) -> ones(x...), d, partition(p, d), p)
+    a = AllocateArray(eltype, ones, false, d, partition(p, d), p)
     return _to_darray(a)
 end
 Base.ones(p::BlocksOrAuto, T::Type, dims::Integer...) = ones(p, T, dims)
@@ -85,7 +98,7 @@ Base.ones(::AutoBlocks, eltype::Type, dims::Dims) =
 
 function Base.zeros(p::Blocks, eltype::Type, dims::Dims)
     d = ArrayDomain(map(x->1:x, dims))
-    a = AllocateArray(eltype, (_, x...) -> zeros(x...), d, partition(p, d), p)
+    a = AllocateArray(eltype, zeros, false, d, partition(p, d), p)
     return _to_darray(a)
 end
 Base.zeros(p::BlocksOrAuto, T::Type, dims::Integer...) = zeros(p, T, dims)

diff --git a/src/array/darray.jl b/src/array/darray.jl
@@ -306,16 +306,12 @@ function Base.isequal(x::ArrayOp, y::ArrayOp)
     x === y
 end
 
-function Base.similar(x::DArray{T,N}) where {T,N}
-    alloc(idx, sz) = Array{T,N}(undef, sz)
-    thunks = [Dagger.@spawn alloc(i, size(x)) for (i, x) in enumerate(x.subdomains)]
-    return DArray(T, x.domain, x.subdomains, thunks, x.partitioning, x.concat)
-end
-
+struct AllocateUndef{S} end
+(::AllocateUndef{S})(T, dims::Dims{N}) where {S,N} = Array{S,N}(undef, dims)
 function Base.similar(A::DArray{T,N} where T, ::Type{S}, dims::Dims{N}) where {S,N}
     d = ArrayDomain(map(x->1:x, dims))
     p = A.partitioning
-    a = AllocateArray(S, (_, _, x...) -> Array{S,N}(undef, x...), d, partition(p, d), p)
+    a = AllocateArray(S, AllocateUndef{S}(), false, d, partition(p, d), p)
     return _to_darray(a)
 end
 

diff --git a/src/array/mul.jl b/src/array/mul.jl
@@ -109,8 +109,8 @@ function gemm_dagger!(
     Bmt, Bnt = size(Bc)
     Cmt, Cnt = size(Cc)
 
-    alpha = _add.alpha
-    beta = _add.beta
+    alpha = T(_add.alpha)
+    beta = T(_add.beta)
 
     if Ant != Bmt
         throw(DimensionMismatch(lazy"A has number of blocks ($Amt,$Ant) but B has number of blocks ($Bmt,$Bnt)"))
@@ -212,8 +212,8 @@ function syrk_dagger!(
     Amt, Ant = size(Ac)
     Cmt, Cnt = size(Cc)
 
-    alpha = _add.alpha
-    beta = _add.beta
+    alpha = T(_add.alpha)
+    beta = T(_add.beta)
 
     uplo = 'U'
     if Ant != Cmt
@@ -233,7 +233,7 @@ function syrk_dagger!(
                         Dagger.@spawn BLAS.herk!(
                             uplo,
                             trans,
-                            alpha,
+                            real(alpha),
                             In(Ac[n, k]),
                             mzone,
                             InOut(Cc[n, n]),