verify_autograd_array_access ignores arrays read or written by tile loads and stores #2314

Bug Description

With wp.config.verify_autograd_array_access = True, Warp marks each kernel array argument as read or written, and warns when an array is written after it was read. wp.tile_load(), wp.tile_load_indexed(), wp.tile_store(), wp.tile_store_indexed(), wp.tile_atomic_add(), and wp.tile_atomic_add_indexed() never set those flags, whereas wp.atomic_add() and array indexing do. Overwriting an array that a tile kernel read, or the reverse, is not reported, while the backward pass uses the overwritten values.

import warp as wp

wp.config.verify_autograd_array_access = True

N = 4


@wp.kernel
def tile_square_kernel(x: wp.array[float], y: wp.array[float]):
    t = wp.tile_load(x, shape=N)
    wp.tile_store(y, t * t)


@wp.kernel
def loss_kernel(y: wp.array[float], loss: wp.array[float]):
    wp.atomic_add(loss, 0, y[wp.tid()])


@wp.kernel
def overwrite_kernel(w: wp.array[float], x: wp.array[float]):
    tid = wp.tid()
    x[tid] = w[tid]


x = wp.array([1.0, 2.0, 3.0, 4.0], dtype=float, requires_grad=True)
y = wp.zeros(N, dtype=float, requires_grad=True)
w = wp.array([-1.0, -2.0, -3.0, -4.0], dtype=float, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch_tiled(tile_square_kernel, dim=[1], inputs=[x, y], block_dim=N)
    wp.launch(loss_kernel, N, inputs=[y, loss])
    wp.launch(overwrite_kernel, N, inputs=[w, x])  # writes x after the tile load read it

tape.backward(loss)
print("x.grad", x.grad.numpy())

Output on main: no warning, and x.grad [-2. -4. -6. -8.]; the correct gradient is [2. 4. 6. 8.]. The same launch order with a regular kernel instead of the tile kernel is reported. Inspecting kernel.adj.args shows is_read and is_write both False for the arrays of a kernel that only uses tile loads and stores.

Proposed fix: mark the arrays in the dispatch functions of these builtins, as atomic_op_dispatch_func() does for atomic operations. Related, and separate: reads in the left operand of a binary operator are not recorded either, which affects wp.tile_load(x, shape=N) * 2.0; that is a different code path (emit_BinOp).

With the check enabled, the fix adds warnings for tile code that was silent before, for example the tile kernels of example_navier_stokes_perturbation.py (ifft_tiled, tiled_transpose) and the tiled wp.sparse.bsr_mv kernel, which reuse an array for loads and stores. The non-tiled bsr_mv_kernel already warns on main.

I have a fix with a test and can open a PR.

Application / Workflow Impact

Found through isolated testing, not in a project. It affects anyone who debugs gradients of tile kernels with the overwrite checker: the check is silent for arrays that tile kernels read or write.

Warp Version

main at 82dab03 (the code in question is unchanged since b0c59f5, where the output above was produced; built from source, CPU only).

System Information

macOS 26.6.2, Apple M5 Pro, Python 3.12.11, CPU. The flags are set during code generation, so the logic does not depend on the device.