Kernel requests too little shared memory when it moves a tile returned by a wp.func to shared storage #2294

Bug Description

A kernel that calls a @wp.func returning a tile, and then passes that tile to an operation that needs it in shared memory such as wp.tile_broadcast(), requests too little dynamic shared memory on CUDA. The tile inside the function is then allocated past the end of the request. Depending on the kernel's other tiles this goes unnoticed or is an illegal memory access.

import numpy as np
import warp as wp

K = 16
N = 20
BLOCK_DIM = 128


@wp.func
def double_tile(a: wp.array[float]):
    t = wp.tile_load(a, shape=K)
    return wp.tile_map(wp.add, t, t)


@wp.kernel(module="unique", enable_backward=False)
def compute(a: wp.array[float], J: wp.array2d[float], out: wp.array2d[float]):
    J_kj = wp.tile_load(J, shape=(K, N), storage="shared")
    d = double_tile(a)
    # tile_broadcast() needs `d` in shared memory
    J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_kj), wp.tile_broadcast(d, shape=(N, K)))
    wp.tile_store(out, wp.tile_matmul(J_ki, J_kj))


wp.config.verify_cuda = True

a = wp.ones(K, dtype=float, device="cuda:0")
J = wp.ones((K, N), dtype=float, device="cuda:0")
out = wp.zeros((N, N), dtype=float, device="cuda:0")

hooks = compute.module.load("cuda:0", BLOCK_DIM).get_kernel_hooks(compute)
# shared tiles: J_kj and J_ki (K * N floats each), d (K floats) and the product (N * N floats),
# plus the K floats that double_tile() returns
print("requested:", hooks.forward_smem_bytes, "allocated:", (2 * K * N + K + N * N + K) * 4)

wp.launch_tiled(compute, dim=[1], inputs=[a, J], outputs=[out], block_dim=BLOCK_DIM, device="cuda:0")
print("ok:", np.allclose(out.numpy(), 2.0 * K))

Output:

requested: 4224 allocated: 4288
Warp CUDA error 700: an illegal memory access was encountered (in function wp_cuda_context_check, /builds/omniverse/warp/warp/native/warp.cu:3023)
RuntimeError: CUDA error 700 detected on device cuda:0: Warp CUDA error 700: an illegal memory access was encountered

Expected: the kernel requests 4288 bytes and the launch succeeds. With K = 13 the request is short by the same 64 bytes but the launch happens to succeed.

The cause is in how the forward shared-memory total is computed. The tile returned by a user function shares its type object with the variable that receives it in the caller. Adjoint.add_call() records func.adj.get_total_required_shared() at the call site, while the function's tile is still a register tile. When the caller later passes the tile to an operation that sets storage = "shared" on its argument, the function's tile becomes a shared tile as well, but the caller's recorded total is not updated. The backward total is not affected, since ModuleBuilder._propagate_backward_shared_memory() reads the callee totals after the module is built.

GH-1901 would remove the in-place storage change that makes this possible.

Application / Workflow Impact

Found in MuJoCo Warp (google-deepmind/mujoco_warp#1879). The two dense Newton Hessian kernels shared a wp.func that returned a weight tile, and the solver crashed with an illegal memory access for some constraint-row limits. We worked around it by not returning the tile from a function.

Warp Version

1.18.0 and 1.20.0.dev20261111 (main at fef20ee)

System Information

Python 3.12 and 3.14, Linux x86_64. CUDA only. The illegal memory access above is from an NVIDIA RTX PRO 4500 Blackwell Server Edition, driver 615.71.09.