add mechanism for spoofing inference work-limiting heuristics

[jrevels]

ref #24676

• add method_for_inference_limit_heuristics field to CodeInfo
• add mechanism for early generator expansion
• create some way to set the flag for early generator expansion
• fix inevitable bugs
• add tests

With what I've got so far, everything seems fine as long as the method_for_inference_heuristics is never inflated. Returning a CodeInfo with an inflated method_for_inference_heuristics field, however, seems to result in a segfault, so I'm pretty sure I've missed some places in src that still need updating. @Keno Any tips here?

[iblislin]

CI timeout seems releated:
https://julia.iblis.cnmc.tw/#/builders/1/builds/4884
https://julia.iblis.cnmc.tw/#/builders/1/builds/4889
and appveyor got timeout as well.
I found there was a process (maybe created by this PR) eating up one CPU core on my worker. kill cannot stop it, only kill -9 makes it.

Here is the backtrace from my worker: https://gist.github.com/iblis17/baf65932991124be3d9aff8760a747d2

[jrevels]

Okay, I've now added a Bool flag to GeneratedFunctionStub that inference can check to decide when to expand the generator. Just to get started, I'm hardcoding the flag to false, but now I need to add a mechanism for actually setting it.

@JeffBezanson What would be the most reasonable way to add this setting mechanism? On the front end, I could add an extra argument to the @generated macro, or add a @generated_early macro, etc. I'm also wondering what the easiest/cleanest way is to propagate that setting to here. I guess I could add a meta node to the expression returned from the @generated macro and then destructure it here?

...I'm a noob when it comes to Julia's parser 😛

[jrevels]

I just did a rebase on top of the changes from #24362. @vtjnash is this PR's strategy still correct in the wake of those changes/are there other cases I need to cover here now that #24362 has been merged?

[jrevels]

Status update:

This still has a bunch of loose odds and ends (mainly the flag-setting thing I asked Jeff about above). At this point, though, I'm waiting on some intermediate review to make sure I'm not on the wrong track.

[jrevels]

@vtjnash recommended offline that the flag-setting strategy should be for downstream users to manually construct and interpolate the GeneratedFunctionStub into their method AST, which works for me as long as I can cook up some representative tests that prevent the feature from getting clobbered by future development 😛

[jrevels]

I'm having trouble actually employing the whole "just interpolate the stub" approach suggested by @vtjnash. Take the following example:

@generated function f(x::NTuple{N,Int}) where {N}
    return Expr(:tuple, (:(x[$i] * x[$j]) for i in 1:N, j in 1:N)...)
end

This essentially expands to

function f(x::NTuple{N, Int}) where N
    if $(Expr(:generated))
        return Expr(:tuple, ((Core._expr)(:call, :*, (Core._expr)(:ref, :x, i), (Core._expr)(:ref, :x, j)) for i = 1:N, j = 1:N)...)
    else
        $(Expr(:meta, :generated_only))
        return nothing
    end
end

The parser will then pick up on the :generated node in the method body, and construct/interpolate the appropriate GeneratedFunctionStub into the body.

I want to emit the equivalent thing as the above example, but manually construct the GeneratedFunctionStub myself. Here's the code I tried:

function _f(x::Type{NTuple{N,Int}}) where {N}
    return Expr(:tuple, (:(x[$i] * x[$j]) for i in 1:N, j in 1:N)...)
end

@eval function f(x::NTuple{N,Int}) where {N}
    $(Expr(:meta, :generated, Core.GeneratedFunctionStub(_f, Any[:f, :x], Any[:N], @__LINE__, Symbol(@__FILE__), false)))
end

...which spits out

ERROR: invalid @generated function; try placing it in global scope

...so I'm assuming my method body for f is malformed. What should I be doing instead?

[maleadt]

The following seems to work (without the additional bool you've added, of course):

function _f(sparam_N, fun, x::Type{NTuple{N,Int}}) where {N}
    return Expr(:tuple, (:(x[$i] * x[$j]) for i in 1:N, j in 1:N)...)
end

@eval function f(x::NTuple{N,Int}) where {N}
    $(Expr(:meta, :generated, Expr(:new, Core.GeneratedFunctionStub, :_f, Any[:f, :x], Any[:N], @__LINE__, QuoteNode(Symbol(@__FILE__)))))
end

[jrevels]

So I added some tests, most of which pass locally. The tests that actually inflate method_for_inference_heuristics, though, cause segfaults. It looks like the addition of that fields breaks jl_compress_ast/jl_uncompress_ast:

julia> f(i) = i
f (generic function with 1 method)

julia> m = methods(f, (Int,)).ms[1]
f(i) in Main at REPL[1]:1

julia> c = @code_lowered f(1)
CodeInfo(:(begin
        nothing
        return i
    end))

julia> ccall(:jl_compress_ast, Any, (Any, Any), m, c); # works fine if `method_for_inference_heuristics` is uninflated

julia> c.method_for_inference_heuristics = m
f(i) in Main at REPL[1]:1

julia> ccall(:jl_compress_ast, Any, (Any, Any), m, c)

signal (11): Segmentation fault: 11
in expression starting at no file:0
module_in_worklist at /Users/jarrettrevels/data/repos/julia-dev/src/dump.c:206 [inlined]
jl_serialize_value_ at /Users/jarrettrevels/data/repos/julia-dev/src/dump.c:638
jl_compress_ast at /Users/jarrettrevels/data/repos/julia-dev/src/dump.c:2160
top-level scope at ./<missing> (unknown line)
jl_call_fptr_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:380 [inlined]
jl_call_method_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:399 [inlined]
jl_toplevel_eval_flex at /Users/jarrettrevels/data/repos/julia-dev/src/toplevel.c:721
jl_toplevel_eval_in at /Users/jarrettrevels/data/repos/julia-dev/src/builtins.c:626
eval at ./repl/REPL.jl:3
jl_call_fptr_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:380 [inlined]
jl_call_method_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:399 [inlined]
jl_apply_generic at /Users/jarrettrevels/data/repos/julia-dev/src/gf.c:2010
eval_user_input at ./repl/REPL.jl:69
jl_call_fptr_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:380 [inlined]
jl_call_method_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:399 [inlined]
jl_apply_generic at /Users/jarrettrevels/data/repos/julia-dev/src/gf.c:2010
macro expansion at ./repl/REPL.jl:100 [inlined]
#1 at ./event.jl:95
jl_call_fptr_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:380 [inlined]
jl_call_method_internal at /Users/jarrettrevels/data/repos/julia-dev/src/./julia_internal.h:399 [inlined]
jl_apply_generic at /Users/jarrettrevels/data/repos/julia-dev/src/gf.c:2010
jl_apply at /Users/jarrettrevels/data/repos/julia-dev/src/./julia.h:1476 [inlined]
start_task at /Users/jarrettrevels/data/repos/julia-dev/src/task.c:268
Allocations: 7940457 (Pool: 7939899; Big: 558); GC: 17
[1]    52261 segmentation fault  ~/data/repos/julia-dev/julia

It looks like the CodeInfo (un)compression strategy is to call jl_serialize_value_ on all the fields besides the last five, which are (de)serialized manually. I had hoped that by placing the new method_for_inference_heuristics field before those last five fields, I'd be able to get away without touching the (un)compression stuff, but it looks like that's not the case...From the stacktrace, it seems that calling jl_serialize_value_ on certain kinds of values (like a Method, here) will cause a segfault when it checks to see if the value's module is in the serializer's worklist. I'm not sure why, though.

Naively, I can work around this by just storing the relevant type signature in method_for_inference_heuristics instead of the entire Method, like:

c.method_for_inference_heuristics = Core.Inference.svec(typeof(f),Int)

(I tried using Tuple{typeof(f), Int} instead, but got the same segfault as before)

Is this workaround valid @vtjnash? I'll push forward assuming so, but let me know if not.

[vtjnash]

it seems that calling jl_serialize_value_ on certain kinds of values (like a Method, here) will cause a segfault

Most sorts of objects can't be serialized. You have to instead push it into the roots array and store the index.

[jrevels]

Ah, make sense.

So far, it seems to me like the signature-based approach is cleaner/more lightweight than the method-based approach anyway, so I'll just stick with that unless it's problematic.

[jrevels]

Argh, looks like jl_(un)compress_ast is still broken here.

This time, the breakage causes inference bootstrap to fail - when calling typeinf_type for Core.Inference.typeinf_ext, the ssavaluetypes field of some CodeInfo is erroneously nothing rather than an Int.

It looks like the (un)compression methods should be (de)serializing ssavaluetypes as part of the automatic (de)serialization step that loops over the first 5 CodeInfo fields, though. Am I reading this code wrong? Maybe the bug is actually happening earlier, e.g. in jl_method_set_source?

EDIT: Fixed now, see my comment in jltypes.c below

[jrevels]

shout out to @maleadt for pointing out the missing comma here that was causing my bootstrap problems

[StefanKarpinski]

been there, buddy

[iblislin]

Got un-killable process on BSD CI worker again:

(lldb) p jl_lineno           
(int) $0 = 1004
(lldb) p jl_filename
(const char *) $1 = 0x000000080fd666b8 "/usr/home/julia/julia-fbsd-buildbot/worker/11rel-amd64/build/test/inference.jl"

Here is the backtrace from lldb
https://gist.github.com/iblis17/cceed54911978ceb80a82534e25038e5#file-bt

[jrevels]

Thanks Iblis! I'll look into that.

I've got Cassette working with this change, and it does seem to help in some places, but it doesn't immediately render overdubbed functions inferable.

AFAICT, this just means that this PR doesn't apply the spoofing mechanism in all the places where it's actually needed, but exploring that problem space can be done after this is merged. Luckily, the spoofing mechanism should make it much easier for inference-knowledgable folks (cough @vtjnash cough) to identify the problem areas, since one can now directly access the underlying signature wherever it is required.

Before merging, it'd be useful to have a test that demonstrates this feature the way that Cassette is using it, but I'm not sure how to accomplish that without basically reimplementing Cassette in Base...

[jrevels]

Seems like the way these new mechanisms are being utilized is incorrect, because it's causing the inference test here to fail. The actual number of specializations on this branch:

julia> map(count_specializations, methods(copy_dims_out))
3-element Array{Int64,1}:
 156
 109
 109

@vtjnash can I get your review on my abstract_call_method changes? I think those might be wrong. Best case scenario, these changes being incorrect might be why I didn't get the performance improvement I was expecting in Cassette...

[jrevels]

A friendly New Year's bump 😛 I'd really like to get this working ASAP so we can start tackling whatever else comes up in Cassette-land.

[jrevels]

Since I'm having trouble getting review here after multiple requests, I've decided to take a different approach and will break up this PR into two parts: the spoofing mechanism implementation, and actually employing the mechanism in inference. The former will now be the sole subject of this PR, while the latter will be the subject of a follow-up PR. I hope this will make the work more amenable to review.

Since the spoofing mechanism itself has been complete for a while now, I plan on merging this today assuming nothing catastrophic happens.

[marius311]

Is there any documentation anywhere on how to use this mechanism? I'm not quite able to follow what Cassette.jl does with it.

I have a package (https://github.com/marius311/SelfFunctions.jl) which does one thing similar to Cassette.jl, basically involving a macro that wraps every single function call in another function call, and unsurprisingly this breaks inference in many cases, and my understanding is the mechanism in this PR / used by Cassette.jl could make it work. (of course maybe I could just use Cassette.jl, which I may look at further down the line) Thanks for any help.

[jrevels]

I have a package (https://github.com/marius311/SelfFunctions.jl) which does one thing similar to Cassette.jl, basically involving a macro that wraps every single function call in another function call, and unsurprisingly this breaks inference in many cases

I'm not sure what's going on here, but just based on what your package seems to be doing from the README, my guess is that this PR is irrelevant to whatever problem you're running into...

Is there any documentation anywhere on how to use this mechanism?

No, it's not really a generally exposed interface; it is very easy to abuse and shouldn't generally be used unless you are writing some new extension to the compiler itself. That being said, this PR just adds two new fields:

The first is the new CodeInfo field method_for_inference_limit_heuristics. If the compiler sees this field is populated with a Method object, it is allowed to use that method when performing limiting heuristics instead of using the actual method represented by the CodeInfo.

The second is the new Generator field expand_early. Before adding this field, generator expansion would always occur after inference limiting. This is usually more efficient, but also obviously prevents inference limiting from exploiting/respecting the method_for_inference_limit_heuristics field of any CodeInfo that would be returned from a generator. To solve this problem, one can set expand_early to true to tell the compiler to expand the generator earlier, so that it can leverage the returned CodeInfo's method_for_inference_limit_heuristics field during inference limiting.

I'm not quite able to follow what Cassette.jl does with it.

Cassette uses this to ensure that compiler does not think overdubbed code is deeply recursing for the sake of limiting inference. In other words, Cassette wants the compiler to treat the call chain overdub(f) -> overdub(g) -> overdub(h) as f -> g -> h for the purpose of limiting inference.

To that end, Cassette sets expand_early to true in the overdub generator and sets the method_for_inference_limit_heuristics field to the method being overdubbed.

[marius311]

Thanks, that helps. I should have been more specific about what aspect of SelfFunctions.jl I think this might be relevant. To explain:

One thing which that macro does is make it so if you call "self" functions from inside other self functions, you don't need explicitly pass the "self" argument. The way this is done is e.g. in that second example

@self Rosenbrock shifted_rosen(x,y) = rosen(x+1, y+1)

is rewritten to

shifted_rosen(self,x,y) = self_call(self, rosen, selfcall(self,+,x,1), selfcall(self,+,y,1))

and then dispatch makes it so that

selfcall(self, normalfunction, args...) = normalfunction(args...)  # this dispatches for `+` above
selfcall(self, selffunction, args...) = selffunction(self, args...) # this dispatches for `rosen`

Doing it this way works even if the inner functions haven't been defined yet (as it should), and preserves type stability. Except that now inference sees tons of recursion into this selfcall function and often I think hits exactly the inference limits discussed here.

I'm not too familiar with how Cassette.jl works but I understood there's alot of this same recursion (the function is overdub instead of selfcall I suppose?) and it seemed like I might be able to copy what you guys did to get around this.

One difference I'm realizing now is that method_for_inference_limit_heuristics is at the level of lowered code. So is there a way to "lower the code in my macro then return it" with the method_for_inference_limit_heuristics applied? I suppose this must be part of the magic that Cassette.jl is doing. I guess I'm closer to thinking I should just figure out how to use Cassette for this?

[jrevels]

Hmm...naively I'm not sure why one would use the @self macro rather than passing an argument explicitly (or equivalently, just defining whatever closure they want)?

Anyway, it might be better to move this discussion to an issue in the SelfFunctions.jl repo.