Skip to content

staticdata: Memoize type_in_worklist query #57917

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 31, 2025
Merged

staticdata: Memoize type_in_worklist query #57917

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9d8ce54
staticdata: Memoize `type_in_worklist` query
topolarity Mar 26, 2025
f0b3460
Fix-up init / destroy logic
topolarity Mar 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 50 additions & 26 deletions src/staticdata.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,22 @@ static const size_t WORLD_AGE_REVALIDATION_SENTINEL = 0x1;
JL_DLLEXPORT size_t jl_require_world = ~(size_t)0;
JL_DLLEXPORT _Atomic(size_t) jl_first_image_replacement_world = ~(size_t)0;

// This structure is used to store hash tables for the memoization
// of queries in staticdata.c (currently only `type_in_worklist`).
typedef struct {
htable_t type_in_worklist;
} jl_query_cache;

static void init_query_cache(jl_query_cache *cache)
{
htable_new(&cache->type_in_worklist, 0);
}

static void destroy_query_cache(jl_query_cache *cache)
{
htable_free(&cache->type_in_worklist);
}

#include "staticdata_utils.c"
#include "precompile_utils.c"

Expand Down Expand Up @@ -552,6 +568,7 @@ typedef struct {
jl_array_t *method_roots_list;
htable_t method_roots_index;
uint64_t worklist_key;
jl_query_cache *query_cache;
jl_ptls_t ptls;
jl_image_t *image;
int8_t incremental;
Expand Down Expand Up @@ -675,14 +692,13 @@ static int jl_needs_serialization(jl_serializer_state *s, jl_value_t *v) JL_NOTS
return 1;
}


static int caching_tag(jl_value_t *v) JL_NOTSAFEPOINT
static int caching_tag(jl_value_t *v, jl_query_cache *query_cache) JL_NOTSAFEPOINT
{
if (jl_is_method_instance(v)) {
jl_method_instance_t *mi = (jl_method_instance_t*)v;
jl_value_t *m = mi->def.value;
if (jl_is_method(m) && jl_object_in_image(m))
return 1 + type_in_worklist(mi->specTypes);
return 1 + type_in_worklist(mi->specTypes, query_cache);
}
if (jl_is_binding(v)) {
jl_globalref_t *gr = ((jl_binding_t*)v)->globalref;
Expand All @@ -697,24 +713,24 @@ static int caching_tag(jl_value_t *v) JL_NOTSAFEPOINT
if (jl_is_tuple_type(dt) ? !dt->isconcretetype : dt->hasfreetypevars)
return 0; // aka !is_cacheable from jltypes.c
if (jl_object_in_image((jl_value_t*)dt->name))
return 1 + type_in_worklist(v);
return 1 + type_in_worklist(v, query_cache);
}
jl_value_t *dtv = jl_typeof(v);
if (jl_is_datatype_singleton((jl_datatype_t*)dtv)) {
return 1 - type_in_worklist(dtv); // these are already recached in the datatype in the image
return 1 - type_in_worklist(dtv, query_cache); // these are already recached in the datatype in the image
}
return 0;
}

static int needs_recaching(jl_value_t *v) JL_NOTSAFEPOINT
static int needs_recaching(jl_value_t *v, jl_query_cache *query_cache) JL_NOTSAFEPOINT
{
return caching_tag(v) == 2;
return caching_tag(v, query_cache) == 2;
}

static int needs_uniquing(jl_value_t *v) JL_NOTSAFEPOINT
static int needs_uniquing(jl_value_t *v, jl_query_cache *query_cache) JL_NOTSAFEPOINT
{
assert(!jl_object_in_image(v));
return caching_tag(v) == 1;
return caching_tag(v, query_cache) == 1;
}

static void record_field_change(jl_value_t **addr, jl_value_t *newval) JL_NOTSAFEPOINT
Expand Down Expand Up @@ -840,7 +856,7 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
jl_datatype_t *dt = (jl_datatype_t*)v;
// ensure all type parameters are recached
jl_queue_for_serialization_(s, (jl_value_t*)dt->parameters, 1, 1);
if (jl_is_datatype_singleton(dt) && needs_uniquing(dt->instance)) {
if (jl_is_datatype_singleton(dt) && needs_uniquing(dt->instance, s->query_cache)) {
assert(jl_needs_serialization(s, dt->instance)); // should be true, since we visited dt
// do not visit dt->instance for our template object as it leads to unwanted cycles here
// (it may get serialized from elsewhere though)
Expand All @@ -851,7 +867,7 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
if (s->incremental && jl_is_method_instance(v)) {
jl_method_instance_t *mi = (jl_method_instance_t*)v;
jl_value_t *def = mi->def.value;
if (needs_uniquing(v)) {
if (needs_uniquing(v, s->query_cache)) {
// we only need 3 specific fields of this (the rest are not used)
jl_queue_for_serialization(s, mi->def.value);
jl_queue_for_serialization(s, mi->specTypes);
Expand All @@ -866,7 +882,7 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
record_field_change((jl_value_t**)&mi->cache, NULL);
}
else {
assert(!needs_recaching(v));
assert(!needs_recaching(v, s->query_cache));
}
// n.b. opaque closures cannot be inspected and relied upon like a
// normal method since they can get improperly introduced by generated
Expand All @@ -876,7 +892,7 @@ static void jl_insert_into_serialization_queue(jl_serializer_state *s, jl_value_
// error now.
}
if (s->incremental && jl_is_binding(v)) {
if (needs_uniquing(v)) {
if (needs_uniquing(v, s->query_cache)) {
jl_binding_t *b = (jl_binding_t*)v;
jl_queue_for_serialization(s, b->globalref->mod);
jl_queue_for_serialization(s, b->globalref->name);
Expand Down Expand Up @@ -1100,9 +1116,9 @@ static void jl_queue_for_serialization_(jl_serializer_state *s, jl_value_t *v, i
// Items that require postorder traversal must visit their children prior to insertion into
// the worklist/serialization_order (and also before their first use)
if (s->incremental && !immediate) {
if (jl_is_datatype(t) && needs_uniquing(v))
if (jl_is_datatype(t) && needs_uniquing(v, s->query_cache))
immediate = 1;
if (jl_is_datatype_singleton((jl_datatype_t*)t) && needs_uniquing(v))
if (jl_is_datatype_singleton((jl_datatype_t*)t) && needs_uniquing(v, s->query_cache))
immediate = 1;
}

Expand Down Expand Up @@ -1265,7 +1281,7 @@ static uintptr_t _backref_id(jl_serializer_state *s, jl_value_t *v, jl_array_t *

static void record_uniquing(jl_serializer_state *s, jl_value_t *fld, uintptr_t offset) JL_NOTSAFEPOINT
{
if (s->incremental && jl_needs_serialization(s, fld) && needs_uniquing(fld)) {
if (s->incremental && jl_needs_serialization(s, fld) && needs_uniquing(fld, s->query_cache)) {
if (jl_is_datatype(fld) || jl_is_datatype_singleton((jl_datatype_t*)jl_typeof(fld)))
arraylist_push(&s->uniquing_types, (void*)(uintptr_t)offset);
else if (jl_is_method_instance(fld) || jl_is_binding(fld))
Expand Down Expand Up @@ -1489,7 +1505,7 @@ static void jl_write_values(jl_serializer_state *s) JL_GC_DISABLED
// write header
if (object_id_expected)
write_uint(f, jl_object_id(v));
if (s->incremental && jl_needs_serialization(s, (jl_value_t*)t) && needs_uniquing((jl_value_t*)t))
if (s->incremental && jl_needs_serialization(s, (jl_value_t*)t) && needs_uniquing((jl_value_t*)t, s->query_cache))
arraylist_push(&s->uniquing_types, (void*)(uintptr_t)(ios_pos(f)|1));
if (f == s->const_data)
write_uint(s->const_data, ((uintptr_t)t->smalltag << 4) | GC_OLD_MARKED | GC_IN_IMAGE);
Expand All @@ -1500,7 +1516,7 @@ static void jl_write_values(jl_serializer_state *s) JL_GC_DISABLED
layout_table.items[item] = (void*)(reloc_offset | (f == s->const_data)); // store the inverse mapping of `serialization_order` (`id` => object-as-streampos)

if (s->incremental) {
if (needs_uniquing(v)) {
if (needs_uniquing(v, s->query_cache)) {
if (jl_typetagis(v, jl_binding_type)) {
jl_binding_t *b = (jl_binding_t*)v;
if (b->globalref == NULL)
Expand Down Expand Up @@ -1529,7 +1545,7 @@ static void jl_write_values(jl_serializer_state *s) JL_GC_DISABLED
assert(jl_is_datatype_singleton(t) && "unreachable");
}
}
else if (needs_recaching(v)) {
else if (needs_recaching(v, s->query_cache)) {
arraylist_push(jl_is_datatype(v) ? &s->fixup_types : &s->fixup_objs, (void*)reloc_offset);
}
}
Expand Down Expand Up @@ -1962,7 +1978,7 @@ static void jl_write_values(jl_serializer_state *s) JL_GC_DISABLED
}
}
void *superidx = ptrhash_get(&serialization_order, dt->super);
if (s->incremental && superidx != HT_NOTFOUND && from_seroder_entry(superidx) > item && needs_uniquing((jl_value_t*)dt->super))
if (s->incremental && superidx != HT_NOTFOUND && from_seroder_entry(superidx) > item && needs_uniquing((jl_value_t*)dt->super, s->query_cache))
arraylist_push(&s->uniquing_super, dt->super);
}
else if (jl_is_typename(v)) {
Expand Down Expand Up @@ -2873,13 +2889,14 @@ JL_DLLEXPORT jl_value_t *jl_as_global_root(jl_value_t *val, int insert)
static void jl_prepare_serialization_data(jl_array_t *mod_array, jl_array_t *newly_inferred,
/* outputs */ jl_array_t **extext_methods JL_REQUIRE_ROOTED_SLOT,
jl_array_t **new_ext_cis JL_REQUIRE_ROOTED_SLOT,
jl_array_t **edges JL_REQUIRE_ROOTED_SLOT)
jl_array_t **edges JL_REQUIRE_ROOTED_SLOT,
jl_query_cache *query_cache)
{
// extext_methods: [method1, ...], worklist-owned "extending external" methods added to functions owned by modules outside the worklist
// edges: [caller1, ext_targets, ...] for worklist-owned methods calling external methods

// Save the inferred code from newly inferred, external methods
*new_ext_cis = queue_external_cis(newly_inferred);
*new_ext_cis = queue_external_cis(newly_inferred, query_cache);

// Collect method extensions and edges data
*extext_methods = jl_alloc_vec_any(0);
Expand Down Expand Up @@ -2909,7 +2926,8 @@ static void jl_prepare_serialization_data(jl_array_t *mod_array, jl_array_t *new
// In addition to the system image (where `worklist = NULL`), this can also save incremental images with external linkage
static void jl_save_system_image_to_stream(ios_t *f, jl_array_t *mod_array,
jl_array_t *worklist, jl_array_t *extext_methods,
jl_array_t *new_ext_cis, jl_array_t *edges)
jl_array_t *new_ext_cis, jl_array_t *edges,
jl_query_cache *query_cache)
{
htable_new(&field_replace, 0);
htable_new(&bits_replace, 0);
Expand Down Expand Up @@ -3016,6 +3034,7 @@ static void jl_save_system_image_to_stream(ios_t *f, jl_array_t *mod_array,
ios_mem(&gvar_record, 0);
ios_mem(&fptr_record, 0);
jl_serializer_state s = {0};
s.query_cache = query_cache;
s.incremental = !(worklist == NULL);
s.s = &sysimg;
s.const_data = &const_data;
Expand Down Expand Up @@ -3373,11 +3392,14 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
int64_t datastartpos = 0;
JL_GC_PUSH4(&mod_array, &extext_methods, &new_ext_cis, &edges);

jl_query_cache query_cache;
init_query_cache(&query_cache);

if (worklist) {
mod_array = jl_get_loaded_modules(); // __toplevel__ modules loaded in this session (from Base.loaded_modules_array)
// Generate _native_data`
if (_native_data != NULL) {
jl_prepare_serialization_data(mod_array, newly_inferred, &extext_methods, &new_ext_cis, NULL);
jl_prepare_serialization_data(mod_array, newly_inferred, &extext_methods, &new_ext_cis, NULL, &query_cache);
jl_precompile_toplevel_module = (jl_module_t*)jl_array_ptr_ref(worklist, jl_array_len(worklist)-1);
*_native_data = jl_precompile_worklist(worklist, extext_methods, new_ext_cis);
jl_precompile_toplevel_module = NULL;
Expand Down Expand Up @@ -3408,7 +3430,7 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
assert((ct->reentrant_timing & 0b1110) == 0);
ct->reentrant_timing |= 0b1000;
if (worklist) {
jl_prepare_serialization_data(mod_array, newly_inferred, &extext_methods, &new_ext_cis, &edges);
jl_prepare_serialization_data(mod_array, newly_inferred, &extext_methods, &new_ext_cis, &edges, &query_cache);
if (!emit_split) {
write_int32(f, 0); // No clone_targets
write_padding(f, LLT_ALIGN(ios_pos(f), JL_CACHE_BYTE_ALIGNMENT) - ios_pos(f));
Expand All @@ -3420,7 +3442,7 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
}
if (_native_data != NULL)
native_functions = *_native_data;
jl_save_system_image_to_stream(ff, mod_array, worklist, extext_methods, new_ext_cis, edges);
jl_save_system_image_to_stream(ff, mod_array, worklist, extext_methods, new_ext_cis, edges, &query_cache);
if (_native_data != NULL)
native_functions = NULL;
// make sure we don't run any Julia code concurrently before this point
Expand Down Expand Up @@ -3449,6 +3471,8 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
}
}

destroy_query_cache(&query_cache);

JL_GC_POP();
*s = f;
if (emit_split)
Expand Down
70 changes: 44 additions & 26 deletions src/staticdata_utils.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,63 +131,81 @@ JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t* ci)
JL_UNLOCK(&newly_inferred_mutex);
}


// compute whether a type references something internal to worklist
// and thus could not have existed before deserialize
// and thus does not need delayed unique-ing
static int type_in_worklist(jl_value_t *v) JL_NOTSAFEPOINT
static int type_in_worklist(jl_value_t *v, jl_query_cache *cache) JL_NOTSAFEPOINT
{
if (jl_object_in_image(v))
return 0; // fast-path for rejection

void *cached = HT_NOTFOUND;
if (cache != NULL)
cached = ptrhash_get(&cache->type_in_worklist, v);

// fast-path for memoized results
if (cached != HT_NOTFOUND)
return cached == v;

int result = 0;
if (jl_is_uniontype(v)) {
jl_uniontype_t *u = (jl_uniontype_t*)v;
return type_in_worklist(u->a) ||
type_in_worklist(u->b);
result = type_in_worklist(u->a, cache) ||
type_in_worklist(u->b, cache);
}
else if (jl_is_unionall(v)) {
jl_unionall_t *ua = (jl_unionall_t*)v;
return type_in_worklist((jl_value_t*)ua->var) ||
type_in_worklist(ua->body);
result = type_in_worklist((jl_value_t*)ua->var, cache) ||
type_in_worklist(ua->body, cache);
}
else if (jl_is_typevar(v)) {
jl_tvar_t *tv = (jl_tvar_t*)v;
return type_in_worklist(tv->lb) ||
type_in_worklist(tv->ub);
result = type_in_worklist(tv->lb, cache) ||
type_in_worklist(tv->ub, cache);
}
else if (jl_is_vararg(v)) {
jl_vararg_t *tv = (jl_vararg_t*)v;
if (tv->T && type_in_worklist(tv->T))
return 1;
if (tv->N && type_in_worklist(tv->N))
return 1;
result = ((tv->T && type_in_worklist(tv->T, cache)) ||
(tv->N && type_in_worklist(tv->N, cache)));
}
else if (jl_is_datatype(v)) {
jl_datatype_t *dt = (jl_datatype_t*)v;
if (!jl_object_in_image((jl_value_t*)dt->name))
return 1;
jl_svec_t *tt = dt->parameters;
size_t i, l = jl_svec_len(tt);
for (i = 0; i < l; i++)
if (type_in_worklist(jl_tparam(dt, i)))
return 1;
if (!jl_object_in_image((jl_value_t*)dt->name)) {
result = 1;
}
else {
jl_svec_t *tt = dt->parameters;
size_t i, l = jl_svec_len(tt);
for (i = 0; i < l; i++) {
if (type_in_worklist(jl_tparam(dt, i), cache)) {
result = 1;
break;
}
}
}
}
else {
return type_in_worklist(jl_typeof(v));
return type_in_worklist(jl_typeof(v), cache);
}
return 0;

// Memoize result
if (cache != NULL)
ptrhash_put(&cache->type_in_worklist, (void*)v, result ? (void*)v : NULL);

return result;
}

// When we infer external method instances, ensure they link back to the
// package. Otherwise they might be, e.g., for external macros.
// Implements Tarjan's SCC (strongly connected components) algorithm, simplified to remove the count variable
static int has_backedge_to_worklist(jl_method_instance_t *mi, htable_t *visited, arraylist_t *stack)
static int has_backedge_to_worklist(jl_method_instance_t *mi, htable_t *visited, arraylist_t *stack, jl_query_cache *query_cache)
{
jl_module_t *mod = mi->def.module;
if (jl_is_method(mod))
mod = ((jl_method_t*)mod)->module;
assert(jl_is_module(mod));
uint8_t is_precompiled = jl_atomic_load_relaxed(&mi->flags) & JL_MI_FLAGS_MASK_PRECOMPILED;
if (is_precompiled || !jl_object_in_image((jl_value_t*)mod) || type_in_worklist(mi->specTypes)) {
if (is_precompiled || !jl_object_in_image((jl_value_t*)mod) || type_in_worklist(mi->specTypes, query_cache)) {
return 1;
}
if (!mi->backedges) {
Expand All @@ -211,7 +229,7 @@ static int has_backedge_to_worklist(jl_method_instance_t *mi, htable_t *visited,
jl_code_instance_t *be;
i = get_next_edge(mi->backedges, i, NULL, &be);
JL_GC_PROMISE_ROOTED(be); // get_next_edge propagates the edge for us here
int child_found = has_backedge_to_worklist(jl_get_ci_mi(be), visited, stack);
int child_found = has_backedge_to_worklist(jl_get_ci_mi(be), visited, stack, query_cache);
if (child_found == 1 || child_found == 2) {
// found what we were looking for, so terminate early
found = 1;
Expand Down Expand Up @@ -243,7 +261,7 @@ static int has_backedge_to_worklist(jl_method_instance_t *mi, htable_t *visited,
// from the worklist or explicitly added by a `precompile` statement, and
// (4) are the most recently computed result for that method.
// These will be preserved in the image.
static jl_array_t *queue_external_cis(jl_array_t *list)
static jl_array_t *queue_external_cis(jl_array_t *list, jl_query_cache *query_cache)
{
if (list == NULL)
return NULL;
Expand All @@ -262,7 +280,7 @@ static jl_array_t *queue_external_cis(jl_array_t *list)
jl_method_instance_t *mi = jl_get_ci_mi(ci);
jl_method_t *m = mi->def.method;
if (ci->owner == jl_nothing && jl_atomic_load_relaxed(&ci->inferred) && jl_is_method(m) && jl_object_in_image((jl_value_t*)m->module)) {
int found = has_backedge_to_worklist(mi, &visited, &stack);
int found = has_backedge_to_worklist(mi, &visited, &stack, query_cache);
assert(found == 0 || found == 1 || found == 2);
assert(stack.len == 0);
if (found == 1 && jl_atomic_load_relaxed(&ci->max_world) == ~(size_t)0) {
Expand Down