datadeps: Support within-DAG deps for AOT scheduler

Author: jpsamaroo

src/datadeps.jl

@@ -96,40 +96,61 @@ struct DAGSpec
                     Dict{Int, Type}(),
                     Dict{Int, Vector{DatadepsArgSpec}}())
 end
-function Base.push!(dspec::DAGSpec, tspec::DTaskSpec, task::DTask)
+function dag_add_task!(dspec::DAGSpec, astate, tspec::DTaskSpec, task::DTask)
+    # Check if this task depends on any other tasks within the DAG,
+    # which we are not yet ready to handle
+    for (idx, (kwpos, arg)) in enumerate(tspec.args)
+        arg, deps = unwrap_inout(arg)
+        pos = kwpos isa Symbol ? kwpos : idx
+        for (dep_mod, readdep, writedep) in deps
+            if arg isa DTask
+                if arg.uid in keys(dspec.uid_to_id)
+                    # Within-DAG dependency, bail out
+                    return false
+                end
+            end
+        end
+    end
+
     add_vertex!(dspec.g)
     id = nv(dspec.g)
 
+    # Record function signature
     dspec.id_to_functype[id] = typeof(tspec.f)
-
-    dspec.id_to_argtypes[id] = DatadepsArgSpec[]
+    argtypes = DatadepsArgSpec[]
     for (idx, (kwpos, arg)) in enumerate(tspec.args)
         arg, deps = unwrap_inout(arg)
         pos = kwpos isa Symbol ? kwpos : idx
         for (dep_mod, readdep, writedep) in deps
             if arg isa DTask
+                #= TODO: Re-enable this when we can handle within-DAG dependencies
                 if arg.uid in keys(dspec.uid_to_id)
                     # Within-DAG dependency
                     arg_id = dspec.uid_to_id[arg.uid]
                     push!(dspec.id_to_argtypes[arg_id], DatadepsArgSpec(pos, DTaskDAGID{arg_id}, dep_mod, UnknownAliasing()))
                     add_edge!(dspec.g, arg_id, id)
                     continue
                 end
+                =#
 
                 # External DTask, so fetch this and track it as a raw value
                 arg = fetch(arg; raw=true)
             end
-            ainfo = aliasing(arg, dep_mod)
-            push!(dspec.id_to_argtypes[id], DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo))
+            ainfo = aliasing(astate, arg, dep_mod)
+            push!(argtypes, DatadepsArgSpec(pos, typeof(arg), dep_mod, ainfo))
         end
     end
+    dspec.id_to_argtypes[id] = argtypes
 
     # FIXME: Also record some portion of options
     # FIXME: Record syncdeps
     dspec.id_to_uid[id] = task.uid
     dspec.uid_to_id[task.uid] = id
 
-    return
+    return true
+end
+function dag_has_task(dspec::DAGSpec, task::DTask)
+    return task.uid in keys(dspec.uid_to_id)
 end
 function Base.:(==)(dspec1::DAGSpec, dspec2::DAGSpec)
     # Are the graphs the same size?
@@ -159,10 +180,7 @@ function Base.:(==)(dspec1::DAGSpec, dspec2::DAGSpec)
             # Are the arguments the same?
             argspec1.value_type === argspec2.value_type || return false
             argspec1.dep_mod === argspec2.dep_mod || return false
-            if !equivalent_structure(argspec1.ainfo, argspec2.ainfo)
-                @show argspec1.ainfo argspec2.ainfo
-                return false
-            end
+            equivalent_structure(argspec1.ainfo, argspec2.ainfo) || return false
         end
     end
 
@@ -454,7 +472,7 @@ function populate_return_info!(state::DataDepsState{DataDepsNonAliasingState}, t
     astate.data_locality[task] = space
     astate.data_origin[task] = space
 end
-function clear_ainfo_owner_readers!(astate::DataDepsAliasingState)
+function reset_ainfo_owner_readers!(astate::DataDepsAliasingState)
     for ainfo in keys(astate.ainfos_owner)
         astate.ainfos_owner[ainfo] = nothing
         empty!(astate.ainfos_readers[ainfo])
@@ -621,24 +639,26 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         @warn "Datadeps support for multi-GPU, multi-worker is currently broken\nPlease be prepared for incorrect results or errors" maxlog=1
     end
 
-    # Round-robin assign tasks to processors
     upper_queue = get_options(:task_queue)
 
     state = DataDepsState(queue.aliasing, all_procs)
     astate = state.alias_state
 
     schedule = Dict{DTask, Processor}()
 
-    if DATADEPS_SCHEDULE_REUSABLE[]
-        # Compute DAG spec
-        for (spec, task) in queue.seen_tasks
-            push!(state.dag_spec, spec, task)
+    # Compute DAG spec
+    for (spec, task) in queue.seen_tasks
+        if !dag_add_task!(state.dag_spec, astate, spec, task)
+            # This task needs to be deferred
+            break
         end
+    end
 
+    if DATADEPS_SCHEDULE_REUSABLE[]
         # Find any matching DAG specs and reuse their schedule
         for (other_spec, spec_schedule) in DAG_SPECS
             if other_spec == state.dag_spec
-                @info "Found matching DAG spec!"
+                @dagdebug nothing :spawn_datadeps "Found matching DAG spec!"
                 #spec_schedule = DAG_SCHEDULE_CACHE[other_spec]
                 schedule = Dict{DTask, Processor}()
                 for (id, proc) in spec_schedule.id_to_proc
@@ -654,13 +674,20 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
 
     # Populate all task dependencies
     write_num = 1
+    task_num = 0
     for (spec, task) in queue.seen_tasks
+        if !dag_has_task(state.dag_spec, task)
+            # This task needs to be deferred
+            break
+        end
         write_num = populate_task_info!(state, spec, task, write_num)
+        task_num += 1
     end
+    @assert task_num > 0
 
     if isempty(schedule)
         # Run AOT scheduling
-        schedule = datadeps_create_schedule(queue.scheduler, state, queue.seen_tasks)::Dict{DTask, Processor}
+        schedule = datadeps_create_schedule(queue.scheduler, state, queue.seen_tasks[1:task_num])::Dict{DTask, Processor}
 
         if DATADEPS_SCHEDULE_REUSABLE[]
             # Cache the schedule
@@ -674,12 +701,17 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         end
     end
 
-    # Clear out ainfo database (will be repopulated during task execution)
-    clear_ainfo_owner_readers!(astate)
+    # Reset ainfo database (will be repopulated during task execution)
+    reset_ainfo_owner_readers!(astate)
 
     # Launch tasks and necessary copies
     write_num = 1
     for (spec, task) in queue.seen_tasks
+        if !dag_has_task(state.dag_spec, task)
+            # This task needs to be deferred
+            break
+        end
+
         our_proc = schedule[task]
         @assert our_proc in all_procs
         our_space = only(memory_spaces(our_proc))
@@ -829,6 +861,9 @@ function distribute_tasks!(queue::DataDepsTaskQueue)
         write_num += 1
     end
 
+    # Remove processed tasks
+    deleteat!(queue.seen_tasks, 1:task_num)
+
     # Copy args from remote to local
     if queue.aliasing
         # We need to replay the writes from all tasks in-order (skipping any
@@ -961,18 +996,25 @@ function spawn_datadeps(f::Base.Callable; static::Bool=true,
     wait_all(; check_errors=true) do
         scheduler = something(scheduler, DATADEPS_SCHEDULER[], RoundRobinScheduler())
         launch_wait = something(launch_wait, DATADEPS_LAUNCH_WAIT[], false)::Bool
+        local result
         if launch_wait
-            result = spawn_bulk() do
+            spawn_bulk() do
                 queue = DataDepsTaskQueue(get_options(:task_queue);
                                           scheduler, aliasing)
-                with_options(f; task_queue=queue)
-                distribute_tasks!(queue)
+                result = with_options(f; task_queue=queue)
+                while !isempty(queue.seen_tasks)
+                    @dagdebug nothing :spawn_datadeps "Entering Datadeps region"
+                    distribute_tasks!(queue)
+                end
             end
         else
             queue = DataDepsTaskQueue(get_options(:task_queue);
                                       scheduler, aliasing)
             result = with_options(f; task_queue=queue)
-            distribute_tasks!(queue)
+            while !isempty(queue.seen_tasks)
+                @dagdebug nothing :spawn_datadeps "Entering Datadeps region"
+                distribute_tasks!(queue)
+            end
         end
         return result
     end