From 8d359fc6133fa49d2d3b7f8bb4b37916e719c344 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 9 Dec 2022 09:09:53 -0800
Subject: Add `diffPair` stdlib function. (#2560)

---
 source/slang/diff.meta.slang               | 22 +++++++++++++++++-----
 tests/autodiff/matrix-arithmetic-fwd.slang |  6 +++---
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index 248112810..f58648657 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -92,19 +92,30 @@ struct DifferentialPair : IDifferentiable
     }
 };
 
-[ForceInline]
+__generic<T: IDifferentiable>
+__intrinsic_op($(kIROp_MakeDifferentialPair))
+DifferentialPair<T> diffPair(T primal, T.Differential diff);
+
+__generic<T: IDifferentiable>
+[__unsafeForceInlineEarly]
+DifferentialPair<T> diffPair(T primal)
+{
+    return diffPair(primal, T.dzero());
+}
+
+[__unsafeForceInlineEarly]
 void updatePrimal<T : IDifferentiable>(inout DifferentialPair<T> p, T newPrimal)
 {
     p = DifferentialPair<T>(newPrimal, p.d);
 }
 
-[ForceInline]
+[__unsafeForceInlineEarly]
 void updateDiff<T : IDifferentiable>(inout DifferentialPair<T> p, T.Differential newDiff)
 {
     p = DifferentialPair<T>(p.p, newDiff);
 }
 
-[ForceInline]
+[__unsafeForceInlineEarly]
 void updatePair<T : IDifferentiable>(inout DifferentialPair<T> p, T newPrimal, T.Differential newDiff)
 {
     p = DifferentialPair<T>(newPrimal, newDiff);
@@ -112,8 +123,8 @@ void updatePair<T : IDifferentiable>(inout DifferentialPair<T> p, T newPrimal, T
 
 // vector-matrix
 __generic<T : __BuiltinFloatingPointType, let N : int, let M : int>
-__target_intrinsic(hlsl)
-__target_intrinsic(glsl, "($1 * $0)")
+[ForceInline]
+[ForwardDerivativeOf(mul)]
 DifferentialPair<vector<T, M>> mul(DifferentialPair<vector<T, N>> left, DifferentialPair<matrix<T, N, M>> right)
 {
     let primal = mul(left.p, right.p);
@@ -135,6 +146,7 @@ DifferentialPair<vector<T,N>> mul(DifferentialPair<matrix<T,N,M>> left, Differen
 
 // matrix-matrix
 __generic<T : __BuiltinFloatingPointType, let R : int, let N : int, let C : int>
+[ForceInline]
 [ForwardDerivativeOf(mul)]
 DifferentialPair<matrix<T,R,C>> mul(DifferentialPair<matrix<T,R,N>> right, DifferentialPair<matrix<T,N,C>> left)
 {
diff --git a/tests/autodiff/matrix-arithmetic-fwd.slang b/tests/autodiff/matrix-arithmetic-fwd.slang
index 7a953cef8..ee909551a 100644
--- a/tests/autodiff/matrix-arithmetic-fwd.slang
+++ b/tests/autodiff/matrix-arithmetic-fwd.slang
@@ -27,8 +27,8 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     float3x3 da = float3x3(1.0);
 
     outputBuffer[0] = __fwd_diff(g)(
-                            DifferentialPair<float3x3>(a, da),
-                            DifferentialPair<float3x3>(b, da)).d._11; // Expect: 8
+                          diffPair(a, da),
+                          diffPair(b, da)).d._11; // Expect: 8
 
     float2x2 l = float2x2(1.0, 2.0, 3.0, 4.0);
     float2x2 r = float2x2(10.0, 11.0, 12.0, 13.0);
@@ -37,5 +37,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     //float2x2 epsilon = d * 0.001f;
     //outputBuffer[1] = (h(l + epsilon, r + epsilon) - h(l - epsilon, r - epsilon)) / (epsilon[0][0] * 2.0));
 
-    outputBuffer[1] = __fwd_diff(h)(DifferentialPair<float2x2>(l, d), DifferentialPair<float2x2>(r, d)).d; // Expect 83.0
+    outputBuffer[1] = __fwd_diff(h)(diffPair(l, d), diffPair(r, d)).d; // Expect 83.0
 }
-- 
cgit v1.2.3