Generators using reflection are now very hard to debug

[maleadt]

#48766 changed generated functions, requiring manual expansion of @generated in order to get access to the world argument that's now mandatory to pass to reflection methods. This is enforced by the current world being set to -1 during generator expansion, triggering an error (code reflection cannot be used from generated functions) when using using world=-1.

However, this change makes generators that do use reflection very hard to debug when they do anything wrong (i.e. not necessarily reflection-related). Take the following example:

function generator(world, source, self, f, tt)
    tt = tt.parameters[1]
    sig = Tuple{f, tt.parameters...}
    mi = Base._which(sig; world)

    error("oh no")

    stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, :ctx, :x, :f), Core.svec())
    stub(world, source, :(nothing))
end

@eval function doit(f, tt)
  $(Expr(:meta, :generated, generator))
  $(Expr(:meta, :generated_only))
end

doit(sin, Tuple{Int})

This generator correctly passes the codegen world to _which, but after that triggers some arbitrary unrelated exception (say, you have a typo in the generator). This results in the generator running again at run time, however, at that point it runs with the world argument also set to -1, resulting in:

ERROR: LoadError: code reflection cannot be used from generated functions
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] _which(tt::Type; method_table::Nothing, world::UInt64, raise::Bool)
   @ Base ./reflection.jl:1644

That makes it very hard to debug generators. In fact, in order to get IRTools and Zygote working again with this change, I've been running with a patch a la:

diff --git a/base/compiler/utilities.jl b/base/compiler/utilities.jl
index 836c370b98..35049b14f8 100644
--- a/base/compiler/utilities.jl
+++ b/base/compiler/utilities.jl
@@ -120,7 +120,9 @@ function get_staged(mi::MethodInstance, world::UInt)
         # user code might throw errors – ignore them
         ci = ccall(:jl_code_for_staged, Any, (Any, UInt), mi, world)::CodeInfo
         return ci
-    catch
+    catch err
+        Core.println(Core.stderr, "Internal error: encountered unexpected ", err, " during expansion of @generated ", mi, ":")
+        ccall(:jl_print_backtrace, Nothing, ())
         return nothing
     end
 end

Internal error: encountered unexpected ErrorException("oh no") during expansion of @generated doit(typeof(Base.sin), Type{Tuple{Int64}}) from doit(Any, Any):
error at ./error.jl:35
generator at /home/tim/Julia/tmp/Zygote.jl/wip3.jl:8
unknown function (ip: 0x7f0aefb00de9)

This at least shows me why the generator failed during compilation, because the run-time exception is useless.

@vtjnash This would have been avoided using a simpler API, #48766 (comment). I get your argument for making the breakage more obvious, but fixing the ecosystem has been a pain.

[vtjnash]

I am not sure how that would help things. The world values you are seeing are all valid expected inputs, and since the compiler handles them, it also expects user code that plugs into the compiler (in the form of @generated functions) to also handle them the same way. Now you could argue that supporting running unspecialized code (aka world=typemax) is only needed for soundness of your implementation when running in the JIT, but not needed for whole program compilation of the sort implemented by the GPUCompiler ecosystem. But how is the system supposed to know that you are not going to run this on the CPU, but rather in the restricted GPU environment, unless your @generated function is implemented to return that information when requested (signaled by world=-1)?

[vtjnash]

tl;dr: this is a misunderstanding of the expected behavior, not affected by the API used to access it

the generator running again at run time, however, at that point it runs with the world argument also set to -1

The world argument is not set because it was running at run time, it is set because it was running at compile time because it found that it couldn't create optimized source code, so it fell back to attempting to compile an unspecialized version of the code.

[maleadt]

attempting to compile an unspecialized version

If world=-1 indicates unspecialized codegen, why doesn't reflection work then, returning nothing or so instead of erroring with code reflection cannot be used from generated functions?

What's the way generators should behave then (to return an unspecialized version of the code)?

And FWIW, this isn't related to GPU code, this is the kind of error that virtually every generated function doing reflection now runs into. I ran into it when debugging IRTools and Zygote.

[vtjnash]

A return of nothing would indeed be more correct. I worried that people would consider that value to be less friendly than an error. And none of these methods are legal to use anyways from @generated, since they fail to return all required information (I think only Core.Compiler.findsup and Core.Compiler.findall actually provide the full interface).

[Keno]

I don't understand why we need the world==-1 behavior at all. Why can't we just throw the error at runtime with whatever world it happened to get called with?