Is there a way to use dr.scatter inside a recorded loop?

[tomas16]

I've developed an iterative algorithm whose structure is something like:

def algo():
    output = dr.zeros(mi.Int, n)

    # a bunch of calculation

    iteration = mi.UInt32(0)
    active = compute_active()
    loop = mi.Loop("Iterations", lambda: (active, output, iteration))
    
    while loop(active):

        # perform calculations involving the creation of several 
        # variables that are only relevant to the current iteration

        dr.scatter(output, ...)
        iteration += 1
        active = compute_active()       # compute_active uses the 'output' variable

    return output

I developed the code first without the loop in place, so I could validate the first iteration and make sure everything's correct. Now that I've added the loop, I get this error:

RuntimeError: jit_var_scatter(): cannot scatter to a placeholder variable!

This refers to one of the variables that gets created where I put the comment inside the loop. Essentially what happens is:

1. I get a SurfaceInteraction3f object, named si
2. From that I spawn new rays and intersect with my mesh again, giving me a 2nd SurfaceInteraction3f object, named si2
3. Finally, I do dr.scatter(si, si2, dr.arange(mi.Int, dr.width(si)), to_update), which is where I get the error.

The to_update variable is also computed from si.
I'm essentially trying to apply a "correction" to a subset of these surface interactions.

So my questions:

1. As per the title: is it at all possible to use dr.scatter inside a recorded loop?
2. If not, is there any alternative for what I'm trying to do here? One way I can see is to make the loop body a drjit kernel, but do the looping and evaluation of the loop condition in plain old python. I'm guessing that would require a call to dr.eval(output) inside the loop.
3. I have additional dr.scatter calls to my output variable. Will these also be broken? Or is this a different case because output is listed as one of the variables that changes through the loop by the lambda?
4. What are the rules around placeholder variables and recorded loops? The documentation only discusses this at a high level.

More conceptually, I'm also a bit puzzled by the difference between placeholder and non-placeholder variables. DrJit is always recording the operations to compile them to a kernel, regardless of whether there's a loop, isn't it? And it's also unclear to me why this scatter isn't allowed/possible, at least in my case where all my variables have the same width that never changes.

[rtabbara]

Hi @tomas16 ,

1. As per the title: is it at all possible to use dr.scatter inside a recorded loop?

Yes, however there are caveats. The first issue is that as you've encountered, you're not allowed to use a placeholder variable as a target. This in particular includes loop state variables and also variables in your body that depend on these state variables

loop = mi.Loop("Iterations", lambda: (active, output, iteration))

while loop(active):
    # If si depends on iteration, active or output then still placeholder
    si = ...

    # Not allowed!
    dr.scatter(si, ...)

So alternatively, you could exclude your scatter target as a loop state variable. However, you cannot attempt to evaluate the target variable if the other scatter parameters (source, index, mask) are placeholders. In fact, attempting to explicitly evaluate any placeholder variable in the loop body isn't allowed

while loop(active):
	# Not allowed
	dr.eval(output)

There's a slight nuance when it comes to dr.scatter however that causes an implicit evaluation. Consider this block of code (no loop)

target = mi.Float(0)
source = mi.Float(1)
index = mi.UInt32(0)
dr.scatter(target, source, index)

target += 1

which with Info logging gives us

jit_eval(): launching 1 kernel.
  -> launching 6831c9701ffc9547 (n=1, in=1, out=0, se=1, ops=7, jit=17.18 us):
jit_eval(): done.

So even though we're not actually evaluating the final target value, a kernel is still launched prior to recording the operation target += 1. The reason for this is that in the context of vectorized types, when we're performing these unordered writes, there needs to be synchronisation across the entire array (not just within a given warp thread) to guarantee all writes have finished before we attempt any subsequent ordered stores/loads. The mechanism that Dr.Jit uses to achieve that is to evaluate the scatter prior to any subsequent operations that depend on the target output.

Going back to the loop then

loop = mi.Loop("Iterations", lambda: (active, iteration))

while loop(active):
    # Not a state variable, so ok
    dr.scatter(output, ...)

    # Trying to perform "ordered" load/store - i.e. we need to perform eval
    # Not allowed!
    output += 1

So the rule of thumb for recorded loops is exclusively scattering (or scatter reduce) to your target will work, but attempting to do any additional arithmetic operations or otherwise will be problematic

2. If not, is there any alternative for what I'm trying to do here? One way I can see is to make the loop body a drjit kernel, but do the looping and evaluation of the loop condition in plain old python. I'm guessing that would require a call to dr.eval(output) inside the loop.

You can alternatively just disable loop recording either globally or just for a block of code using dr.scoped_set_flag

with dr.scoped_set_flag(dr.JitFlag.LoopRecord, False):
   loop = mi.Loop("Iterations", lambda: (active, output, iteration))
   while loop(active):
      ...

Loop state placeholders variables are no longer created when recording the loop body and instead a kernel is launched at each iteration.

3.I have additional dr.scatter calls to my output variable. Will these also be broken? Or is this a different case because output is listed as one of the variables that changes through the loop by the lambda?

As mentioned loop state variables will have corresponding placeholders created so this indeed won't work with recorded loops

More conceptually, I'm also a bit puzzled by the difference between placeholder and non-placeholder variables. DrJit is always recording the operations to compile them to a kernel, regardless of whether there's a loop, isn't it? And it's also unclear to me why this scatter isn't allowed/possible, at least in my case where all my variables have the same width that never changes.

I agree that the differences aren't super clear, but hopefully my explanation has shed some light. For what it's worth, we are in the midst of integrating Nanobind into Dr.Jit and with that includes a lot of improvements and more comprehensive documentation. In particular, I believe with this new version, we gracefully handle the case when a scatter target is a loop state variable. Additionally, the terminology for symbolic, placeholder and recorded-loops has been unified to avoid confusion. This is all still a work in progress, but thought it was nonetheless worth mentioning.

[tomas16]

Thanks, this was super helpful. Had you not mentioned the ongoing work on documentation I would've suggested copying this post into the documentation.

Regarding having placeholder variables in the loop: it seems to me almost every realistic workload will have this problem, because usually everything ultimately has a dependency on the active variable. And in this case, since I'm iteratively updating my output variable, the output itself is a loop state variable.

Thanks for pointing me to the dr.scoped_set_flag trick, I'll try it out. I already did what I suggested myself in my original post and did the looping in regular python. I observed that it launched about 1500 kernels (many of them new and not cached) for only 28 iterations, which was quite puzzling. It worked, but it wasn't very efficient. Perhaps in the way you suggested it'll be better at re-using the same kernel representing the loop body.

[wjakob]

Regarding having placeholder variables in the loop: it seems to me almost every realistic workload will have this problem, because usually everything ultimately has a dependency on the active variable. And in this case, since I'm iteratively updating my output variable, the output itself is a loop state variable.

Let me clarify: You may perform scatters of data that depend on loop state variables, but where the target itself isn't a loop state variable. So what you describe is actually the normaly way that one would use scatters in a loop. We realize that the distinction into which variables are state variables and which aren't can be confusing -- this issue will go away in the future because the framework will just decide this for you automatically.

For now, you can keep the following rule in mind: if a variable occurs as target/source in a dr.gather/dr.scatter operation within a loop, then that is not a state variable.

[tomas16]

Thanks for the clarification. One more thing: in the pseudo-code in my original post, does the output variable have to be a loop state variable? It's being written to (using dr.scatter) in every iteration.

Also, an update on my previous comment: I tested two versions of the code:

The first uses a regular python loop: while np.any(active := compute_active()): # np.any() triggers kernel eval
The second uses @rtabbara 's suggestion of temporarily setting dr.JitFlag.LoopRecord to False, while still using an mi.Loop.

I cleared the kernel cache and ran v1: runtime = 52.8s, and 707 kernels were compiled.
I re-ran v1 with the cache in place: runtime = 7.3s.
I then cleared the cache again and ran v2: it's been running for over 90min and has produced over 10k kernels in the cache. Not sure if it's actually gonna finish or what's going on.

I guess conceptually we're trying to do the same thing in both cases: repeatedly evaluate a kernel that represents the loop body. I probably got something wrong still. But v1 works so I'm happy to leave it at that and move on to the next problem :)