Fix various autodiff crashes related to interface usage. (#2730)

* Fix crash.

* Fix `[ForwradDerivative]` on member functions.

* Update comments.

* Fix crash when [BackwardDerivative] is provided but not [ForwardDerivative].

* Allow calling dynamic dispatched generic method from differentiable func.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

6 files changed, 143 insertions, 0 deletions

diff --git a/tests/autodiff/dynamic-dispatch-generic-member.slang b/tests/autodiff/dynamic-dispatch-generic-member.slang
new file mode 100644
index 000000000..83c3aee7c
--- /dev/null
+++ b/tests/autodiff/dynamic-dispatch-generic-member.slang
@@ -0,0 +1,49 @@
+// Test calling dynamic dispatched generic function from differentiable function.
+
+//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;
+
+interface IFoo
+{
+    float f();
+}
+
+interface IInterface
+{
+    float calc<T:IFoo>(T t, float x);
+}
+
+struct A : IFoo
+{
+    float f() { return 1.0; }
+};
+
+struct B : IInterface
+{
+    float calc<T : IFoo>(T t, float x)
+    {
+        return t.f() * x;
+    }
+};
+
+[BackwardDifferentiable]
+float test(IInterface obj, float x)
+{
+    A objA;
+    return no_diff(obj.calc(objA, x)) * x;
+}
+
+//TEST_INPUT: type_conformance A:IFoo = 0
+//TEST_INPUT: type_conformance B:IInterface = 1
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    var obj = createDynamicObject<IInterface>(dispatchThreadID.x, 1); // B
+    var p = diffPair(3.0);
+    __bwd_diff(test)(obj, p, 1.0);
+    outputBuffer[0] = p.d;
+}
diff --git a/tests/autodiff/dynamic-dispatch-generic-member.slang.expected.txt b/tests/autodiff/dynamic-dispatch-generic-member.slang.expected.txt
new file mode 100644
index 000000000..857cebc03
--- /dev/null
+++ b/tests/autodiff/dynamic-dispatch-generic-member.slang.expected.txt
@@ -0,0 +1,5 @@
+type: float
+3.000000
+0.000000
+0.000000
+0.000000
diff --git a/tests/autodiff/member-func-custom-derivative-2.slang b/tests/autodiff/member-func-custom-derivative-2.slang
new file mode 100644
index 000000000..329f3ade8
--- /dev/null
+++ b/tests/autodiff/member-func-custom-derivative-2.slang
@@ -0,0 +1,49 @@
+//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], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+interface IFoo
+{
+    [BackwardDifferentiable]
+    float3 test(float v, uint offset);
+}
+struct A : IFoo
+{
+    float x;
+
+    float3 f(float v, uint offset)
+    {
+        return v * v;
+    }
+
+    // Provide a backward diff, but leave out forward diff.
+    [BackwardDerivativeOf(f)]
+    [TreatAsDifferentiable]
+    void diff_f(inout DifferentialPair<float> v, uint offset, float3 dOut)
+    {
+        v = diffPair(v.p, 2 * v.p * dOut.x);
+    }
+
+    [BackwardDifferentiable]
+    float3 test(float v, uint offset)
+    {
+        return f(v, 0);
+    }
+}
+
+[BackwardDifferentiable]
+float3 test(IFoo obj, float v)
+{
+    return obj.test(v, 0);
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    A a = {0.0};
+    var p = diffPair(3.0, 0.0);
+    let rs = __bwd_diff(test)(a, p, 1.0);
+    outputBuffer[0] = p.d;
+}
diff --git a/tests/autodiff/member-func-custom-derivative-2.slang.expected.txt b/tests/autodiff/member-func-custom-derivative-2.slang.expected.txt
new file mode 100644
index 000000000..253df0793
--- /dev/null
+++ b/tests/autodiff/member-func-custom-derivative-2.slang.expected.txt
@@ -0,0 +1,2 @@
+type: float
+6.0
diff --git a/tests/autodiff/member-func-custom-derivative.slang b/tests/autodiff/member-func-custom-derivative.slang
new file mode 100644
index 000000000..3ec44e690
--- /dev/null
+++ b/tests/autodiff/member-func-custom-derivative.slang
@@ -0,0 +1,36 @@
+//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], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+struct A
+{
+    float x;
+
+    [ForwardDerivative(diff_f)]
+    float f(float v)
+    {
+        return v * v;
+    }
+
+    DifferentialPair<float> diff_f(DifferentialPair<float> v)
+    {
+        return diffPair(v.p * v.p, v.p * v.d * 2.0);
+    }
+}
+
+[ForwardDifferentiable]
+float test(A obj, float v)
+{
+    return obj.f(v);
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    A a = {0.0};
+    var p = diffPair(3.0, 1.0);
+    let rs = __fwd_diff(test)(a, p);
+    outputBuffer[0] = rs.d;
+}
diff --git a/tests/autodiff/member-func-custom-derivative.slang.expected.txt b/tests/autodiff/member-func-custom-derivative.slang.expected.txt
new file mode 100644
index 000000000..253df0793
--- /dev/null
+++ b/tests/autodiff/member-func-custom-derivative.slang.expected.txt
@@ -0,0 +1,2 @@
+type: float
+6.0