Use correct world in task execution

jpsamaroo · jpsamaroo · commit dc506abad5ff · 2025-08-27T20:53:47.000-07:00
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
@@ -253,13 +253,13 @@ Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Chunk{T}) where {
     Dagger.move(from_proc, to_proc, fetch(x))
 
 # Task execution
-function Dagger.execute!(proc::CuArrayDeviceProc, f, args...; kwargs...)
+function Dagger.execute!(proc::CuArrayDeviceProc, world::UInt64, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()
     task = Threads.@spawn begin
         Dagger.set_tls!(tls)
         with_context!(proc)
-        result = Base.@invokelatest f(args...; kwargs...)
+        result = Base.invoke_in_world(world, f, args...; kwargs...)
         # N.B. Synchronization must be done when accessing result or args
         return result
     end
diff --git a/ext/IntelExt.jl b/ext/IntelExt.jl
@@ -239,13 +239,13 @@ end
 =#
 
 # Task execution
-function Dagger.execute!(proc::oneArrayDeviceProc, f, args...; kwargs...)
+function Dagger.execute!(proc::oneArrayDeviceProc, world::UInt64, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()
     task = Threads.@spawn begin
         Dagger.set_tls!(tls)
         with_context!(proc)
-        result = Base.@invokelatest f(args...; kwargs...)
+        result = Base.invoke_in_world(world, f, args...; kwargs...)
         # N.B. Synchronization must be done when accessing result or args
         return result
     end
diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl
@@ -254,13 +254,13 @@ function Dagger.move_optimized(
 end
 
 # Task execution
-function Dagger.execute!(proc::MtlArrayDeviceProc, f, args...; kwargs...)
+function Dagger.execute!(proc::MtlArrayDeviceProc, world::UInt64, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()
     task = Threads.@spawn begin
         Dagger.set_tls!(tls)
         with_context!(proc)
-        result = Base.@invokelatest f(args...; kwargs...)
+        result = Base.invoke_in_world(world, f, args...; kwargs...)
         # N.B. Synchronization must be done when accessing result or args
         return result
     end
diff --git a/ext/OpenCLExt.jl b/ext/OpenCLExt.jl
@@ -222,13 +222,13 @@ Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Chunk{T}) where {
     Dagger.move(from_proc, to_proc, fetch(x))
 
 # Task execution
-function Dagger.execute!(proc::CLArrayDeviceProc, f, args...; kwargs...)
+function Dagger.execute!(proc::CLArrayDeviceProc, world::UInt64, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()
     task = Threads.@spawn begin
         Dagger.set_tls!(tls)
         with_context!(proc)
-        result = Base.@invokelatest f(args...; kwargs...)
+        result = Base.invoke_in_world(world, f, args...; kwargs...)
         # N.B. Synchronization must be done when accessing result or args
         return result
     end
diff --git a/ext/PythonExt.jl b/ext/PythonExt.jl
@@ -32,7 +32,8 @@ Dagger.move(::CPUProc, ::PythonProcessor, x::Py) = x
 Dagger.move(::CPUProc, ::PythonProcessor, x::PyArray) = x
 # FIXME: Conversion from Python to Julia
 
-function Dagger.execute!(::PythonProcessor, f, args...; kwargs...)
+# N.B. We ignore world here because Python doesn't have world ages
+function Dagger.execute!(::PythonProcessor, world::UInt64, f, args...; kwargs...)
     @assert f isa Py "Function must be a Python object"
     return f(args...; kwargs...)
 end
diff --git a/ext/ROCExt.jl b/ext/ROCExt.jl
@@ -241,13 +241,13 @@ for lib in [BLAS, LAPACK]
 end
 
 # Task execution
-function Dagger.execute!(proc::ROCArrayDeviceProc, f, args...; kwargs...)
+function Dagger.execute!(proc::ROCArrayDeviceProc, world::UInt64, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()
     task = Threads.@spawn begin
         Dagger.set_tls!(tls)
         with_context!(proc)
-        result = Base.@invokelatest f(args...; kwargs...)
+        result = Base.invoke_in_world(world, f, args...; kwargs...)
         # N.B. Synchronization must be done when accessing result or args
         return result
     end
diff --git a/src/processor.jl b/src/processor.jl
@@ -31,11 +31,12 @@ function delete_processor_callback!(name::Symbol)
 end
 
 """
-    execute!(proc::Processor, f, args...; kwargs...) -> Any
+    execute!(proc::Processor, world::UInt64, f, args...; kwargs...) -> Any
 
 Executes the function `f` with arguments `args` and keyword arguments `kwargs`
-on processor `proc`. This function can be overloaded by `Processor` subtypes to
-allow executing function calls differently than normal Julia.
+in inference world `world` on processor `proc`. This function can be overloaded
+by `Processor` subtypes to allow executing function calls differently than
+normal Julia.
 """
 function execute! end
 
diff --git a/src/queue.jl b/src/queue.jl
@@ -1,5 +1,6 @@
 mutable struct DTaskSpec
     fargs::Vector{Argument}
+    world::UInt64
     options::Options
 end
 
diff --git a/src/sch/Sch.jl b/src/sch/Sch.jl
@@ -805,6 +805,7 @@ struct TaskSpec
     scope::Dagger.AbstractScope
     Tf::Type
     data::Vector{Argument}
+    world::UInt64
     options::Options
     ctx_vars::NamedTuple
     sch_handle::SchedulerHandle
@@ -856,7 +857,7 @@ Base.hash(task::TaskSpec, h::UInt) = hash(task.thunk_id, hash(TaskSpec, h))
         push!(to_send, TaskSpec(
             thunk.id,
             task_spec.est_time_util, task_spec.est_alloc_util, task_spec.est_occupancy,
-            task_spec.scope, Tf, args, options,
+            task_spec.scope, Tf, args, thunk.world, options,
             (log_sink=ctx.log_sink, profile=ctx.profile),
             sch_handle, state.uid))
     end
@@ -1521,7 +1522,7 @@ Executes a single task specified by `task` on `to_proc`.
 
         result = Dagger.with_options(propagated) do
             # Execute
-            execute!(to_proc, f, fetched_args...; fetched_kwargs...)
+            execute!(to_proc, task.world, f, fetched_args...; fetched_kwargs...)
         end
 
         # Check if result is safe to store
diff --git a/src/sch/dynamic.jl b/src/sch/dynamic.jl
@@ -235,6 +235,6 @@ function _add_thunk!(ctx, state, task, tid, (f, args, options, future))
             push!(fargs, Dagger.Argument(pos, arg))
         end
     end
-    payload = Dagger.PayloadOne(UInt(0), future, fargs, _options, true)
+    payload = Dagger.PayloadOne(UInt(0), future, fargs, Base.get_world_counter(), _options, true)
     return Dagger.eager_submit_internal!(ctx, state, task, tid, payload)
 end
diff --git a/src/submission.jl b/src/submission.jl
@@ -2,27 +2,31 @@ mutable struct PayloadOne
     uid::UInt
     future::ThunkFuture
     fargs::Vector{Argument}
+    world::UInt64
     options::Options
     reschedule::Bool
 
     PayloadOne() = new()
     PayloadOne(uid::UInt, future::ThunkFuture,
-               fargs::Vector{Argument}, options::Options, reschedule::Bool) =
-        new(uid, future, fargs, options, reschedule)
+               fargs::Vector{Argument}, world::UInt64, options::Options,
+               reschedule::Bool) =
+        new(uid, future, fargs, world, options, reschedule)
 end
 function unset!(p::PayloadOne, _)
     p.uid = 0
     p.future = EMPTY_PAYLOAD_ONE.future
     p.fargs = EMPTY_PAYLOAD_ONE.fargs
+    p.world = EMPTY_PAYLOAD_ONE.world
     p.options = EMPTY_PAYLOAD_ONE.options
     p.reschedule = false
 end
-const EMPTY_PAYLOAD_ONE = PayloadOne(UInt(0), ThunkFuture(), Argument[], Options(), false)
+const EMPTY_PAYLOAD_ONE = PayloadOne(UInt(0), ThunkFuture(), Argument[], UInt64(0), Options(), false)
 mutable struct PayloadMulti
     ntasks::Int
     uid::Vector{UInt}
     future::Vector{ThunkFuture}
     fargs::Vector{Vector{Argument}}
+    world::Vector{UInt64}
     options::Vector{Options}
     reschedule::Bool
 end
@@ -32,6 +36,7 @@ function payload_extract(f, payload::PayloadMulti, i::Integer)
         p1.uid = payload.uid[i]
         p1.future = payload.future[i]
         p1.fargs = payload.fargs[i]
+        p1.world = payload.world[i]
         p1.options = payload.options[i]
         p1.reschedule = true
         return f(p1)
@@ -72,7 +77,7 @@ const UID_TO_TID_CACHE = TaskLocalValue{ReusableCache{Dict{UInt64,Int},Nothing}}
         payload::PayloadOne
 
         uid, future = payload.uid, payload.future
-        fargs, options, reschedule = payload.fargs, payload.options, payload.reschedule
+        fargs, world, options, reschedule = payload.fargs, payload.world, payload.options, payload.reschedule
 
         id = next_id()
 
@@ -169,6 +174,7 @@ const UID_TO_TID_CACHE = TaskLocalValue{ReusableCache{Dict{UInt64,Int},Nothing}}
             thunk = take_or_alloc!(THUNK_SPEC_CACHE[]) do thunk_spec
                 thunk_spec.fargs = fargs
                 thunk_spec.id = id
+                thunk_spec.world = world
                 thunk_spec.options = options
                 return Thunk(thunk_spec)
             end
@@ -329,7 +335,8 @@ function eager_launch!((spec, task)::Pair{DTaskSpec,DTask})
     # Submit the task
     #=FIXME:REALLOC=#
     thunk_id = eager_submit!(PayloadOne(task.uid, task.future,
-                                        spec.fargs, spec.options, true))
+                                        spec.fargs, spec.world,
+                                        spec.options, true))
     task.thunk_ref = thunk_id.ref
 end
 function eager_launch!(specs::Vector{Pair{DTaskSpec,DTask}})
@@ -351,12 +358,14 @@ function eager_launch!(specs::Vector{Pair{DTaskSpec,DTask}})
         eager_process_args_submission_to_local!(id_map, specs)
         [spec.fargs for (spec, _) in specs]
     end
+    all_worlds = UInt64[spec.world for (spec, _) in specs]
     all_options = Options[spec.options for (spec, _) in specs]
 
     # Submit the tasks
     #=FIXME:REALLOC=#
     thunk_ids = eager_submit!(PayloadMulti(ntasks, uids, futures,
-                                           all_fargs, all_options, true))
+                                           all_fargs, all_worlds,
+                                           all_options, true))
     for i in 1:ntasks
         task = specs[i][2]
         task.thunk_ref = thunk_ids[i].ref
diff --git a/src/threadproc.jl b/src/threadproc.jl
@@ -10,7 +10,8 @@ end
 iscompatible(proc::ThreadProc, opts, f, args...) = true
 iscompatible_func(proc::ThreadProc, opts, f) = true
 iscompatible_arg(proc::ThreadProc, opts, x) = true
-function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...); @nospecialize(kwargs...))
+function execute!(proc::ThreadProc, world::UInt64, f, args...; kwargs...)
+    @nospecialize f args kwargs
     tls = get_tls()
     # FIXME: Use return type of the call to specialize container
     result = Ref{Any}()
@@ -19,7 +20,7 @@ function execute!(proc::ThreadProc, @nospecialize(f), @nospecialize(args...); @n
         if task_logging_enabled()
             TimespanLogging.prof_task_put!(tls.sch_handle.thunk_id.id)
         end
-        result[] = @invokelatest f(args...; kwargs...)
+        result[] = Base.invoke_in_world(world, f, args...; kwargs...)
         return
     end
     set_task_tid!(task, proc.tid)
diff --git a/src/thunk.jl b/src/thunk.jl
@@ -7,13 +7,15 @@ const EMPTY_ARGS = Argument[]
 const EMPTY_SYNCDEPS = Set{Any}()
 Base.@kwdef mutable struct ThunkSpec
     fargs::Vector{Argument} = EMPTY_ARGS
+    world::UInt64 = UInt64(0)
     id::Int = 0
     cache_ref::Any = nothing
     affinity::Union{Pair{OSProc,Int}, Nothing} = nothing
     options::Union{Options, Nothing} = nothing
 end
 function unset!(spec::ThunkSpec, _)
     spec.fargs = EMPTY_ARGS
+    spec.world = UInt64(0)
     spec.id = 0
     spec.cache_ref = nothing
     spec.affinity = nothing
@@ -56,6 +58,7 @@ If omitted, options can also be specified by passing key-value pairs as
 """
 mutable struct Thunk
     inputs::Vector{Argument} # TODO: Use `ImmutableArray` in 1.8
+    world::UInt64
     id::Int
     cache_ref::Any
     affinity::Union{Pair{OSProc,Int}, Nothing}
@@ -64,13 +67,14 @@ mutable struct Thunk
     sch_accessible::Bool
     finished::Bool
     function Thunk(spec::ThunkSpec)
-        return new(spec.fargs, spec.id,
+        return new(spec.fargs, spec.world, spec.id,
                    spec.cache_ref, spec.affinity,
                    spec.options,
                    true, true, false)
     end
 end
 function Thunk(f, xs...;
+               world::UInt64=Base.get_world_counter(),
                syncdeps=nothing,
                id::Int=next_id(),
                cache_ref=nothing,
@@ -95,6 +99,7 @@ function Thunk(f, xs...;
             spec.fargs[idx+1] = Argument(something(x.first, idx), x.second)
         end
     end
+    spec.world = world
     if options === nothing
         options = Options()
     end
@@ -561,7 +566,7 @@ function spawn(f, args...; kwargs...)
     unique!(task_options.propagates)
 
     # Construct task spec and handle
-    spec = DTaskSpec(args_kwargs, task_options)
+    spec = DTaskSpec(args_kwargs, Base.get_world_counter(), task_options)
     task = eager_spawn(spec)
 
     # Enqueue the task into the task queue
diff --git a/test/fakeproc.jl b/test/fakeproc.jl
@@ -20,6 +20,6 @@ Dagger.iscompatible_arg(proc::FakeProc, opts, ::Type{<:Integer}) = true
 Dagger.iscompatible_arg(proc::FakeProc, opts, ::Type{<:FakeVal}) = true
 Dagger.move(from_proc::OSProc, to_proc::FakeProc, x::Integer) = FakeVal(x)
 Dagger.move(from_proc::ThreadProc, to_proc::FakeProc, x::Integer) = FakeVal(x)
-Dagger.execute!(proc::FakeProc, func, args...) = FakeVal(42+func(args...).x)
+Dagger.execute!(proc::FakeProc, world, func, args...) = FakeVal(42+func(args...).x)
 
 end
