Add MethodInstance roots to resolve compilation time regression

src/ircode.c

@@ -25,6 +25,7 @@ typedef struct {
     ios_t *s;
     // method we're compressing for
     jl_method_t *method;
+    jl_method_instance_t *method_instance;
     jl_ptls_t ptls;
     uint8_t relocatability;
 } jl_ircode_state;
@@ -68,6 +69,12 @@ static void tagged_root(rle_reference *rr, jl_ircode_state *s, int i)
         s->relocatability = 0;
 }
 
+static void tagged_root_method_instance(rle_reference *rr, jl_ircode_state *s, int i)
+{
+    if (!get_root_reference_method_instance(rr, s->method_instance, i))
+        s->relocatability = 0;
+}
+
 static void literal_val_id(rle_reference *rr, jl_ircode_state *s, jl_value_t *v) JL_GC_DISABLED
 {
     jl_array_t *rs = s->method->roots;
@@ -88,6 +95,19 @@ static void literal_val_id(rle_reference *rr, jl_ircode_state *s, jl_value_t *v)
     return tagged_root(rr, s, jl_array_nrows(rs) - 1);
 }
 
+static void literal_val_id_method_instance(rle_reference *rr, jl_ircode_state *s, jl_value_t *v) JL_GC_DISABLED
+{
+    assert(jl_is_method_instance(v));
+    jl_array_t *rs = s->method_instance->roots;
+    int i, l = jl_array_len(rs);
+    for (i = 0; i < l; i++) {
+        if (jl_egal(jl_array_ptr_ref(rs, i), v))
+            return tagged_root_method_instance(rr, s, i);
+    }
+    jl_add_method_instance_root(s->method_instance, jl_precompile_toplevel_module, v);
+    return tagged_root_method_instance(rr, s, jl_array_len(rs) - 1);
+}
+
 static void jl_encode_int32(jl_ircode_state *s, int32_t x)
 {
     if (x >= INT16_MIN && x <= INT16_MAX) {
@@ -124,6 +144,46 @@ static void jl_encode_as_indexed_root(jl_ircode_state *s, jl_value_t *v)
     }
 }
 
+// Encode a method instance for storing with its method instance, as well as the method name
+// as a method root as a fallback for when the method instance is not available. Intended
+// mainly for storage of method instances associated with linetable entries.
+static void jl_encode_as_indexed_root_method_instance(jl_ircode_state *s, jl_value_t *v)
+{
+    assert(jl_is_method_instance(v));
+
+    rle_reference rr;
+
+    literal_val_id_method_instance(&rr, s, v);
+
+    int id = rr.index;
+    assert(id >= 0);
+    if (rr.key) {
+        write_uint8(s->s, TAG_RELOC_METHODINSTROOT);
+        write_uint64(s->s, rr.key);
+    }
+    if (id <= UINT8_MAX) {
+        write_uint8(s->s, TAG_METHODINSTROOT);
+        write_uint8(s->s, id);
+    }
+    else {
+        assert(id <= UINT32_MAX);
+        write_uint8(s->s, TAG_LONG_METHODINSTROOT);
+        write_uint32(s->s, id);
+    }
+
+    jl_method_instance_t *mi = (jl_method_instance_t*)v;
+    jl_value_t* name = NULL;
+    if (jl_is_module(mi->def.value)) {
+        name = (jl_value_t*)mi->def.module->name;
+    }
+    else if (jl_is_method(mi->def.value)) {
+        name = (jl_value_t*)mi->def.method->name;
+    }
+    assert(name != NULL);
+
+    jl_encode_as_indexed_root(s, name);
+}
+
 static void jl_encode_memory_slice(jl_ircode_state *s, jl_genericmemory_t *mem, size_t offset, size_t len) JL_GC_DISABLED
 {
     jl_datatype_t *t = (jl_datatype_t*)jl_typetagof(mem);
@@ -364,8 +424,15 @@ static void jl_encode_value_(jl_ircode_state *s, jl_value_t *v, int as_literal)
     }
     else if (jl_typetagis(v, jl_lineinfonode_type)) {
         write_uint8(s->s, TAG_LINEINFO);
-        for (i = 0; i < jl_datatype_nfields(jl_lineinfonode_type); i++)
-            jl_encode_value(s, jl_get_nth_field(v, i));
+        for (i = 0; i < jl_datatype_nfields(jl_lineinfonode_type); i++) {
+            jl_value_t *vi = jl_get_nth_field(v, i);
+            if (i == 1 && jl_is_method_instance(vi)) {
+                jl_encode_as_indexed_root_method_instance(s, vi);
+            }
+            else {
+                jl_encode_value(s, vi);
+            }
+        }
     }
     else if (((jl_datatype_t*)jl_typeof(v))->instance == v) {
         write_uint8(s->s, TAG_SINGLETON);
@@ -687,6 +754,40 @@ static jl_value_t *jl_decode_value(jl_ircode_state *s) JL_GC_DISABLED
         return lookup_root(s->method, 0, read_uint8(s->s));
     case TAG_LONG_METHODROOT:
         return lookup_root(s->method, 0, read_uint32(s->s));
+    case TAG_RELOC_METHODINSTROOT:
+        key = read_uint64(s->s);
+        tag = read_uint8(s->s);
+        assert(tag == TAG_METHODINSTROOT || tag == TAG_LONG_METHODINSTROOT);
+        index = -1;
+        if (tag == TAG_METHODINSTROOT)
+            index = read_uint8(s->s);
+        else if (tag == TAG_LONG_METHODINSTROOT)
+            index = read_uint32(s->s);
+        assert(index >= 0);
+        jl_value_t *name = jl_decode_value(s);
+        assert(jl_is_symbol(name));
+        if (s->method_instance != NULL)
+            return lookup_root_method_instance(s->method_instance, key, index);
+        else
+            return name;
+    case TAG_METHODINSTROOT:
+        index = read_uint8(s->s);
+        assert(index >= 0);
+        name = jl_decode_value(s);
+        assert(jl_is_symbol(name));
+        if (s->method_instance != NULL)
+            return lookup_root_method_instance(s->method_instance, 0, index);
+        else
+            return name;
+    case TAG_LONG_METHODINSTROOT:
+        index = read_uint32(s->s);
+        assert(index >= 0);
+        name = jl_decode_value(s);
+        assert(jl_is_symbol(name));
+        if (s->method_instance != NULL)
+            return lookup_root_method_instance(s->method_instance, 0, index);
+        else
+            return name;
     case TAG_SVEC: JL_FALLTHROUGH; case TAG_LONG_SVEC:
         return jl_decode_value_svec(s, tag);
     case TAG_COMMONSYM:
@@ -808,6 +909,10 @@ JL_DLLEXPORT jl_string_t *jl_compress_ir(jl_method_t *m, jl_code_info_t *code)
     JL_LOCK(&m->writelock); // protect the roots array (Might GC)
     assert(jl_is_method(m));
     assert(jl_is_code_info(code));
+    jl_method_instance_t *mi = NULL;
+    if (jl_is_method_instance(code->parent)) {
+        mi = code->parent;
+    }
     ios_t dest;
     ios_mem(&dest, 0);
     int en = jl_gc_enable(0); // Might GC
@@ -817,9 +922,17 @@ JL_DLLEXPORT jl_string_t *jl_compress_ir(jl_method_t *m, jl_code_info_t *code)
         m->roots = jl_alloc_vec_any(0);
         jl_gc_wb(m, m->roots);
     }
+    if (jl_is_method_instance(code->parent)) {
+        jl_method_instance_t *mi = code->parent;
+        if (mi->roots == NULL) {
+            mi->roots = jl_alloc_vec_any(0);
+            jl_gc_wb(mi, mi->roots);
+        }
+    }
     jl_ircode_state s = {
         &dest,
         m,
+        mi,
         jl_current_task->ptls,
         1
     };
@@ -891,6 +1004,9 @@ JL_DLLEXPORT jl_string_t *jl_compress_ir(jl_method_t *m, jl_code_info_t *code)
     if (jl_array_nrows(m->roots) == 0) {
         m->roots = NULL;
     }
+    if (mi != NULL && jl_array_len(mi->roots) == 0) {
+        mi->roots = NULL;
+    }
     JL_GC_PUSH1(&v);
     jl_gc_enable(en);
     JL_UNLOCK(&m->writelock); // Might GC
@@ -907,6 +1023,11 @@ JL_DLLEXPORT jl_code_info_t *jl_uncompress_ir(jl_method_t *m, jl_code_instance_t
     JL_LOCK(&m->writelock); // protect the roots array (Might GC)
     assert(jl_is_method(m));
     assert(jl_is_string(data));
+    assert(metadata == NULL || jl_is_code_instance(metadata));
+    jl_method_instance_t *mi = NULL;
+    if (metadata != NULL && jl_is_method_instance(metadata->def)) {
+        mi = metadata->def;
+    }
     size_t i;
     ios_t src;
     ios_mem(&src, 0);
@@ -916,6 +1037,7 @@ JL_DLLEXPORT jl_code_info_t *jl_uncompress_ir(jl_method_t *m, jl_code_instance_t
     jl_ircode_state s = {
         &src,
         m,
+        mi,
         jl_current_task->ptls,
         1
     };

src/jltypes.c

@@ -3265,7 +3265,7 @@ 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(10,
+                        jl_perm_symsvec(13,
                             "def",
                             "specTypes",
                             "sparam_vals",
@@ -3275,8 +3275,11 @@ void jl_init_types(void) JL_GC_DISABLED
                             "cache",
                             "inInference",
                             "cache_with_orig",
-                            "precompiled"),
-                        jl_svec(10,
+                            "precompiled",
+                            "roots", // !const
+                            "root_blocks", // !const
+                            "nroots_sysimg"),
+                        jl_svec(13,
                             jl_new_struct(jl_uniontype_type, jl_method_type, jl_module_type),
                             jl_any_type,
                             jl_simplevector_type,
@@ -3286,7 +3289,10 @@ void jl_init_types(void) JL_GC_DISABLED
                             jl_any_type,
                             jl_bool_type,
                             jl_bool_type,
-                            jl_bool_type),
+                            jl_bool_type,
+                            jl_array_any_type,
+                            jl_array_uint64_type,
+                            jl_int32_type),
                         jl_emptysvec,
                         0, 1, 3);
     // These fields should be constant, but Serialization wants to mutate them in initialization

src/julia.h

@@ -396,6 +396,12 @@ struct _jl_method_instance_t {
     uint8_t inInference; // flags to tell if inference is running on this object
     uint8_t cache_with_orig; // !cache_with_specTypes
     _Atomic(uint8_t) precompiled; // true if this instance was generated by an explicit `precompile(...)` call
+
+    jl_array_t *roots;  // pointers in generated code (shared to reduce memory), or null
+    // Identify roots by module-of-origin. We only track the module for roots added during incremental compilation.
+    // May be NULL if no external roots have been added, otherwise it's a Vector{UInt64}
+    jl_array_t *root_blocks;   // RLE (build_id.lo, offset) pairs (even/odd indexing)
+    int32_t nroots_sysimg;     // # of roots stored in the system image
 };
 
 // OpaqueClosure

src/julia_internal.h

@@ -689,10 +689,14 @@ int set_next_edge(jl_array_t *list, int i, jl_value_t *invokesig, jl_method_inst
 void push_edge(jl_array_t *list, jl_value_t *invokesig, jl_method_instance_t *caller);
 
 JL_DLLEXPORT void jl_add_method_root(jl_method_t *m, jl_module_t *mod, jl_value_t* root);
+JL_DLLEXPORT void jl_add_method_instance_root(jl_method_instance_t *m, jl_module_t *mod, jl_value_t* root);
 void jl_append_method_roots(jl_method_t *m, uint64_t modid, jl_array_t* roots);
 int get_root_reference(rle_reference *rr, jl_method_t *m, size_t i) JL_NOTSAFEPOINT;
+int get_root_reference_method_instance(rle_reference *rr, jl_method_instance_t *m, size_t i) JL_NOTSAFEPOINT;
 jl_value_t *lookup_root(jl_method_t *m, uint64_t key, int index) JL_NOTSAFEPOINT;
+jl_value_t *lookup_root_method_instance(jl_method_instance_t *m, uint64_t key, int index) JL_NOTSAFEPOINT;
 int nroots_with_key(jl_method_t *m, uint64_t key) JL_NOTSAFEPOINT;
+int nroots_with_key_method_instance(jl_method_instance_t *m, uint64_t key) JL_NOTSAFEPOINT;
 
 int jl_valid_type_param(jl_value_t *v);
 

src/method.c

@@ -493,6 +493,9 @@ JL_DLLEXPORT jl_method_instance_t *jl_new_method_instance_uninit(void)
     mi->inInference = 0;
     mi->cache_with_orig = 0;
     jl_atomic_store_relaxed(&mi->precompiled, 0);
+    mi->roots = NULL;
+    mi->root_blocks = NULL;
+    mi->nroots_sysimg = 0;
     return mi;
 }
 
@@ -1225,6 +1228,18 @@ static void prepare_method_for_roots(jl_method_t *m, uint64_t modid)
     }
 }
 
+static void prepare_method_instance_for_roots(jl_method_instance_t *m, uint64_t modid)
+{
+    if (!m->roots) {
+        m->roots = jl_alloc_vec_any(0);
+        jl_gc_wb(m, m->roots);
+    }
+    if (!m->root_blocks && modid != 0) {
+        m->root_blocks = jl_alloc_array_1d(jl_array_uint64_type, 0);
+        jl_gc_wb(m, m->root_blocks);
+    }
+}
+
 // Add a single root with owner `mod` to a method
 JL_DLLEXPORT void jl_add_method_root(jl_method_t *m, jl_module_t *mod, jl_value_t* root)
 {
@@ -1242,6 +1257,23 @@ JL_DLLEXPORT void jl_add_method_root(jl_method_t *m, jl_module_t *mod, jl_value_
     JL_GC_POP();
 }
 
+// Add a single root with owner `mod` to a method instance
+JL_DLLEXPORT void jl_add_method_instance_root(jl_method_instance_t *m, jl_module_t *mod, jl_value_t* root)
+{
+    JL_GC_PUSH2(&m, &root);
+    uint64_t modid = 0;
+    if (mod) {
+        assert(jl_is_module(mod));
+        modid = mod->build_id.lo;
+    }
+    assert(jl_is_method_instance(m));
+    prepare_method_instance_for_roots(m, modid);
+    if (current_root_id(m->root_blocks) != modid)
+        add_root_block(m->root_blocks, modid, jl_array_len(m->roots));
+    jl_array_ptr_1d_push(m->roots, root);
+    JL_GC_POP();
+}
+
 // Add a list of roots with key `modid` to a method
 void jl_append_method_roots(jl_method_t *m, uint64_t modid, jl_array_t* roots)
 {
@@ -1269,6 +1301,19 @@ int get_root_reference(rle_reference *rr, jl_method_t *m, size_t i)
     return i < m->nroots_sysimg;
 }
 
+int get_root_reference_method_instance(rle_reference *rr, jl_method_instance_t *m, size_t i)
+{
+    if (!m->root_blocks) {
+        rr->key = 0;
+        rr->index = i;
+        return i < m->nroots_sysimg;
+    }
+    rle_index_to_reference(rr, i, (uint64_t*)jl_array_data(m->root_blocks), jl_array_len(m->root_blocks), 0);
+    if (rr->key)
+        return 1;
+    return i < m->nroots_sysimg;
+}
+
 // get a root, given its key and index relative to the key
 // this is the relocatable way to get a root from m->roots
 jl_value_t *lookup_root(jl_method_t *m, uint64_t key, int index)
@@ -1282,6 +1327,17 @@ jl_value_t *lookup_root(jl_method_t *m, uint64_t key, int index)
     return jl_array_ptr_ref(m->roots, i);
 }
 
+jl_value_t *lookup_root_method_instance(jl_method_instance_t *m, uint64_t key, int index)
+{
+    if (!m->root_blocks) {
+        assert(key == 0);
+        return jl_array_ptr_ref(m->roots, index);
+    }
+    rle_reference rr = {key, index};
+    size_t i = rle_reference_to_index(&rr, (uint64_t*)jl_array_data(m->root_blocks), jl_array_len(m->root_blocks), 0);
+    return jl_array_ptr_ref(m->roots, i);
+}
+
 // Count the number of roots added by module with id `key`
 int nroots_with_key(jl_method_t *m, uint64_t key)
 {
@@ -1300,6 +1356,23 @@ int nroots_with_key(jl_method_t *m, uint64_t key)
     return nwithkey;
 }
 
+int nroots_with_key_method_instance(jl_method_instance_t *m, uint64_t key)
+{
+    size_t nroots = 0;
+    if (m->roots)
+        nroots = jl_array_len(m->roots);
+    if (!m->root_blocks)
+        return key == 0 ? nroots : 0;
+    uint64_t *rletable = (uint64_t*)jl_array_data(m->root_blocks);
+    size_t j, nblocks2 = jl_array_len(m->root_blocks);
+    int nwithkey = 0;
+    for (j = 0; j < nblocks2; j+=2) {
+        if (rletable[j] == key)
+            nwithkey += (j+3 < nblocks2 ? rletable[j+3] : nroots) - rletable[j+1];
+    }
+    return nwithkey;
+}
+
 #ifdef __cplusplus
 }
 #endif

src/serialize.h

@@ -65,8 +65,11 @@ extern "C" {
 #define TAG_RELOC_METHODROOT   57
 #define TAG_BINDING            58
 #define TAG_MEMORYT            59
+#define TAG_METHODINSTROOT         60
+#define TAG_LONG_METHODINSTROOT    61
+#define TAG_RELOC_METHODINSTROOT   62
 
-#define LAST_TAG 59
+#define LAST_TAG 62
 
 #define write_uint8(s, n) ios_putc((n), (s))
 #define read_uint8(s) ((uint8_t)ios_getc((s)))