Julia: task switch not allowed from inside staged functions

[damiendr]

I see that task switch is no longer allowed from inside staged functions. This has a number of unfortunate consequences. For instance if a staged function uses println() then it will fail if printing can cause a task switch, as is the case for the IJulia notebook. Is there hope to relax the requirement in the future?

Example notebook:

In[1]

@generated function foo(x)
    println(x)
    :(x)
end

Out[1]

foo (generic function with 1 method)

In[2]

foo(2)

Out[2]

LoadError: task switch not allowed from inside staged nor pure functions
while loading In[2], in expression starting on line 1

 in yieldto(::Task, ::ANY) at ./event.jl:136
 in yieldto(::Task, ::ANY) at /Users/me/julia-dev-HEAD/usr/lib/julia/sys.dylib:?
 in wait() at ./event.jl:169
 in stream_wait(::Task) at ./stream.jl:42
 in uv_write(::Base.PipeEndpoint, ::Ptr{UInt8}, ::UInt64) at ./stream.jl:764
 in unsafe_write(::Base.PipeEndpoint, ::Ptr{UInt8}, ::UInt64) at ./stream.jl:774
 in write(::Base.PipeEndpoint, ::Symbol) at ./io.jl:322
 in print at ./show.jl:3 [inlined]
 in show(::Base.PipeEndpoint, ::TypeName) at ./show.jl:237
 in show_datatype(::Base.PipeEndpoint, ::DataType) at ./show.jl:221
 in show(::Base.PipeEndpoint, ::DataType) at ./show.jl:193
 in print(::Base.PipeEndpoint, ::Type{T}) at ./strings/io.jl:18
 in print(::Base.PipeEndpoint, ::Type{T}, ::Char, ::Vararg{Char,N}) at ./strings/io.jl:29
 in println(::Type{T}) at ./coreio.jl:5
 in foo(...) at ./In[1]:2

Version 0.6.0-dev.666

[vtjnash]

It was never really allowed (for correctness, the compiler algorithms in codegen and type-inference generally assume that this is impossible). The only difference now is that the runtime system will detect this situation and throw an error. I think we could perhaps fix the codegen algorithm, but it would require breaking most user code algorithms for cooperatively-scheduled, green-threaded Tasks, such as the IO system.

However, to assist with the IO needs for debugging generated functions, a simple low-performance STDOUT repr printer is now available as Core.println.

[damiendr]

OK! I think these restrictions are undocumented so far (no hint of it in metaprogramming). There may also be user code out there that will suddenly break when the ban is enforced — I've been task switching and even sometimes invoking type inference within generated functions, and nothing went obviously wrong.

Another useful debug function that's affected is dump().

[vtjnash]

Core.println is a bit more akin to dump in formatting, but you can also use Core.STDOUT as the IO argument to most printing functions to work in that direction also.

OK! I think these restrictions are undocumented so far (no hint of it in metaprogramming).

Unfortunately, I don't think anyone actual knows the limitations of @generated. The only thing definitely safe is very simple template-like code. I agree that should be added to the docs somewhere. Anything else is probably invalid. For example, type-inference might randomly return a wrong answer (or abort) in either of those example cases. Nothing would go obviously wrong, it'll just generate bad code in random places and likely cause unreproducible failures in mysterious places.

[damiendr]

Unfortunately, I don't think anyone actual knows the limitations of @generated.

That's an interesting point! I would shamelessly lobby (;-) to keep the ability to do some reasonable side-effects within staged functions, because there's so many cool things to do with metaprogramming that go beyond simple templating.

Or maybe there is another way to programmatically add methods at runtime that doesn't compromise the integrity of the typeinf/codegen process ?

One idiom I like is essentially a form of multi-stage programming:

• use the Val{x} trick or type parameters to pass data into the staged function (where x is a non-literal value that changes at most a few times during execution).
• in the staged function, do some expensive computation that generates code specialised for x. (_some exotic possibilities: use a solver or theorem prover, compile some code, etc._)
• insert the generated code with no dispatch overhead thanks to @generated.

@generated function foo{x}(::Val{x}, args...)
    code = generate_adhoc_code(x, args) # expensive operation
end

But in the end, I could also do something like this, though it's not as convenient (because method generation isn't automatically triggered for any new x):

function foo
end

x = get_meta_parameter()
gen_foo_method(x)

for i = 1:1000000 # hot loop
    foo(Val{x}, data[i])
end

_EDIT: I guess I haven't explored all the new possibilities opened up by the new, efficient function types yet._

[vtjnash]

That's an interesting point! I would shamelessly lobby (;-) to keep the ability to do some reasonable side-effects within staged functions, because there's so many cool things to do with metaprogramming that go beyond simple templating.

Since that ability isn't present now and has never been present, that would be "add" not "keep".

Or maybe there is another way to programmatically add methods at runtime that doesn't compromise the integrity of the typeinf/codegen process ?

probably not. oddly, this may end up meaning that you probably could programmatically add new types at runtime (although still not during @generated), but you wouldn't be able to call their constructor.

One idiom I like is essentially a form of multi-stage programming:
use the Val{x} trick or type parameters to pass data into the staged function (where x is a non-literal value that changes at most a few times during execution).
in the staged function, do some expensive computation that generates code specialised for x. (some exotic possibilities: use a solver or theorem prover, compile some code, etc.)

a memoization Dict would be equivalent

The system assumes that calling a generated function is cheap and may repeatedly re-call the same function on a handful of possible combinations of arguments, even if they don't make any sense and would never be encountered at runtime.

insert the generated code with no dispatch overhead thanks to @generated

It's not necessarily true that generated functions have no dispatch overhead. In particular, if you construct the Val or Type with a non-constant parameter (e.g. Val{a_local_variable}), the dispatch cost is much higher than for the non-generated function.

[damiendr]

The system assumes that calling a generated function is cheap and may repeatedly re-call the same function on a handful of possible combinations of arguments, even if they don't make any sense and would never be encountered at runtime.

A handful is usually fine in my use cases.

Since that ability isn't present now and has never been present, that would be "add" not "keep".

Fair enough! Do you think the following is already going too far?

func_signatures = []
@generated function foo(args...)
    push!(func_signatures, args)
    quote
        <something safe>
    end
end

In particular, if you construct the Val or Type with a non-constant parameter (e.g. Val{a_local_variable}), the dispatch cost is much higher than for the non-generated function.

Right. Guess I need to benchmark and compare that with a Dict lookup.

[vtjnash]

Fair enough! Do you think the following is already going too far?

It is not safe to assume that func_signatures will contains args while calling foo(args) if that's what you're asking.

[damiendr]

It is not safe to assume that func_signatures will contains args while calling foo(args) if that's what you're asking.

Wow, I have to admit I don't understand how that's possible. How can Julia know which function body to run if it doesn't first go through push!, and how can the side-effect of push! be rolled back by the time the function is executed?

That would be a serious bummer for two of my projects, so I'm really trying to understand the situation here, even though I've used that idiom thousands of time without running into problems:

func_signatures = []
func_ptrs = Ptr{Void}[]

@generated function foo(args...)
    push!(func_signatures, args)
    idx = length(func_signatures)
    quote
        ensure_compiled() # populates func_ptrs from func_signatures
        ccall(func_ptrs[$idx], ...)
    end
end

[vtjnash]

There's actually quite a number of possibilities, including: running inference / a pure interpreter, subprocess / spawn, checkpointing / fork, and memoization / caching.

[damiendr]

OK, makes sense. Thanks for the insight!

[djsegal]

Related to this, Core.println no longer works correctly.

Use Core.print with a \n at the end of the string

// see: https://discourse.julialang.org/t/why-does-println-break-generated-functions/15344

[yuyichao]

What's the code you have trouble with? Your original post there suggests that you were using println instead. Core.println should still work.

[djsegal]

You're right. Scratch that comment 😬