Various auto-diff bug fixes. (#2646)

Co-authored-by: Yong He <yhe@nvidia.com>

5 files changed, 88 insertions, 2 deletions

diff --git a/tests/autodiff/bool-return-val-bwd.slang b/tests/autodiff/bool-return-val-bwd.slang
new file mode 100644
index 000000000..40eb4810a
--- /dev/null
+++ b/tests/autodiff/bool-return-val-bwd.slang
@@ -0,0 +1,28 @@
+//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 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct NonDiff
+{
+    float a;
+}
+
+[BackwardDifferentiable]
+bool myFunc(NonDiff fIn, inout float x)
+{
+    x = pow(x, fIn.a);
+    return x > 100.f;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    float a = 3.0;
+    NonDiff fIn = { a };
+    DifferentialPair<float> dpx = DifferentialPair<float>(4.f, 1.f);
+    __bwd_diff(myFunc)(fIn, dpx);
+
+    outputBuffer[0] = dpx.d;
+}
\ No newline at end of file
diff --git a/tests/autodiff/bool-return-val-bwd.slang.expected.txt b/tests/autodiff/bool-return-val-bwd.slang.expected.txt
new file mode 100644
index 000000000..255e4b2f4
--- /dev/null
+++ b/tests/autodiff/bool-return-val-bwd.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+48.000000
+0.000000
+0.000000
+0.000000
\ No newline at end of file
diff --git a/tests/autodiff/bsdf/bsdf-auto-rev.slang b/tests/autodiff/bsdf/bsdf-auto-rev.slang
index c2ba434f8..7fae5e993 100644
--- a/tests/autodiff/bsdf/bsdf-auto-rev.slang
+++ b/tests/autodiff/bsdf/bsdf-auto-rev.slang
@@ -32,11 +32,16 @@ struct Auto_Bwd_BSDFParameters : IDifferentiable
 };
 
 [BackwardDifferentiable]
-void bsdfGGXSample(in ShadingData sd, in Auto_Bwd_BSDFParameters params, out Auto_Bwd_ScatterSample result)
+bool bsdfGGXSample(in ShadingData sd, in Auto_Bwd_BSDFParameters params, out Auto_Bwd_ScatterSample result)
 {
     float3 wiLocal = no_diff(sd.toLocal(sd.V));
     float2 u = float2(0.8, 0.3);
 
+    if (wiLocal.z < 1e-6)
+    {
+        return false;
+    }
+
     // Taken from Rendering.Materials.Microfacet. Follows the Walter et al. EGSR07 BTDF paper
     float alphaSqr = params.roughness * params.roughness;
     float phi = u.y * (2 * 3.1415926);
@@ -52,6 +57,8 @@ void bsdfGGXSample(in ShadingData sd, in Auto_Bwd_BSDFParameters params, out Aut
     result.wo = no_diff(sd.fromLocal(woLocal)); // wo to world.
     result.pdf = detach(pdf);
     result.weight = evalGGXDivByPDF(wiLocal, woLocal, hLocal, params) * pdf / detach(pdf);
+
+    return woLocal.z > 1e-6;
 }
 
 [BackwardDifferentiable]
@@ -95,4 +102,4 @@ float bsdfGGXPDF(in float3 hLocal, in Auto_Bwd_BSDFParameters params)
     float d = ((cosTheta * a2 - cosTheta) * cosTheta + 1);
 
     return (a2 / (d * d * 3.1415926)) * cosTheta;
-}
+}
\ No newline at end of file
diff --git a/tests/autodiff/reverse-conditional-out-assign.slang b/tests/autodiff/reverse-conditional-out-assign.slang
new file mode 100644
index 000000000..4f6b105f1
--- /dev/null
+++ b/tests/autodiff/reverse-conditional-out-assign.slang
@@ -0,0 +1,40 @@
+//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 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+typedef DifferentialPair<float> dpfloat;
+typedef float.Differential dfloat;
+
+[BackwardDifferentiable]
+bool test_single_branch(float y, out float o)
+{
+    if (y > 0.5)
+    {
+        o = y * 2.0f;
+        return true;
+    }
+
+    o = y + 1.0f;
+
+    return false;    
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    {
+        dpfloat dpa = dpfloat(1.0, 0.0);
+
+        __bwd_diff(test_single_branch)(dpa, 1.0f);
+        outputBuffer[0] = dpa.d; // Expect: 2.0
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+        
+        __bwd_diff(test_single_branch)(dpa, 1.0f);
+        outputBuffer[1] = dpa.d; // Expect: 1.0
+    }
+}
diff --git a/tests/autodiff/reverse-conditional-out-assign.slang.expected.txt b/tests/autodiff/reverse-conditional-out-assign.slang.expected.txt
new file mode 100644
index 000000000..86aa47f11
--- /dev/null
+++ b/tests/autodiff/reverse-conditional-out-assign.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+2.000000
+1.000000
+0.000000
+0.000000
+0.000000