Use jit_optix_overwrite_sbt to share SBT across two scenes.

include/mitsuba/render/endpoint.h

@@ -245,7 +245,7 @@ class MI_EXPORT_LIB Endpoint : public Object {
      *    A direction sampling record, which specifies the query location.
      *
      * \return
-     *    The incident direct radiance/importance accoated with the sample.
+     *    The incident direct radiance/importance associated with the sample.
      */
     virtual Spectrum
     eval_direction(const Interaction3f &ref,

include/mitsuba/render/mesh.h

@@ -184,6 +184,8 @@ class MI_EXPORT_LIB Mesh : public Shape<Float, Spectrum> {
                                                        uint32_t ray_flags = +RayFlags::All,
                                                        Mask active = true) const override;
 
+    void set_scene(Scene<Float, Spectrum> *scene) { m_scene = scene; }
+
     /** \brief Ray-triangle intersection test
      *
      * Uses the algorithm by Moeller and Trumbore discussed at
@@ -460,6 +462,9 @@ class MI_EXPORT_LIB Mesh : public Shape<Float, Spectrum> {
 
     /// Optional: used in eval_parameterization()
     ref<Scene<Float, Spectrum>> m_parameterization;
+
+    /// Pointer to the scene that owns this mesh
+    Scene<Float, Spectrum>* m_scene = nullptr;
 };
 
 MI_EXTERN_CLASS(Mesh)

include/mitsuba/render/optix_api.h

@@ -49,6 +49,8 @@ using OptixDenoiserStructPtr = void*;
 #define OPTIX_COMPILE_OPTIMIZATION_LEVEL_0       0x2340
 #define OPTIX_COMPILE_DEBUG_LEVEL_NONE           0x2350
 #define OPTIX_COMPILE_DEBUG_LEVEL_MINIMAL        0x2351
+#define OPTIX_COMPILE_DEBUG_LEVEL_MODERATE       0x2353
+#define OPTIX_COMPILE_DEBUG_LEVEL_FULL           0x2352
 
 #define OPTIX_BUILD_FLAG_ALLOW_COMPACTION  2
 #define OPTIX_BUILD_FLAG_PREFER_FAST_TRACE 4
@@ -325,7 +327,7 @@ D(optixDenoiserSetup, OptixDenoiserStructPtr, CUstream, unsigned int,
   unsigned int, CUdeviceptr, size_t, CUdeviceptr, size_t);
 D(optixDenoiserInvoke, OptixDenoiserStructPtr, CUstream,
   const OptixDenoiserParams *, CUdeviceptr, size_t,
-  const OptixDenoiserGuideLayer *, const OptixDenoiserLayer *, unsigned int, 
+  const OptixDenoiserGuideLayer *, const OptixDenoiserLayer *, unsigned int,
   unsigned int, unsigned int, CUdeviceptr, size_t);
 D(optixDenoiserComputeIntensity, OptixDenoiserStructPtr, CUstream,
   const OptixImage2D *inputImage, CUdeviceptr, CUdeviceptr, size_t);

include/mitsuba/render/scene.h

@@ -36,7 +36,7 @@ template <typename Float, typename Spectrum>
 class MI_EXPORT_LIB Scene : public Object {
 public:
     MI_IMPORT_TYPES(BSDF, Emitter, EmitterPtr, Film, Sampler, Shape, ShapePtr,
-                     ShapeGroup, Sensor, Integrator, Medium, MediumPtr)
+                    ShapeGroup, Sensor, Integrator, Medium, MediumPtr, Mesh)
 
     /// Instantiate a scene from a \ref Properties object
     Scene(const Properties &props);

src/render/mesh.cpp

@@ -497,6 +497,10 @@ MI_VARIANT void Mesh<Float, Spectrum>::build_parameterization() {
     mesh->initialize();
 
     props.set_object("mesh", mesh.get());
+
+    if (m_scene)
+        props.set_object("parent_scene", m_scene);
+
     m_parameterization = new Scene<Float, Spectrum>(props);
 }
 

src/render/python/shape_v.cpp

@@ -111,6 +111,9 @@ MI_PY_EXPORT(Shape) {
             "it"_a, "sample"_a, "active"_a = true, D(Shape, sample_direction))
         .def("pdf_direction", &Shape::pdf_direction,
             "it"_a, "ps"_a, "active"_a = true, D(Shape, pdf_direction))
+        .def("eval_parameterization", &Shape::eval_parameterization,
+             "uv"_a, "ray_flags"_a = +RayFlags::All, "active"_a = true,
+             D(Shape, eval_parameterization))
         .def("bbox", py::overload_cast<>(
             &Shape::bbox, py::const_), D(Shape, bbox))
         .def("bbox", py::overload_cast<ScalarUInt32>(

src/render/scene.cpp

@@ -2,6 +2,7 @@
 #include <mitsuba/core/plugin.h>
 #include <mitsuba/render/bsdf.h>
 #include <mitsuba/render/medium.h>
+#include <mitsuba/render/mesh.h>
 #include <mitsuba/render/scene.h>
 #include <mitsuba/render/integrator.h>
 
@@ -23,6 +24,7 @@ MI_VARIANT Scene<Float, Spectrum>::Scene(const Properties &props) {
         m_children.push_back(kv.second.get());
 
         Shape *shape           = dynamic_cast<Shape *>(kv.second.get());
+        Mesh *mesh             = dynamic_cast<Mesh *>(kv.second.get());
         Emitter *emitter       = dynamic_cast<Emitter *>(kv.second.get());
         Sensor *sensor         = dynamic_cast<Sensor *>(kv.second.get());
         Integrator *integrator = dynamic_cast<Integrator *>(kv.second.get());
@@ -38,6 +40,8 @@ MI_VARIANT Scene<Float, Spectrum>::Scene(const Properties &props) {
                 m_bbox.expand(shape->bbox());
                 m_shapes.push_back(shape);
             }
+            if (mesh)
+                mesh->set_scene(this);
         } else if (emitter) {
             // Surface emitters will be added to the list when attached to a shape
             if (!has_flag(emitter->flags(), EmitterFlags::Surface))

src/render/scene_optix.inl

@@ -33,6 +33,7 @@ struct OptixSceneState {
     void* ias_buffer = nullptr;
     size_t config_index;
     uint32_t sbt_jit_index;
+    bool own_sbt;
 };
 
 /**
@@ -176,7 +177,7 @@ size_t init_optix_config(bool has_meshes, bool has_others, bool has_instances) {
             optixModuleGetCompilationState(config.module, &compilation_state));
         if (compilation_state != OPTIX_MODULE_COMPILE_STATE_COMPLETED)
             Throw("Optix configuration initialization failed! The OptiX module "
-                  "compilation did not complete succesfully. The module's "
+                  "compilation did not complete successfully. The module's "
                   "compilation state is: %#06x", compilation_state);
 
         // =====================================================
@@ -227,7 +228,7 @@ size_t init_optix_config(bool has_meshes, bool has_others, bool has_instances) {
     return config_index;
 }
 
-MI_VARIANT void Scene<Float, Spectrum>::accel_init_gpu(const Properties &/*props*/) {
+MI_VARIANT void Scene<Float, Spectrum>::accel_init_gpu(const Properties &props) {
     if constexpr (dr::is_cuda_v<Float>) {
         ScopedPhase phase(ProfilerPhase::InitAccel);
         Log(Info, "Building scene in OptiX ..");
@@ -237,61 +238,112 @@ MI_VARIANT void Scene<Float, Spectrum>::accel_init_gpu(const Properties &/*props
         m_accel = new OptixSceneState();
         OptixSceneState &s = *(OptixSceneState *) m_accel;
 
-        // =====================================================
-        //  Initialize OptiX configuration
-        // =====================================================
+        // Check if another scene was passed to the constructor
+        Scene *other_scene = nullptr;
+        for (auto &[k, v] : props.objects()) {
+            other_scene = dynamic_cast<Scene *>(v.get());
+            if (other_scene)
+                break;
+        }
 
-        bool has_meshes = false;
-        bool has_others = false;
-        bool has_instances = false;
+        /* When another scene is passed via props, the new scene should re-use
+           the same configuration, pipeline and update the shader binding table
+           rather than constructing a new one from scratch. This is necessary for
+           two scenes to be ray traced within the same megakernel. */
+        if (other_scene) {
+            Log(Debug, "Re-use OptiX config, pipeline and update SBT ..");
+            OptixSceneState &s2 = *(OptixSceneState *) other_scene->m_accel;
 
-        for (auto& shape : m_shapes) {
-            has_meshes    |= shape->is_mesh();
-            has_others    |= !shape->is_mesh() && !shape->is_instance();
-            has_instances |= shape->is_instance();
-        }
+            const OptixConfig &config = optix_configs[s2.config_index];
 
-        for (auto& shape : m_shapegroups) {
-            has_meshes |= !shape->has_meshes();
-            has_others |= !shape->has_others();
-        }
+            HitGroupSbtRecord* prev_data
+                = (HitGroupSbtRecord*) jit_malloc_migrate(s2.sbt.hitgroupRecordBase, AllocType::Host, 1);
+            dr::sync_thread();
 
-        s.config_index = init_optix_config(has_meshes, has_others, has_instances);
-        const OptixConfig &config = optix_configs[s.config_index];
+            std::vector<HitGroupSbtRecord> hg_sbts;
+            hg_sbts.assign(prev_data, prev_data + s2.sbt.hitgroupRecordCount);
+            jit_free(prev_data);
 
-        // =====================================================
-        //  Shader Binding Table generation
-        // =====================================================
+            fill_hitgroup_records(m_shapes, hg_sbts, config.program_groups);
+            for (auto& shapegroup: m_shapegroups)
+                shapegroup->optix_fill_hitgroup_records(hg_sbts, config.program_groups);
+
+            size_t shapes_count = hg_sbts.size();
+
+            s2.sbt.hitgroupRecordBase = jit_malloc(
+                AllocType::HostPinned, shapes_count * sizeof(HitGroupSbtRecord));
+            s2.sbt.hitgroupRecordCount = (unsigned int) shapes_count;
 
-        std::vector<HitGroupSbtRecord> hg_sbts;
-        fill_hitgroup_records(m_shapes, hg_sbts, config.program_groups);
-        for (auto& shapegroup: m_shapegroups)
-            shapegroup->optix_fill_hitgroup_records(hg_sbts, config.program_groups);
+            jit_memcpy_async(JitBackend::CUDA, s2.sbt.hitgroupRecordBase, hg_sbts.data(),
+                             shapes_count * sizeof(HitGroupSbtRecord));
 
-        size_t shapes_count = hg_sbts.size();
+            s2.sbt.hitgroupRecordBase =
+                jit_malloc_migrate(s2.sbt.hitgroupRecordBase, AllocType::Device, 1);
 
-        s.sbt.missRecordBase =
-            jit_malloc(AllocType::HostPinned, sizeof(MissSbtRecord));
-        s.sbt.missRecordStrideInBytes = sizeof(MissSbtRecord);
-        s.sbt.missRecordCount = 1;
+            jit_optix_update_sbt(s2.sbt_jit_index, &s2.sbt);
 
-        s.sbt.hitgroupRecordBase = jit_malloc(
-            AllocType::HostPinned, shapes_count * sizeof(HitGroupSbtRecord));
-        s.sbt.hitgroupRecordStrideInBytes = sizeof(HitGroupSbtRecord);
-        s.sbt.hitgroupRecordCount = (unsigned int) shapes_count;
+            memcpy(&s.sbt, &s2.sbt, sizeof(OptixShaderBindingTable));
+            s.sbt_jit_index = s2.sbt_jit_index;
+            s.config_index = s2.config_index;
+            s.own_sbt = false;
+        } else {
+            // =====================================================
+            //  Initialize OptiX configuration
+            // =====================================================
 
-        jit_optix_check(optixSbtRecordPackHeader(config.program_groups[0],
-                                                 s.sbt.missRecordBase));
+            bool has_meshes = false;
+            bool has_others = false;
+            bool has_instances = false;
 
-        jit_memcpy_async(JitBackend::CUDA, s.sbt.hitgroupRecordBase, hg_sbts.data(),
-                         shapes_count * sizeof(HitGroupSbtRecord));
+            for (auto& shape : m_shapes) {
+                has_meshes    |= shape->is_mesh();
+                has_others    |= !shape->is_mesh() && !shape->is_instance();
+                has_instances |= shape->is_instance();
+            }
+
+            for (auto& shape : m_shapegroups) {
+                has_meshes |= !shape->has_meshes();
+                has_others |= !shape->has_others();
+            }
 
-        s.sbt.missRecordBase =
-            jit_malloc_migrate(s.sbt.missRecordBase, AllocType::Device, 1);
-        s.sbt.hitgroupRecordBase =
-            jit_malloc_migrate(s.sbt.hitgroupRecordBase, AllocType::Device, 1);
+            s.config_index = init_optix_config(has_meshes, has_others, has_instances);
+            const OptixConfig &config = optix_configs[s.config_index];
 
-        s.sbt_jit_index = jit_optix_configure_sbt(&s.sbt, config.pipeline_jit_index);
+            // =====================================================
+            //  Shader Binding Table generation
+            // =====================================================
+
+            s.sbt.missRecordBase =
+                jit_malloc(AllocType::HostPinned, sizeof(MissSbtRecord));
+            s.sbt.missRecordStrideInBytes = sizeof(MissSbtRecord);
+            s.sbt.missRecordCount = 1;
+
+            jit_optix_check(optixSbtRecordPackHeader(config.program_groups[0],
+                                                     s.sbt.missRecordBase));
+
+            std::vector<HitGroupSbtRecord> hg_sbts;
+            fill_hitgroup_records(m_shapes, hg_sbts, config.program_groups);
+            for (auto& shapegroup: m_shapegroups)
+                shapegroup->optix_fill_hitgroup_records(hg_sbts, config.program_groups);
+
+            size_t shapes_count = hg_sbts.size();
+
+            s.sbt.hitgroupRecordBase = jit_malloc(
+                AllocType::HostPinned, shapes_count * sizeof(HitGroupSbtRecord));
+            s.sbt.hitgroupRecordStrideInBytes = sizeof(HitGroupSbtRecord);
+            s.sbt.hitgroupRecordCount = (unsigned int) shapes_count;
+
+            jit_memcpy_async(JitBackend::CUDA, s.sbt.hitgroupRecordBase, hg_sbts.data(),
+                             shapes_count * sizeof(HitGroupSbtRecord));
+
+            s.sbt.missRecordBase =
+                jit_malloc_migrate(s.sbt.missRecordBase, AllocType::Device, 1);
+            s.sbt.hitgroupRecordBase =
+                jit_malloc_migrate(s.sbt.hitgroupRecordBase, AllocType::Device, 1);
+
+            s.sbt_jit_index = jit_optix_configure_sbt(&s.sbt, config.pipeline_jit_index);
+            s.own_sbt = true;
+        }
 
         // =====================================================
         //  Acceleration data structure building
@@ -319,8 +371,10 @@ MI_VARIANT void Scene<Float, Spectrum>::accel_parameters_changed_gpu() {
             std::vector<OptixInstance> ias;
             prepare_ias(config.context, m_shapes, 0, s.accel, 0u, ScalarTransform4f(), ias);
 
-            // If there is only a single IAS, no need to build the "master" IAS
-            if (ias.size() == 1) {
+            // If we expect only a single IAS, no need to build the "master" IAS
+            if (config.pipeline_compile_options.traversableGraphFlags == OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS) {
+                if (ias.size() != 1)
+                    Throw("OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS used but found multiple IASs.");
                 s.ias_buffer = nullptr;
                 s.ias_handle = ias[0].traversableHandle;
             } else {
@@ -406,6 +460,8 @@ MI_VARIANT void Scene<Float, Spectrum>::accel_parameters_changed_gpu() {
 
 MI_VARIANT void Scene<Float, Spectrum>::accel_release_gpu() {
     if constexpr (dr::is_cuda_v<Float>) {
+        Log(Debug, "Scene GPU acceleration release ..");
+
         // Ensure all ray tracing kernels are terminated before releasing the scene
         dr::sync_thread();
 
@@ -416,11 +472,13 @@ MI_VARIANT void Scene<Float, Spectrum>::accel_release_gpu() {
 
         OptixSceneState *s = (OptixSceneState *) m_accel;
 
-        /* This will decrease the reference count of the shader binding table
-           JIT variable which might trigger the release of the OptiX SBT if no
-           ray tracing calls are pending. */
-        UInt32 handle = UInt32::steal(s->sbt_jit_index);
-        handle = 0;
+        if (s->own_sbt) {
+            /* This will decrease the reference count of the shader binding table
+            JIT variable which might trigger the release of the OptiX SBT if no
+            ray tracing calls are pending. */
+            UInt32 handle = UInt32::steal(s->sbt_jit_index);
+            handle = 0;
+        }
 
         delete s;
 
@@ -431,6 +489,7 @@ MI_VARIANT void Scene<Float, Spectrum>::accel_release_gpu() {
 MI_VARIANT void Scene<Float, Spectrum>::static_accel_initialization_gpu() { }
 MI_VARIANT void Scene<Float, Spectrum>::static_accel_shutdown_gpu() {
     if constexpr (dr::is_cuda_v<Float>) {
+        Log(Debug, "Scene static GPU acceleration shutdown ..");
         for (size_t j = 0; j < OPTIX_CONFIG_COUNT; j++) {
             OptixConfig &config = optix_configs[j];
             if (config.pipeline_jit_index) {

src/render/tests/test_mesh.py

@@ -220,6 +220,12 @@ def test09_eval_parameterization(variants_all_rgb):
     shape = mi.load_dict({
         "type" : "obj",
         "filename" : "resources/data/common/meshes/rectangle.obj",
+        "emitter": {
+            "type": "area",
+            "radiance": {
+                "type": "checkerboard",
+            }
+        }
     })
 
     si = shape.eval_parameterization([-0.01, 0.5])
@@ -234,6 +240,18 @@ def test09_eval_parameterization(variants_all_rgb):
     assert dr.all(si.is_valid())
     assert dr.allclose(si.p, [-.6, -.4, 0])
 
+    # Test with symbolic virtual function call
+    if not 'scalar' in mi.variant():
+        emitter = shape.emitter()
+        N = 4
+        mask = dr.eq(dr.arange(mi.UInt32, N) & 1, 0)
+        emitters = dr.select(mask, mi.EmitterPtr(emitter), dr.zeros(mi.EmitterPtr))
+        it = dr.zeros(mi.Interaction3f, N)
+        it.p = [0, 0, -3]
+        it.t = 0
+        uv = emitters.sample_direction(it, [0.5, 0.5])[0].uv
+        assert dr.allclose(uv, dr.select(mask, mi.Point2f(0.5), mi.Point2f(0.0)))
+
 
 @fresolver_append_path
 def test10_ray_intersect_preliminary(variants_all_rgb):

src/render/tests/test_renders.py

@@ -289,11 +289,12 @@ def render_ref_images(scenes, spp, overwrite, scene=None, variant=None):
             continue
 
         for variant_ in mi.variants():
-            if not variant.split('_')[0] == 'scalar' or variant_.endswith('double'):
-                continue
+            if variant is not None:
+                if not variant.split('_')[0] == 'scalar' or variant_.endswith('double'):
+                    continue
 
-            if variant is not None and variant != variant_:
-                continue
+                if variant != variant_:
+                    continue
 
             if 'polarized' in variant_ and os.path.split(scene_dir)[1] in POLARIZED_EXCLUDE_FOLDERS:
                 continue