Allow for querying of build_id from objects (#53943)

For GPUCompiler we would like to support a native on disk cache of LLVM
IR.
One of the longstanding issues has been the cache invalidation of such
an on disk cache.

With #52233 we now have an integrated cache for the inference results
and we can rely
on `CodeInstance` to be stable across sessions. Due to #52119 we can
also rely on the
`objectid` to be stable.

My inital thought was to key the native disk cache in GPUCompiler on the
objectid of
the corresponding CodeInstance (+ some compilation parameters).

While discussing this with @rayegun yesterday we noted that having a
CodeInstance with
the same objectid might not be enough provenance. E.g we are not
gurantueed that the
CodeInstance is from the same build artifact and the same precise source
code.

For the package images we are tracking this during loading and validate
all contents
at once, and we keep explicitly track of the provenance chain.

This PR adds a lookup up table where we map from "external_blobs" e.g.
loaded images,
to the corresponding top module of each image, and uses this to
determine the
build_id of the package image.

(cherry picked from commit d47cbf6)

base/loading.jl

@@ -3008,6 +3008,14 @@ function module_build_id(m::Module)
     return (UInt128(hi) << 64) | lo
 end
 
+function object_build_id(obj)
+    mod = ccall(:jl_object_top_module, Any, (Any,), obj)
+    if mod === nothing
+        return nothing
+    end
+    return module_build_id(mod::Module)
+end
+
 function isvalid_cache_header(f::IOStream)
     pkgimage = Ref{UInt8}()
     checksum = ccall(:jl_read_verify_header, UInt64, (Ptr{Cvoid}, Ptr{UInt8}, Ptr{Int64}, Ptr{Int64}), f.ios, pkgimage, Ref{Int64}(), Ref{Int64}()) # returns checksum id or zero

src/init.c

@@ -826,6 +826,7 @@ JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel)
 
     arraylist_new(&jl_linkage_blobs, 0);
     arraylist_new(&jl_image_relocs, 0);
+    arraylist_new(&jl_top_mods, 0);
     arraylist_new(&eytzinger_image_tree, 0);
     arraylist_new(&eytzinger_idxs, 0);
     arraylist_push(&eytzinger_idxs, (void*)0);

src/julia.h

@@ -2049,6 +2049,7 @@ JL_DLLEXPORT void jl_create_system_image(void **, jl_array_t *worklist, bool_t e
 JL_DLLEXPORT void jl_restore_system_image(const char *fname);
 JL_DLLEXPORT void jl_restore_system_image_data(const char *buf, size_t len);
 JL_DLLEXPORT jl_value_t *jl_restore_incremental(const char *fname, jl_array_t *depmods, int complete, const char *pkgimage);
+JL_DLLEXPORT jl_value_t *jl_object_top_module(jl_value_t* v) JL_NOTSAFEPOINT;
 
 JL_DLLEXPORT void jl_set_newly_inferred(jl_value_t *newly_inferred);
 JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t *ci);

src/julia_internal.h

@@ -349,6 +349,7 @@ void print_func_loc(JL_STREAM *s, jl_method_t *m);
 extern jl_array_t *_jl_debug_method_invalidation JL_GLOBALLY_ROOTED;
 JL_DLLEXPORT extern arraylist_t jl_linkage_blobs; // external linkage: sysimg/pkgimages
 JL_DLLEXPORT extern arraylist_t jl_image_relocs;  // external linkage: sysimg/pkgimages
+JL_DLLEXPORT extern arraylist_t jl_top_mods;  // external linkage: sysimg/pkgimages
 extern arraylist_t eytzinger_image_tree;
 extern arraylist_t eytzinger_idxs;
 
@@ -1011,7 +1012,8 @@ STATIC_INLINE size_t n_linkage_blobs(void) JL_NOTSAFEPOINT
 
 size_t external_blob_index(jl_value_t *v) JL_NOTSAFEPOINT;
 
-uint8_t jl_object_in_image(jl_value_t* v) JL_NOTSAFEPOINT;
+// Query if this object is perm-allocated in an image.
+JL_DLLEXPORT uint8_t jl_object_in_image(jl_value_t* v) JL_NOTSAFEPOINT;
 
 // the first argument to jl_idtable_rehash is used to return a value
 // make sure it is rooted if it is used after the function returns

src/staticdata.c

@@ -336,6 +336,8 @@ static arraylist_t object_worklist;  // used to mimic recursion by jl_serialize_
 // jl_linkage_blobs.items[2i:2i+1] correspond to build_ids[i]   (0-offset indexing)
 arraylist_t jl_linkage_blobs;
 arraylist_t jl_image_relocs;
+// Keep track of which image corresponds to which top module.
+arraylist_t jl_top_mods;
 
 // Eytzinger tree of images. Used for very fast jl_object_in_image queries
 // See https://algorithmica.org/en/eytzinger
@@ -450,11 +452,23 @@ size_t external_blob_index(jl_value_t *v) JL_NOTSAFEPOINT
     return idx;
 }
 
-uint8_t jl_object_in_image(jl_value_t *obj) JL_NOTSAFEPOINT
+JL_DLLEXPORT uint8_t jl_object_in_image(jl_value_t *obj) JL_NOTSAFEPOINT
 {
     return eyt_obj_in_img(obj);
 }
 
+// Map an object to it's "owning" top module
+JL_DLLEXPORT jl_value_t *jl_object_top_module(jl_value_t* v) JL_NOTSAFEPOINT
+{
+    size_t idx = external_blob_index(v);
+    size_t lbids = n_linkage_blobs();
+    if (idx < lbids) {
+        return (jl_value_t*)jl_top_mods.items[idx];
+    }
+    // The object is runtime allocated
+    return (jl_value_t*)jl_nothing;
+}
+
 // hash of definitions for predefined function pointers
 static htable_t fptr_to_id;
 void *native_functions;   // opaque jl_native_code_desc_t blob used for fetching data from LLVM
@@ -3554,6 +3568,15 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
     arraylist_push(&jl_linkage_blobs, (void*)image_base);
     arraylist_push(&jl_linkage_blobs, (void*)(image_base + sizeof_sysimg));
     arraylist_push(&jl_image_relocs, (void*)relocs_base);
+    if (restored == NULL) {
+        arraylist_push(&jl_top_mods, (void*)jl_top_module);
+    } else {
+        size_t len = jl_array_nrows(*restored);
+        assert(len > 0);
+        jl_module_t *topmod = (jl_module_t*)jl_array_ptr_ref(*restored, len-1);
+        assert(jl_is_module(topmod));
+        arraylist_push(&jl_top_mods, (void*)topmod);
+    }
     jl_timing_counter_inc(JL_TIMING_COUNTER_ImageSize, sizeof_sysimg + sizeof(uintptr_t));
     rebuild_image_blob_tree();
 

test/precompile.jl

@@ -47,6 +47,11 @@ function precompile_test_harness(@nospecialize(f), separate::Bool)
     nothing
 end
 
+@testset "object_build_id" begin
+    @test Base.object_build_id([1]) === nothing
+    @test Base.object_build_id(Base) == Base.module_build_id(Base)
+end
+
 # method root provenance
 
 rootid(m::Module) = Base.module_build_id(Base.parentmodule(m)) % UInt64
@@ -382,6 +387,9 @@ precompile_test_harness(false) do dir
         @test objectid(Foo.a_vec_int) === Foo.oid_vec_int
         @test objectid(Foo.a_mat_int) === Foo.oid_mat_int
         @test Foo.oid_vec_int !== Foo.oid_mat_int
+        @test Base.object_build_id(Foo.a_vec_int) == Base.object_build_id(Foo.a_mat_int)
+        @test Base.object_build_id(Foo) == Base.module_build_id(Foo)
+        @test Base.object_build_id(Foo.a_vec_int) == Base.module_build_id(Foo)
     end
 
     @eval begin function ccallable_test()
@@ -1751,10 +1759,14 @@ let newinterp_path = abspath("compiler/newinterp.jl")
                 @test isdefined(ci, :next)
                 @test ci.owner === nothing
                 @test ci.max_world == typemax(UInt)
+                @test Base.module_build_id(CustomAbstractInterpreterCaching) ==
+                    Base.object_build_id(ci)
                 ci = ci.next
                 @test !isdefined(ci, :next)
                 @test ci.owner === cache_owner
                 @test ci.max_world == typemax(UInt)
+                @test Base.module_build_id(CustomAbstractInterpreterCaching) ==
+                    Base.object_build_id(ci)
             end
             let m = only(methods(sum, (Vector{Float64},)))
                 found = false
@@ -1764,10 +1776,14 @@ let newinterp_path = abspath("compiler/newinterp.jl")
                         @test isdefined(ci, :next)
                         @test ci.owner === cache_owner
                         @test ci.max_world == typemax(UInt)
+                        @test Base.module_build_id(CustomAbstractInterpreterCaching) ==
+                            Base.object_build_id(ci)
                         ci = ci.next
                         @test !isdefined(ci, :next)
                         @test ci.owner === nothing
                         @test ci.max_world == typemax(UInt)
+                        @test Base.module_build_id(CustomAbstractInterpreterCaching) ==
+                            Base.object_build_id(ci)
                         found = true
                         break
                     end