Make ShapeGroup traversable and updatable

include/mitsuba/python/docstr.h

@@ -6926,8 +6926,6 @@ static const char *__doc_mitsuba_ShapeGroup_m_shapes = R"doc()doc";
 
 static const char *__doc_mitsuba_ShapeGroup_m_shapes_registry_ids = R"doc()doc";
 
-static const char *__doc_mitsuba_ShapeGroup_optix_accel_ready = R"doc()doc";
-
 static const char *__doc_mitsuba_ShapeGroup_optix_build_gas = R"doc(Build OptiX geometry acceleration structures)doc";
 
 static const char *__doc_mitsuba_ShapeGroup_optix_fill_hitgroup_records = R"doc()doc";

include/mitsuba/render/shape.h

@@ -572,7 +572,7 @@ class MI_EXPORT_LIB Shape : public Object {
     void* m_optix_data_ptr = nullptr;
 #endif
 
-private:
+protected:
     /// True if the shape's geometry has changed
     bool m_dirty = true;
 };

include/mitsuba/render/shapegroup.h

@@ -17,7 +17,7 @@ NAMESPACE_BEGIN(mitsuba)
 template <typename Float, typename Spectrum>
 class MI_EXPORT_LIB ShapeGroup : public Shape<Float, Spectrum> {
 public:
-    MI_IMPORT_BASE(Shape, m_id)
+    MI_IMPORT_BASE(Shape, m_id, m_dirty)
     MI_IMPORT_TYPES(ShapeKDTree, ShapePtr)
 
     using typename Base::ScalarSize;
@@ -50,9 +50,14 @@ class MI_EXPORT_LIB ShapeGroup : public Shape<Float, Spectrum> {
 
     /// Return whether this shapegroup contains triangle mesh shapes
     bool has_meshes() const { return m_has_meshes; }
+
     /// Return whether this shapegroup contains other type of shapes
     bool has_others() const { return m_has_others; }
 
+    void traverse(TraversalCallback *callback) override;
+    void parameters_changed(const std::vector<std::string> &/*keys*/ = {}) override;
+    bool parameters_grad_enabled() const override;
+
     std::string to_string() const override;
 
 #if defined(MI_ENABLE_CUDA)
@@ -65,12 +70,7 @@ class MI_EXPORT_LIB ShapeGroup : public Shape<Float, Spectrum> {
                                      const OptixProgramGroup *program_groups) override;
 
     /// Build OptiX geometry acceleration structures
-    void optix_build_gas(const OptixDeviceContext& context) {
-        if (!optix_accel_ready) {
-            build_gas(context, m_shapes, m_accel);
-            optix_accel_ready = true;
-        }
-    }
+    void optix_build_gas(const OptixDeviceContext& context) override;
 #endif
 
     MI_DECLARE_CLASS()
@@ -84,13 +84,13 @@ class MI_EXPORT_LIB ShapeGroup : public Shape<Float, Spectrum> {
 
 #if defined(MI_ENABLE_EMBREE)
     RTCScene m_embree_scene = nullptr;
+    std::vector<int> m_embree_geometries;
 #else
     ref<ShapeKDTree> m_kdtree;
 #endif
 
 #if defined(MI_ENABLE_CUDA)
     OptixAccelData m_accel;
-    bool optix_accel_ready = false;
     /// OptiX hitgroup sbt offset
     uint32_t m_sbt_offset;
 #endif

src/render/scene.cpp

@@ -301,9 +301,17 @@ MI_VARIANT void Scene<Float, Spectrum>::parameters_changed(const std::vector<std
 
     bool accel_is_dirty = false;
     for (auto &s : m_shapes) {
-        accel_is_dirty = s->dirty();
-        if (accel_is_dirty)
+        if (s->dirty()) {
+            accel_is_dirty = true;
             break;
+        }
+    }
+
+    for (auto &s : m_shapegroups) {
+        if (s->dirty()) {
+            accel_is_dirty = true;
+            break;
+        }
     }
 
     if (accel_is_dirty) {
@@ -352,9 +360,10 @@ MI_VARIANT void Scene<Float, Spectrum>::static_accel_shutdown() {
 }
 
 MI_VARIANT void Scene<Float, Spectrum>::clear_shapes_dirty() {
-    for (auto &s : m_shapes) {
+    for (auto &s : m_shapes)
+        s->m_dirty = false;
+    for (auto &s : m_shapegroups)
         s->m_dirty = false;
-    }
 }
 
 MI_VARIANT void Scene<Float, Spectrum>::static_accel_initialization_cpu() { }

src/render/shapegroup.cpp

@@ -14,7 +14,7 @@ MI_VARIANT ShapeGroup<Float, Spectrum>::ShapeGroup(const Properties &props) {
     m_has_meshes = false;
     m_has_others = false;
 
-    // Add children to the underlying datastructure
+    // Add children to the underlying data structure
     for (auto &kv : props.objects()) {
         const Class *c_class = kv.second->class_();
         if (c_class->name() == "Instance") {
@@ -71,7 +71,7 @@ MI_VARIANT ShapeGroup<Float, Spectrum>::ShapeGroup(const Properties &props) {
 MI_VARIANT ShapeGroup<Float, Spectrum>::~ShapeGroup() {
 #if defined(MI_ENABLE_EMBREE)
     if constexpr (!dr::is_cuda_v<Float>) {
-        // Ensure all raytracing kernels are terminated before releasing the scene
+        // Ensure all ray tracing kernels are terminated before releasing the scene
         if constexpr (dr::is_llvm_v<Float>)
             dr::sync_thread();
 
@@ -80,49 +80,27 @@ MI_VARIANT ShapeGroup<Float, Spectrum>::~ShapeGroup() {
 #endif
 }
 
-#if defined(MI_ENABLE_EMBREE)
-MI_VARIANT RTCGeometry ShapeGroup<Float, Spectrum>::embree_geometry(RTCDevice device) {
-    DRJIT_MARK_USED(device);
+MI_VARIANT void ShapeGroup<Float, Spectrum>::traverse(TraversalCallback *callback) {
+    for (auto s : m_shapes) {
+        std::string id = s->id();
+        if (id.empty() || string::starts_with(id, "_unnamed_"))
+            id = s->class_()->name();
+        callback->put_object(id, s.get(), +ParamFlags::Differentiable);
+    }
+}
+
+MI_VARIANT void ShapeGroup<Float, Spectrum>::parameters_changed(const std::vector<std::string> &/*keys*/) {
     if constexpr (!dr::is_cuda_v<Float>) {
-        // Construct the BVH only once
-        if (m_embree_scene == nullptr) {
-            m_embree_scene = rtcNewScene(device);
-            for (auto shape : m_shapes) {
-                RTCGeometry geom = shape->embree_geometry(device);
-                rtcAttachGeometry(m_embree_scene, geom);
-                rtcReleaseGeometry(geom);
+        for (auto &s : m_shapes) {
+            if (s->dirty()) {
+                m_dirty = true;
+                break;
             }
-
-            // Ensure shape data pointers are finished evaluating before building
-            if constexpr (dr::is_llvm_v<Float>)
-                dr::sync_thread();
-
-            rtcCommitScene(m_embree_scene);
         }
-
-        RTCGeometry instance = rtcNewGeometry(device, RTC_GEOMETRY_TYPE_INSTANCE);
-        rtcSetGeometryInstancedScene(instance, m_embree_scene);
-        return instance;
-    } else {
-        Throw("embree_geometry() should only be called in CPU mode.");
     }
-}
-#else
-MI_VARIANT
-std::tuple<typename ShapeGroup<Float, Spectrum>::ScalarFloat,
-           typename ShapeGroup<Float, Spectrum>::ScalarPoint2f,
-           typename ShapeGroup<Float, Spectrum>::ScalarUInt32,
-           typename ShapeGroup<Float, Spectrum>::ScalarUInt32>
-ShapeGroup<Float, Spectrum>::ray_intersect_preliminary_scalar(const ScalarRay3f &ray) const {
-    auto pi = m_kdtree->template ray_intersect_scalar<false>(ray);
-    return { pi.t, pi.prim_uv, pi.shape_index, pi.prim_index };
+    Base::parameters_changed();
 }
 
-MI_VARIANT
-bool ShapeGroup<Float, Spectrum>::ray_test_scalar(const ScalarRay3f &ray) const {
-    return m_kdtree->template ray_intersect_scalar<true>(ray).is_valid();
-}
-#endif
 
 MI_VARIANT typename ShapeGroup<Float, Spectrum>::SurfaceInteraction3f
 ShapeGroup<Float, Spectrum>::compute_surface_interaction(const Ray3f &ray,
@@ -178,8 +156,80 @@ MI_VARIANT void ShapeGroup<Float, Spectrum>::optix_fill_hitgroup_records(std::ve
     fill_hitgroup_records(m_shapes, hitgroup_records, program_groups);
 }
 
+MI_VARIANT void ShapeGroup<Float, Spectrum>::optix_build_gas(const OptixDeviceContext& context) {
+    if (dirty()) {
+        build_gas(context, m_shapes, m_accel);
+        for (auto &s : m_shapes)
+            s->m_dirty = false;
+    }
+}
 #endif
 
+#if defined(MI_ENABLE_EMBREE)
+MI_VARIANT RTCGeometry ShapeGroup<Float, Spectrum>::embree_geometry(RTCDevice device) {
+    DRJIT_MARK_USED(device);
+    if constexpr (!dr::is_cuda_v<Float>) {
+        if (dirty()) {
+            if (m_embree_scene == nullptr)
+                m_embree_scene = rtcNewScene(device);
+
+            for (int geo : m_embree_geometries)
+                rtcDetachGeometry(m_embree_scene, geo);
+            m_embree_geometries.clear();
+
+            for (Shape *shape : m_shapes) {
+                RTCGeometry geom = shape->embree_geometry(embree_device);
+                m_embree_geometries.push_back(rtcAttachGeometry(m_embree_scene, geom));
+                rtcReleaseGeometry(geom);
+            }
+
+            // Ensure shape data pointers are finished evaluating before building
+            if constexpr (dr::is_llvm_v<Float>)
+                dr::sync_thread();
+
+            rtcCommitScene(m_embree_scene);
+
+            for (auto &s : m_shapes)
+                s->m_dirty = false;
+
+            // This method is called once per instance, hence make sure we only
+            // rebuild the BVH once per update.
+            m_dirty = false;
+        }
+
+        RTCGeometry instance = rtcNewGeometry(device, RTC_GEOMETRY_TYPE_INSTANCE);
+        rtcSetGeometryInstancedScene(instance, m_embree_scene);
+        return instance;
+    } else {
+        Throw("embree_geometry() should only be called in CPU mode.");
+    }
+}
+#endif
+
+#if !defined(MI_ENABLE_EMBREE)
+MI_VARIANT
+std::tuple<typename ShapeGroup<Float, Spectrum>::ScalarFloat,
+           typename ShapeGroup<Float, Spectrum>::ScalarPoint2f,
+           typename ShapeGroup<Float, Spectrum>::ScalarUInt32,
+           typename ShapeGroup<Float, Spectrum>::ScalarUInt32>
+ShapeGroup<Float, Spectrum>::ray_intersect_preliminary_scalar(const ScalarRay3f &ray) const {
+    auto pi = m_kdtree->template ray_intersect_scalar<false>(ray);
+    return { pi.t, pi.prim_uv, pi.shape_index, pi.prim_index };
+}
+
+MI_VARIANT
+bool ShapeGroup<Float, Spectrum>::ray_test_scalar(const ScalarRay3f &ray) const {
+    return m_kdtree->template ray_intersect_scalar<true>(ray).is_valid();
+}
+#endif
+
+MI_VARIANT bool ShapeGroup<Float, Spectrum>::parameters_grad_enabled() const {
+    for (auto s : m_shapes)
+        if (s->parameters_grad_enabled())
+            return true;
+    return false;
+}
+
 MI_VARIANT std::string ShapeGroup<Float, Spectrum>::to_string() const {
     std::ostringstream oss;
         oss << "ShapeGroup[" << std::endl

src/shapes/shapegroup.cpp

@@ -79,7 +79,7 @@ where only a few distinct types of trees have to be kept in memory. An example i
         # Instantiate the shape group without any kind of transformation
         'first_instance': {
             'type': 'instance',
-            'shapegroup: {
+            'shapegroup': {
                 'type': 'ref',
                 'id': 'my_shape_group'
             }
@@ -88,7 +88,7 @@ where only a few distinct types of trees have to be kept in memory. An example i
         # Create instance of the shape group, but rotated, scaled, and translated
         'second_instance': {
             'to_world': mi.ScalarTransform4f.rotate([1, 0, 0], 45).scale([1.5, 1.5, 1.5]).translate([0, 10, 0])
-            'shapegroup: {
+            'shapegroup': {
                 'type': 'ref',
                 'id': 'my_shape_group'
             }