From f40a3c7c6435578025f123ebe659a325bfc7a98d Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Wed, 19 Mar 2025 05:36:46 -0500 Subject: [PATCH 1/2] Fix package precompilation (PrecompileTools) With the "classic" inference timing disabled, PrecompileTools and SnoopCompile are non-functional on 1.12 & master. #57074 added timing data to the CodeInstances themselves, which should help restore SnoopCompile. However, without the tree structure provided by the old inference timing system, we still need a way to tag MethodInstances that get inferred inside `PrecompileTools.@compile_workload` blocks. This adds a simple mechanism modeled after `@trace_compile` that switching to tagging MethodInstances in `jl_push_newly_inferred`. https://github.com/JuliaLang/PrecompileTools.jl/pull/47 contains the corresponding changes to PrecompileTools.jl. --- src/julia_internal.h | 3 +++ src/staticdata_utils.c | 23 +++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/julia_internal.h b/src/julia_internal.h index 922af5f20e5b7..c5520c92b1e1d 100644 --- a/src/julia_internal.h +++ b/src/julia_internal.h @@ -1281,6 +1281,9 @@ JL_DLLEXPORT void jl_force_trace_compile_timing_disable(void); JL_DLLEXPORT void jl_force_trace_dispatch_enable(void); JL_DLLEXPORT void jl_force_trace_dispatch_disable(void); +JL_DLLEXPORT void jl_tag_newly_inferred_enable(void); +JL_DLLEXPORT void jl_tag_newly_inferred_disable(void); + uint32_t jl_module_next_counter(jl_module_t *m) JL_NOTSAFEPOINT; jl_tupletype_t *arg_type_tuple(jl_value_t *arg1, jl_value_t **args, size_t nargs); diff --git a/src/staticdata_utils.c b/src/staticdata_utils.c index 163eb1ea9cad5..e9f464b64470e 100644 --- a/src/staticdata_utils.c +++ b/src/staticdata_utils.c @@ -86,6 +86,23 @@ static jl_array_t *newly_inferred JL_GLOBALLY_ROOTED /*FIXME*/; // Mutex for newly_inferred jl_mutex_t newly_inferred_mutex; extern jl_mutex_t world_counter_lock; +static _Atomic(uint8_t) jl_tag_newly_inferred_enabled = 0; + +/** + * @brief Enable tagging of all newly inferred CodeInstances. + */ +JL_DLLEXPORT void jl_tag_newly_inferred_enable(void) +{ + jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, 1); // FIXME overflow? +} +/** + * @brief Disable tagging of all newly inferred CodeInstances. + */ +JL_DLLEXPORT void jl_tag_newly_inferred_disable(void) +{ + jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, -1); // FIXME underflow? +} + // Register array of newly-inferred MethodInstances // This gets called as the first step of Base.include_package_for_output @@ -101,6 +118,12 @@ 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); + 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); + jl_atomic_store_relaxed(&mi->flags, miflags | JL_MI_FLAGS_MASK_PRECOMPILED); + } JL_LOCK(&newly_inferred_mutex); size_t end = jl_array_nrows(newly_inferred); jl_array_grow_end(newly_inferred, 1); From ab492b04837a350bb963a96ec8dd8ef2a7913c58 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Wed, 19 Mar 2025 06:47:59 -0500 Subject: [PATCH 2/2] Add a test --- test/precompile.jl | 40 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/test/precompile.jl b/test/precompile.jl index fb845a4274160..9f1aca55d114e 100644 --- a/test/precompile.jl +++ b/test/precompile.jl @@ -738,7 +738,6 @@ end # method root provenance & external code caching precompile_test_harness("code caching") do dir - Bid = rootid(Base) Cache_module = :Cacheb8321416e8a3e2f1 # Note: calling setindex!(::Dict{K,V}, ::Any, ::K) adds both compression and codegen roots write(joinpath(dir, "$Cache_module.jl"), @@ -1068,6 +1067,45 @@ precompile_test_harness("code caching") do dir end end +precompile_test_harness("precompiletools") do dir + PrecompileToolsModule = :PCTb8321416e8a3e2f1 + write(joinpath(dir, "$PrecompileToolsModule.jl"), + """ + module $PrecompileToolsModule + struct MyType + x::Int + end + + function call_findfirst(x, list) + # call a method defined in Base by runtime dispatch + return findfirst(==(Base.inferencebarrier(x)), Base.inferencebarrier(list)) + end + + let + ccall(:jl_tag_newly_inferred_enable, Cvoid, ()) + call_findfirst(MyType(2), [MyType(1), MyType(2), MyType(3)]) + ccall(:jl_tag_newly_inferred_disable, Cvoid, ()) + end + end + """ + ) + pkgid = Base.PkgId(string(PrecompileToolsModule)) + @test !Base.isprecompiled(pkgid) + Base.compilecache(pkgid) + @test Base.isprecompiled(pkgid) + @eval using $PrecompileToolsModule + M = invokelatest(getfield, @__MODULE__, PrecompileToolsModule) + invokelatest() do + m = which(Tuple{typeof(findfirst), Base.Fix2{typeof(==), T}, Vector{T}} where T) + success = 0 + for mi in Base.specializations(m) + sig = Base.unwrap_unionall(mi.specTypes) + success += sig.parameters[3] === Vector{M.MyType} + end + @test success == 1 + end +end + precompile_test_harness("invoke") do dir InvokeModule = :Invoke0x030e7e97c2365aad CallerModule = :Caller0x030e7e97c2365aad