Can we have inferable fetch(task)?

tkf · August 26, 2020, 12:50am

Continuing the discussion in ANN: Parallel `for` loops in FLoops.jl with composable and extensible fold-based API - #18 by c42f, I wonder if we can make fetch(task) inferable using invoke_in_world added in

github.com/JuliaLang/julia

`_apply_in_world` builtin

JuliaLang:master ← JuliaLang:cjf/apply-in-world

opened 12:07AM - 12 May 20 UTC

c42f

+71 -1

This implements a new builtin `Core._apply_in_world` to allow Julia code to be r…un in a frozen world age (when combined with `Base.get_world_counter()`). This is more general than `Core._apply_latest`, so we could remove `_apply_latest` (or replace with a simple shim which calls `_apply_in_world` with `world=typemax(UInt)`) if people think that's a good idea. ### Motivation The original motivation for this was to have a way to freeze the world age of Julia code implementing the Julia parser - see note at https://github.com/JuliaLang/julia/pull/35243#issuecomment-623044749. Using a fixed world should be beneficial for infrastructure code which runs in a user's julia process, but which is otherwise not expected to be modified by the user. There's two benefits: * Users can't accidentally break the basic infrastructure of the language. For example, breaking the Julia parser breaks pretty much everything in the REPL. Likewise, breaking Revise pretty much breaks the user's session. * Method invalidation for infrastructure packages like Pkg will not slow down the user experience when the user loads `$random_package` into their session. Note that world age is dynamically scoped, so fixed world would only apply when package code is entered through an `_apply_in_world` shim. Therefore devs of packages which choose to use this for deployment can still use a Revise-based workflow for developing their packages. ### Possible usage scenarios * When replacing the flisp parser with CSTParser, I'd like to use this to ensure user mistakes don't take down the whole session. * Revise.jl is basic infrastructure but method invalidation slows down load time. @timholy has recently fixed this through heroic efforts, but using a fixed world may be a lot easier and more future proof. Same considerations likely apply to JuliaInterpreter/JuliaDebugger. * Pkg also suffers from method invalidation so this may be of interest for the `pkg>` REPL mode, https://github.com/JuliaLang/Pkg.jl/issues/1816 * Any in-process infrastructure for code editor support, etc, may also benefit from this ### Questions / TODO As public API, I've considered a callable `ApplyInWorld` function wrapper (and maybe an API `Base.freeze_world(f)` to create it) which would capture a function and `Base.get_world_counter()`. There might be other options though? I considered making inference understand the builtin so that the ApplyInWorld shim could be inferred, but in discussions with @Keno and @vtjnash, it seemed there were difficulties with this. For one, Jameson thought capturing the world age in the type parameters of `ApplyInWorld` would break subtyping in some way. For two, I was hoping this would help solve the problem that GeneralizedGenerated.jl solves, but Keno says that's not the case because inference can fundamentally only see older world ages. Actually I'm still a little confused by this because the newer world methods would be part of the global state so in principle accessible during compilation of an older world; but perhaps it just breaks fundamental inference invariants I don't understand yet. A few todos: * [x] Consider the public API. * [x] ~~Replace `_apply_latest` builtin with this~~ Reverted to avoid an extra allocation * [x] As I understand it, the world age isn't unique in precompilation. What complexities does this cause? Does it prevent `ApplyInWorld` wrappers from being saved during precompilation? (Edit: yes) * [x] In general, is it safe to capture the world age `UInt` or does this create an implicit dependency on old methods which may be removed from internal caches? Do we need to lift this implicit dependency into an explicit one? (Edit: Should be ok?) @vtjnash I'd greatly value your input on subtleties I might have missed, especially on the last two questions above.

Here is a quick POC:

module InferableTasks

export @iasync, @ispawn

struct InferableTask{T}
    task::Task
end

Base.fetch(t::InferableTask{T}) where {T} = fetch(t.task)::T
Base.wait(t::InferableTask) = wait(t.task)

macro ispawn(ex)
    inferrable(ex) do ex
        :($Threads.@spawn $ex)
    end |> esc
end

macro iasync(ex)
    inferrable(ex) do ex
        :($Base.@async $ex)
    end |> esc
end

function inferrable(spawn_macro, ex)
    @gensym f T world
    quote
        local $f, $T, $world
        $f() = $ex
        $T = $Core.Compiler.return_type($f, $Tuple{})
        $world = $Base.get_world_counter()
        $InferableTask{$T}($(spawn_macro(:($Base.invoke_in_world($world, $f)))))
    end
end

end

It works?

julia> f() = fetch(@ispawn 1+1)
f (generic function with 1 method)

julia> @code_warntype f()
...

Body::Int64
...

@c42f Do you think the above use of invoke_in_world (_apply_in_world) is unsound? Am I wrong to assume that the call to return_type like above uses the world age that would be obtained via get_world_counter? If not, is passing world to return_type fixes it? The above use of invoke_in_world does not require the inference to understand _apply_in_world, right?

c42f · August 26, 2020, 1:53am

Oh right, I see what you had in mind now. I see that you combined return_type with invoke_in_world to get something similar to what world should Tasks run in? · Issue #35690 · JuliaLang/julia · GitHub immediately. Cool!

It looks like the result of return_type is only used as a hint for efficiency and shouldn’t otherwise affect the semantics of the program. So that’s good.

Looking at the implementation of return_type, we can see that it fetches the current (dynamically scoped) world counter using jl_get_tls_world_age:

github.com

JuliaLang/julia/blob/bf32ea47e940c6b606e8d5d4f412e2d1dcee1422/base/compiler/typeinfer.jl#L676-L679


      
          function return_type(@nospecialize(f), @nospecialize(t))
              world = ccall(:jl_get_tls_world_age, UInt, ())
              return ccall(:jl_call_in_typeinf_world, Any, (Ptr{Ptr{Cvoid}}, Cint), Any[_return_type, f, t, world], 4)
          end

On the other hand, Base.get_world_counter() fetches the global latest world counter which is not what you want. I think you need to replace get_world_counter with a call to jl_get_tls_world_age.

tkf · August 26, 2020, 2:08am

Thanks for reviewing the code! Yeah, #35690 is exactly what I had in mind. I didn’t know that there are different kinds of world age.

By the way, Jameson is mentioning that this has a problem because the inference can’t add the correct edges (?) Slack

tkf · August 26, 2020, 2:15am

I wonder if Jameson’s concern can be workarounded by a hack like this

diff --git a/src/InferableTasks.jl b/src/InferableTasks.jl
index 0a5705a..06dc89d 100644
--- a/src/InferableTasks.jl
+++ b/src/InferableTasks.jl
@@ -21,11 +21,16 @@ macro iasync(ex)
     end |> esc
 end

+const NEVER = Ref(false)
+
 function inferable(spawn_macro, ex)
     @gensym f T world
     quote
         local $f, $T, $world
-        $f() = $ex
+        $Base.@noinline $f() = $ex
+        if NEVER[]
+            $f()
+        end
         $T = $Core.Compiler.return_type($f, $Tuple{})
         $world = $Base.get_world_counter()
         $InferableTask{$T}($(spawn_macro(:($Base.invoke_in_world($world, $f)))))

I’m hoping this would make sure that the caller of @ispawn would be invalidated if the function $f has to be invalidated. If the caller is invalidated, it’ll get a new world age. This in turn, invalidates $f called in the task because it is now called from a new world age (via invoke_in_world).

c42f · August 26, 2020, 2:38am

Ah yes, excellent point about back edges and invalidation. The NEVER hack looks like a clever workaround for that.

tkf · August 26, 2020, 10:47pm

Thanks!

So I put things together in https://github.com/tkf/InferableTasks.jl/blob/master/src/InferableTasks.jl just in case someone wants to try it out later.