Fix ordering of globals in IR

Author: njroussel

src/call.cpp

@@ -679,7 +679,7 @@ void jitc_call_upload(ThreadState *ts) {
         memset(data, 0, call->offset_size);
 
         for (uint32_t i = 0; i < call->n_inst; ++i) {
-            auto it = globals_map.find(GlobalKey(call->inst_hash[i], call->n_inst != 1));
+            auto it = globals_map.find(GlobalKey(call->inst_hash[i], call->n_inst != 1 ? GlobalType::IndirectCallable : GlobalType::Callable));
             if (it == globals_map.end())
                 jitc_fail("jitc_call_upload(): could not find callable!");
 

src/cuda_eval.cpp

@@ -227,7 +227,7 @@ void jitc_cuda_assemble(ThreadState *ts, ScheduledGroup group,
         put('\n');
         put(globals.get() + it.second.start, it.second.length);
         put('\n');
-        if (!it.first.indirect_callable)
+        if (it.first.type != GlobalType::IndirectCallable)
             continue;
         it.second.callable_index = ctr++;
     }
@@ -245,7 +245,7 @@ void jitc_cuda_assemble(ThreadState *ts, ScheduledGroup group,
         fmt("\n.visible .global .align 8 .u64 callables[$u] = {\n",
             indirect_callable_count_unique);
         for (auto const &it : globals_map) {
-            if (!it.first.indirect_callable)
+            if (it.first.type != GlobalType::IndirectCallable)
                 continue;
 
             fmt("    func_$Q$Q$s\n",
@@ -270,8 +270,10 @@ void jitc_cuda_assemble_func(const CallData *call, uint32_t inst,
                         (flags & (uint32_t) JitFlag::PrintIR);
 
     if (call->n_inst == 1)
-        put(".func");
+        // Marked as weak, in case a forward declaration is assembled after this
+        put(".weak .func");
     else
+        // Marked as globally visible for OptiX
         put(".visible .func");
 
     if (out_size)
@@ -1585,7 +1587,7 @@ void jitc_var_call_assemble_cuda(CallData *call, uint32_t call_reg,
     }
 
     if (call->n_inst == 1) {
-        put("            .func ");
+        put("            .weak .func");
         if (out_size)
             fmt("(.param .align $u .b8 result[$u]) ", out_align, out_size);
         put("func_unique_");

src/cuda_scatter.cpp

@@ -67,7 +67,7 @@ void jitc_cuda_render_scatter_reduce_bfly_32(const char *tp, const char *op,
                                              const char *op_ftz, uint32_t shiftamt) {
     fmt_intrinsic(
         ".func reduce_$s_$s(.param .u64 ptr, .param .$s value) {\n"
-        "    .reg .b32 %active, %index, %mask_lt, %mask_gt, %peers, \n"
+        "    .reg .b32 %active, %index, %mask_lt, %mask_gt, %peers,\n"
         "              %peers_lt, %peers_rev, %rank, %rank_bit, %rank_ballot;\n"
         "    .reg .b64 %ptr, %ptr_shift;\n"
         "    .reg .$s %q0, %q1;\n"
@@ -151,7 +151,7 @@ void jitc_cuda_render_scatter_reduce_bfly_64(const char *tp, const char *op,
                                              const char *op_ftz, uint32_t shiftamt) {
     fmt_intrinsic(
         ".func reduce_$s_$s(.param .u64 ptr, .param .$s value) {\n"
-        "    .reg .b32 %active, %index, %mask_lt, %mask_gt, %peers, \n"
+        "    .reg .b32 %active, %index, %mask_lt, %mask_gt, %peers,\n"
         "              %peers_lt, %peers_rev, %rank, %rank_bit, %rank_ballot;\n"
         "    .reg .b64 %ptr, %ptr_shift;\n"
         "    .reg .b32 %q0l, %q0h, %q1l, %q1h;\n"

src/eval.cpp

@@ -914,7 +914,7 @@ XXH128_hash_t jitc_assemble_func(const CallData *call, uint32_t inst,
         kernel_hash.high64 = 0;
     }
 
-    if (globals_map.emplace(GlobalKey(kernel_hash, call->n_inst != 1),
+    if (globals_map.emplace(GlobalKey(kernel_hash, call->n_inst != 1 ? GlobalType::IndirectCallable : GlobalType::Callable),
                             GlobalValue(globals.size(), kernel_length)).second) {
         // Replace '^'s in 'func_^^^..' or '__direct_callable__^^^..' with hash
         size_t hash_offset = strchr(buffer.get() + kernel_offset, '^') - buffer.get(),
@@ -957,7 +957,7 @@ XXH128_hash_t jitc_assemble_func(const CallData *call, uint32_t inst,
 /// Register a global declaration that will be included in the final program
 void jitc_register_global(const char *str) {
     size_t length = strlen(str);
-    if (globals_map.emplace(GlobalKey(XXH128(str, length, 0), false),
+    if (globals_map.emplace(GlobalKey(XXH128(str, length, 0), GlobalType::Global),
                             GlobalValue(globals.size(), length)).second)
         globals.put(str, length);
 }

src/eval.h

@@ -34,21 +34,25 @@ struct ScheduledGroup {
         : size(size), start(start), end(end) { }
 };
 
+enum class GlobalType : uint32_t {
+    IndirectCallable = 0, // Multi-target vcalls, assembled first
+    Callable = 1,         // Single-target vcalls, assembled second
+    Global = 2            // Other globals (intrinsics, etc.), assembled last
+};
+
 struct GlobalKey {
     XXH128_hash_t hash;
-    bool indirect_callable;
+    GlobalType type;
 
-    GlobalKey(XXH128_hash_t hash, bool callable)
-        : hash(hash), indirect_callable(callable) { }
+    GlobalKey(XXH128_hash_t hash, GlobalType type)
+        : hash(hash), type(type) { }
 
     /* Order so that callables are defined before other globals, but don't use
        the callable ID itself for ordering (it can be non-deterministic in
        programs that use Dr.Jit with parallelization) */
     bool operator<(const GlobalKey &v) const {
-        int callable_key_t =   indirect_callable ? 0 : 1,
-            callable_key_v = v.indirect_callable ? 0 : 1;
-        return std::tie(callable_key_t, hash.high64, hash.low64) <
-               std::tie(callable_key_v, v.hash.high64, v.hash.low64);
+        return std::tie(type, hash.high64, hash.low64) <
+               std::tie(v.type, v.hash.high64, v.hash.low64);
     }
 };
 

src/llvm_eval.cpp

@@ -241,7 +241,7 @@ void jitc_llvm_assemble(ThreadState *ts, ScheduledGroup group) {
         put('\n');
         put(globals.get() + it.second.start, it.second.length);
         put('\n');
-        if (!it.first.indirect_callable)
+        if (it.first.type != GlobalType::IndirectCallable)
             continue;
         it.second.callable_index = 1 + ctr++;
     }

src/llvm_mcjit.cpp

@@ -131,7 +131,7 @@ void jitc_llvm_mcjit_compile(void *llvm_module,
         symbols[symbol_pos++] = resolve("callables");
 
         for (auto const &kv: globals_map) {
-            if (!kv.first.indirect_callable)
+            if (kv.first.type != GlobalType::IndirectCallable)
                 continue;
 
             char name_buf[38];

src/llvm_orcv2.cpp

@@ -122,7 +122,7 @@ void jitc_llvm_orcv2_compile(void *llvm_module,
         symbols[symbol_pos++] = resolve("callables");
 
         for (auto const &kv: globals_map) {
-            if (!kv.first.indirect_callable)
+            if (kv.first.type != GlobalType::IndirectCallable)
                 continue;
 
             char name_buf[38];

src/optix_core.cpp

@@ -324,7 +324,7 @@ bool jitc_optix_compile(ThreadState *ts, const char *buf, size_t buf_size,
     bool continuation_callables = jitc_optix_use_continuation_callables();
 
     for (auto const &it : globals_map) {
-        if (!it.first.indirect_callable)
+        if (it.first.type != GlobalType::IndirectCallable)
             continue;
 
         char *name = (char *) malloc_check(58);