Support `out` parameters in backward differentiation. (#2619)

* Support `out` parameters in backward differentiation.

* Fixes.

* Fix cleanup.

---------

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

2 files changed, 61 insertions, 0 deletions

diff --git a/tests/autodiff/reverse-array-out-param.slang b/tests/autodiff/reverse-array-out-param.slang
new file mode 100644
index 000000000..e1c822489
--- /dev/null
+++ b/tests/autodiff/reverse-array-out-param.slang
@@ -0,0 +1,55 @@
+//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;
+
+struct D : IDifferentiable
+{
+    float n;
+    float m;
+}
+
+[BackwardDifferentiable]
+void g(D d, int i, out D outArray[1][2])
+{
+    outArray[0][0].m = d.n;
+    outArray[0][0].n = d.m * 3.0f + 1.0;
+
+    outArray[0][i].m = d.n * d.n;
+    outArray[0][i].n = d.m * 4.0f + 1.0;
+}
+
+[BackwardDifferentiable]
+void f(D d, int i, out D outArray[1][2])
+{
+    g(d, 0, outArray);
+    g(d, i, outArray);
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    D d = { 1.0, 2.0 };
+
+    var dpd = diffPair(d);
+
+    D.Differential tmp[1][2] = { { { 1.0, 1.0 }, { 1.0, 1.0 } } };
+    __bwd_diff(f)(dpd, 1, tmp);
+
+    outputBuffer[0] = dpd.d.m;
+    outputBuffer[1] = dpd.d.n;
+}
+
+[BackwardDifferentiable]
+void inner(float x, out float y)
+{
+    y = 2 * x;
+}
+
+[BackwardDifferentiable]
+void outter(float x, out float y)
+{
+    inner(x, y);
+    inner(x, y);
+}
diff --git a/tests/autodiff/reverse-array-out-param.slang.expected.txt b/tests/autodiff/reverse-array-out-param.slang.expected.txt
new file mode 100644
index 000000000..3370804b6
--- /dev/null
+++ b/tests/autodiff/reverse-array-out-param.slang.expected.txt
@@ -0,0 +1,6 @@
+type: float
+7.000000
+3.000000
+0.000000
+0.000000
+0.000000