Repository navigation
Element assignment and in-place tile solves give wrong results right after the tile is read #2306
Description
Bug Description
A kernel that reads a shared tile and then overwrites it in place gets wrong results on CUDA. Two kinds of in-place write show it: element assignment (y[1, 0] = x) and wp.tile_lower_solve_inplace(). Each kernel below adds y into acc and then overwrites y, four times over, so every block should end with the same acc. Instead, every block comes out wrong.
import numpy as np
import warp as wp
wp.config.enable_mathdx_solver = False
TILE = 32
ITERS = 4
NUM_BLOCKS = 1024
@wp.func
def busy(n: float) -> float:
# loops n times and returns 0
s = float(0.0)
for _ in range(int(n)):
s = s * 0.5 + 1.0
return wp.where(s > 1.0e30, 1.0, 0.0)
@wp.kernel(enable_backward=False)
def assign_after_read(Y: wp.array2d[float], work: wp.array2d[float], out: wp.array3d[float]):
i = wp.tid()
y = wp.tile_load(Y, shape=(TILE, TILE), storage="shared")
w = wp.tile_load(work, shape=(TILE, TILE))
acc = wp.tile_zeros(shape=(TILE, TILE), dtype=float)
for it in range(ITERS):
acc += wp.tile_map(busy, w) # adds 0
acc += y
y[1, 0] = float(it + 100)
wp.tile_store(out[i], acc)
@wp.kernel(enable_backward=False)
def solve_after_read(L: wp.array2d[float], Y: wp.array2d[float], work: wp.array2d[float], out: wp.array3d[float]):
i = wp.tid()
l = wp.tile_load(L, shape=(TILE, TILE), storage="shared")
y = wp.tile_load(Y, shape=(TILE, TILE), storage="shared")
w = wp.tile_load(work, shape=(TILE, TILE))
acc = wp.tile_zeros(shape=(TILE, TILE), dtype=float)
for _ in range(ITERS):
acc += wp.tile_map(busy, w) # adds 0
acc += y
wp.tile_lower_solve_inplace(l, y) # y = y / 2
wp.tile_store(out[i], acc)
rng = np.random.default_rng(0)
Y = rng.integers(0, 8, size=(TILE, TILE)).astype(np.float32)
L = 2.0 * np.eye(TILE, dtype=np.float32)
work = np.zeros((TILE, TILE), dtype=np.float32)
work[1::2] = 2000.0 # more work on odd rows
expected_assign = ITERS * Y
expected_assign[1, 0] = Y[1, 0] + 100 + 101 + 102
expected_solve = Y * (1 + 1 / 2 + 1 / 4 + 1 / 8)
with wp.ScopedDevice("cuda:0"):
out = wp.zeros((NUM_BLOCKS, TILE, TILE), dtype=float)
wp.launch_tiled(
assign_after_read, dim=NUM_BLOCKS, inputs=[wp.array(Y), wp.array(work)], outputs=[out], block_dim=64
)
wrong = np.any(out.numpy() != expected_assign, axis=(1, 2))
print(f"element assignment: blocks with wrong results: {wrong.sum()} of {NUM_BLOCKS}")
out = wp.zeros((NUM_BLOCKS, TILE, TILE), dtype=float)
wp.launch_tiled(
solve_after_read,
dim=NUM_BLOCKS,
inputs=[wp.array(L), wp.array(Y), wp.array(work)],
outputs=[out],
block_dim=64,
)
wrong = np.any(out.numpy() != expected_solve, axis=(1, 2))
print(f"tile_lower_solve_inplace: blocks with wrong results: {wrong.sum()} of {NUM_BLOCKS}")Output on Warp 1.19.0:
element assignment: blocks with wrong results: 1024 of 1024
tile_lower_solve_inplace: blocks with wrong results: 1024 of 1024
Expected output:
element assignment: blocks with wrong results: 0 of 1024
tile_lower_solve_inplace: blocks with wrong results: 0 of 1024
The busy() calls add 0 to acc. They only make the work uneven across the tile: odd rows loop 2000 times and even rows not at all. Without them the failures are rarer but still there. In one run, element assignment failed in 2 of 1024 blocks (559 of 8192 with block_dim=256), and the solve passed.
wp.tile_upper_solve_inplace() and wp.tile_cholesky_solve_inplace() fail the same way in an equivalent kernel. The solves fail only on the fallback path, which the script selects with wp.config.enable_mathdx_solver = False. With cuSolverDx, the solve kernel prints 0 of 1024, while element assignment still fails.
Y holds small integers and L is 2 * I, so every sum and halving is exact in float32. Any mismatch is a wrong result, not rounding.
#2299 reports the same read-then-overwrite pattern for wp.tile_matmul().
Application / Workflow Impact
Found during code review of the fix for #2299 and confirmed with isolated tests, not while working on a project. Both patterns are ordinary tile code: updating a few entries of a shared tile between uses, and repeated in-place solves against the same factor. The results are silently wrong, and nothing raises an error.
Warp Version
1.19.0 (warp-lang from PyPI). It also reproduces on a source build of current main with the pending fix for #2299 applied.
System Information
- Python 3.14.6, Ubuntu 24.04.5 LTS
- CUDA: NVIDIA GeForce RTX 3090, driver 620.70.06
- CPU: both kernels pass with the default
block_dim=1. Withwp.config.enable_cpu_blocks = True, both fail in every block.
Activity