Bug Description
Launching the adjoint pass for a kernel that reads from a differentiable
wp.indexedarray input does not produce correct gradients.
On CUDA, the adjoint launch completes but leaves the source gradient as zeros.
On CPU, the adjoint launch segfaults.
Minimal repro:
import argparse
import faulthandler
import numpy as np
import warp as wp
faulthandler.enable()
@wp.kernel
def weighted_sample_sum(
samples: wp.indexedarray[float],
weights: wp.array[float],
total: wp.array[float],
):
i = wp.tid()
wp.atomic_add(total, 0, samples[i] * weights[i])
def run(device):
base_values = np.linspace(1.0, 6.0, 6, dtype=np.float32)
base = wp.array(base_values, dtype=float, device=device, requires_grad=True)
weights_np = np.array([0.25, 0.5, 1.0], dtype=np.float32)
weights = wp.array(weights_np, dtype=float, device=device)
sample_ids = wp.array([1, 3, 5], dtype=wp.int32, device=device)
samples = wp.indexedarray1d(base, [sample_ids])
total = wp.zeros(1, dtype=float, device=device, requires_grad=True)
wp.launch(weighted_sample_sum, dim=samples.size, inputs=[samples, weights], outputs=[total], device=device)
print(f"{device} forward total:", total.numpy())
total.grad.fill_(1.0)
print(f"{device} launching adjoint")
wp.launch(
weighted_sample_sum,
dim=samples.size,
inputs=[samples, weights],
outputs=[total],
adj_inputs=[base.grad, None],
adj_outputs=[total.grad],
adjoint=True,
device=device,
)
expected = np.zeros_like(base_values)
expected[[1, 3, 5]] = weights_np
print(f"{device} expected grad:", expected)
print(f"{device} actual grad: ", base.grad.numpy())
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--device", default="cpu")
args = parser.parse_args()
run(args.device)Expected gradient on both devices:
Observed on CUDA:
cuda:0 forward total: [8.5]
cuda:0 launching adjoint
cuda:0 expected grad: [0. 0.25 0. 0.5 0. 1. ]
cuda:0 actual grad: [0. 0. 0. 0. 0. 0.]
Observed on CPU:
cpu forward total: [8.5]
cpu launching adjoint
Fatal Python error: Segmentation fault
Current thread:
File ".../warp/_src/context.py", line 8367 in invoke
File ".../warp/_src/context.py", line 8874 in launch
An equivalent kernel that takes a regular wp.array plus explicit index array
and reads base[ids[i]] produces the expected gradient on both CPU and CUDA.
The failure appears specific to wp.indexedarray as the differentiable input.
System Information
Observed with:
Warp 1.15.0.dev0
Python 3.12.13
Linux 6.8.0
CUDA Toolkit 13.0
CUDA Driver 13.0
GPU: NVIDIA RTX PRO 6000 Blackwell Server Edition MIG 1g.24gb
Bug Description
Launching the adjoint pass for a kernel that reads from a differentiable
wp.indexedarrayinput does not produce correct gradients.On CUDA, the adjoint launch completes but leaves the source gradient as zeros.
On CPU, the adjoint launch segfaults.
Minimal repro:
Expected gradient on both devices:
Observed on CUDA:
Observed on CPU:
An equivalent kernel that takes a regular
wp.arrayplus explicit index arrayand reads
base[ids[i]]produces the expected gradient on both CPU and CUDA.The failure appears specific to
wp.indexedarrayas the differentiable input.System Information
Observed with: