Out-of-bounds gradient write in reverse pass of runtime-length while loops #2169

Bug Description

In a runtime-length while loop, the reverse pass writes an array's gradient one element past the end of the array. With the kernel from #2125 and coefficients.requires_grad=True, coefficients.grad stays all zeros, and the gradient lands in the element just past the array. Here coefficients is a 3-element view of a 4-element buffer, so we can see the stray write in storage.grad[3].

import warp as wp

@wp.kernel
def while_max(x: wp.array[float], coefficients: wp.array[float], out: wp.array[float]):
    tid = wp.tid()
    best = float(-1.0e20)
    j = int(0)
    while j < 3:
        c = coefficients[j]
        j += 1
        if c == 0.0:
            continue
        value = x[tid] * c
        if value > best:
            best = value
    out[tid] = best

for device in ["cpu", "cuda:0"]:
    x = wp.array([1.0, 2.0, 3.0, 4.0], dtype=float, requires_grad=True, device=device)
    storage = wp.array([0.0, 0.0, 2.0, 0.0], dtype=float, requires_grad=True, device=device)
    coefficients = storage[:3]
    out = wp.zeros(4, dtype=float, requires_grad=True, device=device)
    with wp.Tape() as tape:
        wp.launch(while_max, dim=4, inputs=[x, coefficients, out], device=device)
    tape.backward(grads={out: wp.ones_like(out)})
    print(device, coefficients.grad.numpy(), storage.grad.numpy()[3])

Output on both devices:

coefficients.grad = [0.0, 0.0, 0.0]   (expected [0.0, 0.0, 10.0])
storage.grad[3]   = 10.0              (outside `coefficients`)

A product-style while loop shows the same thing: it writes 144.0 one element past the end of the input's gradient.

It looks like the reverse pass of a forward-order while loop reads its index variable after the final increment, so the adjoint of coefficients[j] is accumulated at index j + 1. Because this writes outside the array, it can silently corrupt memory next to the array. Gradients through dynamic loops are already tracked in #729, but I couldn't find an existing report of the out-of-bounds write specifically.

System Information

Warp built from main at 0c1ed78. Windows 11, RTX 4090. Reproduces on both cpu and cuda:0.