Incorrect gradients when a kernel uses the value returned by wp.atomic_add() #2157

Bug Description

When a kernel uses the value returned by wp.atomic_add() in a differentiable computation, tape.backward() produces a NaN gradient. The forward result is correct, and Warp reports no error or warning.

import warp as wp


@wp.kernel
def scale_by_previous(x: wp.array[float], total: wp.array[float], y: wp.array[float]):
    previous = wp.atomic_add(total, 0, 1.0)
    y[0] = previous * x[0]


x = wp.array([3.0], dtype=float, requires_grad=True)
total = wp.array([10.0], dtype=float)
y = wp.zeros(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(scale_by_previous, dim=1, inputs=[x, total], outputs=[y])

tape.backward(grads={y: wp.ones(1, dtype=float)})
print("y:", y.numpy())
print("x.grad:", x.grad.numpy())

Output on a CUDA device:

y: [30.]
x.grad: [nan]

Running on the CPU device prints the same output. wp.atomic_sub() also gives a NaN gradient in the same setup, and a kernel that uses the returned value as an array index (slot = wp.atomic_add(count, 0, 1) followed by out[slot] = ...) produces incorrect input gradients.

The differentiability guide covers a related case in its custom replay function example. That example uses the index returned by wp.atomic_add(), shows the incorrect gradient that results, and explains it by saying the index returned by the adjoint of wp.atomic_add() "is always zero". The repro above does not match that explanation: a previous of zero would give an x.grad of [0.], not [nan]. Neither the guide's Limitations and Workarounds section nor the wp.atomic_add() reference page lists this pattern as unsupported.

Expected: x.grad is [10.]. In the forward pass previous is 10.0 and does not depend on x, so the derivative of y with respect to x is 10.0. If Warp cannot differentiate through this pattern, it should raise an error instead of returning NaN.

Related: #1858 covers documenting atomic behavior, including how a consumed return value affects backward replay.

System Information

• Warp 1.20.0.dev0 (current main)
• NVIDIA GeForce RTX 3090, CUDA Toolkit 13.4, Driver 13.4
• Python 3.12.3, Linux x86_64