Fix issue with custom derivatives that mutate differentiable pointer types (#10464)

Fixes: #10195

Enables users to write custom differentiable pointer mutation (like
`getOffset`) and use forward-mode custom derivatives to control how the
derivative counterpart is mutated.

Added a diagnostic to make sure this cannot directly mix with
differentiable value types (i.e. the same function cannot mutate pointer
parameters and also have value parameters that require derivatives).

In the future, it may make more sense to separate such custom
derivatives into their own class as they aren't really forward-mode
derivatives, but for now, this should help unblock some users.

Author: saipraveenb25

source/slang/slang-diagnostics.lua

@@ -1799,6 +1799,13 @@ err(
     span { loc = "expr:Expr", message = "generic constraint for parameter '~param:Type' references type parameter '~referenced:Decl' before it is declared" }
 )
 
+err(
+    "cannot-mix-differentiable-value-and-ptr-outputs",
+    30118,
+    "cannot mix differentiable value types with differentiable pointer outputs",
+    span { loc = "location", message = "function has both IDifferentiable value types and IDifferentiablePtrType outputs, which is not currently supported. Please split the function so that differentiable value parameters and pointer differentiable outputs are in separate functions." }
+)
+
 
 -- Load semantic checking diagnostics (part 3) - Include, Visibility, and Capability
 -- (inlined from slang-diagnostics-semantic-checking-3.lua)

source/slang/slang-ir-autodiff-fwd.cpp

@@ -924,10 +924,75 @@ InstPair ForwardDiffTranscriber::transcribeCall(IRBuilder* builder, IRCall* orig
         diffReturnType = primalReturnType;
     }
 
+    // Determine if this call should use "pointer pair method" mode.
+    // If any output (out/inout/ref parameter, or return value) is a differentiable pointer type
+    // (conforms to IDifferentiablePtrType), we should not mark the call as mixed-differential
+    // so that the transposition logic won't try to reverse it (doesn't make sense for
+    // differential pointers).
+    //
+    // If we detect a differentiable pointer output but there are also any differentiable
+    // value types (either inputs or outputs, i.e. conforming to IDifferentiable), we
+    // emit a diagnostic since we cannot transpose such mixed functions.
+    //
+    bool hasDiffPtrOutput = false;
+    bool hasDiffValueType = false;
+
+    // Check all parameters for differentiable value types (inputs and outputs)
+    // and output parameters for differentiable pointer types.
+    for (UIndex ii = 0; ii < calleeType->getParamCount(); ii++)
+    {
+        auto paramType = calleeType->getParamType(ii);
+        bool isOutput = as<IROutParamTypeBase>(paramType) || as<IRRefParamType>(paramType);
+
+        // For outputs (out/inout/ref), get the value type from the pointer.
+        // For inputs, use the param type directly.
+        IRType* valueType;
+        if (isOutput)
+            valueType = (IRType*)as<IRPtrTypeBase>(paramType)->getValueType();
+        else
+            valueType = paramType;
+
+        // Only outputs can contribute differentiable pointer types.
+        if (isOutput && differentiableTypeConformanceContext.isDifferentiablePtrType(valueType))
+            hasDiffPtrOutput = true;
+
+        // Any parameter (input or output) can contribute differentiable value types.
+        if (differentiableTypeConformanceContext.isDifferentiableValueType(valueType))
+            hasDiffValueType = true;
+    }
+
+    // Check return type.
+    if (auto retType = calleeType->getResultType())
+    {
+        if (differentiableTypeConformanceContext.isDifferentiablePtrType((IRType*)retType))
+            hasDiffPtrOutput = true;
+        if (differentiableTypeConformanceContext.isDifferentiableValueType((IRType*)retType))
+            hasDiffValueType = true;
+    }
+
+    // If there are pointer differentiable outputs AND any differentiable value types
+    // (inputs or outputs), emit a diagnostic.
+    if (hasDiffPtrOutput && hasDiffValueType)
+    {
+        getSink()->diagnose(Diagnostics::CannotMixDifferentiableValueAndPtrOutputs{
+            .location = origCall->sourceLoc});
+
+        // We'll continue transcribing instead of erroring out right away since we
+        // can still generate code (it just won't be correct). The diagnostic
+        // will still disallow any final code from being generated.
+    }
+
+    bool usePointerPairMode = hasDiffPtrOutput;
+
     auto callInst = argBuilder.emitCallInst(diffReturnType, diffCallee, args);
     placeholderCall->removeAndDeallocate();
 
-    argBuilder.markInstAsMixedDifferential(callInst, diffReturnType);
+    // In pointer pair mode, do not mark as mixed differential since
+    // the transposition logic should not try to reverse this call.
+    if (!usePointerPairMode)
+    {
+        argBuilder.markInstAsMixedDifferential(callInst, diffReturnType);
+    }
     argBuilder.addAutoDiffOriginalValueDecoration(callInst, primalCallee);
 
     *builder = afterBuilder;

tests/autodiff/diff-ptr-mixed-value-and-ptr-error.slang

@@ -0,0 +1,40 @@
+//DIAGNOSTIC_TEST:SIMPLE(diag=CHECK,non-exhaustive):-target spirv
+
+// Test that mixing IDifferentiable value types with IDifferentiablePtrType outputs
+// in a differentiable function produces an error diagnostic.
+
+struct MyPtr : IDifferentiablePtrType
+{
+    typealias Differential = MyPtr;
+    uint offset;
+}
+
+// This function has both an IDifferentiablePtrType output (inout MyPtr)
+// and an IDifferentiable value input (float x), which is not supported.
+[Differentiable]
+[ForwardDerivative(fwd_mixed)]
+float mixedFunc(inout MyPtr ptr, float x)
+{
+    return x * ptr.offset;
+}
+
+DifferentialPair<float> fwd_mixed(inout DifferentialPtrPair<MyPtr> ptr, DifferentialPair<float> x)
+{
+    return diffPair(x.p * ptr.p.offset, x.d * ptr.p.offset);
+}
+
+[Differentiable]
+float caller(float x)
+{
+    var ptr = MyPtr(2.0);
+    let r = mixedFunc(ptr, x);
+//CHECK:             ^ cannot mix differentiable value types with differentiable pointer outputs
+    return r;
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    fwd_diff(caller)(diffPair(1.0, 1.0));
+}

tests/autodiff/diff-ptr-type-custom.slang

@@ -0,0 +1,67 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUFFER): -vk -shaderobj -output-using-type -emit-spirv-directly
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=BUFFER):-cuda -compute -shaderobj -output-using-type
+
+struct MyAddress : IDifferentiablePtrType
+{
+    typealias Differential = MyAddress;
+
+    RWStructuredBuffer<float>.Handle handle;
+    uint baseIndex;
+
+    __init(RWStructuredBuffer<float>.Handle h)
+    {
+        handle = h;
+        baseIndex = 0;
+    }
+
+    __subscript(uint index) -> float
+    {
+        [nonmutating]
+        get { return handle[baseIndex + index]; }
+    }
+
+    [Differentiable]
+    [ForwardDerivative(fwd_getOffset)]
+    This getOffset(int elements)
+    {
+        This addr = This(handle);
+        addr.baseIndex = baseIndex + uint(elements);
+        return addr;
+    }
+
+    static DifferentialPtrPair<This> fwd_getOffset(DifferentialPtrPair<This> self, int elements)
+    {
+        return DifferentialPtrPair<This>(
+            self.p.getOffset(elements),
+            self.d.getOffset(elements));
+    }
+}
+
+//TEST_INPUT: set buf = ubuffer(data=[1.0 2.0 3.0 4.0], stride=4)
+uniform RWStructuredBuffer<float>.Handle buf;
+
+//TEST_INPUT: set gradBuf = ubuffer(data=[0.0 0.0 0.0 0.0], stride=4)
+uniform RWStructuredBuffer<float>.Handle gradBuf;
+
+//TEST_INPUT: ubuffer(data=[0], stride=4):out,name=result
+RWStructuredBuffer<uint> result;
+
+[Differentiable]
+float compute(MyAddress addr, float x)
+{
+    let addr2 = addr.getOffset(2);
+    return x * no_diff addr2[0];
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    let addr = MyAddress(buf);
+    let gradAddr = MyAddress(gradBuf);
+    var addrPair = DifferentialPtrPair<MyAddress>(addr, gradAddr);
+    var xPair = diffPair(1.0f);
+    bwd_diff(compute)(addrPair, xPair, 1.0f);
+    result[0] = xPair.d == 3.0 ? 1 : 0;
+    // BUFFER: 1
+}