Major Dr.Jit update

This PR updates the revision of Dr.Jit to incorporate a few major
improvements:

- Variables are now consistently ordered, which improves the
  effectiveness of the function de-duplication optimization

- The 4 billion variable limit is gone

- Scene loading can now occur in parallel in JIT modes

For details, see

- mitsuba-renderer/drjit-core#38
- mitsuba-renderer/drjit-core#39
- mitsuba-renderer/drjit-core#40
- mitsuba-renderer/drjit-core#41

src/core/plugin.cpp

@@ -177,7 +177,21 @@ ref<Object> PluginManager::create_object(const Properties &props,
     if (class_->name() == "Scene")
        return class_->construct(props);
 
-    const Class *plugin_class = get_plugin_class(props.plugin_name(), class_->variant());
+    std::string variant = class_->variant();
+    const Class *plugin_class = get_plugin_class(props.plugin_name(), variant);
+
+    /* Construct each plugin in its own scope to isolate them from each other.
+     * This is important when plugins are created in parallel. */
+
+#if defined(MI_ENABLE_LLVM)
+    if (string::starts_with(variant, "llvm_"))
+        jit_new_scope(JitBackend::LLVM);
+#endif
+
+#if defined(MI_ENABLE_CUDA)
+    if (string::starts_with(variant, "cuda_"))
+        jit_new_scope(JitBackend::CUDA);
+#endif
 
     Assert(plugin_class != nullptr);
     ref<Object> object = plugin_class->construct(props);

src/core/xml.cpp

@@ -246,21 +246,28 @@ ColorMode variant_to_color_mode() {
 }
 
 struct XMLParseContext {
+    std::string variant;
+    bool parallel;
+
     std::unordered_map<std::string, XMLObject> instances;
     Transform4f transform;
     size_t id_counter = 0;
     ColorMode color_mode;
-    std::string variant;
-    bool parallel;
+    uint32_t backend = 0;
 
     XMLParseContext(const std::string &variant, bool parallel)
         : variant(variant), parallel(parallel) {
         color_mode = MI_INVOKE_VARIANT(variant, variant_to_color_mode);
+
+#if defined(MI_ENABLE_CUDA) || defined(MI_ENABLE_LLVM)
+        if (string::starts_with(variant, "cuda_"))
+            backend = (uint32_t) JitBackend::CUDA;
+        else if (string::starts_with(variant, "llvm_"))
+            backend = (uint32_t) JitBackend::LLVM;
+#endif
     }
 
-    bool is_cuda() const { return string::starts_with(variant, "cuda_"); }
-    bool is_llvm() const { return string::starts_with(variant, "llvm_"); }
-    bool is_jit()  const { return is_cuda() || is_llvm(); }
+    bool is_jit()  const { return backend != 0; }
 };
 
 /// Helper function to check if attributes are fully specified
@@ -1166,6 +1173,10 @@ static ref<Object> instantiate_top_node(XMLParseContext &ctx, const std::string
     ThreadEnvironment env;
     std::unordered_map<std::string, Task*> task_map;
     instantiate_node(ctx, id, env, task_map, true);
+    #if defined(MI_ENABLE_CUDA) || defined(MI_ENABLE_LLVM)
+    if (ctx.is_jit() && ctx.parallel)
+        jit_new_scope((JitBackend) ctx.backend);
+    #endif
     return ctx.instances.find(id)->second.object;
 }
 

src/mitsuba/mitsuba.cpp

@@ -86,11 +86,6 @@ Usage: mitsuba [options] <One or more scene XML files>
         supported width is legal and causes arithmetic operations to be
         replicated multiple times.
 
-    -P
-        Force parallel scene loading, which is disabled by default
-        in JIT modes since interferes with the ability to reuse
-        cached compiled kernels across separate runs of the renderer.
-
 )";
 }
 
@@ -177,7 +172,6 @@ int main(int argc, char *argv[]) {
 
     // Specialized flags for the JIT compiler
     auto arg_optim_lev = parser.add(StringVec{ "-O" }, true);
-    auto arg_load_par  = parser.add(StringVec{ "-P" });
     auto arg_wavefront = parser.add(StringVec{ "-W" });
     auto arg_source    = parser.add(StringVec{ "-S" });
     auto arg_vec_width = parser.add(StringVec{ "-V" }, true);
@@ -316,8 +310,6 @@ int main(int argc, char *argv[]) {
 
         size_t sensor_i  = (*arg_sensor_i ? arg_sensor_i->as_int() : 0);
 
-        bool parallel_loading = !(llvm || cuda) || (*arg_load_par);
-
         // Append the mitsuba directory to the FileResolver search path list
         ref<Thread> thread = Thread::thread();
         ref<FileResolver> fr = thread->file_resolver();
@@ -362,7 +354,7 @@ int main(int argc, char *argv[]) {
             // Try and parse a scene from the passed file.
             std::vector<ref<Object>> parsed =
                 xml::load_file(arg_extra->as_string(), mode, params,
-                               *arg_update, parallel_loading);
+                               *arg_update, true);
 
             if (parsed.size() != 1)
                 Throw("Root element of the input file is expanded into "