Merge pull request #546 from pszufe/psz/daggerrand

Add rand! implementation

Author: jpsamaroo

Diff for: Project.toml

@@ -28,11 +28,13 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [weakdeps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
+DistributionsExt = "Distributions"
 GraphVizExt = "GraphViz"
 GraphVizSimpleExt = "Colors"
 JSON3Ext = "JSON3"
@@ -43,6 +45,7 @@ Adapt = "4.0.4"
 Colors = "0.12"
 DataFrames = "1"
 DataStructures = "0.18"
+Distributions = "0.25"
 GraphViz = "0.2"
 Graphs = "1"
 MacroTools = "0.5"
@@ -61,6 +64,7 @@ julia = "1.8"
 [extras]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

Diff for: docs/src/darray.md

@@ -431,6 +431,10 @@ From `Base`:
 - `map`/`reduce`/`mapreduce`
 - `sum`/`prod`
 - `minimum`/`maximum`/`extrema`
+- `map!`
+
+From `Random`:
+- `rand!`/`randn!`
 
 From `Statistics`:
 - `mean`

Diff for: ext/DistributionsExt.jl

@@ -0,0 +1,15 @@
+module DistributionsExt
+
+if isdefined(Base, :get_extension)
+    using Distributions
+else
+    using ..Distributions
+end
+
+using Dagger, Random
+
+Random.rand!(s::Sampleable, A::DArray{T}) where T = map!(_ -> rand(s), A)
+Random.rand!(rng::AbstractRNG, s::Sampleable{Univariate}, A::DArray{T}) where T = map!(_ -> rand(rng, s), A)
+Random.rand!(rng::AbstractRNG, s::Sampleable{ArrayLikeVariate{M}}, A::DArray{T}) where {M, T} = map!(_ -> rand(rng, s), A)
+
+end # module DistributionsExt

Diff for: src/Dagger.jl

@@ -13,6 +13,8 @@ import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remot
 
 import LinearAlgebra
 import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, diagind, ishermitian, issymmetric
+import Random
+import Random: AbstractRNG
 
 import UUIDs: UUID, uuid4
 
@@ -82,6 +84,7 @@ include("array/sort.jl")
 include("array/linalg.jl")
 include("array/mul.jl")
 include("array/cholesky.jl")
+include("array/random.jl")
 
 # Visualization
 include("visualization.jl")
@@ -101,6 +104,9 @@ function __init__()
     system_uuid()
 
     @static if !isdefined(Base, :get_extension)
+        @require Distributions="31c24e10-a181-5473-b8eb-7969acd0382f" begin
+            include(joinpath(dirname(@__DIR__), "ext", "DistributionsExt.jl"))
+        end
         @require Graphs="86223c79-3864-5bf0-83f7-82e725a168b6" begin
             @require GraphViz="f526b714-d49f-11e8-06ff-31ed36ee7ee0" begin
                 include(joinpath(dirname(@__DIR__), "ext", "GraphVizExt.jl"))

Diff for: src/array/cholesky.jl

@@ -22,7 +22,7 @@ function LinearAlgebra._chol!(A::DArray{T,2}, ::Type{UpperTriangular}) where T
 
     info = [convert(LinearAlgebra.BlasInt, 0)]
     try
-        Dagger.spawn_datadeps(;aliasing=true) do
+        Dagger.spawn_datadeps() do
             for k in range(1, mt)
                 Dagger.@spawn potrf_checked!(uplo, InOut(Ac[k, k]), Out(info))
                 for n in range(k+1, nt)

Diff for: src/array/operators.jl

@@ -111,3 +111,24 @@ end
 
 Base.first(A::DArray) = A[begin]
 Base.last(A::DArray) = A[end]
+
+# In-place operations
+
+function Base.map!(f, a::DArray{T}) where T
+    Dagger.spawn_datadeps() do
+        for ca in chunks(a)
+            Dagger.@spawn map!(f, InOut(ca), ca)
+        end
+    end
+    return a
+end
+
+function Base.map!(f, a::DArray{T}, b::AbstractArray{U}) where {T, U}
+    b2 = view(b, a.partitioning)
+    Dagger.spawn_datadeps() do
+        for (c_a, c_b2) in zip(chunks(a), chunks(b2))
+            Dagger.@spawn map!(f, InOut(c_a), c_b2)
+        end
+    end
+    return a
+end

Diff for: src/array/random.jl

@@ -0,0 +1,26 @@
+Random.rand!(dm::DArray{T}) where T = map!(_ -> rand(T), dm)
+Random.randn!(dm::DArray{T}) where T = map!(_ -> randn(T), dm)
+
+randfork(rng::R, n::Integer) where {R<:AbstractRNG} =
+    R(abs(rand(copy(rng), Int) + n))
+
+function Random.rand!(rng::AbstractRNG, A::DArray{T}) where T
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            Dagger.@spawn map!(_->rand(rng, T), InOut(Ac), Ac)
+        end
+    end
+    return A
+end
+function Random.randn!(rng::AbstractRNG, A::DArray{T}) where T
+    part_sz = prod(map(length, first(A.subdomains).indexes))
+    Dagger.spawn_datadeps() do
+        for Ac in chunks(A)
+            rng = randfork(rng, part_sz)
+            Dagger.@spawn map!(_->randn(rng, T), InOut(Ac), Ac)
+        end
+    end
+    return A
+end

Diff for: src/datadeps.jl

@@ -98,6 +98,9 @@ struct DataDepsAliasingState
     ainfos_readers::Dict{AbstractAliasing,Vector{Pair{DTask,Int}}}
     ainfos_overlaps::Dict{AbstractAliasing,Set{AbstractAliasing}}
 
+    # Cache ainfo lookups
+    ainfo_cache::Dict{Tuple{Any,Any},AbstractAliasing}
+
     function DataDepsAliasingState()
         data_origin = Dict{AbstractAliasing,MemorySpace}()
         data_locality = Dict{AbstractAliasing,MemorySpace}()
@@ -106,8 +109,11 @@ struct DataDepsAliasingState
         ainfos_readers = Dict{AbstractAliasing,Vector{Pair{DTask,Int}}}()
         ainfos_overlaps = Dict{AbstractAliasing,Set{AbstractAliasing}}()
 
+        ainfo_cache = Dict{Tuple{Any,Any},AbstractAliasing}()
+
         return new(data_origin, data_locality,
-                   ainfos_owner, ainfos_readers, ainfos_overlaps)
+                   ainfos_owner, ainfos_readers, ainfos_overlaps,
+                   ainfo_cache)
     end
 end
 struct DataDepsNonAliasingState
@@ -156,6 +162,12 @@ struct DataDepsState{State<:Union{DataDepsAliasingState,DataDepsNonAliasingState
     end
 end
 
+function aliasing(astate::DataDepsAliasingState, arg, dep_mod)
+    return get!(astate.ainfo_cache, (arg, dep_mod)) do
+        return aliasing(arg, dep_mod)
+    end
+end
+
 # Determine which arguments could be written to, and thus need tracking
 
 "Whether `arg` has any writedep in this datadeps region."
@@ -190,7 +202,7 @@ function has_writedep(state::DataDepsState, arg, deps, task::DTask)
             for (readdep, writedep, other_ainfo, _, _) in other_taskdeps
                 writedep || continue
                 for (dep_mod, _, _) in deps
-                    ainfo = aliasing(arg, dep_mod)
+                    ainfo = aliasing(state.alias_state, arg, dep_mod)
                     if will_alias(ainfo, other_ainfo)
                         return true
                     end
@@ -221,7 +233,7 @@ function is_writedep(arg, deps, task::DTask)
 end
 
 # Aliasing state setup
-function populate_task_info!(state::DataDepsState, spec, task)
+function populate_task_info!(state::DataDepsState, spec::DTaskSpec, task::DTask)
     # Populate task dependencies
     dependencies_to_add = Vector{Tuple{Bool,Bool,AbstractAliasing,<:Any,<:Any}}()
 
@@ -233,13 +245,13 @@ function populate_task_info!(state::DataDepsState, spec, task)
         # Unwrap the Chunk underlying any DTask arguments
         arg = arg isa DTask ? fetch(arg; raw=true) : arg
 
-        # Skip non-mutable arguments
-        Base.datatype_pointerfree(typeof(arg)) && continue
+        # Skip non-aliasing arguments
+        type_may_alias(typeof(arg)) || continue
 
         # Add all aliasing dependencies
         for (dep_mod, readdep, writedep) in deps
             if state.aliasing
-                ainfo = aliasing(arg, dep_mod)
+                ainfo = aliasing(state.alias_state, arg, dep_mod)
             else
                 ainfo = UnknownAliasing()
             end
@@ -251,15 +263,16 @@ function populate_task_info!(state::DataDepsState, spec, task)
     end
 
     # Track the task result too
-    push!(dependencies_to_add, (true, true, UnknownAliasing(), identity, task))
+    # N.B. We state no readdep/writedep because, while we can't model the aliasing info for the task result yet, we don't want to synchronize because of this
+    push!(dependencies_to_add, (false, false, UnknownAliasing(), identity, task))
 
     # Record argument/result dependencies
     push!(state.dependencies, task => dependencies_to_add)
 end
 function populate_argument_info!(state::DataDepsState{DataDepsAliasingState}, arg, deps)
     astate = state.alias_state
     for (dep_mod, readdep, writedep) in deps
-        ainfo = aliasing(arg, dep_mod)
+        ainfo = aliasing(astate, arg, dep_mod)
 
         # Initialize owner and readers
         if !haskey(astate.ainfos_owner, ainfo)
@@ -405,7 +418,7 @@ function generate_slot!(state::DataDepsState, dest_space, data)
             data_converted = move(from_proc, to_proc, data)
             data_chunk = tochunk(data_converted, to_proc)
             @assert processor(data_chunk) in processors(dest_space)
-            @assert memory_space(data_chunk) == memory_space(data_converted)
+            @assert memory_space(data_converted) == memory_space(data_chunk) "space mismatch! $(memory_space(data_converted)) != $(memory_space(data_chunk)) ($(typeof(data_converted)) vs. $(typeof(data_chunk))), spaces ($orig_space -> $dest_space)"
             @assert orig_space != memory_space(data_chunk) "space preserved! $orig_space != $(memory_space(data_chunk)) ($(typeof(data)) vs. $(typeof(data_chunk))), spaces ($orig_space -> $dest_space)"
             return data_chunk
         end
@@ -664,7 +677,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
             # Is the data written previously or now?
             arg, deps = unwrap_inout(arg)
             arg = arg isa DTask ? fetch(arg; raw=true) : arg
-            if Base.datatype_pointerfree(typeof(arg)) || !has_writedep(state, arg, deps, task)
+            if !type_may_alias(typeof(arg)) || !has_writedep(state, arg, deps, task)
                 @dagdebug nothing :spawn_datadeps "($(repr(spec.f)))[$idx] Skipped copy-to (unwritten)"
                 spec.args[idx] = pos => arg
                 continue
@@ -676,7 +689,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
             end
             if queue.aliasing
                 for (dep_mod, _, _) in deps
-                    ainfo = aliasing(arg, dep_mod)
+                    ainfo = aliasing(astate, arg, dep_mod)
                     data_space = astate.data_locality[ainfo]
                     nonlocal = our_space != data_space
                     if nonlocal
@@ -736,10 +749,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         for (idx, (_, arg)) in enumerate(task_args)
             arg, deps = unwrap_inout(arg)
             arg = arg isa DTask ? fetch(arg; raw=true) : arg
-            Base.datatype_pointerfree(typeof(arg)) && continue
+            type_may_alias(typeof(arg)) || continue
             if queue.aliasing
                 for (dep_mod, _, writedep) in deps
-                    ainfo = aliasing(arg, dep_mod)
+                    ainfo = aliasing(astate, arg, dep_mod)
                     if writedep
                         @dagdebug nothing :spawn_datadeps "($(repr(spec.f)))[$idx][$dep_mod] Syncing as writer"
                         get_write_deps!(state, ainfo, task, write_num, syncdeps)
@@ -769,10 +782,10 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         for (idx, (_, arg)) in enumerate(task_args)
             arg, deps = unwrap_inout(arg)
             arg = arg isa DTask ? fetch(arg; raw=true) : arg
-            Base.datatype_pointerfree(typeof(arg)) && continue
+            type_may_alias(typeof(arg)) || continue
             if queue.aliasing
                 for (dep_mod, _, writedep) in deps
-                    ainfo = aliasing(arg, dep_mod)
+                    ainfo = aliasing(astate, arg, dep_mod)
                     if writedep
                         @dagdebug nothing :spawn_datadeps "($(repr(spec.f)))[$idx][$dep_mod] Set as owner"
                         add_writer!(state, ainfo, task, write_num)
@@ -864,7 +877,7 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         for arg in keys(astate.data_origin)
             # Is the data previously written?
             arg, deps = unwrap_inout(arg)
-            if Base.datatype_pointerfree(typeof(arg)) || !has_writedep(state, arg, deps)
+            if !type_may_alias(typeof(arg)) || !has_writedep(state, arg, deps)
                 @dagdebug nothing :spawn_datadeps "Skipped copy-from (unwritten)"
             end
 
@@ -935,7 +948,7 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true,
         throw(ArgumentError("Dynamic scheduling is no longer available"))
     end
     wait_all(; check_errors=true) do
-        scheduler = something(scheduler, DATADEPS_SCHEDULER[], :naive)::Symbol
+        scheduler = something(scheduler, DATADEPS_SCHEDULER[], :roundrobin)::Symbol
         launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool
         if launch_wait
             result = spawn_bulk() do