[GPUCodegen] Re-infer deleted methods if needed (#35655)

Author: vchuravy

src/aotcompile.cpp

@@ -250,9 +250,12 @@ static void makeSafeName(GlobalObject &G)
 
 
 // takes the running content that has collected in the shadow module and dump it to disk
-// this builds the object file portion of the sysimage files for fast startup
+// this builds the object file portion of the sysimage files for fast startup, and can
+// also be used be extern consumers like GPUCompiler.jl to obtain a module containing
+// all reachable & inferrrable functions. The `policy` flag switches between the defaul
+// mode `0` and the extern mode `1`.
 extern "C" JL_DLLEXPORT
-void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams)
+void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams, int _policy)
 {
     jl_native_code_desc_t *data = new jl_native_code_desc_t;
     jl_codegen_params_t params;
@@ -263,12 +266,17 @@ void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams)
     JL_GC_PUSH1(&src);
     JL_LOCK(&codegen_lock);
 
+    CompilationPolicy policy = (CompilationPolicy) _policy;
+
     // compile all methods for the current world and type-inference world
     size_t compile_for[] = { jl_typeinf_world, jl_world_counter };
     for (int worlds = 0; worlds < 2; worlds++) {
         params.world = compile_for[worlds];
         if (!params.world)
             continue;
+        // Don't emit methods for the typeinf_world with extern policy
+        if (policy == CompilationPolicy::Extern && params.world == jl_typeinf_world)
+            continue;
         size_t i, l;
         for (i = 0, l = jl_array_len(methods); i < l; i++) {
             // each item in this list is either a MethodInstance indicating something
@@ -311,8 +319,9 @@ void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams)
                 }
             }
         }
+
         // finally, make sure all referenced methods also get compiled or fixed up
-        jl_compile_workqueue(emitted, params);
+        jl_compile_workqueue(emitted, params, policy);
     }
     JL_GC_POP();
 

src/codegen.cpp

@@ -6534,8 +6534,10 @@ jl_compile_result_t jl_emit_codeinst(
 
 void jl_compile_workqueue(
     std::map<jl_code_instance_t*, jl_compile_result_t> &emitted,
-    jl_codegen_params_t &params)
+    jl_codegen_params_t &params, CompilationPolicy policy)
 {
+    jl_code_info_t *src = NULL;
+    JL_GC_PUSH1(&src);
     while (!params.workqueue.empty()) {
         jl_code_instance_t *codeinst;
         Function *protodecl;
@@ -6565,7 +6567,17 @@ void jl_compile_workqueue(
                 decls = &std::get<1>(result);
             }
             else {
-                result = jl_emit_codeinst(codeinst, NULL, params);
+                // Reinfer the function. The JIT came along and removed the inferred
+                // method body. See #34993
+                if (policy == CompilationPolicy::Extern &&
+                    codeinst->inferred && codeinst->inferred == jl_nothing) {
+                    src = jl_type_infer(codeinst->def, jl_world_counter, 0);
+                    if (src)
+                        result = jl_emit_code(codeinst->def, src, src->rettype, params);
+                }
+                else {
+                    result = jl_emit_codeinst(codeinst, NULL, params);
+                }
                 if (std::get<0>(result))
                     decls = &std::get<1>(result);
                 else
@@ -6619,6 +6631,7 @@ void jl_compile_workqueue(
             }
         }
     }
+    JL_GC_POP();
 }
 
 

src/jitlayers.cpp

@@ -97,7 +97,7 @@ static jl_callptr_t _jl_compile_codeinst(
         jl_compile_result_t result = jl_emit_codeinst(codeinst, src, params);
         if (std::get<0>(result))
             emitted[codeinst] = std::move(result);
-        jl_compile_workqueue(emitted, params);
+        jl_compile_workqueue(emitted, params, CompilationPolicy::Default);
 
         jl_add_to_ee();
         StringMap<std::unique_ptr<Module>*> NewExports;

src/jitlayers.h

@@ -75,9 +75,15 @@ jl_compile_result_t jl_emit_codeinst(
         jl_code_info_t *src,
         jl_codegen_params_t &params);
 
+enum CompilationPolicy {
+    Default = 0,
+    Extern = 1
+};
+
 void jl_compile_workqueue(
     std::map<jl_code_instance_t*, jl_compile_result_t> &emitted,
-    jl_codegen_params_t &params);
+    jl_codegen_params_t &params,
+    CompilationPolicy policy);
 
 Function *jl_cfunction_object(jl_function_t *f, jl_value_t *rt, jl_tupletype_t *argt,
     jl_codegen_params_t &params);

src/julia_internal.h

@@ -615,7 +615,7 @@ JL_DLLEXPORT jl_value_t *jl_dump_fptr_asm(uint64_t fptr, int raw_mc, const char*
 JL_DLLEXPORT jl_value_t *jl_dump_llvm_asm(void *F, const char* asm_variant, const char *debuginfo);
 JL_DLLEXPORT jl_value_t *jl_dump_function_ir(void *f, char strip_ir_metadata, char dump_module, const char *debuginfo);
 
-void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams);
+void *jl_create_native(jl_array_t *methods, const jl_cgparams_t cgparams, int policy);
 void jl_dump_native(void *native_code,
         const char *bc_fname, const char *unopt_bc_fname, const char *obj_fname, const char *asm_fname,
         const char *sysimg_data, size_t sysimg_len);

src/precompile.c

@@ -405,7 +405,7 @@ void *jl_precompile(int all)
         }
     }
     m = NULL;
-    void *native_code = jl_create_native(m2, jl_default_cgparams);
+    void *native_code = jl_create_native(m2, jl_default_cgparams, 0);
     JL_GC_POP();
     return native_code;
 }