Add checking for differentiability of the primal substitute function. (#6277)

Co-authored-by: Yong He <yonghe@outlook.com>
Co-authored-by: Ellie Hermaszewska <ellieh@nvidia.com>

1 files changed, 21 insertions, 0 deletions

diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 2ce4f81f1..76074f551 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -12114,6 +12114,27 @@ static void checkDerivativeAttribute(
         imaginaryArguments.directions,
         imaginaryArguments.thisArg,
         imaginaryArguments.thisArgDirection);
+
+    // For primal-substitute we'd also want to make sure that the differentiability
+    // level of the target is as high as the funcDecl itself
+    //
+    if (auto declRefExpr = as<DeclRefExpr>(attr->funcExpr))
+    {
+        if (auto declRef = declRefExpr->declRef)
+        {
+            auto targetDiffLevel = visitor->getShared()->getFuncDifferentiableLevel(
+                declRef.as<FunctionDeclBase>().getDecl());
+            auto currDiffLevel = visitor->getShared()->getFuncDifferentiableLevel(funcDecl);
+            if (targetDiffLevel < currDiffLevel)
+            {
+                visitor->getSink()->diagnose(
+                    attr->loc,
+                    Diagnostics::primalSubstituteTargetMustHaveHigherDifferentiabilityLevel,
+                    declRefExpr->declRef.getDecl(),
+                    funcDecl);
+            }
+        }
+    }
 }
 
 static void checkCudaKernelAttribute(