Merge pull request #558 from JuliaParallel/jps/optimizations

Various performance and memory optimizations

Author: jpsamaroo

docs/src/api-dagger/types.md

@@ -16,7 +16,6 @@ DTask
 ## Task Options Types
 ```@docs
 Options
-Sch.ThunkOptions
 Sch.SchedulerOptions
 ```
 

docs/src/checkpointing.md

@@ -71,7 +71,7 @@ z = collect(Z)
 ```
 
 Two changes were made: first, we `enumerate(X.chunks)` so that we can get a
-unique index to identify each `chunk`; second, we specify a `ThunkOptions` to
+unique index to identify each `chunk`; second, we specify options to
 `delayed` with a `checkpoint` and `restore` function that is specialized to
 write or read the given chunk to or from a file on disk, respectively. Notice
 the usage of `collect` in the `checkpoint` function, and the use of

docs/src/task-spawning.md

@@ -12,9 +12,9 @@ or `spawn` if it's more convenient:
 
 `Dagger.spawn(f, Dagger.Options(options), args...; kwargs...)`
 
-When called, it creates an [`DTask`](@ref) (also known as a "thunk" or
-"task") object representing a call to function `f` with the arguments `args` and
-keyword arguments `kwargs`. If it is called with other thunks as args/kwargs,
+When called, it creates an [`DTask`](@ref) (also known as a "task" or
+"thunk") object representing a call to function `f` with the arguments `args` and
+keyword arguments `kwargs`. If it is called with other tasks as args/kwargs,
 such as in `Dagger.@spawn f(Dagger.@spawn g())`, then, in this example, the
 function `f` gets passed the results of executing `g()`, once that result is
 available. If `g()` isn't yet finished executing, then the execution of `f`
@@ -29,9 +29,17 @@ it'll be passed as-is to the function `f` (with some exceptions).
 
 !!! note "Task / thread occupancy"
     By default, `Dagger` assumes that tasks saturate the thread they are running on and does not try to schedule other tasks on the thread.
-    This default can be controlled by specifying [`Sch.ThunkOptions`](@ref) (more details can be found under [Scheduler and Thunk options](@ref)).
+    This default can be controlled by specifying [`Options`](@ref) (more details can be found under [Task and Scheduler options](@ref)).
     The section [Changing the thread occupancy](@ref) shows a runnable example of how to achieve this.
 
+## Options
+
+The [`Options`](@ref Dagger.Options) struct in the second argument position is
+optional; if provided, it is passed to the scheduler to control its
+behavior. [`Options`](@ref Dagger.Options) contains option
+key-value pairs, which can be any field in [`Options`](@ref)
+(see [Task and Scheduler options](@ref)).
+
 ## Simple example
 
 Let's see a very simple directed acyclic graph (or DAG) constructed with Dagger:
@@ -51,7 +59,7 @@ s = Dagger.@spawn combine(p, q, r)
 @assert fetch(s) == 16
 ```
 
-The thunks `p`, `q`, `r`, and `s` have the following structure:
+The tasks `p`, `q`, `r`, and `s` have the following structure:
 
 ![graph](https://user-images.githubusercontent.com/25916/26920104-7b9b5fa4-4c55-11e7-97fb-fe5b9e73cae6.png)
 
@@ -108,15 +116,16 @@ x::DTask
 @assert fetch(x) == 3 # fetch the result of `@spawn`
 ```
 
-This is useful for nested execution, where an `@spawn`'d thunk calls `@spawn`. This is detailed further in [Dynamic Scheduler Control](@ref).
+This is useful for nested execution, where an `@spawn`'d task calls `@spawn`.
+This is detailed further in [Dynamic Scheduler Control](@ref).
 
 ## Options
 
 The [`Options`](@ref Dagger.Options) struct in the second argument position is
 optional; if provided, it is passed to the scheduler to control its
 behavior. [`Options`](@ref Dagger.Options) contains a `NamedTuple` of option
 key-value pairs, which can be any of:
-- Any field in [`Sch.ThunkOptions`](@ref) (see [Scheduler and Thunk options](@ref))
+- Any field in [`Options`](@ref) (see [Task and Scheduler options](@ref))
 - `meta::Bool` -- Pass the input [`Chunk`](@ref) objects themselves to `f` and
   not the value contained in them.
 
@@ -127,19 +136,19 @@ There are also some extra options that can be passed, although they're considere
 
 ## Errors
 
-If a thunk errors while running under the eager scheduler, it will be marked as
-having failed, all dependent (downstream) thunks will be marked as failed, and
-any future thunks that use a failed thunk as input will fail. Failure can be
+If a task errors while running under the eager scheduler, it will be marked as
+having failed, all dependent (downstream) tasks will be marked as failed, and
+any future tasks that use a failed task as input will fail. Failure can be
 determined with `fetch`, which will re-throw the error that the
-originally-failing thunk threw. `wait` and `isready` will *not* check whether a
-thunk or its upstream failed; they only check if the thunk has completed, error
+originally-failing task threw. `wait` and `isready` will *not* check whether a
+task or its upstream failed; they only check if the task has completed, error
 or not.
 
 This failure behavior is not the default for lazy scheduling ([Lazy API](@ref)),
-but can be enabled by setting the scheduler/thunk option ([Scheduler and Thunk options](@ref))
+but can be enabled by setting the scheduler/task option ([Task and Scheduler options](@ref))
 `allow_error` to `true`.  However, this option isn't terribly useful for
-non-dynamic usecases, since any thunk failure will propagate down to the output
-thunk regardless of where it occurs.
+non-dynamic usecases, since any task failure will propagate down to the output
+task regardless of where it occurs.
 
 ## Cancellation
 
@@ -198,7 +207,7 @@ end
 ```
 
 Alternatively, if you want to compute but not fetch the result of a lazy
-operation, you can call `compute` on the thunk. This will return a `Chunk`
+operation, you can call `compute` on the task. This will return a `Chunk`
 object which references the result (see [Chunks](@ref) for more details):
 
 ```julia
@@ -215,16 +224,14 @@ Note that, as a legacy API, usage of the lazy API is generally discouraged for m
 - Distinct schedulers don't share runtime metrics or learned parameters, thus causing the scheduler to act less intelligently
 - Distinct schedulers can't share work or data directly
 
-## Scheduler and Thunk options
+## Task and Scheduler options
 
 While Dagger generally "just works", sometimes one needs to exert some more
 fine-grained control over how the scheduler allocates work. There are two
-parallel mechanisms to achieve this: Scheduler options (from
-[`Sch.SchedulerOptions`](@ref)) and Thunk options (from
-[`Sch.ThunkOptions`](@ref)). These two options structs contain many shared
-options, with the difference being that Scheduler options operate
-globally across an entire DAG, and Thunk options operate on a thunk-by-thunk
-basis.
+parallel mechanisms to achieve this: Task options (from [`Options`](@ref)) and
+Scheduler options (from [`Sch.SchedulerOptions`](@ref)). Scheduler
+options operate globally across an entire DAG, and Task options operate on a
+task-by-task basis.
 
 Scheduler options can be constructed and passed to `collect()` or `compute()`
 as the keyword argument `options` for lazy API usage:
@@ -238,7 +245,7 @@ compute(t; options=opts)
 collect(t; options=opts)
 ```
 
-Thunk options can be passed to `@spawn/spawn`, `@par`, and `delayed` similarly:
+Task options can be passed to `@spawn/spawn`, `@par`, and `delayed` similarly:
 
 ```julia
 # Execute on worker 1
@@ -251,8 +258,9 @@ delayed(+; single=1)(1, 2)
 
 ## Changing the thread occupancy
 
-One of the supported [`Sch.ThunkOptions`](@ref) is the `occupancy` keyword.
-This keyword can be used to communicate that a task is not expected to fully saturate a CPU core (e.g. due to being IO-bound).
+One of the supported [`Options`](@ref) is the `occupancy` keyword.
+This keyword can be used to communicate that a task is not expected to fully
+saturate a CPU core (e.g. due to being IO-bound).
 The basic usage looks like this:
 
 ```julia

ext/GraphVizExt.jl

@@ -21,19 +21,22 @@ Requires the `all_task_deps` event enabled in `enable_logging!`
 
 Options:
 - `disconnected`: If `true`, render disconnected vertices (tasks or arguments without upstream/downstream dependencies)
+- `show_data`: If `true`, show the data dependencies in the graph
 - `color_by`: How to color tasks; if `:fn`, then color by unique function name, if `:proc`, then color by unique processor
 - `layout_engine`: The layout engine to use for GraphViz rendering
 - `times`: If `true`, annotate each task with its start and finish times
 - `times_digits`: Number of digits to display in the time annotations
 - `colors`: A list of colors to use for coloring tasks
 - `name_to_color`: A function that maps task names to colors
 """
-function Dagger.render_logs(logs::Dict, ::Val{:graphviz}; disconnected=false,
+function Dagger.render_logs(logs::Dict, ::Val{:graphviz};
+                            disconnected=false, show_data::Bool=true,
                             color_by=:fn, layout_engine="dot",
                             times::Bool=true, times_digits::Integer=3,
                             colors=Dagger.Viz.default_colors,
                             name_to_color=Dagger.Viz.name_to_color)
-    dot = Dagger.Viz.logs_to_dot(logs; disconnected, times, times_digits,
+    dot = Dagger.Viz.logs_to_dot(logs; disconnected, show_data,
+                                 times, times_digits,
                                  color_by, colors, name_to_color)
     gv = GraphViz.Graph(dot)
     GraphViz.layout!(gv; engine=layout_engine)

src/Dagger.jl

@@ -17,9 +17,9 @@ import Random: AbstractRNG
 import UUIDs: UUID, uuid4
 
 if !isdefined(Base, :ScopedValues)
-    import ScopedValues: ScopedValue, with
+    import ScopedValues: ScopedValue, @with, with
 else
-    import Base.ScopedValues: ScopedValue, with
+    import Base.ScopedValues: ScopedValue, @with, with
 end
 import TaskLocalValues: TaskLocalValue
 
@@ -32,7 +32,6 @@ import TimespanLogging: timespan_start, timespan_finish
 
 import Adapt
 
-# Preferences
 import Preferences: @load_preference, @set_preferences!
 
 if @load_preference("distributed-package") == "DistributedNext"
@@ -43,29 +42,35 @@ else
     import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remotecall, remotecall_wait, remotecall_fetch, check_same_host
 end
 
+import MacroTools: @capture, prewalk
+
 include("lib/util.jl")
 include("utils/dagdebug.jl")
 
 # Distributed data
 include("utils/locked-object.jl")
 include("utils/tasks.jl")
-
-import MacroTools: @capture, prewalk
-
-include("options.jl")
+include("utils/reuse.jl")
 include("processor.jl")
 include("threadproc.jl")
+include("sch_options.jl")
 include("context.jl")
 include("utils/processors.jl")
+include("scopes.jl")
+include("utils/scopes.jl")
+include("chunks.jl")
+include("utils/signature.jl")
+include("options.jl")
 include("dtask.jl")
 include("cancellation.jl")
 include("task-tls.jl")
-include("scopes.jl")
-include("utils/scopes.jl")
+include("argument.jl")
 include("queue.jl")
 include("thunk.jl")
+include("utils/fetch.jl")
+include("utils/chunks.jl")
+include("utils/logging.jl")
 include("submission.jl")
-include("chunks.jl")
 include("memory-spaces.jl")
 
 # Task scheduling
@@ -85,15 +90,15 @@ include("stream.jl")
 include("stream-buffers.jl")
 include("stream-transfer.jl")
 
+# File IO
+include("file-io.jl")
+
 # Array computations
 include("array/darray.jl")
 include("array/alloc.jl")
 include("array/map-reduce.jl")
 include("array/copy.jl")
-
-# File IO
-include("file-io.jl")
-
+include("array/random.jl")
 include("array/operators.jl")
 include("array/indexing.jl")
 include("array/setindex.jl")
@@ -104,19 +109,19 @@ include("array/linalg.jl")
 include("array/mul.jl")
 include("array/cholesky.jl")
 include("array/lu.jl")
-include("array/random.jl")
 
 import KernelAbstractions, Adapt
 
 # GPU
 include("gpu.jl")
 
-# Logging and Visualization
+# Logging
+include("utils/logging-events.jl")
+
+# Visualization
 include("visualization.jl")
 include("ui/gantt-common.jl")
 include("ui/gantt-text.jl")
-include("utils/logging-events.jl")
-include("utils/logging.jl")
 include("utils/viz.jl")
 
 """

src/argument.jl

@@ -0,0 +1,46 @@
+mutable struct ArgPosition
+    positional::Bool
+    idx::Int
+    kw::Symbol
+end
+ArgPosition() = ArgPosition(true, 0, :NULL)
+ArgPosition(pos::ArgPosition) = ArgPosition(pos.positional, pos.idx, pos.kw)
+ispositional(pos::ArgPosition) = pos.positional
+iskw(pos::ArgPosition) = !pos.positional
+raw_position(pos::ArgPosition) = ispositional(pos) ? pos.idx : pos.kw
+function pos_idx(pos::ArgPosition)
+    @assert pos.positional
+    @assert pos.idx > 0
+    @assert pos.kw == :NULL
+    return pos.idx
+end
+function pos_kw(pos::ArgPosition)
+    @assert !pos.positional
+    @assert pos.idx == 0
+    @assert pos.kw != :NULL
+    return pos.kw
+end
+mutable struct Argument
+    pos::ArgPosition
+    value
+end
+Argument(pos::Integer, value) = Argument(ArgPosition(true, pos, :NULL), value)
+Argument(kw::Symbol, value) = Argument(ArgPosition(false, 0, kw), value)
+ispositional(arg::Argument) = ispositional(arg.pos)
+iskw(arg::Argument) = iskw(arg.pos)
+pos_idx(arg::Argument) = pos_idx(arg.pos)
+pos_kw(arg::Argument) = pos_kw(arg.pos)
+raw_position(arg::Argument) = raw_position(arg.pos)
+value(arg::Argument) = arg.value
+valuetype(arg::Argument) = typeof(arg.value)
+Base.iterate(arg::Argument) = (arg.pos, true)
+function Base.iterate(arg::Argument, state::Bool)
+    if state
+        return (arg.value, false)
+    else
+        return nothing
+    end
+end
+
+Base.copy(arg::Argument) = Argument(ArgPosition(arg.pos), arg.value)
+chunktype(arg::Argument) = chunktype(value(arg))

src/array/alloc.jl

@@ -86,7 +86,7 @@ function stage(ctx, a::AllocateArray)
         args = a.want_index ? (i, size(x)) : (size(x),)
 
         if isnothing(a.procgrid)
-            scope = get_options(:compute_scope, get_options(:scope, DefaultScope()))
+            scope = get_compute_scope()
         else
             scope = ExactScope(a.procgrid[CartesianIndex(mod1.(Tuple(I), size(a.procgrid))...)])
         end

src/array/darray.jl

@@ -173,7 +173,7 @@ domainchunks(d::DArray) = d.subdomains
 size(x::DArray) = size(domain(x))
 stage(ctx, c::DArray) = c
 
-function Base.collect(d::DArray; tree=false)
+function Base.collect(d::DArray{T,N}; tree=false, copyto=false) where {T,N}
     a = fetch(d)
     if isempty(d.chunks)
         return Array{eltype(d)}(undef, size(d)...)
@@ -183,6 +183,13 @@ function Base.collect(d::DArray; tree=false)
         return fetch(a.chunks[1])
     end
 
+    if copyto
+        C = Array{T,N}(undef, size(a))
+        DC = view(C, Blocks(size(a)...))
+        copyto!(DC, a)
+        return C
+    end
+
     dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)]
     if tree
         collect(fetch(treereduce_nd(map(x -> ((args...,) -> Dagger.@spawn x(args...)) , dimcatfuncs), a.chunks)))
@@ -458,7 +465,7 @@ function stage(ctx::Context, d::Distribute)
             # TODO: fix hashing
             #hash = uhash(idx, Base.hash(Distribute, Base.hash(d.data)))
             if isnothing(d.procgrid)
-                scope = get_options(:compute_scope, get_options(:scope, DefaultScope()))
+                scope = get_compute_scope()
             else
                 scope = ExactScope(d.procgrid[CartesianIndex(mod1.(Tuple(I), size(d.procgrid))...)])
             end
@@ -478,7 +485,7 @@ function stage(ctx::Context, d::Distribute)
             #hash = uhash(c, Base.hash(Distribute, Base.hash(d.data)))
             c = d.domainchunks[I]
             if isnothing(d.procgrid)
-                scope = get_options(:compute_scope, get_options(:scope, DefaultScope()))
+                scope = get_compute_scope()
             else
                 scope = ExactScope(d.procgrid[CartesianIndex(mod1.(Tuple(I), size(d.procgrid))...)])
             end