mitsuba-renderer / drjit

Dr.Jit — A Just-In-Time-Compiler for Differentiable Rendering
BSD 3-Clause "New" or "Revised" License
593 stars 43 forks source link

[Question] How to convert static arrays like Vector3f to tensors inside recorded loops? #90

Closed uguuuuuu closed 11 months ago

uguuuuuu commented 2 years ago

I want to convert static arrays like Vector3f to TensorXf inside recorded loops by something like mi.TensorXf(a), but when I tried, the following errors occured:

jit_var_eval(): variable r51154 remains dirty after evaluation!
jit_var_eval(): variable r51170 remains dirty after evaluation!
TypeError: TensorXf: expect an array that implements the array interface protocol!

And I couldn't even ravel them

Critical Dr.Jit compiler failure: jit_var_new_scatter(): variable remains dirty after evaluation!

Is there a way to do that?

Speierers commented 2 years ago

Hi @uguuuuuu,

Could you provide a minimal reproducer for this? It will be hard to help you otherwise.

uguuuuuu commented 2 years ago

Hi @uguuuuuu,

Could you provide a minimal reproducer for this? It will be hard to help you otherwise.

Sorry I forgot to attach a reproducer and thank you so much for your reply. The problem was in fact about dr.scatter() and dr.gather(). Below is a reproducer:

import drjit as dr
import mitsuba as mi
mi.set_variant('cuda_ad_rgb')

i = mi.UInt32(0)
loop = mi.Loop(name="array to tensor",
            state=lambda: (i))

while loop(i < 10):
    v = dr.zeros(mi.Vector3f, 20)
    t = mi.TensorXf(dr.ravel(v), [20,3])

    # dr.gather() won't work
    t = t[None]

    # Neither will dr.scatter
    a = dr.empty(mi.TensorXf, [20, 3])
    dr.scatter(a.array, t.array, dr.arange(mi.UInt32, 60))
Critical Dr.Jit compiler failure: jit_var_new_gather(): variable remains dirty after evaluation!
Aborted (core dumped)

OR

Critical Dr.Jit compiler failure: jit_var_new_scatter(): variable remains dirty after evaluation!
Aborted (core dumped)
Speierers commented 2 years ago

In Dr.Jit the rule is that all gather calls happening inside of a symbolic loop must use as a source an array that is evaluated before the loop body. E.g. in your code you will need to take t = ... outside of the loop body.

Similarly, the target of scatter operations happening inside of a symbolic loop must be evaluated before the loop body.

uguuuuuu commented 2 years ago

In Dr.Jit the rule is that all gather calls happening inside of a symbolic loop must use as a source an array that is evaluated before the loop body. E.g. in your code you will need to take t = ... outside of the loop body.

Similarly, the target of scatter operations happening inside of a symbolic loop must be evaluated before the loop body.

I'm not sure if I understand you right, but after I moved t out of the loop, the error message became RuntimeError: jit_var_new_gather(): cannot gather from a placeholder variable!

import drjit as dr
import mitsuba as mi
mi.set_variant('cuda_ad_rgb')

i = mi.UInt32(0)
v = dr.zeros(mi.Vector3f, 20)
t = mi.TensorXf(dr.ravel(v), [20,3])
loop = mi.Loop(name="array to tensor",
            state=lambda: (i, t))

while loop(i < 10):
    # dr.gather() won't work
    t = t[None]

    # Neither will dr.scatter
    a = dr.empty(mi.TensorXf, [20, 3])
    dr.scatter(a.array, t.array, dr.arange(mi.UInt32, 60))

Output:

RuntimeError: jit_var_new_gather(): cannot gather from a placeholder variable!
drjit-autodiff: scope leak detected (1 scopes remain in use)!
drjit-autodiff: scope leak detected (1 scopes remain in use)!
Speierers commented 2 years ago

Unfortunately, it is not possible to gather and scatter from/in the same array in a symbolic loop. This is because in order for the scatter operation of the previous iteration to take place, we would need to evaluate a kernel, which would break the symbolic loop.

You will either have to find another formulation of your algorithm that works with symbolic loops, or switch to wavefront loops (which will launch a kernel at the end of every iteration) using dr.set_flag(dr.JitFlag.LoopRecord, False).