diff --git a/slangpy/tests/slangpy_tests/test_torchintegration.py b/slangpy/tests/slangpy_tests/test_torchintegration.py index a8d6d8afc..1944f6b10 100644 --- a/slangpy/tests/slangpy_tests/test_torchintegration.py +++ b/slangpy/tests/slangpy_tests/test_torchintegration.py @@ -649,6 +649,55 @@ def test_null_grad_idifftensor(device_type: DeviceType): loss.backward() +@pytest.mark.parametrize("device_type", DEVICE_TYPES) +@pytest.mark.parametrize("a_grad,b_grad", [(True, True), (True, False), (False, True)]) +def test_mixed_requires_grad_idifftensor(device_type: DeviceType, a_grad: bool, b_grad: bool): + """ + Regression for shader-slang/slangpy#1056. + + A [Differentiable] function with two IDiffTensor inputs failed during backward + when only some inputs had requires_grad=True. The compiled backward kernel + scatters _grad_out unconditionally for every IDiffTensor param, but dispatch only + bound a grad buffer for requires_grad=True inputs. On CUDA the other input's + _grad_out was a dangling device pointer the atomic scatter faulted on, aborting + the process with CUDA_ERROR_ILLEGAL_ADDRESS; on other backends type resolution + rejected the unbound gradient with a TypeError. yes/yes worked; yes/no and no/yes + failed. The autouse torch_bridge_mode fixture exercises both bridge paths. + """ + src = """ +import slangpy; + +[Differentiable] +void mul2(uint index, IDiffTensor a, IDiffTensor b, IWDiffTensor out) +{ + out[index] = a[index] * b[index]; +} +""" + import torch + + device = helpers.get_torch_device(device_type) + module = helpers.create_module(device, src) + + a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device="cuda", requires_grad=a_grad) + b = torch.tensor([5.0, 6.0, 7.0, 8.0], dtype=torch.float32, device="cuda", requires_grad=b_grad) + out = torch.zeros(size=(4,), dtype=torch.float32, device="cuda", requires_grad=True) + + module.mul2(index=grid(shape=(4,)), a=a, b=b, out=out) + out.sum().backward() + + # sum(a * b): d/da = b, d/db = a. torch reports .grad only for requires_grad leaves. + if a_grad: + assert a.grad is not None + compare_tensors(b.detach(), a.grad) + else: + assert a.grad is None + if b_grad: + assert b.grad is not None + compare_tensors(a.detach(), b.grad) + else: + assert b.grad is None + + @pytest.mark.parametrize("device_type", DEVICE_TYPES) def test_nn_parameter_as_input(device_type: DeviceType): """ diff --git a/src/slangpy_ext/utils/slangpy.cpp b/src/slangpy_ext/utils/slangpy.cpp index ca4be81d0..ac048ff23 100644 --- a/src/slangpy_ext/utils/slangpy.cpp +++ b/src/slangpy_ext/utils/slangpy.cpp @@ -572,7 +572,16 @@ nb::tuple NativeCallData::autograd_backward( pair->grad = bridge.create_zeros_like_tensor(pair->primal); input_grads.append(pair->grad); } else { - pair->grad = nb::none(); + // An IDiffTensor input always binds as a DiffTensor in the backward pass, so it + // needs an output-gradient buffer regardless of whether this leaf wants a + // gradient: on CUDA the compiled kernel scatters into _grad_out through a raw + // device pointer, so an unbound buffer is a dangling pointer the atomic write + // faults on (CUDA_ERROR_ILLEGAL_ADDRESS, #1056); on other backends type + // resolution rejects a diff tensor with no associated output gradient outright + // (tensorcommon.py). Bind a throwaway zeroed buffer so both paths are satisfied; + // it is discarded - torch still sees no gradient for this leaf via the None + // appended below. + pair->grad = bridge.create_zeros_like_tensor(pair->primal); input_grads.append(nb::none()); } } else {