From 1c9b33157322751c456bf7abbd386edccf4413c3 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Sat, 14 Jan 2023 22:50:57 -0800
Subject: Support custom backward derivative attribute. (#2594)

---
 source/slang/slang-ir-autodiff.cpp | 60 ++++++++++++++++++++++++--------------
 1 file changed, 38 insertions(+), 22 deletions(-)

(limited to 'source/slang/slang-ir-autodiff.cpp')

diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp
index 74afa4002..7182375de 100644
--- a/source/slang/slang-ir-autodiff.cpp
+++ b/source/slang/slang-ir-autodiff.cpp
@@ -6,6 +6,21 @@
 
 namespace Slang
 {
+
+bool isBackwardDifferentiableFunc(IRInst* func)
+{
+    for (auto decorations : func->getDecorations())
+    {
+        switch (decorations->getOp())
+        {
+        case kIROp_BackwardDifferentiableDecoration:
+        case kIROp_UserDefinedBackwardDerivativeDecoration:
+            return true;
+        }
+    }
+    return false;
+}
+
 static IRInst* _lookupWitness(IRBuilder* builder, IRInst* witness, IRInst* requirementKey)
 {
     if (auto witnessTable = as<IRWitnessTable>(witness))
@@ -388,7 +403,7 @@ void DifferentiableTypeConformanceContext::buildGlobalWitnessDictionary()
     {
         if (auto pairType = as<IRDifferentialPairType>(globalInst))
         {
-            differentiableWitnessDictionary.Add(pairType->getValueType(), pairType->getWitness());
+            differentiableWitnessDictionary.AddIfNotExists(pairType->getValueType(), pairType->getWitness());
         }
     }
 }
@@ -406,6 +421,8 @@ void stripDerivativeDecorations(IRInst* inst)
         case kIROp_BackwardDerivativeIntermediateTypeDecoration:
         case kIROp_BackwardDerivativePropagateDecoration:
         case kIROp_BackwardDerivativePrimalDecoration:
+        case kIROp_UserDefinedBackwardDerivativeDecoration:
+        case kIROp_AutoDiffOriginalValueDecoration:
             decor->removeAndDeallocate();
             break;
         default:
@@ -435,6 +452,8 @@ void stripAutoDiffDecorationsFromChildren(IRInst* parent)
             case kIROp_BackwardDerivativePrimalDecoration:
             case kIROp_BackwardDerivativePrimalContextDecoration:
             case kIROp_BackwardDerivativePrimalReturnDecoration:
+            case kIROp_AutoDiffOriginalValueDecoration:
+            case kIROp_UserDefinedBackwardDerivativeDecoration:
                 decor->removeAndDeallocate();
                 break;
             default:
@@ -456,27 +475,26 @@ void stripAutoDiffDecorations(IRModule* module)
 }
 
 
-void stripBlockTypeDecorations(IRFunc* func)
+void stripTempDecorations(IRInst* inst)
 {
-    for (auto child : func->getChildren())
+    for (auto decor = inst->getFirstDecoration(); decor; )
     {
-        if (auto block = as<IRBlock>(child))
+        auto next = decor->getNextDecoration();
+        switch (decor->getOp())
         {
-            for (auto decor = block->getFirstDecoration(); decor; )
-            {
-                auto next = decor->getNextDecoration();
-                switch (decor->getOp())
-                {
-                case kIROp_DifferentialInstDecoration:
-                case kIROp_MixedDifferentialInstDecoration:
-                    decor->removeAndDeallocate();
-                    break;
-                default:
-                    break;
-                }
-                decor = next;
-            }
+        case kIROp_DifferentialInstDecoration:
+        case kIROp_MixedDifferentialInstDecoration:
+        case kIROp_AutoDiffOriginalValueDecoration:
+            decor->removeAndDeallocate();
+            break;
+        default:
+            break;
         }
+        decor = next;
+    }
+    for (auto child : inst->getChildren())
+    {
+        stripTempDecorations(child);
     }
 }
 
@@ -554,9 +572,7 @@ struct AutoDiffPass : public InstPassBase
             auto inner = findGenericReturnVal(baseGeneric);
             if (auto typeDecor = inner->findDecoration<IRBackwardDerivativeIntermediateTypeDecoration>())
             {
-                auto typeSpec = cast<IRSpecialize>(typeDecor->getBackwardDerivativeIntermediateType());
-                auto typeSpecBase = typeSpec->getBase();
-                return typeSpecBase;
+                return typeDecor->getBackwardDerivativeIntermediateType();
             }
         }
         else if (auto func = as<IRFunc>(base))
@@ -742,7 +758,7 @@ struct AutoDiffPass : public InstPassBase
             // passes since they don't expect decorations in blocks.
             // 
             for (auto diffFunc : autodiffCleanupList)
-                stripBlockTypeDecorations(diffFunc);
+                stripTempDecorations(diffFunc);
 
             autodiffCleanupList.clear();
 
-- 
cgit v1.2.3