Add scope for union of processor type

Adds ProcessorTypeScope(T), which matches processors that are a subtype
of T. In the process, also expands the scope system to support
lazily-evaluated scoping behavior, such as doing a subtype check or
checking for `default_enabled`, via "taints".

Author: jpsamaroo

src/sch/Sch.jl

@@ -9,7 +9,7 @@ import Random: randperm
 import Base: @invokelatest
 
 import ..Dagger
-import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, WeakChunk, OSProc, AnyScope
+import ..Dagger: Context, Processor, Thunk, WeakThunk, ThunkFuture, ThunkFailedException, Chunk, WeakChunk, OSProc, DefaultScope
 import ..Dagger: order, dependents, noffspring, istask, inputs, unwrap_weak_checked, affinity, tochunk, timespan_start, timespan_finish, procs, move, chunktype, processor, default_enabled, get_processors, get_parent, execute!, rmprocs!, addprocs!, thunk_processor, constrain, cputhreadtime
 
 const OneToMany = Dict{Thunk, Set{Thunk}}
@@ -634,7 +634,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
         scope = if task.f isa Chunk
             task.f.scope
         else
-            AnyScope()
+            DefaultScope()
         end
         for input in task.inputs
             input = unwrap_weak_checked(input)
@@ -693,7 +693,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
                 end
             end
         end
-        state.cache[task] = SchedulingException("No processors available, try making proclist more liberal")
+        state.cache[task] = SchedulingException("No processors available, try widening scope")
         state.errored[task] = true
         set_failed!(state, task)
         @goto pop_task
@@ -723,7 +723,7 @@ function schedule!(ctx, state, procs=procs_to_use(ctx))
                 if procs_found
                     push!(failed_scheduling, task)
                 else
-                    state.cache[task] = SchedulingException("No processors available, try making proclist more liberal")
+                    state.cache[task] = SchedulingException("No processors available, try widening scope")
                     state.errored[task] = true
                     set_failed!(state, task)
                 end

src/sch/util.jl

@@ -14,7 +14,7 @@ function get_propagated_options(thunk)
     nt = NamedTuple()
     for key in thunk.propagates
         value = if key == :scope
-            isa(thunk.f, Chunk) ? thunk.f.scope : AnyScope()
+            isa(thunk.f, Chunk) ? thunk.f.scope : DefaultScope()
         elseif key == :processor
             isa(thunk.f, Chunk) ? thunk.f.processor : OSProc()
         elseif key in fieldnames(Thunk)
@@ -285,18 +285,20 @@ function can_use_proc(task, gproc, proc, opts, scope)
     # Check against single
     if opts.single !== nothing
         if gproc.pid != opts.single
-            @debug "Rejected $proc: gproc.pid != single"
+            @debug "[$(task.id)] Rejected $proc: gproc.pid ($(gproc.pid)) != single ($(opts.single))"
             return false
         end
     end
 
     # Check scope
     proc_scope = Dagger.ExactScope(proc)
     if constrain(scope, proc_scope) isa Dagger.InvalidScope
-        @debug "Rejected $proc: Task scope ($scope) vs. processor scope ($proc_scope)"
+        @debug "[$(task.id)] Rejected $proc: Not contained in task scope ($scope)"
         return false
     end
 
+    @debug "[$(task.id)] Accepted $proc"
+
     return true
 end
 

src/scopes.jl

@@ -1,10 +1,28 @@
-export AnyScope, UnionScope, NodeScope, ProcessScope, ExactScope
+export AnyScope, DefaultScope, UnionScope, NodeScope, ProcessScope, ExactScope, ProcessorTypeScope
 
 abstract type AbstractScope end
 
-"Default scope that is unconstrained."
+"Widest scope that contains all processors."
 struct AnyScope <: AbstractScope end
 
+abstract type AbstractScopeTaint end
+
+"Taints a scope for later evaluation."
+struct TaintScope <: AbstractScope
+    scope::AbstractScope
+    taints::Set{AbstractScopeTaint}
+end
+Base.:(==)(ts1::TaintScope, ts2::TaintScope) =
+    ts1.scope == ts2.scope &&
+    length(ts1.taints) == length(ts2.taints) &&
+    all(collect(ts1.taints) .== collect(ts2.taints))
+
+struct DefaultEnabledTaint <: AbstractScopeTaint end
+
+"Default scope that contains the set of `default_enabled` processors."
+DefaultScope() = TaintScope(AnyScope(),
+                            Set{AbstractScopeTaint}([DefaultEnabledTaint()]))
+
 "Union of two or more scopes."
 struct UnionScope <: AbstractScope
     scopes::Tuple
@@ -35,6 +53,13 @@ end
 ProcessScope(p::OSProc) = ProcessScope(p.pid)
 ProcessScope() = ProcessScope(myid())
 
+"Scoped to any processor with a given supertype."
+struct ProcessorTypeTaint{T} <: AbstractScopeTaint end
+
+ProcessorTypeScope(T) =
+    TaintScope(AnyScope(),
+               Set{AbstractScopeTaint}([ProcessorTypeTaint{T}()]))
+
 "Scoped to a specific processor."
 struct ExactScope <: AbstractScope
     parent::ProcessScope
@@ -58,7 +83,46 @@ Base.isless(x, ::AnyScope) = true
 constrain(::AnyScope, ::AnyScope) = AnyScope()
 constrain(::AnyScope, y) = y
 
+# N.B. TaintScope taints constraining (until encountering an `ExactScope`) to
+# allow lazy evaluation of the taint matching on the final processor
+taint_match(::DefaultEnabledTaint, x::Processor) = default_enabled(x)
+taint_match(::ProcessorTypeTaint{T}, x::Processor) where T = x isa T
+Base.isless(::TaintScope, ::TaintScope) = false
+Base.isless(::TaintScope, ::AnyScope) = true
+Base.isless(::TaintScope, x) = false
+Base.isless(x, ::TaintScope) = true
+function constrain(x::TaintScope, y::TaintScope)
+    scope = constrain(x.scope, y.scope)
+    if scope isa InvalidScope
+        return scope
+    end
+    taints = Set{AbstractScopeTaint}()
+    for tx in x.taints
+        push!(taints, tx)
+    end
+    for ty in y.taints
+        push!(taints, ty)
+    end
+    return TaintScope(scope, taints)
+end
+function constrain(x::TaintScope, y)
+    scope = constrain(x.scope, y)
+    if scope isa InvalidScope
+        return scope
+    end
+    return TaintScope(scope, x.taints)
+end
+function constrain(x::TaintScope, y::ExactScope)
+    for taint in x.taints
+        if !taint_match(taint, y.processor)
+            return InvalidScope(x, y)
+        end
+    end
+    return constrain(x.scope, y)
+end
+
 Base.isless(::UnionScope, ::UnionScope) = false
+Base.isless(::UnionScope, ::TaintScope) = true
 Base.isless(::UnionScope, ::AnyScope) = true
 function constrain(x::UnionScope, y::UnionScope)
     zs = Vector{AbstractScope}()
@@ -75,12 +139,14 @@ end
 constrain(x::UnionScope, y) = constrain(x, UnionScope((y,)))
 
 Base.isless(::NodeScope, ::NodeScope) = false
+Base.isless(::NodeScope, ::TaintScope) = true
 Base.isless(::NodeScope, ::AnyScope) = true
 constrain(x::NodeScope, y::NodeScope) =
     x == y ? y : InvalidScope(x, y)
 
 Base.isless(::ProcessScope, ::ProcessScope) = false
 Base.isless(::ProcessScope, ::NodeScope) = true
+Base.isless(::ProcessScope, ::TaintScope) = true
 Base.isless(::ProcessScope, ::AnyScope) = true
 constrain(x::ProcessScope, y::ProcessScope) =
     x == y ? y : InvalidScope(x, y)
@@ -90,6 +156,7 @@ constrain(x::NodeScope, y::ProcessScope) =
 Base.isless(::ExactScope, ::ExactScope) = false
 Base.isless(::ExactScope, ::ProcessScope) = true
 Base.isless(::ExactScope, ::NodeScope) = true
+Base.isless(::ExactScope, ::TaintScope) = true
 Base.isless(::ExactScope, ::AnyScope) = true
 constrain(x::ExactScope, y::ExactScope) =
     x == y ? y : InvalidScope(x, y)

src/thunk.jl

@@ -39,7 +39,7 @@ arguments are still passed as-is.
 - `processor::Processor=OSProc()` - The processor associated with `f`. Useful if
 `f` is a callable struct that exists on a given processor and should be
 transferred appropriately.
-- `scope::Dagger.AbstractScope=AnyScope()` - The scope associated with `f`.
+- `scope::Dagger.AbstractScope=DefaultScope()` - The scope associated with `f`.
 Useful if `f` is a function or callable struct that may only be transferred to,
 and executed within, the specified scope.
 
@@ -80,7 +80,7 @@ mutable struct Thunk
         if !isa(f, Chunk) && (!isnothing(processor) || !isnothing(scope))
             f = tochunk(f,
                         something(processor, OSProc()),
-                        something(scope, AnyScope()))
+                        something(scope, DefaultScope()))
         end
         xs = Any[xs...]
         if options !== nothing
@@ -315,7 +315,7 @@ function spawn(f, args...; processor=nothing, scope=nothing, kwargs...)
     if !isnothing(processor) || !isnothing(scope)
         f = tochunk(f,
                     something(processor, get_options(:processor, OSProc())),
-                    something(scope, get_options(:scope, AnyScope())))
+                    something(scope, get_options(:scope, DefaultScope())))
     end
     uid, future, finalizer_ref, thunk_ref = if myid() == 1
         _spawn(f, args...; options, kwargs...)

test/scopes.jl

@@ -31,6 +31,17 @@
     es1_ch = Dagger.tochunk(nothing, OSProc(), es1)
     es2_ch = Dagger.tochunk(nothing, OSProc(), es2)
 
+    os1 = ExactScope(OSProc(1))
+
+    @testset "Default Scope" begin
+        ds = DefaultScope()
+        for (s1, s2) in ((ds, es1), (es1, ds))
+            @test Dagger.constrain(s1, s2) == es1
+        end
+        for (s1, s2) in ((ds, os1), (os1, ds))
+            @test Dagger.constrain(s1, s2) isa Dagger.InvalidScope
+        end
+    end
     @testset "Node Scope" begin
         @everywhere node_scope_test(ch...) = Dagger.system_uuid()
 
@@ -128,6 +139,26 @@
         @test es1 in us_res.scopes
         @test !(es2 in us_res.scopes)
     end
+    @testset "Processor Type Scope" begin
+        pts_th = ProcessorTypeScope(Dagger.ThreadProc)
+        pts_os = ProcessorTypeScope(Dagger.OSProc)
+
+        @test Dagger.constrain(pts_th, es1) == es1
+        @test Dagger.constrain(pts_th, os1) isa Dagger.InvalidScope
+
+        @test Dagger.constrain(pts_os, es1) isa Dagger.InvalidScope
+        @test Dagger.constrain(pts_os, os1) == os1
+
+        # Duplicate
+        pts_th_dup = Dagger.constrain(pts_th, pts_th)
+        @test Dagger.constrain(pts_th_dup, es1) == es1
+        @test Dagger.constrain(pts_th_dup, os1) isa Dagger.InvalidScope
+
+        # Empty intersection
+        pts_all = Dagger.constrain(pts_th, pts_os)
+        @test Dagger.constrain(pts_all, es1) isa Dagger.InvalidScope
+        @test Dagger.constrain(pts_all, os1) isa Dagger.InvalidScope
+    end
     # TODO: Test scope propagation
 
     rmprocs([wid1, wid2])

test/thunk.jl

@@ -21,6 +21,7 @@ import Dagger: Chunk
     end
     MulProc() = MulProc(myid())
     Dagger.get_parent(mp::MulProc) = OSProc(mp.owner)
+    Dagger.move(src::MulProc, dest::Dagger.OSProc, x::Function) = Base.:*
     Dagger.move(src::MulProc, dest::Dagger.ThreadProc, x::Function) = Base.:*
 end
 
@@ -202,7 +203,7 @@ end
             a = delayed(+; processor=Dagger.ThreadProc(1,1))(1,2)
             @test a.f isa Chunk
             @test a.f.processor isa Dagger.ThreadProc
-            @test a.f.scope isa AnyScope
+            @test a.f.scope == DefaultScope()
 
             a = delayed(+; processor=Dagger.ThreadProc(1,1), scope=NodeScope())(1,2)
             @test a.f isa Chunk
@@ -236,7 +237,7 @@ end
             a = Dagger.Sch._find_thunk(_a)
             @test a.f isa Chunk
             @test a.f.processor isa Dagger.ThreadProc
-            @test a.f.scope isa AnyScope
+            @test a.f.scope == DefaultScope()
 
             _a = Dagger.spawn(+, 1, 2; processor=Dagger.ThreadProc(1,1), scope=NodeScope())
             a = Dagger.Sch._find_thunk(_a)