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.

1 files changed, 27 insertions, 11 deletions

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;
     }
 }