Frozen Functions #336

DoeringChristian · 2025-02-03T08:58:00Z

This PR implements the frozen function feature in drjit.
It adds a new FrozenFunction struct that is a wrapper over a Python callable.
When calling the operator() of the object, it traverses the input and determines a unique layout/key.
This key is then used to determine if the function has to be recorded or if it's possible to replay it.
The PR relies on functionality implemented in Frozen ThreadState, and #288.

This replaces #317 and changes the source branch.

njroussel

This is looking good so far, I still have to go through some of the main bits like freeze.cpp/freeze.h. I'll come back to this tomorrow.

njroussel · 2025-02-04T17:10:32Z

include/drjit/texture.h

+
+    DR_TRAVERSE_CB(drjit::TraversableBase, m_value, m_shape_opaque,
+                   m_inv_resolution);


This is more of a note (to myself):

This will most likely need to be updated with #329 before this gets merged, or we do it in #329 if we merge this fist.

src/python/common.h

tests/test_custom_type_ext.py

src/python/apply.cpp

include/drjit/array_traverse.h

include/drjit/python.h

njroussel

I've gone through the last few files I had missed.

Ignoring the few new comments, I'm overall pretty happy with the current state of the source code. As for the documentation, the only part that truly matters is the dr.freeze and I think it's good too - I don't really have the full picture yet, so take that with a pinch of salt. For the C++ documentation, it was plenty enough for me to understand how these mechanisms work and should be used 👍

include/drjit/traversable_base.h

src/python/common.h

src/python/detail.cpp

src/python/freeze.cpp

include/drjit/traversable_base.h

merlinND

Very partial review, I only looked at the files that had < ~150 lines of diff.

drjit/__init__.py

include/drjit/autodiff.h

include/drjit/python.h

include/drjit/texture.h

drjit/__init__.py

include/drjit/traversable_base.h

src/python/detail.cpp

include/drjit/array_traverse.h

wjakob

Some comments from me on this PR

include/drjit/python.h

include/drjit/array_traverse.h

include/drjit/texture.h

src/python/freeze.cpp

wjakob · 2025-06-11T14:48:30Z

src/python/freeze.cpp

+ * assigning to the input. The function takes an optional python-type if
+ * it is known.
+ */
+void FlatVariables::traverse_ad_index(uint64_t index, TraverseContext &ctx,


There is something basic I do not understand. My assumption was that inputs/outputs to @dr.freeze-annotated functions cannot be AD-attached. Yet there is lots of handling related to AD variables.

The documentation (docs/freeze.rst) you added does not discuss AD variables. It would be good to explain somewhere what behavior is legal and what isn't.

The Auto Opaque PR #361 adds the documentation in freeze.rst. It should explain that only propagating gradients from inside to the inputs and their dependencies is possible from a frozen function. These traverse_ad_index functions are required to handle this case.

wjakob · 2025-06-11T14:51:39Z

src/python/freeze.cpp

@@ -0,0 +1,1551 @@
+#include "freeze.h"


Like in the Dr.Jit-Core implementation of kernel freezing, this file needs a detailed comment describing the high level "algorithm". The various functions are individually documented, but the big picture does not easily emerge when going through them.

It should cover the high level ideas like: how are the arguments to a function flattened, how do we detect misses/hits? You also implemented an optimization related to literal variables, that would be good to briefly mention too.

Will ad that in the auto opaque PR #361 if ok. Since this changes quite a bit.

Added comment explaining traverse with `rw` arguments Added comment about c++ traversal Added recursion guards Simplified comment on TraverseCallback operator() Removed redundant profiling loop Fixed recursion_guard by wrapping it in namespace Fixed typing import Improved handling of nullptr variant/domains Removed warn log Added JitFlag FreezingTraverseScope Using try_cast in get_traversable_base Renamed to EnableObjectTraversal flag Added comment for traverse_py_cb_ro_impl Added rw value for traversal of trampolines Removed outdated function Excluding texture from traversal outside of frozen function Testing fix for windows compilation Fixing windows compilation bugs Added base value test for custom_type_ext traversal Added nested traversal test Exposing frozen function related flags to python Improved warning about non-drjit types Added option to specify backend if no input variable is specified Allow for changes in the last dimension of tensor shapes Marked payload and fn as used Reverted comment Formatting Removed printing of keys Fixed backend test Added method to detect recording frozen functions with the wrong backend Added comments for arguments of dr.freeze decorator Added flag test to DR_TRAVERSE_CB_RO/RW macros Removed test for `EnableObjectTraversal` flag in `DR_TRAVERSE_CB` macro Fixed tensor freezing indexing issue Removed deprecated logging code Added comment about tensor shape inference Using Wenzel's documentation, and renamed arguments Added comment regarding `frozen_function_tp_traverse` and `frozen_function_clear` Fixed typo Added warning text to documentation Fixed freezing drjit optimizers Added optimizer freezing tests Small refactor Suggestion from Wenzel

Small fix

DoeringChristian mentioned this pull request Feb 3, 2025

Frozen Functions #317

Closed

DoeringChristian force-pushed the frozen-functions branch from a62bdf7 to d47afff Compare February 3, 2025 18:02

njroussel mentioned this pull request Feb 4, 2025

TraversableBase intrusive interface #288

Closed

njroussel reviewed Feb 4, 2025

View reviewed changes

DoeringChristian requested review from wjakob and merlinND February 7, 2025 16:52

DoeringChristian force-pushed the frozen-functions branch from bfbd9b5 to bfb7674 Compare February 7, 2025 18:17

njroussel reviewed Feb 10, 2025

View reviewed changes

merlinND reviewed Feb 14, 2025

View reviewed changes

DoeringChristian force-pushed the frozen-functions branch from 1ae57b6 to c12c89e Compare February 21, 2025 13:26

DoeringChristian force-pushed the frozen-functions branch 6 times, most recently from 5f7efff to 3af2b69 Compare March 4, 2025 10:25

merlinND reviewed Mar 6, 2025

View reviewed changes

src/python/detail.cpp Outdated Show resolved Hide resolved

DoeringChristian force-pushed the frozen-functions branch 10 times, most recently from 15220b6 to 3d80462 Compare March 16, 2025 15:03

DoeringChristian force-pushed the frozen-functions branch 3 times, most recently from 971b79f to 2232018 Compare March 23, 2025 11:49

DoeringChristian force-pushed the frozen-functions branch 2 times, most recently from 1f08b94 to f57a954 Compare April 12, 2025 11:23

wjakob force-pushed the master branch 4 times, most recently from c821068 to dcb5217 Compare April 14, 2025 16:01

DoeringChristian force-pushed the frozen-functions branch 4 times, most recently from 7de6d1c to bec6588 Compare April 18, 2025 11:42

DoeringChristian force-pushed the frozen-functions branch from 48daace to 843f668 Compare April 24, 2025 20:56

DoeringChristian force-pushed the frozen-functions branch 2 times, most recently from fd8cf08 to 8ef5ddf Compare June 4, 2025 13:15

DoeringChristian force-pushed the frozen-functions branch from 8ef5ddf to 60a9a4a Compare June 11, 2025 08:15

wjakob reviewed Jun 11, 2025

View reviewed changes

include/drjit/array_traverse.h Outdated Show resolved Hide resolved

wjakob reviewed Jun 11, 2025

View reviewed changes

DoeringChristian force-pushed the frozen-functions branch 5 times, most recently from 382c8b9 to f31870a Compare June 18, 2025 13:11

DoeringChristian force-pushed the frozen-functions branch from f31870a to 09ff550 Compare June 18, 2025 13:27

DoeringChristian added 5 commits June 19, 2025 21:05

Fixed recursion guard

4bb2674

Added bindings to opaque_width

4416301

Implemented auto_opaque feature

1192616

Small fix

Added and updated tests for auto opaque

4672954

Added documentation for dr.freeze

cac9a9b

wjakob merged commit e1f88e2 into master Jun 19, 2025

wjakob deleted the frozen-functions branch June 19, 2025 12:06


		DR_TRAVERSE_CB(drjit::TraversableBase, m_value, m_shape_opaque,
		m_inv_resolution);

Frozen Functions #336

Frozen Functions #336

Conversation

DoeringChristian commented Feb 3, 2025

Uh oh!

njroussel left a comment

Choose a reason for hiding this comment

Uh oh!

njroussel Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

njroussel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

merlinND left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wjakob left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wjakob Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

DoeringChristian Jun 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wjakob Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

DoeringChristian Jun 11, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

DoeringChristian Jun 11, 2025 •

edited

Loading