Getting the first non-specular bounce

[Mephisto405]

Hello,

I'm genuinely enjoying programming with Mitsuba 3, which is easier than Mitsuba 0.6 + 2.0.
Now, I'm implementing a code to get the first non-specular bounce for some purpose.
That is, not just tracing a ray on a single bounce; I'd like to query the first non-specular bounce by tracing a sequence of rays (e.g., to get G-buffers).

Given the bsdf and si_bsdf, I implemented such function as follow:

bsdf = ... # bsdf at a surface
si_bsdf = ... # surfaceinteraction3f at the surface

active = mi.Bool(si_bsdf.is_valid())
loop = mi.Loop(name="BSDF ray tracing", state=lambda: (si_bsdf))
loop.set_max_iterations(max_depth)

while loop(active):
    bsdf_sample, bsdf_weight = bsdf.sample(bsdf_ctx, si_bsdf, sampler.next_float32(), mi.Point2f(sampler.next_float32(), sampler.next_float32()), True)
    ray = si_bsdf.spawn_ray(si_bsdf.to_world(bsdf_sample.wo))
    
    si_bsdf = scene.ray_intersect(ray, ray_flags=mi.RayFlags.All, coherent=False)
    bsdf = si_bsdf.bsdf(ray)

    active = si_bsdf.is_valid() & (mi.has_flag(bsdf.flags(), mi.BSDFFlags.Delta) | mi.has_flag(bsdf.flags(), mi.BSDFFlags.Delta1D))

But I got the error below at the while loop(active): line:

RuntimeError: jit_var_loop(): loop condition does not depend on any of the loop variables

I thought my loop condition active indeed depends on the loop variable si_bsdf, but the error message claims the opposite.
So I guess I understand Dr.Jit in a wrong way.

Any help would be appreciated!!

[njroussel]

Hi!

loop = mi.Loop(name="BSDF ray tracing", state=lambda: (si_bsdf))

Can you try with this instead:

loop = mi.Loop(name="BSDF ray tracing", state=lambda: (active, si_bsdf))

The error message is a bit misleasing here. I'd encourage you to take a look at the Dr.Jit documentation on recorded loops. As a rule of thumb, your loop state variables should usually be all variables which need to be read before written at any iteration. Here, you need to read active (in the loop condition) at every iteration before you write to it in active = si_bsdf.is_valid() ...

I believe that should be enough 😄

[Mephisto405]

Very thankful! That works!

Besides, I saw unintuitive behavior now.
I wanted to see how many bounces are required to reach a non-specular shape.
And this code works fine:

active = mi.Bool(si_bsdf.is_valid())
i = dr.llvm.UInt32(0)
loop = mi.Loop(name="BSDF ray tracing", state=lambda: (i, active, si_bsdf))
loop.set_max_iterations(max_depth)
#path = [ps.p]

while loop(active):
    bsdf_sample, bsdf_weight = bsdf.sample(bsdf_ctx, si_bsdf, sampler.next_float32(), mi.Point2f(sampler.next_float32(), sampler.next_float32()), True)
    ray = si_bsdf.spawn_ray(si_bsdf.to_world(bsdf_sample.wo))
    
    si_bsdf = scene.ray_intersect(ray, ray_flags=mi.RayFlags.All, coherent=False)
    bsdf = si_bsdf.bsdf(ray)

    active = si_bsdf.is_valid() & (mi.has_flag(bsdf.flags(), mi.BSDFFlags.Delta) | mi.has_flag(bsdf.flags(), mi.BSDFFlags.Delta1D))
    i += 1
    #path.append(dr.select(si_bsdf.is_valid(), si_bsdf.p, dr.fma(ray.d, 10., ray.o)))
print(i)

However, when I change i = dr.llvm.UInt32(0) to i=0 (and also remove i in the mi.Loop instantiation), it turns out to be weird. Every print(i) prints 2.

In fact, I tried your suggestion before starting this discussion, but I tested my code with integer i, and it seemed to fail (since it seems that every query requires exactly two bounces to reach a non-spec shape) at the moment due to the aforementioned behavior.

[njroussel]

There's a few things to keep in mind when writing Mitsuba code.

Firstly, when using a JIT variant (CUDA/LLVVM), instructions are not actually evaluated/executed they are simply recorded. In fact, the entire rendering loop is typically only recorded/traced by the just-in-time compiler (Dr.Jit). It is only when writing the samples to an output, that the recorded computation is then compiled into a CUDA/LLVM program which is finally executed. Obviously, this only works for Dr.Jit/Mitsuba types. Any other computation (not involving Dr.Jit/Mitsuba types) will not be recorded but executed instantly (this is just the normal Python behavior).
In your example, when using a Python integer i, the increment is evaluated instantly and not recorded. It just so happens that the body of a mi.Loop is recorded twice for technical reason. So any code that doesn't involve Mitsuba/DrJit will be executed exactly twice no matter what the loop condition is. This is why you'd be seeing i == 2 when the variable is not part of the recorded loop.

Secondly, practically everything in Mitsuba/Drjit is vectorized. For example, a mi.Ray object is actually a N-wide array of rays. In your snippet, i will also be N-wide and when you're incrementing it i += 1 you're actually incrementing all N "instances". What you would actually want to do is only increment the indices which are not done yet:

i += dr.select(active & ~glossy, 1, 0)

Now, at then of your loop, i would have different values corresponding to the depth of each ray. Whereas before you were technically computing the max depth among all your rays.

[Mephisto405]

OMG! I didn't keep in mind these principles of Mitsuba 3.0. So thankful your generous and detailed explanation.
Many thanks again thesedays.

[njroussel]

No problem :)