Issues porting existing code to 3.6.0

[tomas16]

dr.scatter issues

Hi, I'm porting an existing algorithm to work with 3.6.0 and running into some issues. Not sure if these may be bugs or if I'm doing things wrong. Most of the issues are related to using dr.scatter to write results to output arrays inside a symbolic loop. I can't share the full algorithm, but it's mostly straightforward raytracing. At every bounce, I want to record various bits of information such as intersection location, costi/costt, and some others.

The general structure of the algorithm is

@dr.syntax
def algo(scene: mi.Scene, ray: mi.Ray3f, max_bounces: int):
    num_rays = dr.width(ray)
    numel = num_rays * max_bounces      # for each ray, we want to record information per bounce

    # output arrays (write-only)
    out_xyz = dr.zeros(mi.Float, numel * 3)
    out_costi = dr.zeros(mi.Float, numel)

    # these 2 output arrays get written to multiple times per iteration, with different value/index/active
    out_active = dr.zeros(mi.Bool, numel)           # one per bounce, per ray
    out_raystatus = dr.zeros(mi.UInt, num_rays)     # one per ray

    ray_idx = dr.arange(mi.UInt, num_rays)

    active = mi.Bool(True)                          # expands to num_rays inside the loop body (not shown)
    bounce = dr.zeros(mi.UInt, num_rays)
    while dr.hint(active & (bounce < max_bounces), mode='symbolic'):
    # or simply and equivalently:
    # while active & (bounce < max_bounces):

        # index vector into the output arrays for the current bounce/iteration
        idx = ray_idx + bounce * num_rays

        # Do raytracing stuff (si = scene.ray_intersect() etc)
        # ... a bunch of variables are created here

        # this works
        dr.scatter(out_costi, si.wi.z, idx, active)

        # this doesn't work, but used to before 3.6.0
        dr.scatter(out_xyz, si.p, idx, active)

        # this works (in real algo there's other stuff in between)
        dr.scatter(out_active, False, idx, mask)            # mask is computed and handles special cases only
        dr.scatter(out_active, active, idx, active)         # main update

        # this doesn't work (in real algo there's 4 lines like this with other stuff in between)
        # -> can only scatter once to out_raystatus inside the loop body
        dr.scatter(out_raystatus, 1, ray_idx, mask2)
        dr.scatter(out_raystatus, 2, ray_idx, mask3)

out_costi

The simplest case, out_costi, performs a single scatter of a scalar value per ray and works fine.

out_xyz

The 2nd case, out_xyz, is very similar, only this time I want to record a Point3f per ray. The implementation shown above worked pre-3.6.0. Now I get

    dr.scatter(out_xyz, si.p, idx, active)
RuntimeError: Variable r476 remains dirty despite an attempt to evaluate it.
This normally indicates the following situation:
1. You executed an operation causing a *side effect*, e.g., a
...

I don't understand why out_xyz would have to get evaluated. Like out_costi, it's never read from, and is only ever used as the target of dr.scatter.

If I initialize out_xyz as out_xyz = dr.zeros(mi.Point3f, numel), I don't get an error. Running this with toy data and converting to a numpy array at the end, I'm expecting a 7x5x3 ndarray, representing a max of 7 bounces and 5 rays. out_xyz.ravel() seems to give me:

• x coordinate of bounce 0, all 5 rays
• x coordinate of bounce 1, all 5 rays
• ...
• after the x coordinates of all bounces and all rays, the remainder of the array is all zeros. The Y and Z coordinates aren't present.

I also tried initializing it to other array types like Array3f, which resulted in the error TypeError: drjit.scatter(): 'target' argument must be a dynamic 1D array!

Given this used to work pre-3.6.0 and it doesn't result in an error when using Point3f, I wonder if this may be a bug?

out_active

Doing multiple scatter operations with different value/active args to out_active works.

out_raystatus

Multiple scatter operations to out_raystatus result in the same error as encountered by out_xyz: RuntimeError: Variable r372 remains dirty despite an attempt to evaluate it..

The only difference with the previous case is that out_active is a larger array of all bounces and all rays, whereas out_raystatus is one value per ray. And of course mi.Bool vs mi.UInt.

Other issues

On a few occasions, I've encountered this message:

Dr.Jit encountered an unrecoverable error and will now shut
down. Please re-run your program in debug mode to check for
out-of-bounds reads, writes, and other sources of undefined
behavior. You can do so by calling

   dr.set_flag(dr.JitFlag.Debug, True)

at the beginning of the program. If these additional checks
fail to pinpoint the problem, then you have likely found a
bug. We are happy to help investigate and fix the problem if
you can you create a self-contained reproducer and submit it
at https://github.com/mitsuba-renderer/drjit.

The error message of this specific failure is as follows:
>>> jit_var(r0): unknown variable!

When I enable the debug mode, I don't get any additional info output to stdout/stderr. Instead, I get this same error message, even when I wasn't getting it before turning on the debug mode. Letting you guys know in case it may be a bug as well.

Running in the debugger also prints out the following on startup:

PYDEV DEBUGGER WARNING:
sys.settrace() should not be used when the debugger is being used.
This may cause the debugger to stop working correctly.
If this is needed, please check: 
http://pydev.blogspot.com/2007/06/why-cant-pydev-debugger-work-with.html
to see how to restore the debug tracing back correctly.

Versions

• mitsuba 3.6.0
• drjit 1.0.1
• python 3.12.7
• llvm 16

[rtabbara]

Hi @tomas16 ,

Multiple scatter operations to out_raystatus result in the same error as encountered by out_xyz

I'm able to reproduce this with a local toy example within a loop, and given that this used to work previously, I would be inclined to say that this indeed is a legitimate bug. With respect to out_xyz, I think they are one and the same issue because when you do

dr.scatter(out_xyz, si.p, idx, active)

because out_xyz is of type Float this is in fact syntactic sugar for

dr.scatter(out_xyz, si.p.x, idx*3 + 0)
dr.scatter(out_xyz, si.p.y, idx*3 + 1)
dr.scatter(out_xyz, si.p.z, idx*3 + 2)

and so it's the same issue as with out_raystatus — namely, that performing multiple scatters on the the same target within the loop raises this error. We'll take a look at that.

When I enable the debug mode, I don't get any additional info output to stdout/stderr.

I don't think debug mode necessarily will find anything so I believe this is ok. What it is useful for is for tracking things such as out-of-bounds gathers/scatters and should then hopefully report something.

[tomas16]

Hi @rtabbara, that makes sense. How would you explain that the two scatters to out_active do work?

[rtabbara]

How would you explain that the two scatters to out_active do work?

I had a quick look and had a similar setup your example of out_active. By setting dr.set_log_level(dr.LogLevel.Debug) I can see the trace and that for the first scatter we have

jit_var_scatter(): r3[r16] = r18 (type=bool, mask=r17, ptr=r0, mode=auto) [skipped, target/source are value variables with the same value]

so we have an optimisation where the first scatter isn't even performed. If I instead change it to

# changed from False to True
dr.scatter(out_active, True, idx, mask)

then I do indeed get the same issues as reported with the other variables. So I think this was the reason.

[rtabbara]

Hi @tomas16 ,

An update on this — the latest Mitsuba master should have incorporated fixes to support multiple scatter calls to the same target within a symbolic loop. If you have a chance to pull the latest and try it out that would be great.

As an aside, just to highlight a potential pitfall, in evaluated mode, Dr.Jit tries to be conservative with respect to sequences of scatters to the same target to avoid write conflicts. So in particular, by default, it will evaluate the target prior to a subsequent scatter (which is why you initially encountered this bug in the symbolic scope). If you know that your write indices are independent across scatters, you use the argument mode=drjit.ReduceMode.NoConflicts which can avoid further kernel splits.

All of this is to say, that when we're in a symbolic loop body, obviously these conservative checks can no longer be used because we can no longer evaluate the target within the loop body. So whether you have one or multiple scatters is irrelevant — there is still an implicit assumption that scatters to the same target within the loop body have to have no write conflicts across iterations (and within the body itself). You're probably aware of this, but maybe it's good to highlight to other users too.

Anyway, I'll mark this discussion as answered, but if you encounter any other issues with regards to the port, please add independent issues and we'll track them.

[tomas16]

Hi @rtabbara ,

Awesome, I tried the latest commit and my code works now!

I'm not sure I fully understand your explanation. Perhaps you can define what exactly you mean by write conflict?

I have many variations on the theme of multiple scatter ops to the same target inside a loop body. In some cases, I'm using the same target and index, but changing the values and/or active mask. I'm certainly sometimes overwriting the same element of the target array. My understanding of the limitation is that when you're inside of the symbolic loop, you can't scatter (write), then read from that array before performing another scatter to it. In my code I'm scattering to output arrays that aren't read until after mitsuba has finished executing and I convert the outputs to numpy arrays.

Thank you and the whole mitsuba team for developing this awesome tool!

And an extra shout out for enabling this (didn't used to work last time I tried it):

pip install git+ssh://[email protected]/mitsuba-renderer/drjit.git
pip install git+ssh://[email protected]/mitsuba-renderer/mitsuba3.git

It makes testing unreleased versions a little easier :)

[rtabbara]

Hi @tomas16 ,

Glad you got it working!

I'm not sure I fully understand your explanation. Perhaps you can define what exactly you mean by write conflict?

As a super contrived example consider this:

import drjit as dr

@dr.syntax
def foo():
    n = 10000000
    y = dr.zeros(dr.llvm.Float, n)

    idx = dr.arange(dr.llvm.UInt32, n)
    idx2 = n - 1 - idx

    a = dr.full(dr.llvm.Float, 1, n)
    b = dr.full(dr.llvm.Float, 2, n)

    j = dr.zeros(dr.llvm.UInt32, n)
    while dr.hint(j < 1, mode='symbolic'):
        j += 1
        dr.scatter(target=y, index=idx2, value=a)

        # If everything was in lock-step we would expect y == b
        dr.scatter(target=y, index=idx, value=b)

    return y

y = foo()

# We would expect min == max == 2
print(f'Min: {dr.min(y)} Max: {dr.max(y)}')

Running this on my machine I erroneously get

Min: [1] Max: [2]

however if I instead change the mode to evaluated (where we evaluate the target before the second scatter) , I do indeed get the expected result of

Min: [2] Max: [2]

so this is what I meant with regards to write conflicts. If you have multiple scatters in your kernel to the same region of memory of the target, there's no guarantee that those writes will happen in sequence because they're happening across multiple threads. In practice, you may be fine if the indices don't change across invocations, but in general this is undefined behaviour.

[tomas16]

Thanks for this explanation @rtabbara

I'm definitely doing something like this in my code, only I don't change the index argument to scatter, but the active mask. The index is the same arange in all these cases. I haven't seen any issues yet (including with pre-3.6.0), but I'll use an evaluated loop instead of a symbolic one just in case.

I'm still a little confused as to why it is the way it is. You're saying these writes are happening across multiple threads. However, my mental model is that the whole loop is compiled into a single kernel. That kernel runs in parallel over the data width using SIMD and threads. So for most operations, if they're sequential in the kernel they should be executed sequentially for each "lane" of data. I can see how scatter may be different due to its random access nature, where these parallel lanes of execution now write to each other's data. But then I would think the same issue would happen with any (non-repeated) scatter inside a symbolic loop, as one lane of execution could be a full iteration ahead when it performs its scatter.

Perhaps this also explains why I haven't run into any issues. Since my index is always an arange, each lane of execution only writes to its own data. I'm modulating whether the write occurs or not through the active mask, but this happens sequentially within each lane. Does that make sense?

[rtabbara]

Hi @tomas16 ,

I can see how scatter may be different due to its random access nature, where these parallel lanes of execution now write to each other's data. But then I would think the same issue would happen with any (non-repeated) scatter inside a symbolic loop, as one lane of execution could be a full iteration ahead when it performs its scatter.

You're exactly right. That's precisely the potential danger. In my contrived example, I used multiple scatters in the loop body just to make it a bit easier to convey, but whether you, for example, have a loop of iteration length 10 with a single scatter in the loop body or a loop of iteration length 5 with two scatters in the loop body, it's kind of the same thing in terms of the actual operations that are performed. And as you alluded to, even with a single scatter in your loop body, you have to worry about write conflicts across iterations.

Since my index is always an arange, each lane of execution only writes to its own data

I think this is probably fine then because it sounds like you could also express your write as

x = dr.select(active, new_value, old_val)