Fix CUDA illegal memory access in backward with mixed requires_grad (#1056)

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<float, 1> a, IDiffTensor<float, 1> b, IWDiffTensor<float, 1> 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):
     """

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 {