Keep track of variable-sized data in a JIT loop

[tomas16]

Hi, I'm using mitsuba for geometric raytracing. The idea is to spawn some rays, have them intersect with the scene, and when they intersect, refract and spawn a new ray. At every intersection I need to keep track of some data, so that I can use this as input to a separate physics simulation.

What I have now is something like:

# origin and direction are Nx3 numpy arrays
# scene is mitsuba.Scene

def raytrace(scene, origin, direction):
    ray = mi.Ray3f(o=mi.Point3f(origin), d=mi.Vector3f(direction))
    opl = mi.Float(0)
    active_mask = mi.Bool(True)
    iteration = mi.Int(0)           # for debugging
    loop = mi.Loop("Raytracing", lambda: (active_mask, ray, opl, i))
    while loop(active_mask):
        si = scene.ray_intersect(ray, mi.RayFlags.ShadingFrame, coherent=False, active=active_mask)

        # Refract and generate ray for next iteration
        cos_theta_i = si.wi.z
        F, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, 1.5)       # TODO: don't hardcode refractive index
        raydir = si.to_world(mi.refract(si.wi, cos_theta_t, eta_ti))
        ray = si.spawn_ray(raydir)

        # Update active_mask and save results from this iteration
        active_mask = ~dr.isinf(si.t)
        opl[active_mask] += si.t[active_mask]

        i += 1

    return np.array(opl), np.array(i)

In above example, I'm adding up the path length, and I think this should run as a single kernel.

But what I really want is something more akin to a python list and append to it on every iteration. Can this be done this way, or does it require having a python loop around a drjit kernel? Is there any way to have a growing array/matrix/... in drjit?

I should note I'm not too worried about performance as I have ~1000 rays and probably <30 iterations. But I would like to learn to use mitsuba "the right way".

Side question:

opl[active_mask] += si.t[active_mask]

Is this the right idiom to update a variable only for the still active lanes? I thought I read somewhere in the docs that this type of indexing should be avoided.

[njroussel]

Hi @tomas16

Is there any way to have a growing array/matrix/... in drjit?

In short, no.

However, most of times you can now the the maxmium length/size of your data that you're trying to record. For example, the maxmium number of bounces in the scene in your setup. In that case you can create a flat array of the appropriate size and write to it at each iteration. Here's a quick and dirty example:

import mitsuba as mi
import drjit as dr

dr.set_log_level(dr.LogLevel.Info)
mi.set_variant('cuda_ad_rgb')

W = 4 # wavefront width
N = 5 # max iterations

scene = mi.load_dict(mi.cornell_box())

origin = dr.zeros(mi.Point3f, W)
origin.z = 4
direction = mi.Vector3f(0, 0, -1)

output_t = dr.zeros(mi.Float, W * N) # Flat buffer to hold all distances

iteration = mi.Int(0)
ray = mi.Ray3f(o=origin, d=direction)
active = mi.Bool(True)

loop = mi.Loop("", lambda: (iteration, ray, active))
while loop(active):
    si = scene.ray_intersect(ray, active=active)
    active &= si.is_valid()

    idx = dr.arange(mi.UInt32, W) + iteration * W
    dr.scatter(output_t, si.t, idx, active)

    ctx = mi.BSDFContext()
    bsdf_sample, _ = si.bsdf().sample(ctx, si, 0.5, mi.Point2f(0.1, 0.7), active)
    ray = si.spawn_ray(si.to_world(bsdf_sample.wo))

    iteration += 1
    active &= iteration < N

print(output_t)

This launches a single kernel. However, it obviously now requires some "smart" indexing into output_t or whatever your flat buffer is to retrieve data.

---

As for your side question, the mask on the right hand side is unnecessary. You can just write opl[active_mask] += si.t. This is the appropriate to do such an operation.
Generally, I'd recommend looking at src/python/python/ad/integrators/prb.py for what is "correct/appropriate" Dr.Jit practices.

[tomas16]

Got it, makes sense. Thanks for sharing that code snippet and pointing me to the non-trivial example of prb.py.

[tomas16]

Actually one more question: is there a way to get the "wavefront width" from a given Ray3f, Point3f or Vector3f variable? I'm passing numpy arrays into my function because I can get their shape.

[tomas16]

Answered my question: dr.width(value)

[tomas16]

@njroussel if I could trouble you for a related question.

I need to keep track of which shape is being intersected.

I tried the following but it results in a RuntimeError:

# before the loop:
shape = dr.zeros(mi.ShapePtr, num_rays * max_bounces)

# inside the loop body:
dr.scatter(shape, si.shape, idx, active)

I've also been playing around with creating a python dictionary to map each Mesh in the scene to an index, but then the issue is I'm not sure how to do the dictionary lookup inside the vectorized JIT loop. And it doesn't really seem like the right approach.

Another option would be to attach some attribute to each Mesh, but that also doesn't work, probably because Mesh is C++ class and has no __dict__. The meshes do have a name, but I don't see a way to store those names in a manner similar to output_t in your example. That's why I'm trying to map it to some number (which can either be some index I come up with, or the actual ShapePtr pointer value that I can map to something else afterwards).

[tomas16]

Based on the question in #882, turns out I can use dr.reinterpret_array_v(mi.UInt32, si.shape)

[njroussel]

Sorry for the late reply.

We've added functionnality for this in this commit: 0b2a032
Unless you're building the project yourself, you'll need to wait for the next official PyPI release (later today).