From 551bbb5fbd61b53253de8f6ba3303bb4d29f8c86 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Tue, 11 Feb 2025 03:08:27 -0800
Subject: 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>
---
 source/slang/slang-check-decl.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'source/slang/slang-check-decl.cpp')

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(
-- 
cgit v1.2.3