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, 18 insertions, 4 deletions

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