Data flow validation pass for diagnosing derivative loss. (#2523)

1 files changed, 12 insertions, 0 deletions

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 4c3f4d646..61b6fcb76 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -11,6 +11,7 @@
 #include "slang-ir-diff-jvp.h"
 #include "slang-ir-inline.h"
 #include "slang-ir-insts.h"
+#include "slang-ir-check-differentiability.h"
 #include "slang-ir-missing-return.h"
 #include "slang-ir-sccp.h"
 #include "slang-ir-ssa.h"
@@ -3113,6 +3114,14 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
                 baseVal.val));
     }
 
+    LoweredValInfo visitTreatAsDifferentiableExpr(TreatAsDifferentiableExpr* expr)
+    {
+        auto baseVal = lowerSubExpr(expr->innerExpr);
+        SLANG_ASSERT(baseVal.flavor == LoweredValInfo::Flavor::Simple);
+        getBuilder()->addDecoration(baseVal.val, kIROp_TreatAsDifferentiableCallDecoration);
+        return baseVal;
+    }
+
     // Emit IR to denote the forward-mode derivative
     // of the inner func-expr. This will be resolved 
     // to a concrete function during the derivative 
@@ -8998,6 +9007,9 @@ RefPtr<IRModule> generateIRForTranslationUnit(
 
     checkForMissingReturns(module, compileRequest->getSink());
 
+    // Check for invalid differentiable function body.
+    checkAutoDiffUsages(module, compileRequest->getSink());
+
     // The "mandatory" optimization passes may make use of the
     // `IRHighLevelDeclDecoration` type to relate IR instructions
     // back to AST-level code in order to improve the quality