Skip to content

Commit c89b1ff

Browse files
authored
Fix package precompilation (PrecompileTools) (#57828)
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 switches to tagging MethodInstances in `jl_push_newly_inferred`. JuliaLang/PrecompileTools.jl#47 contains (or will contain) the corresponding changes to PrecompileTools.jl.
1 parent f8eaeb4 commit c89b1ff

File tree

3 files changed

+65
-1
lines changed

3 files changed

+65
-1
lines changed

src/julia_internal.h

+3
Original file line numberDiff line numberDiff line change
@@ -1280,6 +1280,9 @@ JL_DLLEXPORT void jl_force_trace_compile_timing_disable(void);
12801280
JL_DLLEXPORT void jl_force_trace_dispatch_enable(void);
12811281
JL_DLLEXPORT void jl_force_trace_dispatch_disable(void);
12821282

1283+
JL_DLLEXPORT void jl_tag_newly_inferred_enable(void);
1284+
JL_DLLEXPORT void jl_tag_newly_inferred_disable(void);
1285+
12831286
uint32_t jl_module_next_counter(jl_module_t *m) JL_NOTSAFEPOINT;
12841287
jl_tupletype_t *arg_type_tuple(jl_value_t *arg1, jl_value_t **args, size_t nargs);
12851288

src/staticdata_utils.c

+23
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,23 @@ static jl_array_t *newly_inferred JL_GLOBALLY_ROOTED /*FIXME*/;
8686
// Mutex for newly_inferred
8787
jl_mutex_t newly_inferred_mutex;
8888
extern jl_mutex_t world_counter_lock;
89+
static _Atomic(uint8_t) jl_tag_newly_inferred_enabled = 0;
90+
91+
/**
92+
* @brief Enable tagging of all newly inferred CodeInstances.
93+
*/
94+
JL_DLLEXPORT void jl_tag_newly_inferred_enable(void)
95+
{
96+
jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, 1); // FIXME overflow?
97+
}
98+
/**
99+
* @brief Disable tagging of all newly inferred CodeInstances.
100+
*/
101+
JL_DLLEXPORT void jl_tag_newly_inferred_disable(void)
102+
{
103+
jl_atomic_fetch_add(&jl_tag_newly_inferred_enabled, -1); // FIXME underflow?
104+
}
105+
89106

90107
// Register array of newly-inferred MethodInstances
91108
// 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)
101118
{
102119
if (!newly_inferred)
103120
return;
121+
uint8_t tag_newly_inferred = jl_atomic_load_relaxed(&jl_tag_newly_inferred_enabled);
122+
if (tag_newly_inferred) {
123+
jl_method_instance_t *mi = jl_get_ci_mi((jl_code_instance_t*)ci);
124+
uint8_t miflags = jl_atomic_load_relaxed(&mi->flags);
125+
jl_atomic_store_relaxed(&mi->flags, miflags | JL_MI_FLAGS_MASK_PRECOMPILED);
126+
}
104127
JL_LOCK(&newly_inferred_mutex);
105128
size_t end = jl_array_nrows(newly_inferred);
106129
jl_array_grow_end(newly_inferred, 1);

test/precompile.jl

+39-1
Original file line numberDiff line numberDiff line change
@@ -738,7 +738,6 @@ end
738738

739739
# method root provenance & external code caching
740740
precompile_test_harness("code caching") do dir
741-
Bid = rootid(Base)
742741
Cache_module = :Cacheb8321416e8a3e2f1
743742
# Note: calling setindex!(::Dict{K,V}, ::Any, ::K) adds both compression and codegen roots
744743
write(joinpath(dir, "$Cache_module.jl"),
@@ -1068,6 +1067,45 @@ precompile_test_harness("code caching") do dir
10681067
end
10691068
end
10701069

1070+
precompile_test_harness("precompiletools") do dir
1071+
PrecompileToolsModule = :PCTb8321416e8a3e2f1
1072+
write(joinpath(dir, "$PrecompileToolsModule.jl"),
1073+
"""
1074+
module $PrecompileToolsModule
1075+
struct MyType
1076+
x::Int
1077+
end
1078+
1079+
function call_findfirst(x, list)
1080+
# call a method defined in Base by runtime dispatch
1081+
return findfirst(==(Base.inferencebarrier(x)), Base.inferencebarrier(list))
1082+
end
1083+
1084+
let
1085+
ccall(:jl_tag_newly_inferred_enable, Cvoid, ())
1086+
call_findfirst(MyType(2), [MyType(1), MyType(2), MyType(3)])
1087+
ccall(:jl_tag_newly_inferred_disable, Cvoid, ())
1088+
end
1089+
end
1090+
"""
1091+
)
1092+
pkgid = Base.PkgId(string(PrecompileToolsModule))
1093+
@test !Base.isprecompiled(pkgid)
1094+
Base.compilecache(pkgid)
1095+
@test Base.isprecompiled(pkgid)
1096+
@eval using $PrecompileToolsModule
1097+
M = invokelatest(getfield, @__MODULE__, PrecompileToolsModule)
1098+
invokelatest() do
1099+
m = which(Tuple{typeof(findfirst), Base.Fix2{typeof(==), T}, Vector{T}} where T)
1100+
success = 0
1101+
for mi in Base.specializations(m)
1102+
sig = Base.unwrap_unionall(mi.specTypes)
1103+
success += sig.parameters[3] === Vector{M.MyType}
1104+
end
1105+
@test success == 1
1106+
end
1107+
end
1108+
10711109
precompile_test_harness("invoke") do dir
10721110
InvokeModule = :Invoke0x030e7e97c2365aad
10731111
CallerModule = :Caller0x030e7e97c2365aad

0 commit comments

Comments
 (0)