Need `@force_interpret`? (PrecompileTools still broken)

[timholy]

In writing tests for the revamp of SnoopCompile, I'm noticing that with more complex workloads, PrecompileTools isn't working. Upon tracing the problem, it seems to be due to premature compilation of the full expression returned by @compile_workload. Specifically, in

julia> using PrecompileTools

julia> macroexpand(Main, quote
       @compile_workload begin
           InvalidA.callscallsf(1)                      # resolved callee
           InvalidA.alsocallsf(1)                       # resolved callee (different branch)
           InvalidA.invokesfs(1)                        # invoked callee
           InvalidA.callscallsfrta(1)                   # runtime-dispatched callee
           InvalidA.callsfrtr(1)
           InvalidA.callsfrts(1)
       end
       end)
quote
    #= REPL[2]:2 =#
    begin
        #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:85 =#
        if ccall(:jl_generating_output, PrecompileTools.Cint, ()) == 1 && (PrecompileTools).workload_enabled(Main) || (PrecompileTools).verbose[]
            #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:86 =#
            begin
                #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:77 =#
                (PrecompileTools).tag_newly_inferred_enable()
                #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:78 =#
                try
                    #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:79 =#
                    begin
                        #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:71 =#
                        begin
                            #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:72 =#
                            $(Expr(:meta, :force_compile))
                            #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:73 =#
                            begin
                                #= REPL[2]:3 =#
                                InvalidA.callscallsf(1)
                                #= REPL[2]:4 =#
                                InvalidA.alsocallsf(1)
                                #= REPL[2]:5 =#
                                InvalidA.invokesfs(1)
                                #= REPL[2]:6 =#
                                InvalidA.callscallsfrta(1)
                                #= REPL[2]:7 =#
                                InvalidA.callsfrtr(1)
                                #= REPL[2]:8 =#
                                InvalidA.callsfrts(1)
                            end
                        end
                    end
                finally
                    #= /home/tim/.julia/dev/PrecompileTools/src/workloads.jl:81 =#
                    (PrecompileTools).tag_newly_inferred_disable()
                end
            end
        end
    end
end

the intent is to call PrecompileTools.tag_newly_inferred_enable() before the block inside the try/finally gets compiled. Unfortunately, with the following Julia diff:

$ git diff
diff --git a/src/staticdata_utils.c b/src/staticdata_utils.c
index e9f464b644..7607736fb2 100644
--- a/src/staticdata_utils.c
+++ b/src/staticdata_utils.c
@@ -93,6 +93,7 @@ static _Atomic(uint8_t) jl_tag_newly_inferred_enabled = 0;
  */
 JL_DLLEXPORT void jl_tag_newly_inferred_enable(void)
 {
+    jl_printf(JL_STDOUT, "Enabling tagging of newly inferred CodeInstances.\n");
     jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, 1);  // FIXME overflow?
 }
 /**
@@ -100,6 +101,7 @@ JL_DLLEXPORT void jl_tag_newly_inferred_enable(void)
  */
 JL_DLLEXPORT void jl_tag_newly_inferred_disable(void)
 {
+    jl_printf(JL_STDOUT, "Disabling tagging of newly inferred CodeInstances.\n");
     jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, -1);  // FIXME underflow?
 }

@@ -119,6 +121,9 @@ JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t* ci)
     if (!newly_inferred)
         return;
     uint8_t tag_newly_inferred = jl_atomic_load_relaxed(&jl_tag_newly_inferred_enabled);
+    jl_printf(JL_STDOUT, "Tagging newly inferred CodeInstance: %d\n",
+             (int)(tag_newly_inferred ? 1 : 0));
+    jl_(ci);
     if (tag_newly_inferred) {
         jl_method_instance_t *mi = jl_get_ci_mi((jl_code_instance_t*)ci);
         uint8_t miflags = jl_atomic_load_relaxed(&mi->flags);

it's quite clear from the log that this isn't happening:

...
Tagging newly inferred CodeInstance: 0
Core.CodeInstance(def=f(Signed) from f(Integer), owner=nothing, next=#<null>, min_world=0x0000000000000001, max_world=0x0000000000009f5e, rettype=Int64, exctype=Union{}, rettype_const=1, inferred=nothing, debuginfo=#<null>, edges=svec(), analysis_results=Base.Compiler.AnalysisResults(result=nothing, next=#<null>), ipo_purity_bits=0x000040e0, time_infer_total=0x0000, time_infer_cache_saved=0x0000, time_infer_self=0x0000, time_compile=0x0000, specsigflags=0x00, precompile=false, invoke=0x00007f8e136ea4a0, specptr=0x0000000000000000)
Tagging newly inferred CodeInstance: 0
Core.CodeInstance(def=callsfrts(Int64) from callsfrts(Any), owner=nothing, next=#<null>, min_world=0x0000000000009f57, max_world=0x0000000000009f5e, rettype=Int64, exctype=Any, rettype_const=1, inferred=nothing, debuginfo=#<null>, edges=svec(), analysis_results=Base.Compiler.AnalysisResults(result=nothing, next=#<null>), ipo_purity_bits=0x000040c1, time_infer_total=0x0000, time_infer_cache_saved=0x0000, time_infer_self=0x0000, time_compile=0x0000, specsigflags=0x00, precompile=false, invoke=0x0000000000000000, specptr=0x0000000000000000)
Enabling tagging of newly inferred CodeInstances.
Disabling tagging of newly inferred CodeInstances.
...

(note that callfrts is the last entry in the workload)

One possible fix would be to have a way to specify that the block should be run in the interpreter until it reaches the @force_compile. I'd be interested in any alternative proposals.

[KristofferC]

The docstring for @time has

In some cases the system will look inside the @time expression and compile some of the called code before execution of the top-level expression begins. When that happens, some compilation time will not
be counted. To include this time you can run @time @eval ....

Could such an @eval "shield" be used here?

[timholy]

I don't think so, because if you use

@setup_workload begin   # ideally this would trigger a `let`
    x = 22
    @compile_workload begin
        foo(x)
    end
end

won't that @eval occur outside the scope defined by @setup_workload?

The only reason that @setup_workload isn't currently implemented with let is because that, too, sometimes seems to trigger "premature" compilation. (You may recall JuliaLang/PrecompileTools.jl#17)

[Keno]

Something like

fw = ()->$(workload)
(PrecompileTools).tag_newly_inferred_enable()
invokelatest(fw)
(PrecompileTools).tag_newly_inferred_disable()

may also do the trick.

[timholy]

It's a good idea, but it won't work for examples like this one which have to be run at toplevel. (IIUC, their goal was to precompile the code that gets run by the macro.)

[KristofferC]

Is there any way to split the setup (tag_newly_inferred_enable) and the workload into different top-level expressions?

if ccall(:jl_generating_output, PrecompileTools.Cint, ()) == 1 && (PrecompileTools).workload_enabled(Main) || (PrecompileTools).verbose[]
    (PrecompileTools).tag_newly_inferred_enable()
end

if ccall(:jl_generating_output, PrecompileTools.Cint, ()) == 1 && (PrecompileTools).workload_enabled(Main) || (PrecompileTools).verbose[]
    try 
        # force_compile
        # workload
    finally
        (PrecompileTools).tag_newly_inferred_disable()
end

[timholy]

I didn't try with an explicit :toplevel Expr, but with a standard quote it seems the interpreter's heuristic analyzes the entire block returned by the macro, thus defeating this strategy. I can play with that more, however.

[vtjnash]

The forced ccall and meta will override other heuristics and prevent this from working as well. Getting rid of both of those should help, and then you can also add Core.@latestworld to forbid inference from continuing past that point in the IR

[vtjnash]

Note also that RelationalAI added a replacement macro @trace_dispatch to avoid the problems with newly-inferred reliability. In theory it is a mostly drop-in replacement macro.

[timholy]

Core.@latestworld is the bomb. @trace_dispatch isn't very useful because it prints rather than returning the MethodInstances.

Seemingly fixed in JuliaLang/PrecompileTools.jl#53, I'll let that close this when it merges.