From 47715e625337d489f3c0131bbc2b849378b48a5a Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 20 Feb 2023 14:42:50 -0800
Subject: Miscellaneous backward autodiff fixes. (#2665)

* Fix differentiable type registration

* Fix use of non-differentiable return value in a differentiable func.

* Fix use of primal inst that does not dominate the diff block.

* Fix primal inst hoisting, and add missing type legalization logic.

* Make `detach` defined on all differentiable T.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/diff.meta.slang | 24 ++++++++----------------
 1 file changed, 8 insertions(+), 16 deletions(-)

(limited to 'source/slang/diff.meta.slang')

diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index a60a77cc3..859b8a488 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -319,7 +319,13 @@ void mul(inout DifferentialPair<matrix<T, R, N>> left, inout DifferentialPair<ma
 
 // Detach and set derivatives to zero
 
-__generic<T : __BuiltinFloatingPointType>
+__generic<T : IDifferentiable>
+T detach(T x)
+{
+    return x;
+}
+
+__generic<T : IDifferentiable>
 [ForwardDerivativeOf(detach)]
 DifferentialPair<T> __d_detach(DifferentialPair<T> dpx)
 {
@@ -329,27 +335,13 @@ DifferentialPair<T> __d_detach(DifferentialPair<T> dpx)
     );
 }
 
-__generic<T : __BuiltinFloatingPointType, let N : int>
-[ForwardDerivativeOf(detach)]
-DifferentialPair<vector<T, N>> __d_detach_vector(DifferentialPair<vector<T, N>> dpx)
-{
-    VECTOR_MAP_D_UNARY(T, N, __d_detach, dpx);
-}
-
-__generic<T : __BuiltinFloatingPointType>
+__generic<T : IDifferentiable>
 [BackwardDerivativeOf(detach)]
 void __d_detach(inout DifferentialPair<T> dpx, T.Differential dOut)
 {
     dpx = diffPair(dpx.p, T.dzero());
 }
 
-__generic<T : __BuiltinFloatingPointType, let N : int>
-[BackwardDerivativeOf(detach)]
-void __d_detach_vector(inout DifferentialPair<vector<T, N>> dpx, vector<T, N>.Differential dOut)
-{
-    dpx = diffPair(dpx.p, vector<T, N>.dzero());
-}
-
 // Natural Exponent
 
 __generic<T : __BuiltinFloatingPointType>
-- 
cgit v1.2.3