datadeps: Properly skip copying non-written args

Author: jpsamaroo

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)

src/datadeps.jl

@@ -221,7 +221,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,8 +233,8 @@ 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
@@ -251,7 +251,8 @@ 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)
@@ -664,7 +665,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
@@ -736,7 +737,7 @@ 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)
@@ -769,7 +770,7 @@ 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)
@@ -864,7 +865,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
 

src/memory-spaces.jl

@@ -33,6 +33,8 @@ function unwrap(x::Chunk)
     @assert root_worker_id(x.processor) == myid()
     MemPool.poolget(x.handle)
 end
+move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::T, from::F) where {T,F} =
+    throw(ArgumentError("No `move!` implementation defined for $F -> $T"))
 function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Chunk, from::Chunk)
     to_w = root_worker_id(to_space)
     remotecall_wait(to_w, dep_mod, to_space, from_space, to, from) do dep_mod, to_space, from_space, to, from
@@ -44,6 +46,10 @@ function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Chun
     end
     return
 end
+function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::Base.RefValue{T}, from::Base.RefValue{T}) where {T}
+    to[] = from[]
+    return
+end
 function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::AbstractArray{T,N}, from::AbstractArray{T,N}) where {T,N}
     move!(to_space, from_space, dep_mod(to), dep_mod(from))
 end
@@ -66,6 +72,23 @@ end
 
 ### Aliasing and Memory Spans
 
+type_may_alias(::Type{String}) = false
+type_may_alias(::Type{Symbol}) = false
+type_may_alias(::Type{<:Type}) = false
+type_may_alias(::Type{C}) where C<:Chunk{T} where T = type_may_alias(T)
+function type_may_alias(::Type{T}) where T
+    if isbitstype(T)
+        return false
+    elseif ismutabletype(T)
+        return true
+    elseif isstructtype(T)
+        for FT in fieldtypes(T)
+            type_may_alias(FT) && return true
+        end
+    end
+    return false
+end
+
 may_alias(::MemorySpace, ::MemorySpace) = true
 may_alias(space1::CPURAMMemorySpace, space2::CPURAMMemorySpace) = space1.owner == space2.owner
 
@@ -104,8 +127,71 @@ memory_spans(::NoAliasing) = MemorySpan{CPURAMMemorySpace}[]
 struct UnknownAliasing <: AbstractAliasing end
 memory_spans(::UnknownAliasing) = [MemorySpan{CPURAMMemorySpace}(C_NULL, typemax(UInt))]
 
+warn_unknown_aliasing(T) =
+    @warn "Cannot resolve aliasing for object of type $T\nExecution may become sequential"
+
+struct CombinedAliasing <: AbstractAliasing
+    sub_ainfos::Vector{AbstractAliasing}
+end
+function memory_spans(ca::CombinedAliasing)
+    # FIXME: Don't hardcode CPURAMMemorySpace
+    all_spans = MemorySpan{CPURAMMemorySpace}[]
+    for sub_a in ca.sub_ainfos
+        append!(all_spans, memory_spans(sub_a))
+    end
+    return all_spans
+end
+Base.:(==)(ca1::CombinedAliasing, ca2::CombinedAliasing) =
+    ca1.sub_ainfos == ca2.sub_ainfos
+Base.hash(ca1::CombinedAliasing, h::UInt) =
+    hash(ca1.sub_ainfos, hash(CombinedAliasing, h))
+
+struct ObjectAliasing <: AbstractAliasing
+    ptr::Ptr{Cvoid}
+    sz::UInt
+end
+function ObjectAliasing(x::T) where T
+    @nospecialize x
+    ptr = pointer_from_objref(x)
+    sz = sizeof(T)
+    return ObjectAliasing(ptr, sz)
+end
+function memory_spans(oa::ObjectAliasing)
+    rptr = RemotePtr{Cvoid}(oa.ptr)
+    span = MemorySpan{CPURAMMemorySpace}(rptr, oa.sz)
+    return [span]
+end
+
 aliasing(x, T) = aliasing(T(x))
-aliasing(x) = isbits(x) ? NoAliasing() : UnknownAliasing()
+function aliasing(x::T) where T
+    if isbits(x)
+        return NoAliasing()
+    elseif isstructtype(T)
+        as = AbstractAliasing[]
+        # If the object itself is mutable, it can alias
+        if ismutabletype(T)
+            push!(as, ObjectAliasing(x))
+        end
+        # Check all object fields (recursive)
+        for field in fieldnames(T)
+            sub_as = aliasing(getfield(x, field))
+            if sub_as isa NoAliasing
+                continue
+            elseif sub_as isa CombinedAliasing
+                append!(as, sub_as.sub_ainfos)
+            else
+                push!(as, sub_as)
+            end
+        end
+        return CombinedAliasing(as)
+    else
+        warn_unknown_aliasing(T)
+        return UnknownAliasing()
+    end
+end
+aliasing(::String) = NoAliasing() # FIXME: Not necessarily true
+aliasing(::Symbol) = NoAliasing()
+aliasing(::Type) = NoAliasing()
 aliasing(x::Chunk, T) = remotecall_fetch(root_worker_id(x.processor), x, T) do x, T
     aliasing(unwrap(x), T)
 end
@@ -129,6 +215,7 @@ function aliasing(x::Array{T}) where T
     else
         # FIXME: Also ContiguousAliasing of container
         #return IteratedAliasing(x)
+        warn_unknown_aliasing(T)
         return UnknownAliasing()
     end
 end
@@ -173,6 +260,7 @@ function aliasing(x::SubArray{T,N,A}) where {T,N,A<:Array}
     else
         # FIXME: Also ContiguousAliasing of container
         #return IteratedAliasing(x)
+        warn_unknown_aliasing(T)
         return UnknownAliasing()
     end
 end

test/datadeps.jl

@@ -1,7 +1,21 @@
 using LinearAlgebra, Graphs
 
-empty!(Dagger.DAGDEBUG_CATEGORIES)
-push!(Dagger.DAGDEBUG_CATEGORIES, :spawn_datadeps)
+@testset "Memory Aliasing" begin
+    A = rand(4)
+    a = Dagger.aliasing(A)
+    @test a isa Dagger.ContiguousAliasing
+    @test a.span.ptr.addr == UInt(pointer(A))
+    @test a.span.len == sizeof(Float64) * length(A)
+
+    r = Ref(3)
+    a = Dagger.aliasing(r)
+    @test a isa Dagger.CombinedAliasing
+    @test length(a.sub_ainfos) == 1
+    s = only(a.sub_ainfos)
+    @test s isa Dagger.ObjectAliasing
+    @test s.ptr == pointer_from_objref(r)
+    @test s.sz == sizeof(3)
+end
 
 function with_logs(f)
     Dagger.enable_logging!(;taskdeps=true, taskargs=true)