From bbd1e1786401bb88c34802b987d4da72e2364503 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 1 Feb 2023 14:18:57 -0800
Subject: Support `out` parameters in backward differentiation. (#2619)

* Support `out` parameters in backward differentiation.

* Fixes.

* Fix cleanup.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 tests/autodiff/reverse-array-out-param.slang       | 55 ++++++++++++++++++++++
 .../reverse-array-out-param.slang.expected.txt     |  6 +++
 2 files changed, 61 insertions(+)
 create mode 100644 tests/autodiff/reverse-array-out-param.slang
 create mode 100644 tests/autodiff/reverse-array-out-param.slang.expected.txt

(limited to 'tests')

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
-- 
cgit v1.2.3