Improve the behavior of control flow operations (dr.if_stmt(), drjit.while_loop()) in AD-suspended scopes
#299
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
wjakob
force-pushed
the
ad-suspend-fixes
branch
2 times, most recently
from
32128f8
to
ce7b4d8
Compare
njroussel
approved these changes
Oct 21, 2024
The AD layer exposes a function named ``ad_var_inc_ref()`` that
increases the reference count of a variable analogous to
``jit_var_inc_ref()``. However, one difference between the two is that
the former detaches AD variables when the underlying index has
derivative tracking disabled.
For example, this ensures that code like
```python
x = Float(0)
dr.enable_grad(x)
with dr.suspend_grad():
y = Float(y)
```
creates a non-differentiable copy.
However, since there are many other operations throughout the Dr.Jit
codebase that require reference counting, there were quite a few places
that exhibited this detaching behavior, which is not always wanted.
(see issue #253).
This commit provides two reference counting functions:
- ``ad_var_inc_ref()`` which increases the reference count *without*
detaching, and
- ``ad_var_copy_ref()``, which detaches (i.e., reproducing the former
behavior)
Following this split, only the constructor of AD arrays uses the
detaching ``ad_var_copy_ref()``, while all other operations use
the new ``ad_var_inc_ref()``.
…pended mode Dr.Jit control flow operations (``dr.if_stmt(), drjit.while_loop()``) disable gradient tracking of all variable state when the operation takes place within an AD-disabled scope. This can be surprising when a ``@dr.syntax`` transformation silently passes local variables to such an operation, which then become non-differentiable. This commit carves out an exception: when variables aren't actually modified by the control flow operation, they can retain their AD identity. This is part #1 of the fix for issue #253 reported by @dvicini and targets ``dr.if_stmt()`` only. The next commit will also fix the same problem for while loops.
…AD-suspended mode Dr.Jit control flow operations (``dr.if_stmt(), drjit.while_loop()``) disable gradient tracking of all variable state when the operation takes place within an AD-disabled scope. This can be surprising when a ``@dr.syntax`` transformation silently passes local variables to such an operation, which then become non-differentiable. This commit carves out an exception: when variables aren't actually modified by the control flow operation, they can retain their AD identity. This is part #2 of the fix for issue #253 reported by @dvicini and targets ``dr.while_lop()`` only. The previous commit fixed the same problem for ``if`` statements.
njroussel
force-pushed
the
ad-suspend-fixes
branch
from
ce7b4d8
to
030f584
Compare
|
Hi, may I ask if this fix can relate to this issue? Thanks! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Dr.Jit control flow operations (
dr.if_stmt(), drjit.while_loop()) currently disable gradient tracking of all variable state when the operation takes place within an AD-disabled scope.This can be surprising when a
@dr.syntaxtransformation silently passes local variables to such an operation, which then become non-differentiable. @dvicini reported this in issue #253.This commit carves out an exception: when variables aren't actually modified by the control flow operation, they can retain their AD identity.
The PR has 3 parts:
ad_var_inc_refso that it only does reference counting. The old behavior (drop ref if variable is currently non-differentiable) is moved to a new functionad_var_copy_ref()(see the description of this commit for details). This already fixes 80% of the problems.dr.if_stmt()anddr.while_loop(), respectively.The PR depends on a Dr.Jit-Core PR: mitsuba-renderer/drjit-core#104