Initial support for differentiating existential types (#3111)

* Merge

* WIP: Complete auto-diff logic for existential types

* Revert "Add compiler option for generating representative hash"

This reverts commit 13b09ef4621e73844c96d64d9c111a8ed0d45aae.

* More fixes for fwd-mode AD on existential types

* Add anyValueSize inference pass

* Fix checking of `Differential.Differential==Differential`

* In-progress: infer any-value-size for existential types

* Existentials now work in forward-mode

* Overhaul handling of existential AD types. Fwd-mode works, reverse-mode requires front-end changes

* Reverse-mode now works on existentials

* Cleanup

* Remove diff rules for create existential object for now

* Revert treat-as-differentiable changes

* Fixes

* More fixes

* Cleanup

* more cleanup

* signed/unsigned

* Revert "Cleanup"

This reverts commit e4f7d71f07bb207736f90708961eeecd09a1b652.

* Cleanup (again)

* Remove public/export/keep-alive on null differential after AD pass

* Minor fix

* Update dictionary accessors

* Keep export decoration

* More fixes + Support for `kIROp_PackAnyValue`

* Merge upstream

* Update expected-failure.txt

8 files changed, 236 insertions, 1 deletions

diff --git a/tests/autodiff/existential-1.slang b/tests/autodiff/existential-1.slang
new file mode 100644
index 000000000..4f93dd04d
--- /dev/null
+++ b/tests/autodiff/existential-1.slang
@@ -0,0 +1,76 @@
+// Test calling differentiable function through dynamic dispatch.
+
+//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 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+[anyValueSize(16)]
+interface IInterface : IDifferentiable
+{
+    [Differentiable]
+    float calc(float x);
+}
+
+struct A : IInterface
+{
+    int data1;
+
+    __init(int data1) { this.data1 = data1; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * x * data1; }
+};
+
+struct B : IInterface
+{
+    int data1;
+    int data2;
+
+    __init(int data1, int data2) { this.data1 = data1; this.data2 = data2; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * data1 * data2; }
+};
+
+[Differentiable]
+float doThing(IInterface obj, float x)
+{
+    return obj.calc(x);
+}
+
+[Differentiable]
+float f(uint id, float x)
+{
+    IInterface obj;
+
+    if (id == 0)
+        obj = no_diff(A(2));
+    else
+        obj = no_diff(B(2, 3));
+
+    return doThing(obj, x);
+}
+
+//TEST_INPUT: type_conformance A:IInterface = 0
+//TEST_INPUT: type_conformance B:IInterface = 1
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    outputBuffer[0] = fwd_diff(f)(dispatchThreadID.x, DifferentialPair<float>(1.0, 2.0)).d; // A.calc, expect 12
+    outputBuffer[1] = fwd_diff(f)(dispatchThreadID.x + 1, DifferentialPair<float>(1.5, 1.0)).d; // B.calc, expect 18
+
+    {
+        var dpx = diffPair(1.0);
+        bwd_diff(f)(dispatchThreadID.x, dpx, 2.0); // A.calc, expect 12
+        outputBuffer[2] = dpx.d;
+    }
+
+    {
+        var dpx = diffPair(1.5);
+        bwd_diff(f)(dispatchThreadID.x + 1, dpx, 1.0); // B.calc, expect 18
+        outputBuffer[3] = dpx.d;
+    }
+}
diff --git a/tests/autodiff/existential-1.slang.expected.txt b/tests/autodiff/existential-1.slang.expected.txt
new file mode 100644
index 000000000..c5fce63b1
--- /dev/null
+++ b/tests/autodiff/existential-1.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+12.000000
+18.000000
+12.000000
+18.000000
+0.000000
diff --git a/tests/autodiff/existential-2.slang b/tests/autodiff/existential-2.slang
new file mode 100644
index 000000000..dd68c2972
--- /dev/null
+++ b/tests/autodiff/existential-2.slang
@@ -0,0 +1,75 @@
+// Test calling differentiable function through dynamic dispatch.
+
+//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 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+[anyValueSize(16)]
+interface IInterface : IDifferentiable
+{
+    [Differentiable]
+    float calc(float x);
+}
+
+struct A : IInterface
+{
+    float data1;
+
+    [Differentiable]
+    __init(float data1) { this.data1 = data1; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * x * data1; }
+};
+
+struct B : IInterface
+{
+    float data1;
+    float data2;
+
+    [Differentiable]
+    __init(float data1, float data2) { this.data1 = data1; this.data2 = data2; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * data1 * data2; }
+};
+
+[Differentiable]
+float doThing(IInterface obj, float x)
+{
+    return obj.calc(x);
+}
+ 
+[Differentiable]
+float f(uint id, float x)
+{
+    IInterface obj;
+
+    if (id == 0)
+        obj = A(x); // x^4
+    else
+        obj = B(x, x); // x^4
+
+    return doThing(obj, x) + doThing(obj, x); // 2 * x^4
+}
+
+//TEST_INPUT: type_conformance A:IInterface = 0
+//TEST_INPUT: type_conformance B:IInterface = 1
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    {
+        var dpx = diffPair(1.0);
+        bwd_diff(f)(dispatchThreadID.x, dpx, 2.0); 
+        outputBuffer[0] = dpx.d; // Expect: 2 * 4 * x^3 * dx = 8 * x^3 * dx = 16
+    }
+
+    {
+        var dpx = diffPair(1.5);
+        bwd_diff(f)(dispatchThreadID.x + 1, dpx, 1.0);
+        outputBuffer[1] = dpx.d; // Expect: 2 * 4 * x^3 * dx = 8 * (1.5)^3 * 1.0 = 27
+    }
+}
diff --git a/tests/autodiff/existential-2.slang.expected.txt b/tests/autodiff/existential-2.slang.expected.txt
new file mode 100644
index 000000000..98cc988a3
--- /dev/null
+++ b/tests/autodiff/existential-2.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+16.000000
+27.000000
+0.000000
+0.000000
+0.000000
diff --git a/tests/autodiff/null-differential.slang b/tests/autodiff/null-differential.slang
new file mode 100644
index 000000000..71e0f1046
--- /dev/null
+++ b/tests/autodiff/null-differential.slang
@@ -0,0 +1,68 @@
+// Test calling differentiable function through dynamic dispatch.
+
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+//DISABLE_TEST(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;
+
+[anyValueSize(16)]
+interface IInterface : IDifferentiable
+{
+    [Differentiable]
+    float calc(float x);
+}
+
+struct A : IInterface
+{
+    int data1;
+
+    __init(int data1) { this.data1 = data1; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * x * data1; }
+};
+
+struct B : IInterface
+{
+    int data1;
+    int data2;
+
+    __init(int data1, int data2) { this.data1 = data1; this.data2 = data2; }
+
+    [Differentiable]
+    float calc(float x) { return x * x * data1 * data2; }
+};
+
+[Differentiable]
+float doThing(IInterface obj, IInterface obj2, float x)
+{
+    return obj.calc(x);
+}
+
+//TEST_INPUT: type_conformance A:IInterface = 0
+//TEST_INPUT: type_conformance B:IInterface = 1
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    {
+        IInterface obj;
+        if (dispatchThreadID.x == 0)
+            obj = A(2);
+        else
+            obj = B(2, 3);
+
+        IInterface obj2 = A(3);
+
+        var dpx = diffPair(1.0);
+        DifferentialPair<IInterface> dpobj = diffPair<IInterface>(obj);
+        DifferentialPair<IInterface> dpobj2 = diffPair<IInterface>(obj2);
+
+        bwd_diff(doThing)(dpobj, dpobj2, dpx, 2.0); // A.calc, expect 12
+
+        IDifferentiable objd = dpobj.d;
+
+        outputBuffer[0] = (isNullDifferential((objd)) ? 1.f : 0.f);
+    }
+}
diff --git a/tests/autodiff/treat-as-differentiable.slang b/tests/autodiff/treat-as-differentiable-1.slang
index 95423d978..95423d978 100644
--- a/tests/autodiff/treat-as-differentiable.slang
+++ b/tests/autodiff/treat-as-differentiable-1.slang
diff --git a/tests/autodiff/treat-as-differentiable.slang.expected.txt b/tests/autodiff/treat-as-differentiable-1.slang.expected.txt
index 9d11e5c94..9d11e5c94 100644
--- a/tests/autodiff/treat-as-differentiable.slang.expected.txt
+++ b/tests/autodiff/treat-as-differentiable-1.slang.expected.txt
diff --git a/tests/expected-failure.txt b/tests/expected-failure.txt
index ad1901ca3..a04eb90c9 100644
--- a/tests/expected-failure.txt
+++ b/tests/expected-failure.txt
@@ -173,4 +173,8 @@ tests/slang-extension/cas-int64-byte-address-buffer.slang.4 (vk)
 tests/slang-extension/exchange-int64-byte-address-buffer.slang.4 (vk)
 tests/slang-extension/realtime-clock.slang.2 (vk)
 tests/spirv/spirv-instruction.slang (vk)
-tests/type/texture-sampler/texture-sampler-2d.slang (vk)
\ No newline at end of file
+tests/type/texture-sampler/texture-sampler-2d.slang (vk)
+tests/autodiff/existential-1.slang.1 (vk)
+tests/autodiff/existential-2.slang.1 (vk)
+tests/ir/loop-inversion.slang.4 (vk)
+tests/spirv/direct-spirv-control-flow-2.slang (vk)
\ No newline at end of file