Added a new differential type system and various improvements (#2343)

* 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()'

* Refactored variable names

* Added initial IR defs for new type system

* Added necessary logic for semantic checking

* Overhauled type system to use builtin pair types and conform to the IDifferentiable interface

* Automatically replace IRDifferentiablePairType to a custom IRStructType

* Added generics handling by expanding the conformance context functionality and allowing for type parameters

* Minor fix: early return in processPairTypes()

* Minor fixes to differentiable resolution on generic types

* Added new instructions for differential pairs. Basic tests work now.

Looking into generic types.

* Adjusted most tests to the new type system. OutType and InOutType are still not properly working.

* Updated __jvp to produce both primal and differential output

* Moved autodiff related declarations to diff.meta.slang

* Refactored variable names

* Added initial IR defs for new type system

* Added necessary logic for semantic checking

* Overhauled type system to use builtin pair types and conform to the IDifferentiable interface

* Automatically replace IRDifferentiablePairType to a custom IRStructType

* Added generics handling by expanding the conformance context functionality and allowing for type parameters

* Minor fix: early return in processPairTypes()

* Minor fixes to differentiable resolution on generic types

* Added new instructions for differential pairs. Basic tests work now.

Looking into generic types.

* Adjusted most tests to the new type system. OutType and InOutType are still not properly working.

* Updated __jvp to produce both primal and differential output

* Moved autodiff related declarations to diff.meta.slang

* Removed external changes

* Cleanup the transcription logic: each case returns a pair of insts for the primal and differential computation.

10 files changed, 142 insertions, 50 deletions

diff --git a/tests/autodiff/arithmetic-jvp.slang b/tests/autodiff/arithmetic-jvp.slang
index 4e4d200e1..ddd1a4aa9 100644
--- a/tests/autodiff/arithmetic-jvp.slang
+++ b/tests/autodiff/arithmetic-jvp.slang
@@ -4,14 +4,17 @@
 //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;
+
 __differentiate_jvp float f(float x)
 {
     return x;
 }
 
-float g_jvp_(float x, float dx)
+dpfloat g_jvp_(dpfloat dpx)
 {
-    return 2 * dx;
+    return dpfloat(dpx.p(), 2 * dpx.d());
 }
 
 [__custom_jvp(g_jvp_)]
@@ -37,15 +40,13 @@ __differentiate_jvp float j(float x, float y)
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
     {
-        float a = 2.0;
-        float b = 1.5;
-        float da = 1.0;
-        float db = 1.0;
-
-        outputBuffer[0] = __jvp(f)(a, da);        // Expect: 1
-        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
+        dpfloat dpa = dpfloat(2.0, 1.0);
+        dpfloat dpb = dpfloat(1.5, 1.0);
+
+        outputBuffer[0] = __jvp(f)(dpa).d();                    // Expect: 1
+        outputBuffer[1] = __jvp(f)(dpfloat(dpa.p(), 0.0)).d();  // Expect: 0
+        outputBuffer[2] = __jvp(g)(dpa).d();                    // Expect: 2
+        outputBuffer[3] = __jvp(h)(dpa, dpb).d();               // Expect: 8
+        outputBuffer[4] = __jvp(j)(dpa, dpb).d();               // Expect: 1
     }
 }
diff --git a/tests/autodiff/generic-jvp.slang b/tests/autodiff/generic-jvp.slang
new file mode 100644
index 000000000..48993c21c
--- /dev/null
+++ b/tests/autodiff/generic-jvp.slang
@@ -0,0 +1,30 @@
+//TEST_IGNORE_FILE:(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//TEST_IGNORE_FILE:(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 __DifferentialPair<double> dpdouble;
+typedef __DifferentialPair<float3> dpfloat3;
+
+__generic<T:__BuiltinArithmeticType>
+__differentiate_jvp T g(T x)
+{
+    return x + x;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    {
+        dpfloat dpa = dpfloat(2.0, 1.0);
+        dpdouble dpb = dpdouble(1.5, 2.0);
+        dpfloat3 dpf3 = dpfloat3(float3(1.0, 3.0, 5.0), float3(0.5, 1.5, 2.5));
+
+        outputBuffer[0] = f(dpa.p());                               // Expect: 1
+        outputBuffer[1] = __jvp(f)(dpfloat(2.0, 0.0)).d();          // Expect: 0
+        outputBuffer[2] = (float)__jvp(f)(dpb).d();                 // Expect: 2
+        outputBuffer[3] = __jvp(f)(dpf3).d().y;              // Expect: 1.5
+    }
+}
diff --git a/tests/autodiff/inout-parameters-jvp.slang b/tests/autodiff/inout-parameters-jvp.slang
index 40e9d30ca..ba04c6b65 100644
--- a/tests/autodiff/inout-parameters-jvp.slang
+++ b/tests/autodiff/inout-parameters-jvp.slang
@@ -4,6 +4,7 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float> dpfloat;
 
 __differentiate_jvp void g(float x, float y, inout float z)
 {
@@ -30,14 +31,16 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     float dy = 0.5;
     float dz = 2.5;
 
-    __jvp(h)(x, y, z, dx, dy, dz);
+    dpfloat dpz = dpfloat(z, dz);
 
-    outputBuffer[0] = dz; // Expect: 12.0
-    outputBuffer[1] = z;  // Expect: 1.0
+    __jvp(h)(dpfloat(x, dx), dpfloat(y, dy), dpz);
 
-    __jvp(g)(x, y, z, dx, dy, dz);
+    outputBuffer[0] = dpz.d(); // Expect: 12.0
+    outputBuffer[1] = dpz.p(); // Expect: 6.75
 
-    outputBuffer[2] = dz; // Expect: 21.5
-    outputBuffer[3] = z;  // Expect: 1.0
+    __jvp(g)(dpfloat(x, dx), dpfloat(y, dy), dpz);
+
+    outputBuffer[2] = dpz.d(); // Expect: 21.5
+    outputBuffer[3] = dpz.p(); // Expect: 12.5
 
 }
\ 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 c48ef7bf6..324de53ca 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
+6.75
 21.5
-1.0
\ No newline at end of file
+12.5
\ No newline at end of file
diff --git a/tests/autodiff/local-redecl-custom-jvp.slang b/tests/autodiff/local-redecl-custom-jvp.slang
index 2bc7cd582..6241a8bf5 100644
--- a/tests/autodiff/local-redecl-custom-jvp.slang
+++ b/tests/autodiff/local-redecl-custom-jvp.slang
@@ -3,11 +3,16 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float> dpfloat;
+typedef float.Differential dfloat;
+
 import test_intrinsics;
 
-float my_pow_jvp(float x, float n, float dx, float dn)
+dpfloat my_pow_jvp(dpfloat x, dpfloat n)
 {
-    return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x);
+    return dpfloat(
+        pow(x.p(), n.p()),
+        x.d() * n.p() * pow(x.p(), n.p()-1) + n.d() * pow(x.p(), n.p()) * log(x.p()));
 }
 
 [__custom_jvp(my_pow_jvp)]
@@ -17,12 +22,12 @@ float _pow(float, float);
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
      {
-        float a = 5.0;
-        float n = 2;
-        float da = 1.0;
-        float dn = 0;
+        dpfloat dpa = dpfloat(5.0, 1.0);
+        dpfloat dpn = dpfloat(2, 0.0);
 
-        outputBuffer[0] = __jvp(_pow)(a, n, da, dn);        // Expect: 10.0
-        outputBuffer[1] = __jvp(_pow)(a, n, 0.0, 1.0);      // Expect: 40.23595
+        outputBuffer[0] = __jvp(_pow)(dpa, dpn).d();        // Expect: 10.0
+        outputBuffer[1] = __jvp(_pow)(
+            dpfloat(dpa.p(), 0.0),
+            dpfloat(dpn.p(), 1.0)).d();                     // Expect: 40.23595
     }
 }
diff --git a/tests/autodiff/nested-jvp.slang b/tests/autodiff/nested-jvp.slang
index 222396ec8..baebeee56 100644
--- a/tests/autodiff/nested-jvp.slang
+++ b/tests/autodiff/nested-jvp.slang
@@ -4,29 +4,34 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float> dpfloat;
+typedef __DifferentialPair<float3> dpfloat3;
+
 [__custom_jvp(pow_jvp)]
 float pow_(float x, float n)
 {
     return pow<float>(x, n);
 }
 
-
 [__custom_jvp(max_jvp)]
 float max_(float x, float y)
 {
     return max<float>(x, y);
 }
 
-
-float pow_jvp(float x, float n, float dx, float dn)
+dpfloat pow_jvp(dpfloat x, dpfloat n)
 {
-    return dx * n * pow(x, n-1) + ((dn != 0.0) ? (dn * pow(x, n) * log(x)) : 0.0);
+    return dpfloat(
+        pow(x.p(), n.p()),
+        x.d() * n.p() * pow(x.p(), n.p()-1) + 
+            ((n.d() != 0.0) ? (n.d() * pow(x.p(), n.p()) * log(x.p())) : 0.0));
 }
 
-
-float max_jvp(float x, float y, float dx, float dy)
+dpfloat max_jvp(dpfloat x, dpfloat y)
 {
-    return (x > y) ? dx : dy;
+    return dpfloat(
+        max(x.p(), y.p()),
+        (x.p() > y.p()) ? x.d() : y.d());
 }
 
 
@@ -53,7 +58,10 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float3 d_f90 = float3(0.9, 0.9, 0.9);
         float d_cosTheta = 1.0;
 
-        outputBuffer[0] = __jvp(fresnel)(f0, f90, cosTheta, d_f0, d_f90, d_cosTheta).y; // Expect: -0.031250
+        outputBuffer[0] = __jvp(fresnel)(
+            dpfloat3(f0, d_f0),
+            dpfloat3(f90, d_f90),
+            dpfloat(cosTheta, d_cosTheta)).d().y; // Expect: -0.031250
 
         float a = 1.0;
         float b = -0.4;
@@ -63,8 +71,16 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         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;
+        outputBuffer[1] = __jvp(g)(
+            dpfloat(a, da),
+            dpfloat(b, db),
+            dpfloat(c, dc)).d();                 // Expect: -0.24375
+
+        outputBuffer[2] = g(a, b, c);       // Expect: 0.95625
+
+        outputBuffer[3] = __jvp(g)(
+            dpfloat(a, da),
+            dpfloat(b, db),
+            dpfloat(3.0, dc)).d();               // Expect: -0.4;
     }
 }
diff --git a/tests/autodiff/out-parameters-jvp.slang b/tests/autodiff/out-parameters-jvp.slang
index 58c6cfeb0..b243d4fb5 100644
--- a/tests/autodiff/out-parameters-jvp.slang
+++ b/tests/autodiff/out-parameters-jvp.slang
@@ -4,6 +4,8 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float> dpfloat;
+
 __differentiate_jvp void h(float x, float y, out float result)
 {
     float m = x + y;
@@ -20,9 +22,9 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     float dx = 1.0;
     float dy = 0.5;
 
-    float dresult = 0.0f;
-    __jvp(h)(x, y, dx, dy, dresult);
+    dpfloat dresult;
+    __jvp(h)(dpfloat(x, dx), dpfloat(y, dy), dresult);
 
-    outputBuffer[0] = dresult; // Expect: 9.5
+    outputBuffer[0] = dresult.d(); // Expect: 9.5
 
 }
\ No newline at end of file
diff --git a/tests/autodiff/test-intrinsics-jvp.slang b/tests/autodiff/test-intrinsics-jvp.slang
new file mode 100644
index 000000000..333c89189
--- /dev/null
+++ b/tests/autodiff/test-intrinsics-jvp.slang
@@ -0,0 +1,17 @@
+//TEST_IGNORE_FILE:
+
+__exported import test_intrinsics;
+
+[__custom_jvp(pow_jvp)]
+float pow_(float x, float n);
+float pow_jvp(float x, float n, float dx, float dn)
+{
+    return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x);
+}
+
+[__custom_jvp(max_jvp)]
+float max_(float x, float y);
+float max_jvp(float x, float y, float dx, float dy)
+{
+    return (x > y) ? dx : dy;
+}
\ No newline at end of file
diff --git a/tests/autodiff/vector-arithmetic-jvp.slang b/tests/autodiff/vector-arithmetic-jvp.slang
index 2b43f1752..393cc18ec 100644
--- a/tests/autodiff/vector-arithmetic-jvp.slang
+++ b/tests/autodiff/vector-arithmetic-jvp.slang
@@ -4,6 +4,10 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float2> dpfloat2;
+typedef __DifferentialPair<float3> dpfloat3;
+typedef __DifferentialPair<float4> dpfloat4;
+
 __differentiate_jvp float3 f(float3 x)
 {
     return x;
@@ -37,6 +41,7 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         float3 a = float3(2.0, 2.0, 2.0);
         float3 b = float3(1.5, 1.5, 1.5);
         float3 da = float3(1.0, 1.0, 1.0);
+        //dpfloat3 dpa = dpfloat3(a, da);
 
         float2 a2 = float2(2.0, 1.0);
         float2 b2 = float2(1.5, -2.0);
@@ -44,9 +49,18 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         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
+        outputBuffer[0] = __jvp(f)(dpfloat3(a, da)).d().z;    // Expect: 1
+
+        outputBuffer[1] = __jvp(g)(
+            dpfloat3(a, da),
+            dpfloat3(b, float3(2.0, 1.0, 0.0))).d().y;        // Expect: 8
+
+        outputBuffer[2] = __jvp(h)(
+            dpfloat2(a2, float2(1.0, 0.0)),
+            dpfloat2(b2, float2(1.0, 1.0))).d().x;            // Expect: 8
+
+        outputBuffer[3] = __jvp(j)(
+            dpfloat4(a4, float4(1.0)),
+            dpfloat4(b4, float4(2.0))).d().w;                 // Expect: 9
     }
 }
diff --git a/tests/autodiff/vector-swizzle-jvp.slang b/tests/autodiff/vector-swizzle-jvp.slang
index 6722b54dc..775c0140e 100644
--- a/tests/autodiff/vector-swizzle-jvp.slang
+++ b/tests/autodiff/vector-swizzle-jvp.slang
@@ -4,6 +4,10 @@
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
 
+typedef __DifferentialPair<float2> dpfloat2;
+typedef __DifferentialPair<float3> dpfloat3;
+typedef __DifferentialPair<float4> dpfloat4;
+
 __differentiate_jvp float2 f(float3 x)
 {
     return x.zy;
@@ -23,16 +27,16 @@ 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
+        outputBuffer[0] = __jvp(f)(dpfloat3(a, da)).d().x;  // Expect: 1
+        outputBuffer[1] = __jvp(f)(dpfloat3(a, da)).d().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
+        outputBuffer[2] = __jvp(g)(dpfloat3(x, dx), dpfloat4(y, dy)).d().x; // Expect: -2.25
+        outputBuffer[3] = __jvp(g)(dpfloat3(x, dx), dpfloat4(y, dy)).d().y; // Expect: 0.5
         
     }
 }