Frozen Function Support #1477

DoeringChristian · 2025-01-27T13:23:09Z

Description

This PR adds TRAVERSE_CB macros to most classes in Mitsuba. This allows the frozen function feature to traverse a scene and enables frozen function support (see here).
Additionally, it expands on #1326 to pass UInt32 type seeds, allowing rendering scenes in frozen functions with different seeds without re-tracing.

The following is a more concise list of changes:

Object is now inheriting from drjit::TraversableBase and implements the traverse_1_cb_ro and traverse_1_cb_rw functions.
Classes that can be added to the registry now use drjit::registry_put instead of jit_registry_put, enforcing that they inherit from drjit::TraversableBase at compile time.
Most classes now implement traverse_1_cb_ro and traverse_cb_rw from the drjit::TraversableBase class using either a combination of the macros DR_TRAVERSE_CB, MI_DECLARE_TRAVERSE_CB, and MI_IMPLEMENT_TRAVERSE_CB.
Implements the latter two macros in Mitsuba.
Makes the optimizer classes traversable by adding a DRJIT_STRUCT property
Adds special traversal functions for Scenes and it's subtypes (GPU and CPU)

Testing

Added tests for shapes, bsdfs, emitters, integrators, optimizers, meidia and samplers as well as tests verifying gradient descent in frozen functions (pose and material parameter optimization).

Checklist

My code follows the style guidelines of this project
My changes generate no new warnings
My code also compiles for cuda_* and llvm_* variants. If you can't test this, please leave below
I have commented my code
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
I cleaned the commit history and removed any "Merge" commits
I give permission that the Mitsuba 3 project may redistribute my contributions under the terms of its license

merlinND

Very nice PR! I like that the tests systematically cover a lot of the plugins.
Given the number of new tests, we may have to check that they don't slow down the CI too much.

Mostly small comments & consistency fixes, the bigger points are about:

How to name the traverse macros
Ways to improve the tests a bit

merlinND · 2025-02-14T10:46:26Z

include/mitsuba/core/field.h

+public:
+    void traverse_1_cb_ro(void *payload,
+                          drjit::detail::traverse_callback_ro fn) const {
+


Suggested change

merlinND · 2025-02-14T10:48:15Z

include/mitsuba/core/fwd.h

@@ -368,6 +368,30 @@ extern "C" {
                            })
 #endif

+#define MI_DECLARE_TRAVERSE_CB()                                           \


Alignment of \ (also below)

merlinND · 2025-02-14T10:49:27Z

include/mitsuba/core/fwd.h

+        void Type<Float, Spectrum>::traverse_1_cb_ro(                          \
+            void *payload, drjit::detail::traverse_callback_ro fn) const {     \
+            if constexpr (!std ::is_same_v<Base, drjit ::TraversableBase>)     \
+                Base ::traverse_1_cb_ro(payload, fn);                          \


Suggested change

Base ::traverse_1_cb_ro(payload, fn); \

Base::traverse_1_cb_ro(payload, fn); \

and also below.

merlinND · 2025-02-14T10:55:18Z

include/mitsuba/render/bsdf.h

@@ -609,6 +609,8 @@ class MI_EXPORT_LIB BSDF : public Object {

    /// Identifier (if available)
    std::string m_id;
+
+    DR_TRAVERSE_CB(Object);


Would it be possible to use Base consistently for most of these?

merlinND · 2025-02-14T12:26:29Z

include/mitsuba/render/scene.h

@@ -636,6 +636,13 @@ class MI_EXPORT_LIB Scene : public Object {
    std::unique_ptr<DiscreteDistribution<Float>> m_silhouette_distr = nullptr;

    bool m_shapes_grad_enabled;
+
+    void traverse_1_cb_ro_cpu(void *payload,


Add a comment to explain why these two are needed on top of the normal MI_DECLARE_TRAVERSE_CB() macro.

merlinND · 2025-02-14T14:20:24Z