Fix --image-codegen (#49631)

JuliaLang · May 17, 2023 · 0b599ce · 0b599ce · aviatesk · May 22, 2023
1 parent 10dc33e
commit 0b599ce
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 27 deletions.
diff --git a/src/cgutils.cpp b/src/cgutils.cpp
@@ -333,7 +333,8 @@ static Constant *julia_pgv(jl_codectx_t &ctx, const char *cname, void *addr)
     StringRef localname;
     std::string gvname;
     if (!gv) {
-        raw_string_ostream(gvname) << cname << ctx.global_targets.size();
+        uint64_t id = ctx.emission_context.imaging ? jl_atomic_fetch_add(&globalUniqueGeneratedNames, 1) : ctx.global_targets.size();
+        raw_string_ostream(gvname) << cname << id;
         localname = StringRef(gvname);
     }
     else {

diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -1284,7 +1284,7 @@ static const auto &builtin_func_map() {
 
 static const auto jl_new_opaque_closure_jlcall_func = new JuliaFunction<>{XSTR(jl_new_opaque_closure_jlcall), get_func_sig, get_func_attrs};
 
-static _Atomic(int) globalUniqueGeneratedNames{1};
+static _Atomic(uint64_t) globalUniqueGeneratedNames{1};
 
 // --- code generation ---
 extern "C" {

diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp
@@ -212,35 +212,46 @@ static jl_callptr_t _jl_compile_codeinst(
 
         if (params._shared_module)
             jl_ExecutionEngine->addModule(orc::ThreadSafeModule(std::move(params._shared_module), params.tsctx));
-        StringMap<orc::ThreadSafeModule*> NewExports;
-        StringMap<void*> NewGlobals;
-        for (auto &global : params.globals) {
-            NewGlobals[global.second->getName()] = global.first;
-        }
-        for (auto &def : emitted) {
-            orc::ThreadSafeModule &TSM = std::get<0>(def.second);
-            //The underlying context object is still locked because params is not destroyed yet
-            auto M = TSM.getModuleUnlocked();
-            for (auto &F : M->global_objects()) {
-                if (!F.isDeclaration() && F.getLinkage() == GlobalValue::ExternalLinkage) {
-                    NewExports[F.getName()] = &TSM;
+        if (!params.imaging) {
+            StringMap<orc::ThreadSafeModule*> NewExports;
+            StringMap<void*> NewGlobals;
+            for (auto &global : params.globals) {
+                NewGlobals[global.second->getName()] = global.first;
+            }
+            for (auto &def : emitted) {
+                orc::ThreadSafeModule &TSM = std::get<0>(def.second);
+                //The underlying context object is still locked because params is not destroyed yet
+                auto M = TSM.getModuleUnlocked();
+                for (auto &F : M->global_objects()) {
+                    if (!F.isDeclaration() && F.getLinkage() == GlobalValue::ExternalLinkage) {
+                        NewExports[F.getName()] = &TSM;
+                    }
+                }
+                // Let's link all globals here also (for now)
+                for (auto &GV : M->globals()) {
+                    auto InitValue = NewGlobals.find(GV.getName());
+                    if (InitValue != NewGlobals.end()) {
+                        jl_link_global(&GV, InitValue->second);
+                    }
                 }
             }
-            // Let's link all globals here also (for now)
-            for (auto &GV : M->globals()) {
-                auto InitValue = NewGlobals.find(GV.getName());
-                if (InitValue != NewGlobals.end()) {
-                    jl_link_global(&GV, InitValue->second);
+            DenseMap<orc::ThreadSafeModule*, int> Queued;
+            std::vector<orc::ThreadSafeModule*> Stack;
+            for (auto &def : emitted) {
+                // Add the results to the execution engine now
+                orc::ThreadSafeModule &M = std::get<0>(def.second);
+                jl_add_to_ee(M, NewExports, Queued, Stack);
+                assert(Queued.empty() && Stack.empty() && !M);
+            }
+        } else {
+            jl_jit_globals(params.globals);
+            auto main = std::move(emitted[codeinst].first);
+            for (auto &def : emitted) {
+                if (def.first != codeinst) {
+                    jl_merge_module(main, std::move(def.second.first));
                 }
             }
-        }
-        DenseMap<orc::ThreadSafeModule*, int> Queued;
-        std::vector<orc::ThreadSafeModule*> Stack;
-        for (auto &def : emitted) {
-            // Add the results to the execution engine now
-            orc::ThreadSafeModule &M = std::get<0>(def.second);
-            jl_add_to_ee(M, NewExports, Queued, Stack);
-            assert(Queued.empty() && Stack.empty() && !M);
+            jl_ExecutionEngine->addModule(std::move(main));
         }
         ++CompiledCodeinsts;
         MaxWorkqueueSize.updateMax(emitted.size());

diff --git a/test/llvmpasses/image-codegen.jl b/test/llvmpasses/image-codegen.jl
@@ -0,0 +1,22 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+# RUN: export JULIA_LLVM_ARGS="--print-before=loop-vectorize --print-module-scope"
+# RUN: rm -rf %t
+# RUN: mkdir %t
+# RUN: julia --image-codegen --startup-file=no %s 2> %t/output.txt
+# RUN: FileCheck %s < %t/output.txt
+
+# COM: checks that global variables compiled in imaging codegen
+# COM: are marked as external and not internal
+# COM: Also makes sure that --imaging-codegen doesn't crash
+
+# CHECK: *** IR Dump Before
+# CHECK-NOT: internal global
+# CHECK-NOT: private global
+# CHECK: jl_global
+# CHECK-SAME: = global
+# CHECK: julia_f_
+# CHECK-NOT: internal global
+# CHECK-NOT: private global
+
+f() = "abcd"
+f()