Added TraversableBase class for c++ type traversal

Improved traversal of c++ iterators and references

Using TraversableBase for Texture traversal

Added traversal helpers for trampoline classes

Fixed double definition of traverse_py_cb_* methods

Changed traverse function, allowing assigment of c++ objects

Added traversal tests

Added ad_scope_postponed function

Added freezing feature

Added freezing tests

Improved logging

Cleanup of freeze.cpp

Added declarations in freeze.h

Reordered traversal functions

Cleanup and added clear function

Test ordering

Removed outdated fix comment

Improved exceptions

Minor performance improvements

Using nb::hash for key hashing

Using the FreezingScope flag for passthrough recording

Removed kw_only attribute from dataclasses in freeze tests

Changed traversable base casting

Wrapped DRJIT_EXPORT in ifndef

Added warning for traversable base casting

Temporary fix for registry traversal

First freeze documentation attempt

Improved registry traversal and insertion

Cleanup freezing tests

Updated drjit-core submodule

Enabled retreiving the variant of c++ objects

Added variance mismatch test

Attempt at variant traversal in traverse_cb_ro

Added registry domain traversal

Moved declaration of `freeze` to python and added `traverse_callback_*` aliases

Formatting and removed print in tests

Removed default impls for Variant and Domain

Fixed leak of FrozenFunctions by custom currying class

Implemented domain traversal

Added comments for frozen function tests

Removed prints from tests

Added comments for frozen function tests

Fixed typo

Fixed traverse test

Remove b.value()

Formatting

Added comments for traversable_base.h and fixed some include patterns

Put freezing profiler phase behind ndebug flag

Added traversable_base.h comments

Handled comments from Merlin

Added comments in apply.h

Moved freeze documentation to __init__.py

Added suggestions from Merlin

Added default value to resize for registry traversal

Fixed unnecesarry format change

Changed docstring comments for custom_type_ext tests

Changed include style

Comments

Fixed tensors getting leaked

Removed tracing log levels

Fixed tensor assignment leak

Improved profiling helper

Using move only Layout

Fixed variable missmatch

Fixed out of bounds write

Added fast callback traversal

Using info from jit_set_backend`

Improved traversal of opaque C++ objects

Improved layout for ad/jit indices

Reduced jit_var_inc_ref overhead

Improved traversal api

Removed redundant logging calls

Added Profile flags behind DRJIT_ENABLE_NVTX

Removed indices vector

Added domain and variant to rw traversal functions

Fixed warnings

Fixed logging of key difference

Fixed segfault with logging

Reverted on std::function helpers

Fixed double traversal of C++ objects

Added heuristic for reserving layout and maps in FlatVariables

Added back-assignment

Fixed variant domain not being added at traversal

Refactor to use VarLayout

Deferred traversal

Removed input eval

Removed deep_make_opaque and deep_eval

Comments

Removed RecordingKey

Moved assignment into record path

Fixed vairable leaks on failures

Removed fix comments

Removed redundant code

Renamed Hasher and Equal and added comments

Comments

Comments

Cleanup

Author: DoeringChristian

drjit/__init__.py

@@ -1878,6 +1878,102 @@ def binary_search(start, end, pred):
 
     return start
 
+def freeze(f):
+    '''
+    Decorator to freeze a function for replaying kernels without compilation.
+
+    This decorator wraps a function, and enables replaying it in order remove the
+    need for tracing python operations. The first time a frozen function is called,
+    it is executed regularly, while all operations performed are recorded. When the
+    frozen function is called again, with compatible arguments, the recorded
+    operations are replayed instead of launching the Python function. This saves on
+    the overhead that is inherent when tracing and compiling python operations into
+    kernels. The frozen function keeps a record of all previously recorded calls,
+    and associates them with the input layout. Since the kernels might be compiled
+    differently when changing certain attributes of variables, the frozen function
+    will be re-recorded if the inputs are no longer compatible with previous
+    recordings. We also track the layout of any container and the values of Python
+    types of the inputs to the frozen function. The following are the input
+    attributes which are tracked and might cause the function to be re-recorded if
+    changed.
+
+    - Python type of any variable or container
+    - Number of members in a container, such as the length of a list
+    - Key or field names, such as dictionary keys or dataclass field names
+    - The `Dr.Jit` type of any variable
+    - Whether a variable has size $1$
+    - Whether the memory, referenced by a variable is unaligned (only applies to
+      mapped NumPy arrays)
+    - Whether the variable has gradients enabled
+    - The sets of variables that have the same size
+    - The hash of any Python variable, that is not a Dr.Jit variable or a container
+
+    The width of a variable itself is not tracked, as we want to allow replaying of
+    kernels with different sizes. However, if multiple variables are scheduled and
+    evaluated at the same time, they can get compiled into the same kernel if their
+    size matches. Therefore, we track the size class of the variable, and re-record
+    the kernel if it changed. Because the size is not tracked it is leaked into the
+    frozen function context and we infer the size of kernels using a heuristic based
+    on the sizes of the input variables. This allows gathering from variables with
+    their width as an argument. However, using the size of a variable in more
+    complicated computations might lead to undefined behavior.
+
+    ```python
+    y = dr.gather(type(x), x, dr.width(x)//2)
+    ```
+
+    Similarly, calculating the mean of a variable relies on the number of entries,
+    which will be baked into the frozen function. To avoid this, we suggest
+    supplying the number of entries as a Dr.Jit literal in the arguments to the
+    function.
+    '''
+    import functools
+
+    class FrozenFunction:
+        def __init__(self, f) -> None:
+            self.ff = detail.FrozenFunction(f)
+
+        def __call__(self, *args, **kwargs):
+            return self.ff(*args, **kwargs)
+
+        @property
+        def n_recordings(self):
+            return self.ff.n_recordings
+
+        @property
+        def n_cached_recordings(self):
+            return self.ff.n_cached_recordings
+
+        def clear(self):
+            return self.ff.clear()
+
+        def __get__(self, obj, type=None):
+            if obj is None:
+                return self
+            else:
+                return FrozenMethod(self.ff, obj)
+
+    class FrozenMethod(FrozenFunction):
+        """
+        A FrozenMethod currying the object into the __call__ method.
+
+        If the ``freeze`` decorator is applied to a method of some class, it has
+        to call the internal frozen function with the ``self`` argument. To this
+        end we implement the ``__get__`` method of the frozen function, to
+        return a ``FrozenMethod``, which holds a reference to the object.
+        The ``__call__`` method of the ``FrozenMethod`` then supplies the object
+        in addition to the arguments to the internal function.
+        """
+        def __init__(self, ff, obj) -> None:
+            super().__init__(ff)
+            self.obj = obj
+
+        def __call__(self, *args, **kwargs):
+            return self.ff(self.obj, *args, **kwargs)
+
+
+
+    return functools.wraps(f)(FrozenFunction(f))
 
 def assert_true(
     cond,

ext/drjit-core

@@ -15,6 +15,7 @@
 
 #pragma once
 
+#include <type_traits>
 #define DRJIT_STRUCT_NODEF(Name, ...)                                          \
     Name(const Name &) = default;                                              \
     Name(Name &&) = default;                                                   \
@@ -140,6 +141,18 @@ namespace detail {
     using det_traverse_1_cb_rw =
         decltype(T(nullptr)->traverse_1_cb_rw(nullptr, nullptr));
 
+    template <typename T>
+    using det_get = decltype(std::declval<T&>().get());
+
+    template <typename T>
+    using det_const_get = decltype(std::declval<const T &>().get());
+
+    template<typename T>
+    using det_begin = decltype(std::declval<T &>().begin());
+
+    template<typename T>
+    using det_end = decltype(std::declval<T &>().begin());
+
     inline drjit::string get_label(const char *s, size_t i) {
         auto skip = [](char c) {
             return c == ' ' || c == '\r' || c == '\n' || c == '\t' || c == ',';
@@ -180,10 +193,17 @@ template <typename T> auto labels(const T &v) {
 }
 
 template <typename Value>
-void traverse_1_fn_ro(const Value &value, void *payload, void (*fn)(void *, uint64_t)) {
-    (void) payload; (void) fn;
+void traverse_1_fn_ro(const Value &value, void *payload,
+                      void (*fn)(void *, uint64_t, const char *,
+                                 const char *)) {
+    (void) payload;
+    (void) fn;
     if constexpr (is_jit_v<Value> && depth_v<Value> == 1) {
-        fn(payload, value.index_combined());
+        if constexpr(Value::IsClass)
+            fn(payload, value.index_combined(), Value::CallSupport::Variant,
+               Value::CallSupport::Domain);
+        else
+            fn(payload, value.index_combined(), "", "");
     } else if constexpr (is_traversable_v<Value>) {
         traverse_1(fields(value), [payload, fn](auto &x) {
             traverse_1_fn_ro(x, payload, fn);
@@ -198,14 +218,36 @@ void traverse_1_fn_ro(const Value &value, void *payload, void (*fn)(void *, uint
                          is_detected_v<detail::det_traverse_1_cb_ro, Value>) {
         if (value)
             value->traverse_1_cb_ro(payload, fn);
+
+    } else if constexpr (is_detected_v<detail::det_begin, Value> &&
+                         is_detected_v<detail::det_end, Value>) {
+        for (auto elem : value) {
+            traverse_1_fn_ro(elem, payload, fn);
+        }
+    } else if constexpr (is_detected_v<detail::det_const_get, Value>) {
+        const auto *tmp = value.get();
+        traverse_1_fn_ro(tmp, payload, fn);
+    } else if constexpr (is_detected_v<detail::det_traverse_1_cb_ro, Value *>) {
+        value.traverse_1_cb_ro(payload, fn);
+    } else {
+        // static_assert(false, "Failed to traverse field!");
     }
 }
 
 template <typename Value>
-void traverse_1_fn_rw(Value &value, void *payload, uint64_t (*fn)(void *, uint64_t)) {
-    (void) payload; (void) fn;
+void traverse_1_fn_rw(Value &value, void *payload,
+                      uint64_t (*fn)(void *, uint64_t, const char *,
+                                     const char *)) {
+    (void) payload;
+    (void) fn;
     if constexpr (is_jit_v<Value> && depth_v<Value> == 1) {
-        value = Value::borrow((typename Value::Index) fn(payload, value.index_combined()));
+        if constexpr(Value::IsClass)
+            value = Value::borrow((typename Value::Index) fn(
+                payload, value.index_combined(), Value::CallSupport::Variant,
+                Value::CallSupport::Domain));
+        else
+            value = Value::borrow((typename Value::Index) fn(
+                payload, value.index_combined(), "", ""));
     } else if constexpr (is_traversable_v<Value>) {
         traverse_1(fields(value), [payload, fn](auto &x) {
             traverse_1_fn_rw(x, payload, fn);
@@ -220,6 +262,18 @@ void traverse_1_fn_rw(Value &value, void *payload, uint64_t (*fn)(void *, uint64
                          is_detected_v<detail::det_traverse_1_cb_rw, Value>) {
         if (value)
             value->traverse_1_cb_rw(payload, fn);
+    } else if constexpr (is_detected_v<detail::det_begin, Value> &&
+                         is_detected_v<detail::det_end, Value>) {
+        for (auto elem : value) {
+            traverse_1_fn_rw(elem, payload, fn);
+        }
+    } else if constexpr (is_detected_v<detail::det_get, Value>) {
+        auto *tmp = value.get();
+        traverse_1_fn_rw(tmp, payload, fn);
+    } else if constexpr (is_detected_v<detail::det_traverse_1_cb_rw, Value *>) {
+        value.traverse_1_cb_rw(payload, fn);
+    } else {
+        // static_assert(false, "Failed to traverse field!");
     }
 }
 

include/drjit/array_traverse.h

@@ -971,7 +971,8 @@ NAMESPACE_BEGIN(detail)
 /// Internal operations for traversing nested data structures and fetching or
 /// storing indices. Used in ``call.h`` and ``loop.h``.
 
-template <bool IncRef> void collect_indices_fn(void *p, uint64_t index) {
+template <bool IncRef>
+void collect_indices_fn(void *p, uint64_t index, const char *, const char *) {
     vector<uint64_t> &indices = *(vector<uint64_t> *) p;
     if constexpr (IncRef)
         index = ad_var_inc_ref(index);
@@ -983,7 +984,8 @@ struct update_indices_payload {
     size_t &pos;
 };
 
-inline uint64_t update_indices_fn(void *p, uint64_t) {
+inline uint64_t update_indices_fn(void *p, uint64_t, const char *,
+                                  const char *) {
     update_indices_payload &payload = *(update_indices_payload *) p;
     return payload.indices[payload.pos++];
 }

include/drjit/autodiff.h

@@ -247,6 +247,11 @@ extern DRJIT_EXTRA_EXPORT bool ad_release_one_output(drjit::detail::CustomOpBase
 extern DRJIT_EXTRA_EXPORT void ad_copy_implicit_deps(drjit::vector<uint32_t> &,
                                                      bool input);
 
+/// Retrieve a list of ad indices, that are the target of edges, that have been
+/// postponed by the current scope
+extern DRJIT_EXTRA_EXPORT void ad_scope_postponed(drjit::vector<uint32_t> &);
+
+
 /// Kahan-compensated floating point atomic scatter-addition
 extern DRJIT_EXTRA_EXPORT void
 ad_var_scatter_add_kahan(uint64_t *target_1, uint64_t *target_2, uint64_t value,

include/drjit/extra.h

@@ -54,6 +54,9 @@
 #include <drjit/math.h>
 #include <drjit-core/python.h>
 #include <nanobind/stl/array.h>
+#include <nanobind/intrusive/counter.h>
+#include <nanobind/nanobind.h>
+#include <drjit/traversable_base.h>
 
 NAMESPACE_BEGIN(drjit)
 struct ArrayBinding;
@@ -1057,25 +1060,74 @@ template <typename T> void bind_all(ArrayBinding &b) {
 // Expose already existing object tree traversal callbacks (T::traverse_1_..) in Python.
 // This functionality is needed to traverse custom/opaque C++ classes and correctly
 // update their members when they are used in vectorized loops, function calls, etc.
-template <typename T, typename... Args> auto& bind_traverse(nanobind::class_<T, Args...> &cls) {
+template <typename T, typename... Args> auto &bind_traverse(nanobind::class_<T, Args...> &cls)
+{
     namespace nb = nanobind;
-    struct Payload { nb::callable c; };
+    struct Payload {
+        nb::callable c;
+    };
+
+    static_assert(std::is_base_of_v<TraversableBase, T>);
 
-    cls.def("_traverse_1_cb_ro", [](const T *self, nb::callable c) {
+    cls.def("_traverse_1_cb_ro", [](T *self, nb::callable c) {
         Payload payload{ std::move(c) };
-        self->traverse_1_cb_ro((void *) &payload, [](void *p, uint64_t index) {
-            ((Payload *) p)->c(index);
-        });
+        self->traverse_1_cb_ro((void *) &payload,
+                               [](void *p, uint64_t index, const char *variant,
+                                  const char *domain) {
+                                   ((Payload *) p)->c(index, variant, domain);
+                               });
     });
 
     cls.def("_traverse_1_cb_rw", [](T *self, nb::callable c) {
         Payload payload{ std::move(c) };
-        self->traverse_1_cb_rw((void *) &payload, [](void *p, uint64_t index) {
-            return nb::cast<uint64_t>(((Payload *) p)->c(index));
+        self->traverse_1_cb_rw((void *) &payload, [](void *p, uint64_t index,
+                                                     const char *variant,
+                                                     const char *domain) {
+            return nb::cast<uint64_t>(
+                ((Payload *) p)->c(index, variant, domain));
         });
     });
 
     return cls;
 }
 
+inline void traverse_py_cb_ro(const TraversableBase *base, void *payload,
+                              void (*fn)(void *, uint64_t, const char *variant,
+                                         const char *domain)) {
+    namespace nb    = nanobind;
+    nb::handle self = base->self_py();
+    if (!self)
+        return;
+
+    auto detail = nb::module_::import_("drjit.detail");
+    nb::callable traverse_py_cb_ro =
+        nb::borrow<nb::callable>(nb::getattr(detail, "traverse_py_cb_ro"));
+
+    traverse_py_cb_ro(self,
+                      nb::cpp_function([&](uint64_t index, const char *variant,
+                                           const char *domain) {
+                          fn(payload, index, variant, domain);
+                      }));
+}
+
+inline void traverse_py_cb_rw(TraversableBase *base, void *payload,
+                              uint64_t (*fn)(void *, uint64_t, const char *,
+                                             const char *)) {
+
+    namespace nb    = nanobind;
+    nb::handle self = base->self_py();
+    if (!self)
+        return;
+
+    auto detail = nb::module_::import_("drjit.detail");
+    nb::callable traverse_py_cb_rw =
+        nb::borrow<nb::callable>(nb::getattr(detail, "traverse_py_cb_rw"));
+
+    traverse_py_cb_rw(self,
+                      nb::cpp_function([&](uint64_t index, const char *variant,
+                                           const char *domain) {
+                          return fn(payload, index, variant, domain);
+                      }));
+}
+
 NAMESPACE_END(drjit)

include/drjit/python.h

@@ -17,6 +17,7 @@
 #include <drjit/idiv.h>
 #include <drjit/jit.h>
 #include <drjit/tensor.h>
+#include <drjit/traversable_base.h>
 
 #pragma once
 
@@ -41,7 +42,7 @@ enum class CudaTextureFormat : uint32_t {
     Float16 = 1, /// Half precision storage format
 };
 
-template <typename _Storage, size_t Dimension> class Texture {
+template <typename _Storage, size_t Dimension> class Texture : TraversableBase {
 public:
     static constexpr bool IsCUDA = is_cuda_v<_Storage>;
     static constexpr bool IsDiff = is_diff_v<_Storage>;
@@ -1386,6 +1387,9 @@ template <typename _Storage, size_t Dimension> class Texture {
     WrapMode m_wrap_mode;
     bool m_use_accel = false;
     mutable bool m_migrated = false;
+
+    DR_TRAVERSE_CB(drjit::TraversableBase, m_value, m_shape_opaque,
+                   m_inv_resolution);
 };
 
 NAMESPACE_END(drjit)