From 1717c6c507043a5cc960c2344dd556524804ce59 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Tue, 25 Apr 2023 18:32:45 -0700
Subject: Support recomputing phi params in bwd prop func. (#2841)

---
 source/slang/slang-ir.cpp | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

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

diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 9225faf6b..e74a57424 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -7478,7 +7478,7 @@ namespace Slang
         return nullptr;
     }
 
-    IRInst* getResolvedInstForDecorations(IRInst* inst)
+    IRInst* getResolvedInstForDecorations(IRInst* inst, bool resolveThroughDifferentiation)
     {
         IRInst* candidate = inst;
         for(;;)
@@ -7488,6 +7488,20 @@ namespace Slang
                 candidate = specInst->getBase();
                 continue;
             }
+            if (resolveThroughDifferentiation)
+            {
+                switch (candidate->getOp())
+                {
+                case kIROp_ForwardDifferentiate:
+                case kIROp_BackwardDifferentiate:
+                case kIROp_BackwardDifferentiatePrimal:
+                case kIROp_BackwardDifferentiatePropagate:
+                    candidate = candidate->getOperand(0);
+                    continue;
+                default:
+                    break;
+                }
+            }
             if( auto genericInst = as<IRGeneric>(candidate) )
             {
                 if( auto returnVal = findGenericReturnVal(genericInst) )
-- 
cgit v1.2.3