Add support for `[PrimalSubstitute]` and `[PrimalSubstituteOf]`. (#2691)

* Add support for `[PrimalSubstitute]` and `[PrimalSubstituteOf]`.

* Fix

* Fix.

* Cleanup.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

1 files changed, 13 insertions, 2 deletions

diff --git a/source/slang/slang-ir-check-differentiability.cpp b/source/slang/slang-ir-check-differentiability.cpp
index 186b0cc03..14f6394e2 100644
--- a/source/slang/slang-ir-check-differentiability.cpp
+++ b/source/slang/slang-ir-check-differentiability.cpp
@@ -39,12 +39,17 @@ public:
         return false;
     }
 
-
     bool _isDifferentiableFuncImpl(IRInst* func, DifferentiableLevel level)
     {
         func = getResolvedInstForDecorations(func);
         if (!func)
             return false;
+        if (auto substDecor = func->findDecoration<IRPrimalSubstituteDecoration>())
+        {
+            func = getResolvedInstForDecorations(substDecor->getPrimalSubstituteFunc());
+            if (!func)
+                return false;
+        }
 
         for (auto decorations : func->getDecorations())
         {
@@ -84,7 +89,13 @@ public:
         if (!func)
             return false;
 
-        
+        if (auto substDecor = func->findDecoration<IRPrimalSubstituteDecoration>())
+        {
+            func = getResolvedInstForDecorations(substDecor->getPrimalSubstituteFunc());
+            if (!func)
+                return false;
+        }
+
         if (auto existingLevel = differentiableFunctions.TryGetValue(func))
             return *existingLevel >= level;