Address review comments

- Rename `precompiled` field in `jl_method_instance_t` to `flags` and use bit 0 as `precompiled` and bit 1 as `dispatched`. - Use relaxed atomic operations to access `flags`. - Add setting of `dispatched` to `jl_gf_invoke_by_method`.
JuliaLang · Sep 24, 2024 · e3fab0a · e3fab0a
1 parent 8728e7e
commit e3fab0a
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 15 deletions.
diff --git a/src/gf.c b/src/gf.c
@@ -3102,7 +3102,8 @@ static void jl_compile_now(jl_method_instance_t *mi)
 JL_DLLEXPORT void jl_compile_method_instance(jl_method_instance_t *mi, jl_tupletype_t *types, size_t world)
 {
     size_t tworld = jl_typeinf_world;
-    jl_atomic_store_relaxed(&mi->precompiled, 1);
+    uint8_t miflags = jl_atomic_load_relaxed(&mi->flags) | JL_MI_FLAGS_MASK_PRECOMPILED;
+    jl_atomic_store_relaxed(&mi->flags, miflags);
     if (jl_generating_output()) {
         jl_compile_now(mi);
         // In addition to full compilation of the compilation-signature, if `types` is more specific (e.g. due to nospecialize),
@@ -3117,7 +3118,8 @@ JL_DLLEXPORT void jl_compile_method_instance(jl_method_instance_t *mi, jl_tuplet
             types2 = jl_type_intersection_env((jl_value_t*)types, (jl_value_t*)mi->def.method->sig, &tpenv2);
             jl_method_instance_t *mi2 = jl_specializations_get_linfo(mi->def.method, (jl_value_t*)types2, tpenv2);
             JL_GC_POP();
-            jl_atomic_store_relaxed(&mi2->precompiled, 1);
+            miflags = jl_atomic_load_relaxed(&mi2->flags) | JL_MI_FLAGS_MASK_PRECOMPILED;
+            jl_atomic_store_relaxed(&mi2->flags, miflags);
             if (jl_rettype_inferred_native(mi2, world, world) == jl_nothing)
                 (void)jl_type_infer(mi2, world, SOURCE_MODE_NOT_REQUIRED);
             if (jl_typeinf_func && jl_atomic_load_relaxed(&mi->def.method->primary_world) <= tworld) {
@@ -3394,8 +3396,11 @@ JL_DLLEXPORT jl_value_t *jl_apply_generic(jl_value_t *F, jl_value_t **args, uint
                                                      jl_int32hash_fast(jl_return_address()),
                                                      world);
     JL_GC_PROMISE_ROOTED(mfunc);
-    if (!mfunc->dispatched) {
-        mfunc->dispatched = 1;
+    uint8_t miflags = jl_atomic_load_relaxed(&mfunc->flags);
+    uint8_t was_dispatched = miflags & JL_MI_FLAGS_MASK_DISPATCHED;
+    if (!was_dispatched) {
+        miflags |= JL_MI_FLAGS_MASK_DISPATCHED;
+        jl_atomic_store_relaxed(&mfunc->flags, miflags);
         record_dispatch_statement(mfunc);
     }
     return _jl_invoke(F, args, nargs, mfunc, world);
@@ -3504,6 +3509,13 @@ jl_value_t *jl_gf_invoke_by_method(jl_method_t *method, jl_value_t *gf, jl_value
             jl_gc_sync_total_bytes(last_alloc); // discard allocation count from compilation
     }
     JL_GC_PROMISE_ROOTED(mfunc);
+    uint8_t miflags = jl_atomic_load_relaxed(&mfunc->flags);
+    uint8_t was_dispatched = miflags & JL_MI_FLAGS_MASK_DISPATCHED;
+    if (!was_dispatched) {
+        miflags |= JL_MI_FLAGS_MASK_DISPATCHED;
+        jl_atomic_store_relaxed(&mfunc->flags, miflags);
+        record_dispatch_statement(mfunc);
+    }
     size_t world = jl_current_task->world_age;
     return _jl_invoke(gf, args, nargs - 1, mfunc, world);
 }

diff --git a/src/jltypes.c b/src/jltypes.c
@@ -3610,23 +3610,21 @@ void jl_init_types(void) JL_GC_DISABLED
     jl_method_instance_type =
         jl_new_datatype(jl_symbol("MethodInstance"), core,
                         jl_any_type, jl_emptysvec,
-                        jl_perm_symsvec(8,
+                        jl_perm_symsvec(7,
                             "def",
                             "specTypes",
                             "sparam_vals",
                             "backedges",
                             "cache",
                             "cache_with_orig",
-                            "precompiled",
-                            "dispatched"),
-                        jl_svec(8,
+                            "flags"),
+                        jl_svec(7,
                             jl_new_struct(jl_uniontype_type, jl_method_type, jl_module_type),
                             jl_any_type,
                             jl_simplevector_type,
                             jl_array_any_type,
                             jl_any_type/*jl_code_instance_type*/,
                             jl_bool_type,
-                            jl_bool_type,
                             jl_bool_type),
                         jl_emptysvec,
                         0, 1, 3);

diff --git a/src/julia.h b/src/julia.h
@@ -410,9 +410,14 @@ struct _jl_method_instance_t {
     jl_array_t *backedges; // list of method-instances which call this method-instance; `invoke` records (invokesig, caller) pairs
     _Atomic(struct _jl_code_instance_t*) cache;
     uint8_t cache_with_orig; // !cache_with_specTypes
-    _Atomic(uint8_t) precompiled; // true if this instance was generated by an explicit `precompile(...)` call
-    uint8_t dispatched; // true if this instance was dispatched to
+
+    // flags for this method instance
+    //   bit 0: generated by an explicit `precompile(...)`
+    //   bit 1: dispatched
+    _Atomic(uint8_t) flags;
 };
+#define JL_MI_FLAGS_MASK_PRECOMPILED    0x01
+#define JL_MI_FLAGS_MASK_DISPATCHED     0x02
 
 // OpaqueClosure
 typedef struct _jl_opaque_closure_t {

diff --git a/src/method.c b/src/method.c
@@ -629,8 +629,7 @@ JL_DLLEXPORT jl_method_instance_t *jl_new_method_instance_uninit(void)
     mi->backedges = NULL;
     jl_atomic_store_relaxed(&mi->cache, NULL);
     mi->cache_with_orig = 0;
-    jl_atomic_store_relaxed(&mi->precompiled, 0);
-    mi->dispatched = 0;
+    jl_atomic_store_relaxed(&mi->flags, 0);
     return mi;
 }
 

diff --git a/src/staticdata.c b/src/staticdata.c
@@ -1678,7 +1678,7 @@ static void jl_write_values(jl_serializer_state *s) JL_GC_DISABLED
             else if (jl_is_method_instance(v)) {
                 assert(f == s->s);
                 jl_method_instance_t *newmi = (jl_method_instance_t*)&f->buf[reloc_offset];
-                jl_atomic_store_relaxed(&newmi->precompiled, 0);
+                jl_atomic_store_relaxed(&newmi->flags, 0);
             }
             else if (jl_is_code_instance(v)) {
                 assert(f == s->s);

diff --git a/src/staticdata_utils.c b/src/staticdata_utils.c
@@ -159,7 +159,8 @@ static int has_backedge_to_worklist(jl_method_instance_t *mi, htable_t *visited,
     if (jl_is_method(mod))
         mod = ((jl_method_t*)mod)->module;
     assert(jl_is_module(mod));
-    if (jl_atomic_load_relaxed(&mi->precompiled) || !jl_object_in_image((jl_value_t*)mod) || type_in_worklist(mi->specTypes)) {
+    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)) {
         return 1;
     }
     if (!mi->backedges) {