From 1817e9530989072ac34ff16d11fcc570f7862998 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Tue, 1 Nov 2022 12:55:36 -0400
Subject: Added a vector intrinsic definition for exp (to serve as template for
 other vector intrinsics) (#2481)

* Added vector exp definition

* Naming
---
 source/slang/diff.meta.slang                     | 82 ++++++++++--------------
 tests/autodiff/dstdlib-vector.slang              | 26 ++++++++
 tests/autodiff/dstdlib-vector.slang.expected.txt |  7 ++
 3 files changed, 66 insertions(+), 49 deletions(-)
 create mode 100644 tests/autodiff/dstdlib-vector.slang
 create mode 100644 tests/autodiff/dstdlib-vector.slang.expected.txt

diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index 674531048..1c3066e1d 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -59,60 +59,15 @@ extension float : IDifferentiable
     }
 }
 
-extension vector<float, 3> : IDifferentiable
+__generic<let N:int>
+extension vector<float, N> : IDifferentiable
 {
-    typedef vector<float, 3> Differential;
+    typedef vector<float, N> Differential;
 
     [__unsafeForceInlineEarly]
     static Differential dzero()
     {
-        return vector<float, 3>(0.f);
-    }
-
-    [__unsafeForceInlineEarly]
-    static Differential dadd(Differential a, Differential b)
-    {
-        return a + b;
-    }
-
-    [__unsafeForceInlineEarly]
-    static Differential dmul(This a, Differential b)
-    {
-        return a * b;
-    }
-}
-
-extension vector<float, 2> : IDifferentiable
-{
-    typedef vector<float, 2> Differential;
-
-    [__unsafeForceInlineEarly]
-    static Differential dzero()
-    {
-        return vector<float, 2>(0.f);
-    }
-
-    [__unsafeForceInlineEarly]
-    static Differential dadd(Differential a, Differential b)
-    {
-        return a + b;
-    }
-
-    [__unsafeForceInlineEarly]
-    static Differential dmul(This a, Differential b)
-    {
-        return a * b;
-    }
-}
-
-extension vector<float, 4> : IDifferentiable
-{
-    typedef vector<float, 4> Differential;
-
-    [__unsafeForceInlineEarly]
-    static Differential dzero()
-    {
-        return vector<float, 4>(0.f);
+        return vector<float, N>(0.f);
     }
 
     [__unsafeForceInlineEarly]
@@ -207,6 +162,9 @@ struct DifferentialPair : IDifferentiable
 
 typealias IDFloat = IFloat & IDifferentiable;
 
+#define VECTOR_MAP_UNARY(TYPE, COUNT, FUNC, VALUE) \
+    vector<TYPE,COUNT> result; for(int i = 0; i < COUNT; ++i) { result[i] = FUNC(VALUE[i]); } return result
+
 namespace dstd
 {
     // Natural Exponent
@@ -262,4 +220,30 @@ namespace dstd
             cos(dpx.p()),
             T.dmul(-sin(dpx.p()), dpx.d()));
     }
+
+    __generic<let N : int>
+    __target_intrinsic(hlsl)
+    __target_intrinsic(glsl)
+    __target_intrinsic(spirv_direct, "12 resultType resultId glsl450 27 _0")
+    [ForwardDerivative(d_exp_vector)]
+    vector<float, N> exp(vector<float, N> x)
+    {
+        VECTOR_MAP_UNARY(float, N, dstd.exp, x);
+    }
+
+    __generic<let N : int>
+    DifferentialPair<vector<float, N>> d_exp_vector(DifferentialPair<vector<float, N>> dpx)
+    {
+        vector<float, N> result;
+        vector<float, N>.Differential d_result;
+        for(int i = 0; i < N; ++i)
+        { 
+            DifferentialPair<float> dpexp = dstd.d_exp(DifferentialPair<float>(dpx.p()[i], dpx.d()[i]));
+            result[i] = dpexp.p(); 
+            d_result[i] = dpexp.d(); 
+        }
+
+        return DifferentialPair<vector<float, N>>(result, d_result);
+    }
+
 };
diff --git a/tests/autodiff/dstdlib-vector.slang b/tests/autodiff/dstdlib-vector.slang
new file mode 100644
index 000000000..44315daf6
--- /dev/null
+++ b/tests/autodiff/dstdlib-vector.slang
@@ -0,0 +1,26 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+typedef DifferentialPair<float> dpfloat;
+
+[ForwardDifferentiable]
+float f(float x)
+{
+    float3 vx = float3(x, 2*x, 3*x);
+    float3 vexpx = dstd.exp(vx);
+    return vexpx.x + vexpx.y + vexpx.z;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
+{
+    {
+        dpfloat dpa = dpfloat(2.0, 1.0);
+
+        outputBuffer[0] = f(dpa.p());        // Expect: 465.415999
+        outputBuffer[1] = __fwd_diff(f)(dpa).d(); // Expect: 1326.871736
+    }
+}
\ No newline at end of file
diff --git a/tests/autodiff/dstdlib-vector.slang.expected.txt b/tests/autodiff/dstdlib-vector.slang.expected.txt
new file mode 100644
index 000000000..3ab3a05e5
--- /dev/null
+++ b/tests/autodiff/dstdlib-vector.slang.expected.txt
@@ -0,0 +1,7 @@
+type: float
+465.415999
+1326.871736
+0.000000
+0.000000
+0.000000
+0.000000
-- 
cgit v1.2.3