Sch: Use W1T1 if profitable

Author: jpsamaroo

src/sch/Sch.jl

@@ -626,16 +626,17 @@ end
         resize!(sorted_procs, length(input_procs))
         costs = @reusable_dict :schedule!_costs Processor Float64 OSProc() 0.0 32
         costs_cleanup = @reuse_defer_cleanup empty!(costs)
-        estimate_task_costs!(sorted_procs, costs, state, input_procs, task)
-        costs_cleanup() # We don't use costs here
+        estimate_task_costs!(sorted_procs, costs, state, input_procs, task; sig)
         input_procs_cleanup()
         scheduled = false
 
-        # Move our corresponding ThreadProc to be the last considered
+        # Move our corresponding ThreadProc to be the last considered,
+        # if the task is expected to run for longer than the time it takes to
+        # schedule it onto another worker (estimated at 1ms).
         if length(sorted_procs) > 1
             sch_threadproc = Dagger.ThreadProc(myid(), Threads.threadid())
             sch_thread_idx = findfirst(proc->proc==sch_threadproc, sorted_procs)
-            if sch_thread_idx !== nothing
+            if sch_thread_idx !== nothing && costs[sch_threadproc] > 1_000_000 # 1ms
                 deleteat!(sorted_procs, sch_thread_idx)
                 push!(sorted_procs, sch_threadproc)
             end

src/sch/util.jl

@@ -551,13 +551,13 @@ current estimated per-processor compute pressure, and transfer costs for each
 `Chunk` argument to `task`. Returns `(procs, costs)`, with `procs` sorted in
 order of ascending cost.
 """
-function estimate_task_costs(state, procs, task)
+function estimate_task_costs(state, procs, task; sig=nothing)
     sorted_procs = Vector{Processor}(undef, length(procs))
     costs = Dict{Processor,Float64}()
-    estimate_task_costs!(sorted_procs, costs, state, procs, task)
+    estimate_task_costs!(sorted_procs, costs, state, procs, task; sig)
     return sorted_procs, costs
 end
-@reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task)
+@reuse_scope function estimate_task_costs!(sorted_procs, costs, state, procs, task; sig=nothing)
     tx_rate = state.transfer_rate[]
 
     # Find all Chunks
@@ -569,6 +569,12 @@ end
         end
     end
 
+    # Estimate the cost of executing the task itself
+    if sig === nothing
+        sig = signature(task.f, task.inputs)
+    end
+    est_time_util = get(state.signature_time_cost, sig, 1000^3)
+
     # Estimate total cost for executing this task on each candidate processor
     for proc in procs
         gproc = get_parent(proc)
@@ -582,8 +588,14 @@ end
         tx_cost = impute_sum(affinity(chunk)[2] for chunk in chunks_filt)
 
         # Estimate total cost to move data and get task running after currently-scheduled tasks
-        est_time_util = get(state.worker_time_pressure[gproc.pid], proc, 0)
-        costs[proc] = est_time_util + (tx_cost/tx_rate)
+        est_business = get(state.worker_time_pressure[get_parent(proc).pid], proc, 0)
+
+        # Add fixed cost for cross-worker task transfer (esimated at 1ms)
+        # TODO: Actually estimate/benchmark this
+        task_xfer_cost = gproc.pid != myid() ? 1_000_000 : 0 # 1ms
+
+        # Compute final cost
+        costs[proc] = est_time_util + est_business + (tx_cost/tx_rate) + task_xfer_cost
     end
     chunks_cleanup()