Robustness fixes around reverse-mode differentiation of variables & inout parameters (#2985)

* Add a new test case for checking loop in the reverse mode

* Create duplicate var for primal value to avoid inconsistent inputs to backward call

Fixes an issue with inout parameters where the backprop call may use the post-call value of the var instead of the pre-call value.

* `IRStore`s transpose to `IRLoad` and an `IRStore(0)` to clear differential.

Fixes some subtle issues around transposing

* Simplify test

* Delete out-edited.hlsl

---------

Co-authored-by: Lifan Wu <lifanw@nvidia.com>
Co-authored-by: Yong He <yonghe@outlook.com>

4 files changed, 142 insertions, 1 deletions

diff --git a/source/slang/slang-ir-autodiff-transpose.h b/source/slang/slang-ir-autodiff-transpose.h
index 843612df9..739ce6553 100644
--- a/source/slang/slang-ir-autodiff-transpose.h
+++ b/source/slang/slang-ir-autodiff-transpose.h
@@ -1560,6 +1560,14 @@ struct DiffTransposePass
     TranspositionResult transposeStore(IRBuilder* builder, IRStore* fwdStore, IRInst*)
     {
         IRInst* revVal = builder->emitLoad(fwdStore->getPtr());
+
+        auto primalType = tryGetPrimalTypeFromDiffInst(fwdStore->getVal());
+        SLANG_ASSERT(primalType);
+
+        // Clear the value at the differential address, by setting to 0.
+        IRInst* emptyVal = emitDZeroOfDiffInstType(builder, primalType);
+        builder->emitStore(fwdStore->getPtr(), emptyVal);
+
         if (auto diffPairType = as<IRDifferentialPairType>(revVal->getDataType()))
         {
             revVal = builder->emitDifferentialPairGetDifferential(
diff --git a/source/slang/slang-ir-autodiff-unzip.h b/source/slang/slang-ir-autodiff-unzip.h
index 1d05a8081..2a5d73a01 100644
--- a/source/slang/slang-ir-autodiff-unzip.h
+++ b/source/slang/slang-ir-autodiff-unzip.h
@@ -340,7 +340,21 @@ struct DiffUnzipPass
                     // the remerged diff pair.
                     auto diffPairType = as<IRDifferentialPairType>(as<IRPtrTypeBase>(arg->getDataType())->getValueType());
                     auto primalValueType = diffPairType->getValueType();
-                    auto diffPairRef = diffBuilder->emitReverseGradientDiffPairRef(arg->getDataType(), primalArg, diffArg);
+                    
+                    // We can't simply reuse primalArg for an inout parameter since this will represent the value 
+                    // after the primal call which can potentially alter primalArg. Therefore, we will find the
+                    // first store into primalArg, and create a temp var holding that value (i.e. value prior to primal call)
+                    // 
+                    auto storeUse = findUniqueStoredVal(cast<IRVar>(primalArg));
+                    auto storeInst = cast<IRStore>(storeUse->getUser());
+
+                    auto storedVal = storeInst->getVal();
+
+                    // Emit the temp var into the primal blocks since it's holding a primal value.
+                    auto tempPrimalVar = primalBuilder->emitVar(primalValueType);
+                    primalBuilder->emitStore(tempPrimalVar, storedVal);
+                    
+                    auto diffPairRef = diffBuilder->emitReverseGradientDiffPairRef(arg->getDataType(), tempPrimalVar, diffArg);
                     diffBuilder->markInstAsDifferential(diffPairRef, primalValueType);
                     diffArgs.add(diffPairRef);
                 }
diff --git a/tests/autodiff/path-tracer/pt-loop.slang b/tests/autodiff/path-tracer/pt-loop.slang
new file mode 100644
index 000000000..2e865b80a
--- /dev/null
+++ b/tests/autodiff/path-tracer/pt-loop.slang
@@ -0,0 +1,115 @@
+//Tests automatic synthesis of Differential type requirement.
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct PathData : IDifferentiable
+{
+    float3 thp;    
+    uint length;
+    bool terminated;
+    bool isHit;
+
+    [BackwardDifferentiable]
+    __init()
+    {
+        this.thp = float3(1.f);
+        this.length = 0;
+        this.terminated = false;
+        this.isHit = false;
+    }
+}
+
+bool traceRayInline(uint length)
+{
+    if (length < 2) return true;
+    else return false;
+}
+
+float3 getAlbedo(uint length)
+{
+    return float3(0.9f, 1.f, 1.f);
+}
+
+float3 getAlbedoDerivative(uint length)
+{
+    return float3(1.f, 0.f, 0.f);
+}
+
+[ForwardDerivativeOf(getAlbedo)]
+[TreatAsDifferentiable]
+DifferentialPair<float3> __fwd_d_getAlbedo(uint length)
+{
+    float3 primalValue = getAlbedo(length);
+    float3 derivativeValue = no_diff getAlbedoDerivative(length);
+    return DifferentialPair<float3>(primalValue, derivativeValue);
+}
+
+[BackwardDerivativeOf(getAlbedo)]
+[TreatAsDifferentiable]
+void __bwd_d_getAlbedo(uint length, float3 dOut)
+{
+    outputBuffer[2] += dOut.x;
+}
+
+[BackwardDifferentiable]
+void handleHit(inout PathData pathData)
+{
+    if (pathData.length >= 2)
+    {
+        pathData.terminated = true;
+        return;
+    }
+
+    float3 albedo = getAlbedo(pathData.length);
+    pathData.thp *= albedo;
+    pathData.length++;
+}
+
+[BackwardDifferentiable]
+[PreferRecompute]
+float3 tracePath()
+{
+    PathData pathData = PathData();
+
+    if (traceRayInline(pathData.length))
+    {
+        pathData.isHit = true;
+    }
+    else
+    {
+        pathData.terminated = true;
+        pathData.isHit = false;
+    }
+
+    [MaxIters(4)]
+    while (!pathData.terminated)
+    {
+        if (pathData.isHit)
+        {
+            handleHit(pathData);
+
+            //pathData.isHit = traceRayInline(pathData.length);
+            if (!traceRayInline(pathData.length)) pathData.isHit = false;
+            else pathData.isHit = true;
+        }
+        else
+        {
+            pathData.terminated = true;
+        }
+    }
+    return pathData.thp;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispathThreadID: SV_DispatchThreadID)
+{
+    DifferentialPair<float3> dpThp = __fwd_diff(tracePath)();
+    outputBuffer[0] = dpThp.p.x;
+    outputBuffer[1] = dpThp.d.x;
+
+    __bwd_diff(tracePath)(float3(1.f, 0.f, 0.f));
+}
diff --git a/tests/autodiff/path-tracer/pt-loop.slang.expected.txt b/tests/autodiff/path-tracer/pt-loop.slang.expected.txt
new file mode 100644
index 000000000..5731f4ce2
--- /dev/null
+++ b/tests/autodiff/path-tracer/pt-loop.slang.expected.txt
@@ -0,0 +1,4 @@
+type: float
+0.81
+1.8
+1.8
\ No newline at end of file