Revenge of the precompile statement generator (JuliaLang#28371)

Third times the charm?

Makefile

@@ -206,7 +206,7 @@ define sysimg_builder
 $$(build_private_libdir)/sys$1-o.a $$(build_private_libdir)/sys$1-bc.a : $$(build_private_libdir)/sys$1-%.a : $$(build_private_libdir)/sys.ji
 	@$$(call PRINT_JULIA, cd $$(JULIAHOME)/base && \
 	if ! $$(call spawn,$3) $2 -C "$$(JULIA_CPU_TARGET)" --output-$$* $$(call cygpath_w,$$@).tmp $$(JULIA_SYSIMG_BUILD_FLAGS) \
-		--startup-file=no --warn-overwrite=yes --sysimage $$(call cygpath_w,$$<) -e nothing; then \
+		--startup-file=no --warn-overwrite=yes --sysimage $$(call cygpath_w,$$<) $$(call cygpath_w,$$(JULIAHOME)/contrib/generate_precompile.jl) $$(call cygpath_w,$$<); then \
 		echo '*** This error is usually fixed by running `make clean`. If the error persists$$(COMMA) try `make cleanall`. ***'; \
 		false; \
 	fi )

base/options.jl

@@ -26,6 +26,7 @@ struct JLOptions
     warn_overwrite::Int8
     can_inline::Int8
     polly::Int8
+    trace_compile::Int8
     fast_math::Int8
     worker::Int8
     cookie::Ptr{UInt8}

base/sysimg.jl

@@ -922,11 +922,10 @@ Base.init_load_path() # want to be able to find external packages in userimg.jl
 
 let
 tot_time_userimg = @elapsed (Base.isfile("userimg.jl") && Base.include(Main, "userimg.jl"))
-tot_time_precompile = Base.is_primary_base_module ? (@elapsed Base.include(Base, "precompile.jl")) : 0.0
 
 
 tot_time_base = (Base.end_base_include - Base.start_base_include) * 10.0^(-9)
-tot_time = tot_time_base + Base.tot_time_stdlib[] + tot_time_userimg + tot_time_precompile
+tot_time = tot_time_base + Base.tot_time_stdlib[] + tot_time_userimg
 
 println("Sysimage built. Summary:")
 print("Total ─────── "); Base.time_print(tot_time               * 10^9); print(" \n");
@@ -935,7 +934,6 @@ print("Stdlibs: ──── "); Base.time_print(Base.tot_time_stdlib[] * 10^9);
 if isfile("userimg.jl")
 print("Userimg: ──── "); Base.time_print(tot_time_userimg       * 10^9); print(" "); show(IOContext(stdout, :compact=>true), (tot_time_userimg       / tot_time) * 100); println("%")
 end
-print("Precompile: ─ "); Base.time_print(tot_time_precompile    * 10^9); print(" "); show(IOContext(stdout, :compact=>true), (tot_time_precompile    / tot_time) * 100); println("%")
 end
 
 empty!(LOAD_PATH)

contrib/add_license_to_files.jl

@@ -30,7 +30,6 @@ const excludedirs = [
 
 const skipfiles = [
     "../contrib/add_license_to_files.jl",
-    "../contrib/fixup_precompile.jl",
     # files to check - already copyright
     # see: https://github.com/JuliaLang/julia/pull/11073#issuecomment-98099389
     "../base/special/trig.jl",

contrib/generate_precompile.jl

@@ -0,0 +1,156 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+# Prevent this from being put into the Main namespace
+let
+M = Module()
+@eval M begin
+if !isdefined(Base, :uv_eventloop)
+    Base.reinit_stdio()
+end
+Base.include(@__MODULE__, joinpath(Sys.BINDIR, "..", "share", "julia", "test", "testhelpers", "FakePTYs.jl"))
+import .FakePTYs: with_fake_pty
+
+CTRL_C = '\x03'
+UP_ARROW = "\e[A"
+DOWN_ARROW = "\e[B"
+
+precompile_script = """
+2+2
+print("")
+@time 1+1
+; pwd
+? reinterpret
+using Ra\t$CTRL_C
+\\alpha\t$CTRL_C
+\e[200~paste here ;)\e[201~"$CTRL_C
+$UP_ARROW$DOWN_ARROW$CTRL_C
+123\b\b\b$CTRL_C
+\b\b$CTRL_C
+f(x) = x03
+f(1,2)
+[][1]
+cd("complet_path\t\t$CTRL_C
+"""
+
+julia_cmd() = (julia = joinpath(Sys.BINDIR, Base.julia_exename()); `$julia`)
+have_repl =  haskey(Base.loaded_modules,
+                    Base.PkgId(Base.UUID("3fa0cd96-eef1-5676-8a61-b3b8758bbffb"), "REPL"))
+have_pkg =  haskey(Base.loaded_modules,
+                    Base.PkgId(Base.UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f"), "Pkg"))
+
+if have_pkg
+    precompile_script *= """
+    tmp = mktempdir()
+    cd(tmp)
+    touch("Project.toml")
+    ] activate .
+    st
+    $CTRL_C
+    rm(tmp; recursive=true)
+    """
+end
+
+function generate_precompile_statements()
+    start_time = time()
+
+    # Precompile a package
+    mktempdir() do prec_path
+        push!(DEPOT_PATH, prec_path)
+        push!(LOAD_PATH, prec_path)
+        pkgname = "__PackagePrecompilationStatementModule"
+        mkpath(joinpath(prec_path, pkgname, "src"))
+        write(joinpath(prec_path, pkgname, "src", "$pkgname.jl"),
+              """
+              module $pkgname
+              end
+              """)
+        @eval using __PackagePrecompilationStatementModule
+        empty!(LOAD_PATH)
+        empty!(DEPOT_PATH)
+    end
+
+    # Create a staging area where all the loaded packages are available
+    PrecompileStagingArea = Module()
+    for (_pkgid, _mod) in Base.loaded_modules
+        if !(_pkgid.name in ("Main", "Core", "Base"))
+            eval(PrecompileStagingArea, :($(Symbol(_mod)) = $_mod))
+        end
+    end
+
+    # TODO: Implement REPL replayer for Windows
+    @static if !Sys.iswindows()
+        print("Generating precompile statements...")
+        sysimg = isempty(ARGS) ? joinpath(dirname(Sys.BINDIR), "lib", "julia", "sys.ji") : ARGS[1]
+
+        # Run a repl process and replay our script
+        stdout_accumulator, stderr_accumulator = IOBuffer(), IOBuffer()
+        with_fake_pty() do slave, master
+            with_fake_pty() do slave_err, master_err
+                done = false
+                withenv("JULIA_HISTORY" => tempname(), "JULIA_PROJECT" => nothing,
+                        "TERM" => "") do
+                    p = run(`$(julia_cmd()) -O0 --trace-compile=yes --sysimage $sysimg
+                                           --startup-file=no --color=yes`,
+                            slave, slave, slave_err; wait=false)
+                    readuntil(master, "julia>", keep=true)
+                    for (tty, accumulator) in (master     => stdout_accumulator,
+                                               master_err => stderr_accumulator)
+                        @async begin
+                            while true
+                                done && break
+                                write(accumulator, readavailable(tty))
+                            end
+                        end
+                    end
+                    if have_repl
+                        for l in split(precompile_script, '\n'; keepempty=false)
+                            write(master, l, '\n')
+                        end
+                    end
+                    write(master, "exit()\n")
+                    wait(p)
+                    done = true
+                end
+            end
+        end
+
+        stderr_output = String(take!(stderr_accumulator))
+        # println(stderr_output)
+        # stdout_output = String(take!(stdout_accumulator))
+        # println(stdout_output)
+
+        # Extract the precompile statements from stderr
+        statements = Set{String}()
+        for statement in split(stderr_output, '\n')
+            m = match(r"(precompile\(Tuple{.*)", statement)
+            m === nothing && continue
+            statement = m.captures[1]
+            occursin(r"Main.", statement) && continue
+            push!(statements, statement)
+        end
+
+        # Load the precompile statements
+        statements_ordered = join(sort(collect(statements)), '\n')
+        # println(statements_ordered)
+        if have_repl
+            # Seems like a reasonable number right now, adjust as needed
+            @assert length(statements) > 700
+        end
+
+        Base.include_string(PrecompileStagingArea, statements_ordered)
+        print(" $(length(statements)) generated in ")
+        Base.time_print((time() - start_time) * 10^9)
+        println()
+    end
+
+    # Fall back to explicit list on Windows, might as well include them
+    # for everyone though
+    Base.include(PrecompileStagingArea, "precompile_explicit.jl")
+
+    return
+end
+
+generate_precompile_statements()
+
+end # @eval
+end # let

base/precompile.jl → contrib/precompile_explicit.jl

@@ -3,20 +3,11 @@
 # Steps to regenerate this file:
 # 1. Remove all `precompile` calls
 # 2. Rebuild system image
-# 3. Enable TRACE_COMPILE in options.h and rebuild
-# 4. Run `./julia 2> precompiles.txt` and do various things.
-# 5. Run `./julia contrib/fixup_precompile.jl precompiles.txt to overwrite `precompile.jl`
-#    or ./julia contrib/fixup_precompile.jl --merge precompiles.txt to merge into existing
-#    `precompile.jl`
+# 3. Start julia with `--trace-compile=yes and do some stuff
+# 5. Run `grep -v '#[0-9]' precompiles.txt >> contrib/precompile_explicit.jl`
+# (filters out closures, which might have different generated names in different environments)
+# This list is only used on Windows, otherwise precompile statements are generated dynamically.
 
-let
-PrecompileStagingArea = Module()
-for (_pkgid, _mod) in Base.loaded_modules
-    if !(_pkgid.name in ("Main", "Core", "Base"))
-        @eval PrecompileStagingArea $(Symbol(_mod)) = $_mod
-    end
-end
-@eval PrecompileStagingArea begin
 precompile(Tuple{Type{Array{Base.StackTraces.StackFrame, 1}}, UndefInitializer, Int64})
 precompile(Tuple{Type{Array{Union{Nothing, String}, 1}}, UndefInitializer, Int64})
 precompile(Tuple{Type{Base.CoreLogging.LogState}, Logging.ConsoleLogger})
@@ -742,5 +733,3 @@ precompile(Tuple{typeof(REPL.setup_interface), REPL.LineEditREPL})
 precompile(Tuple{typeof(REPL.start_repl_backend), Base.Channel{Any}, Base.Channel{Any}})
 precompile(Tuple{typeof(Random.__init__)})
 precompile(Tuple{typeof(eval), Module, Expr})
-end
-end

src/gf.c

@@ -1088,13 +1088,13 @@ static jl_method_instance_t *jl_mt_assoc_by_type(jl_methtable_t *mt, jl_datatype
     if (entry != NULL) {
         jl_method_t *m = entry->func.method;
         if (!jl_has_call_ambiguities((jl_value_t*)tt, m)) {
-#ifdef TRACE_COMPILE
-            if (!jl_has_free_typevars((jl_value_t*)tt)) {
-                jl_printf(JL_STDERR, "precompile(");
-                jl_static_show(JL_STDERR, (jl_value_t*)tt);
-                jl_printf(JL_STDERR, ")\n");
+            if (jl_options.trace_compile) {
+                if (!jl_has_free_typevars((jl_value_t*)tt)) {
+                    jl_printf(JL_STDERR, "precompile(");
+                    jl_static_show(JL_STDERR, (jl_value_t*)tt);
+                    jl_printf(JL_STDERR, ")\n");
+                }
             }
-#endif
             if (!mt_cache) {
                 intptr_t nspec = (mt == jl_type_type_mt ? m->nargs + 1 : mt->max_args + 2);
                 jl_compilation_sig(tt, env, m, nspec, &newparams);

src/jloptions.c

@@ -55,6 +55,11 @@ jl_options_t jl_options = { 0,    // quiet
                             0,    // method overwrite warning
                             1,    // can_inline
                             JL_OPTIONS_POLLY_ON, // polly
+#ifdef TRACE_COMPILE
+                            1, // trace_compile
+#else
+                            0, // trace_compile
+#endif
                             JL_OPTIONS_FAST_MATH_DEFAULT,
                             0,    // worker
                             NULL, // cookie
@@ -159,6 +164,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
            opt_warn_overwrite,
            opt_inline,
            opt_polly,
+           opt_trace_compile,
            opt_math_mode,
            opt_worker,
            opt_bind_to,
@@ -215,6 +221,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
         { "warn-overwrite",  required_argument, 0, opt_warn_overwrite },
         { "inline",          required_argument, 0, opt_inline },
         { "polly",           required_argument, 0, opt_polly },
+        { "trace-compile",   required_argument, 0, opt_trace_compile },
         { "math-mode",       required_argument, 0, opt_math_mode },
         { "handle-signals",  required_argument, 0, opt_handle_signals },
         // hidden command line options
@@ -567,6 +574,15 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
                 jl_errorf("julia: invalid argument to --polly (%s)", optarg);
             }
             break;
+         case opt_trace_compile:
+            if (!strcmp(optarg,"yes"))
+                jl_options.trace_compile = 1;
+            else if (!strcmp(optarg,"no"))
+                jl_options.trace_compile = 0;
+            else {
+                jl_errorf("julia: invalid argument to --trace-compile (%s)", optarg);
+            }
+            break;
         case opt_math_mode:
             if (!strcmp(optarg,"ieee"))
                 jl_options.fast_math = JL_OPTIONS_FAST_MATH_OFF;

src/julia.h

@@ -1799,6 +1799,7 @@ typedef struct {
     int8_t warn_overwrite;
     int8_t can_inline;
     int8_t polly;
+    int8_t trace_compile;
     int8_t fast_math;
     int8_t worker;
     const char *cookie;

stdlib/REPL/src/REPL.jl

@@ -183,7 +183,7 @@ struct REPLBackendRef
     response_channel::Channel
 end
 
-function run_repl(repl::AbstractREPL, consumer::Function = x->nothing)
+function run_repl(repl::AbstractREPL, @nospecialize(consumer = x -> nothing))
     repl_channel = Channel(1)
     response_channel = Channel(1)
     backend = start_repl_backend(repl_channel, response_channel)

test/precompile.jl

@@ -235,7 +235,12 @@ try
                  :Future, :Libdl, :LinearAlgebra, :Logging, :Mmap, :Printf,
                  :Profile, :Random, :Serialization, :SharedArrays, :SparseArrays, :SuiteSparse, :Test,
                  :Unicode, :REPL, :InteractiveUtils, :OldPkg, :Pkg, :LibGit2, :SHA, :UUIDs, :Sockets,
-                 :Statistics, ]))
+                 :Statistics, ]),
+                # Plus precompilation module generated at build time
+                let id = Base.PkgId("__PackagePrecompilationStatementModule")
+                    Dict(id => Base.module_build_id(Base.root_module(id)))
+                end
+           )
         @test discard_module.(deps) == deps1
 
         @test current_task()(0x01, 0x4000, 0x30031234) == 2