Add stdlib tests for `clamp` derivatives which also checks `max` and `min` derivatives (#4136)

* Add stdlib tests for `clamp` derivatives which also checks `max` and `min` derivatives

* Extend test

2 files changed, 212 insertions, 0 deletions

diff --git a/tests/autodiff-dstdlib/dstdlib-clamp.slang b/tests/autodiff-dstdlib/dstdlib-clamp.slang
new file mode 100644
index 000000000..32b1cc8eb
--- /dev/null
+++ b/tests/autodiff-dstdlib/dstdlib-clamp.slang
@@ -0,0 +1,181 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+typedef DifferentialPair<float> dpfloat;
+typedef DifferentialPair<float2> dpfloat2;
+typedef DifferentialPair<float3> dpfloat3;
+
+[Differentiable]
+float _clamp(float x, float min, float max)
+{
+    return clamp(x, min, max);
+}
+
+[Differentiable]
+float3 _clamp3(float3 x, float3 min, float3 max)
+{
+    return clamp(x, min, max);
+}
+
+[Differentiable]
+float _clamp_equiv(float x, float _min, float _max)
+{
+    return max(_min, min(_max, x));
+}
+
+[Differentiable]
+float3 _clamp_equiv(float3 x, float3 _min, float3 _max)
+{
+    return max(_min, min(_max, x));
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    // x in between max and min
+    {
+        dpfloat dpx = dpfloat(2.0, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp)(dpx, dpmin, dpmax);
+        outputBuffer[0] = res.d; // Expected: 0.1
+    }
+
+    // x less than min
+    {
+        dpfloat dpx = dpfloat(0.5, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp)(dpx, dpmin, dpmax);
+        outputBuffer[1] = res.d; // Expected: 0.3
+    }
+
+    // x greater than max
+    {
+        dpfloat dpx = dpfloat(4.0, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp)(dpx, dpmin, dpmax);
+        outputBuffer[2] = res.d; // Expected: 0.2
+    }
+
+    // float3 version with one in between, one below min and one above max.
+    {
+        dpfloat3 dpx = dpfloat3(float3(2.0, 0.5, 4.0), float3(0.1, 0.1, 0.1));
+        dpfloat3 dpmax = dpfloat3(float3(3.0, 3.0, 3.0), float3(0.2, 0.2, 0.2));
+        dpfloat3 dpmin = dpfloat3(float3(1.0, 1.0, 1.0), float3(0.3, 0.3, 0.3));
+
+        dpfloat3 res = fwd_diff(_clamp3)(dpx, dpmin, dpmax);
+        outputBuffer[3] = res.d.x; // Expected: 0.1
+        outputBuffer[4] = res.d.y; // Expected: 0.3
+        outputBuffer[5] = res.d.z; // Expected: 0.2
+    }
+
+    // Equivalent to the first test, but with a different implementation of clamp
+    {
+        dpfloat dpx = dpfloat(2.0, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp_equiv)(dpx, dpmin, dpmax);
+        outputBuffer[6] = res.d; // Expected: 0.1
+    }
+
+    // Equivalent to the second test, but with a different implementation of clamp
+    {
+        dpfloat dpx = dpfloat(0.5, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp_equiv)(dpx, dpmin, dpmax);
+        outputBuffer[7] = res.d; // Expected: 0.3
+    }
+
+    // Equivalent to the third test, but with a different implementation of clamp
+    {
+        dpfloat dpx = dpfloat(4.0, 0.1);
+        dpfloat dpmax = dpfloat(3.0, 0.2);
+        dpfloat dpmin = dpfloat(1.0, 0.3);
+
+        dpfloat res = fwd_diff(_clamp_equiv)(dpx, dpmin, dpmax);
+        outputBuffer[8] = res.d; // Expected: 0.2
+    }
+
+    // Equivalent to the fourth test, but with a different implementation of clamp
+    {
+        dpfloat3 dpx = dpfloat3(float3(2.0, 0.5, 4.0), float3(0.1, 0.1, 0.1));
+        dpfloat3 dpmax = dpfloat3(float3(3.0, 3.0, 3.0), float3(0.2, 0.2, 0.2));
+        dpfloat3 dpmin = dpfloat3(float3(1.0, 1.0, 1.0), float3(0.3, 0.3, 0.3));
+
+        dpfloat3 res = fwd_diff(_clamp_equiv)(dpx, dpmin, dpmax);
+        outputBuffer[9] = res.d.x; // Expected: 0.1
+        outputBuffer[10] = res.d.y; // Expected: 0.3
+        outputBuffer[11] = res.d.z; // Expected: 0.2
+    }
+
+    // Reverse-mode tests.
+    
+    // x in between max and min
+    {
+        dpfloat dpx = dpfloat(2.0, 0.0);
+        dpfloat dpmax = dpfloat(3.0, 0.0);
+        dpfloat dpmin = dpfloat(1.0, 0.0);
+
+        bwd_diff(_clamp)(dpx, dpmin, dpmax, 1.0);
+        
+        outputBuffer[12] = dpx.d; // Expected: 1.0
+        outputBuffer[13] = dpmin.d; // Expected: 0.0
+        outputBuffer[14] = dpmax.d; // Expected: 0.0
+    }
+
+    // x less than min
+    {
+        dpfloat dpx = dpfloat(0.5, 0.0);
+        dpfloat dpmax = dpfloat(3.0, 0.0);
+        dpfloat dpmin = dpfloat(1.0, 0.0);
+
+        bwd_diff(_clamp)(dpx, dpmin, dpmax, 1.0);
+        
+        outputBuffer[15] = dpx.d; // Expected: 0.0
+        outputBuffer[16] = dpmin.d; // Expected: 1.0
+        outputBuffer[17] = dpmax.d; // Expected: 0.0
+    }
+
+    // x greater than max
+    {
+        dpfloat dpx = dpfloat(4.0, 0.0);
+        dpfloat dpmax = dpfloat(3.0, 0.0);
+        dpfloat dpmin = dpfloat(1.0, 0.0);
+
+        bwd_diff(_clamp)(dpx, dpmin, dpmax, 1.0);
+        
+        outputBuffer[18] = dpx.d; // Expected: 0.0
+        outputBuffer[19] = dpmin.d; // Expected: 0.0
+        outputBuffer[20] = dpmax.d; // Expected: 1.0
+    }
+
+    // float3 version with one in between, one below min and one above max.
+    {
+        dpfloat3 dpx = dpfloat3(float3(2.0, 0.5, 4.0), float3(0.0, 0.0, 0.0));
+        dpfloat3 dpmax = dpfloat3(float3(3.0, 3.0, 3.0), float3(0.0, 0.0, 0.0));
+        dpfloat3 dpmin = dpfloat3(float3(1.0, 1.0, 1.0), float3(0.0, 0.0, 0.0));
+
+        bwd_diff(_clamp3)(dpx, dpmin, dpmax, float3(0.1, 0.2, 0.3));
+        
+        outputBuffer[21] = dpx.d.x; // Expected: 0.1
+        outputBuffer[22] = dpx.d.y; // Expected: 0.0
+        outputBuffer[23] = dpx.d.z; // Expected: 0.0
+        outputBuffer[24] = dpmin.d.x; // Expected: 0.0
+        outputBuffer[25] = dpmin.d.y; // Expected: 0.2
+        outputBuffer[26] = dpmin.d.z; // Expected: 0.0
+        outputBuffer[27] = dpmax.d.x; // Expected: 0.0
+        outputBuffer[28] = dpmax.d.y; // Expected: 0.0
+        outputBuffer[29] = dpmax.d.z; // Expected: 0.3
+    }
+}
diff --git a/tests/autodiff-dstdlib/dstdlib-clamp.slang.expected.txt b/tests/autodiff-dstdlib/dstdlib-clamp.slang.expected.txt
new file mode 100644
index 000000000..b00b0060b
--- /dev/null
+++ b/tests/autodiff-dstdlib/dstdlib-clamp.slang.expected.txt
@@ -0,0 +1,31 @@
+type: float
+0.100000
+0.300000
+0.200000
+0.100000
+0.300000
+0.200000
+0.100000
+0.300000
+0.200000
+0.100000
+0.300000
+0.200000
+1.000000
+0.000000
+0.000000
+0.000000
+1.000000
+0.000000
+0.000000
+0.000000
+1.000000
+0.100000
+0.000000
+0.000000
+0.000000
+0.200000
+0.000000
+0.000000
+0.000000
+0.300000