Added forward-mode autodiff support for more instructions (#2331)

* Merge slang-ir-diff-jvp.cpp

* Added support and tests for other float vector types

* Added swizzle test and code to handle it (tests failing currently)

* Fixed one test, the other is still pending

* Fixed instruction cloning logic to avoid modifying original function

* Fixed an issue with custom 'pow_jvp' and added support for vector contructor

* Minor update to comments

* Fixed support for division

* Fixed an issue with uninitialized diagnostic sink

* Moved derivative processing to after mandatory inlining.

Skip instructions that don't have side-effects and aren't used by anything.

* WIP: Handling unconditional control flow and multi-block functions

* Support for unconditional multi-block functions

* Added a dead code elimination step to the derivative pass

* Changed name of 'hasNoSideEffects()'

10 files changed, 114 insertions, 13 deletions

diff --git a/tests/autodiff/arithmetic-jvp.slang b/tests/autodiff/arithmetic-jvp.slang
index 8632aa332..4e4d200e1 100644
--- a/tests/autodiff/arithmetic-jvp.slang
+++ b/tests/autodiff/arithmetic-jvp.slang
@@ -1,7 +1,7 @@
 //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
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
 __differentiate_jvp float f(float x)
@@ -27,11 +27,16 @@ __differentiate_jvp float h(float x, float y)
     return m * n + 2 * x * y;
 }
 
+__differentiate_jvp float j(float x, float y)
+{
+    float m = x / y;
+    return m * y;
+}
 
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
-     {
+    {
         float a = 2.0;
         float b = 1.5;
         float da = 1.0;
@@ -41,5 +46,6 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         outputBuffer[1] = __jvp(f)(a, 0.0);       // Expect: 0
         outputBuffer[2] = __jvp(g)(a, da);        // Expect: 2
         outputBuffer[3] = __jvp(h)(a, b, da, db); // Expect: 8
+        outputBuffer[4] = __jvp(j)(a, b, da, db); // Expect: 1
     }
 }
diff --git a/tests/autodiff/arithmetic-jvp.slang.expected.txt b/tests/autodiff/arithmetic-jvp.slang.expected.txt
index 0545c08a1..597fe4035 100644
--- a/tests/autodiff/arithmetic-jvp.slang.expected.txt
+++ b/tests/autodiff/arithmetic-jvp.slang.expected.txt
@@ -2,4 +2,5 @@ type: float
 1.0
 0.0
 2.0
-8.0
\ No newline at end of file
+8.0
+1.0
\ No newline at end of file
diff --git a/tests/autodiff/inout-parameters-jvp.slang b/tests/autodiff/inout-parameters-jvp.slang
index 989e56c02..40e9d30ca 100644
--- a/tests/autodiff/inout-parameters-jvp.slang
+++ b/tests/autodiff/inout-parameters-jvp.slang
@@ -4,6 +4,14 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+
+__differentiate_jvp void g(float x, float y, inout float z)
+{
+    float m = x + y;
+    float n = x - y;
+    z += m * n + 2 * x * y;
+}
+
 __differentiate_jvp void h(float x, float y, inout float z)
 {
     float m = x + y;
@@ -27,4 +35,9 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     outputBuffer[0] = dz; // Expect: 12.0
     outputBuffer[1] = z;  // Expect: 1.0
 
+    __jvp(g)(x, y, z, dx, dy, dz);
+
+    outputBuffer[2] = dz; // Expect: 21.5
+    outputBuffer[3] = z;  // Expect: 1.0
+
 }
\ No newline at end of file
diff --git a/tests/autodiff/inout-parameters-jvp.slang.expected.txt b/tests/autodiff/inout-parameters-jvp.slang.expected.txt
index d8a590c0e..c48ef7bf6 100644
--- a/tests/autodiff/inout-parameters-jvp.slang.expected.txt
+++ b/tests/autodiff/inout-parameters-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 12.0
 1.0
-0.0
-0.0
\ No newline at end of file
+21.5
+1.0
\ No newline at end of file
diff --git a/tests/autodiff/nested-jvp.slang b/tests/autodiff/nested-jvp.slang
index c37641c91..222396ec8 100644
--- a/tests/autodiff/nested-jvp.slang
+++ b/tests/autodiff/nested-jvp.slang
@@ -20,7 +20,7 @@ float max_(float x, float y)
 
 float pow_jvp(float x, float n, float dx, float dn)
 {
-    return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x);
+    return dx * n * pow(x, n-1) + ((dn != 0.0) ? (dn * pow(x, n) * log(x)) : 0.0);
 }
 
 
@@ -36,6 +36,10 @@ __differentiate_jvp float3 fresnel(float3 f0, float3 f90, float cosTheta)
     return f0 + (f90 - f0) * pow_(max_(1 - cosTheta, 0.0), 5);
 }
 
+__differentiate_jvp float g(float a, float b, float c)
+{
+    return fresnel(float3(a), float3(b), 2 * c * c).y;
+}
 
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
@@ -50,5 +54,17 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float d_cosTheta = 1.0;
 
         outputBuffer[0] = __jvp(fresnel)(f0, f90, cosTheta, d_f0, d_f90, d_cosTheta).y; // Expect: -0.031250
+
+        float a = 1.0;
+        float b = -0.4;
+        float c = 0.5;
+
+        float da = -0.4;
+        float db = -1.0;
+        float dc = 0.2;
+
+        outputBuffer[1] = __jvp(g)(a, b, c, da, db, dc);    // Expect: -0.24375
+        outputBuffer[2] = g(a, b, c);                       // Expect: 0.95625
+        outputBuffer[3] = __jvp(g)(a, b, 3.0, da, db, dc);  // Expect: -0.4;
     }
 }
diff --git a/tests/autodiff/nested-jvp.slang.expected.txt b/tests/autodiff/nested-jvp.slang.expected.txt
index 107153351..f200c0a88 100644
--- a/tests/autodiff/nested-jvp.slang.expected.txt
+++ b/tests/autodiff/nested-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 -0.031250
-0.0
-0.0
-0.0
\ No newline at end of file
+-0.243750
+0.95625
+-0.4
\ No newline at end of file
diff --git a/tests/autodiff/vector-arithmetic-jvp.slang b/tests/autodiff/vector-arithmetic-jvp.slang
index 5852b49a7..2b43f1752 100644
--- a/tests/autodiff/vector-arithmetic-jvp.slang
+++ b/tests/autodiff/vector-arithmetic-jvp.slang
@@ -16,6 +16,20 @@ __differentiate_jvp float3 g(float3 x, float3 y)
     return a * b + 2 * x * y;
 }
 
+__differentiate_jvp float2 h(float2 x, float2 y)
+{
+    float2 a = x + y;
+    float2 b = x - y;
+    return a * b + 2 * x * y;
+}
+
+__differentiate_jvp float4 j(float4 x, float4 y)
+{
+    float4 a = x + y;
+    float4 b = x - y;
+    return a * b + 2 * x * y;
+}
+
 [numthreads(1, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
@@ -24,7 +38,15 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float3 b = float3(1.5, 1.5, 1.5);
         float3 da = float3(1.0, 1.0, 1.0);
 
-        outputBuffer[0] = __jvp(f)(a, da).z;                          // Expect: 1
-        outputBuffer[1] = __jvp(g)(a, b, da, float3(2.0, 1.0, 0.0)).y;       // Expect: 8
+        float2 a2 = float2(2.0, 1.0);
+        float2 b2 = float2(1.5, -2.0);
+
+        float4 a4 = float4(2.0, 1.0, 0.0, 2.0);
+        float4 b4 = float4(1.5, -2.0, 1.0, 1.5);
+
+        outputBuffer[0] = __jvp(f)(a, da).z;                                            // Expect: 1
+        outputBuffer[1] = __jvp(g)(a, b, da, float3(2.0, 1.0, 0.0)).y;                  // Expect: 8
+        outputBuffer[2] = __jvp(h)(a2, b2, float2(1.0, 0.0), float2(1.0, 1.0)).x;       // Expect: 8
+        outputBuffer[3] = __jvp(j)(a4, b4, float4(1.0), float4(2.0)).w;                 // Expect: 9
     }
 }
diff --git a/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
index adbb9a448..9202deda2 100644
--- a/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
+++ b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt
@@ -1,5 +1,5 @@
 type: float
 1.0
 8.0
-0.0
-0.0
\ No newline at end of file
+8.0
+9.0
\ No newline at end of file
diff --git a/tests/autodiff/vector-swizzle-jvp.slang b/tests/autodiff/vector-swizzle-jvp.slang
new file mode 100644
index 000000000..6722b54dc
--- /dev/null
+++ b/tests/autodiff/vector-swizzle-jvp.slang
@@ -0,0 +1,38 @@
+//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;
+
+__differentiate_jvp float2 f(float3 x)
+{
+    return x.zy;
+}
+
+__differentiate_jvp float2 g(float3 x, float4 y)
+{
+    float3 a = x + y.zyx;
+    float2 b = x.zx - y.yw;
+    return a.yz * b + 2 * x.xy * y.yz;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    {
+        float3 a = float3(2.0, 2.0, 2.0);
+        float3 da = float3(1.0, 0.5, 1.0);
+
+        outputBuffer[0] = __jvp(f)(a, da).x;  // Expect: 1
+        outputBuffer[1] = __jvp(f)(a, da).y;  // Expect: 0.5
+        
+        float3 x = float3(0.5, 2.0, 0.5);
+        float4 y = float4(-1.5, 1.0, 4.0, 2.0);
+        float3 dx = float3(1.0, 0.0, -1.0);
+        float4 dy = float4(0.0, 0.5, -0.25, 1.0);
+
+        outputBuffer[2] = __jvp(g)(x, y, dx, dy).x; // Expect: -2.25
+        outputBuffer[3] = __jvp(g)(x, y, dx, dy).y; // Expect: 0.5
+        
+    }
+}
diff --git a/tests/autodiff/vector-swizzle-jvp.slang.expected.txt b/tests/autodiff/vector-swizzle-jvp.slang.expected.txt
new file mode 100644
index 000000000..3a49054c9
--- /dev/null
+++ b/tests/autodiff/vector-swizzle-jvp.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+1.0
+0.5
+-2.25
+0.5
\ No newline at end of file