ADmode flags

[saeedhd96]

Hello,

I see the use of ADFlags in this tutorial. I am a little unclear on the ADFlags. A few questions:

1. Introducing changes/bias in any scenario? Based on the documentation for dr.traverse(), I understand that setting and unsetting these flags would have no effect on the values of the gradients, but it only affects the performance. Is that correct?

2. What are edges, vertices, interior, input in this tree an d what is stored in them (on a high level)? For example, how clearing an edge vs. its vertices could be different in the resulting gradients/performance?

3. Do any of these flags matter when calling dr.backward() also? because I assume backward has only one traversal if there is one output, while forward might need multiple traversals if there are many inputs.

Thanks for you time in advance

[wjakob]

Hi Saeed -- please take a look at the example of an AD computation graph in the Dr.Jit paper (Figure 10) and related discussion: https://rgl.s3.eu-central-1.amazonaws.com/media/papers/Jakob2022DrJit.pdf. This hopefully clarifies what edges and vertices are.

There are two aspects to this feature:

• Lowering memory usage: by destroying the graph during traversal, Dr.Jit uses less memory, which is useful when using the system in wavefront model. For megakernels, there is no difference since Dr.Jit variables represent GPU registers that do not consume device global memory.
• Correctness: At its core, the AD traversal of a graph is a very simple mechanical operation. You go through vertices in forward or reverse order, pick up gradient values present at the vertex, and propagate it along edges connecting variables. Whenever a gradient value is propagated along an edge, it must be scaled by a sensitivity value associated with the edge. This can all matter quite a bit when you want to reuse an AD graph and propagate gradient values through it multiple times. If there are gradients at interior vertices left from the previous pass, you will end up computing the wrong thing. That's why these flags exist -- they let you control very precisely what is done to the graph during the traversal.