Add diagnostic for calling non-bwd-diff func from bwd-diff func. (shader-slang#2602)

csyonghe · web-flow · commit 6fae15cd1210 · 2023-01-19T08:58:20.000-08:00
diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
@@ -6894,38 +6894,34 @@ namespace Slang
 
     bool SharedSemanticsContext::isDifferentiableFunc(FunctionDeclBase* func)
     {
-        // A function is differentiable if it is marked as differentiable, or it
-        // has an associated derivative function.
-        if (func->findModifier<DifferentiableAttribute>())
-            return true;
-        for (auto assocDecl : getAssociatedDeclsForDecl(func))
-        {
-            switch (assocDecl.kind)
-            {
-            case DeclAssociationKind::ForwardDerivativeFunc:
-            case DeclAssociationKind::BackwardDerivativeFunc:
-                return true;
-            default:
-                break;
-            }
-        }
-        return false;
+        return getFuncDifferentiableLevel(func) != FunctionDifferentiableLevel::None;
     }
 
     bool SharedSemanticsContext::isBackwardDifferentiableFunc(FunctionDeclBase* func)
     {
-        // A function is differentiable if it is marked as differentiable, or it
-        // has an associated derivative function.
+        return getFuncDifferentiableLevel(func) == FunctionDifferentiableLevel::Backward;
+    }
+
+    FunctionDifferentiableLevel SharedSemanticsContext::getFuncDifferentiableLevel(FunctionDeclBase* func)
+    {
         if (func->findModifier<BackwardDifferentiableAttribute>())
-            return true;
+            return FunctionDifferentiableLevel::Backward;
         if (func->findModifier<BackwardDerivativeAttribute>())
-            return true;
+            return FunctionDifferentiableLevel::Backward;
+
+        FunctionDifferentiableLevel diffLevel = FunctionDifferentiableLevel::None;
+        if (func->findModifier<DifferentiableAttribute>())
+            diffLevel = FunctionDifferentiableLevel::Forward;
+
         for (auto assocDecl : getAssociatedDeclsForDecl(func))
         {
             switch (assocDecl.kind)
             {
             case DeclAssociationKind::BackwardDerivativeFunc:
-                return true;
+                return FunctionDifferentiableLevel::Backward;
+            case DeclAssociationKind::ForwardDerivativeFunc:
+                diffLevel = FunctionDifferentiableLevel::Forward;
+                break;
             default:
                 break;
             }
@@ -6937,12 +6933,12 @@ namespace Slang
             case BuiltinRequirementKind::DAddFunc:
             case BuiltinRequirementKind::DMulFunc:
             case BuiltinRequirementKind::DZeroFunc:
-                return true;
+                return FunctionDifferentiableLevel::Backward;
             default:
                 break;
             }
         }
-        return false;
+        return diffLevel;
     }
 
     List<ExtensionDecl*> const& getCandidateExtensions(
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
@@ -1967,6 +1967,10 @@ namespace Slang
 
         if (m_parentDifferentiableAttr)
         {
+            FunctionDifferentiableLevel callerDiffLevel = FunctionDifferentiableLevel::None;
+            if (m_parentFunc)
+                callerDiffLevel = getShared()->getFuncDifferentiableLevel(m_parentFunc);
+
             if (auto checkedInvokeExpr = as<InvokeExpr>(checkedExpr))
             {
                 // Register types for final resolved invoke arguments again.
@@ -1978,7 +1982,8 @@ namespace Slang
                 {
                     if (auto calleeDecl = as<FunctionDeclBase>(calleeExpr->declRef.getDecl()))
                     {
-                        if (getShared()->isDifferentiableFunc(calleeDecl))
+                        auto calleeDiffLevel = getShared()->getFuncDifferentiableLevel(calleeDecl);
+                        if (calleeDiffLevel >= callerDiffLevel)
                         {
                             if (!m_treatAsDifferentiableExpr)
                             {
diff --git a/source/slang/slang-check-impl.h b/source/slang/slang-check-impl.h
@@ -11,7 +11,12 @@
 
 namespace Slang
 {
-    
+    enum class FunctionDifferentiableLevel
+    {
+        None,
+        Forward,
+        Backward
+    };
         /// Should the given `decl` be treated as a static rather than instance declaration?
     bool isEffectivelyStatic(
         Decl*           decl);
@@ -292,6 +297,7 @@ namespace Slang
 
         bool isDifferentiableFunc(FunctionDeclBase* func);
         bool isBackwardDifferentiableFunc(FunctionDeclBase* func);
+        FunctionDifferentiableLevel getFuncDifferentiableLevel(FunctionDeclBase* func);
 
     private:
             /// Mapping from type declarations to the known extensiosn that apply to them
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
@@ -576,7 +576,7 @@ DIAGNOSTIC(41010, Warning, missingReturn, "control flow may reach end of non-'vo
 DIAGNOSTIC(41011, Error, typeDoesNotFitAnyValueSize, "type '$0' does not fit in the size required by its conforming interface.")
 DIAGNOSTIC(41012, Note, typeAndLimit, "sizeof($0) is $1, limit is $2")
 DIAGNOSTIC(41012, Error, typeCannotBePackedIntoAnyValue, "type '$0' contains fields that cannot be packed into an AnyValue.")
-DIAGNOSTIC(41020, Error, lossOfDerivativeDueToCallOfNonDifferentiableFunction, "derivative cannot be propagated through call to non-differentiable function `$0`, use 'no_diff' to clarify intention.")
+DIAGNOSTIC(41020, Error, lossOfDerivativeDueToCallOfNonDifferentiableFunction, "derivative cannot be propagated through call to non-$1-differentiable function `$0`, use 'no_diff' to clarify intention.")
 DIAGNOSTIC(41021, Error, differentiableFuncMustHaveOutput, "a differentiable function must have at least one differentiable output.")
 DIAGNOSTIC(41022, Error, differentiableFuncMustHaveInput, "a differentiable function must have at least one differentiable input.")
 DIAGNOSTIC(41023, Error, getStringHashMustBeOnStringLiteral, "getStringHash can only be called when argument is statically resolvable to a string literal")
diff --git a/source/slang/slang-ir-check-differentiability.cpp b/source/slang/slang-ir-check-differentiability.cpp
@@ -12,7 +12,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
     DiagnosticSink* sink;
     AutoDiffSharedContext sharedContext;
 
-    HashSet<IRInst*> differentiableFunctions;
+    enum DifferentiableLevel
+    {
+        Forward, Backward
+    };
+    Dictionary<IRInst*, DifferentiableLevel> differentiableFunctions;
 
     CheckDifferentiabilityPassContext(IRModule* inModule, DiagnosticSink* inSink)
         : InstPassBase(inModule), sink(inSink), sharedContext(inModule->getModuleInst())
@@ -59,7 +63,7 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
     }
 
 
-    bool _isDifferentiableFuncImpl(IRInst* func)
+    bool _isDifferentiableFuncImpl(IRInst* func, DifferentiableLevel level)
     {
         func = getLeafFunc(func);
         if (!func)
@@ -71,32 +75,41 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
             {
             case kIROp_ForwardDerivativeDecoration:
             case kIROp_ForwardDifferentiableDecoration:
+                if (level == DifferentiableLevel::Forward)
+                    return true;
+                break;
             case kIROp_UserDefinedBackwardDerivativeDecoration:
             case kIROp_BackwardDerivativeDecoration:
             case kIROp_BackwardDifferentiableDecoration:
                 return true;
+            default:
+                break;
             }
         }
         return false;
     }
 
-    bool isDifferentiableFunc(IRInst* func)
+    bool isDifferentiableFunc(IRInst* func, DifferentiableLevel level)
     {
-        switch (func->getOp())
+        if (level == DifferentiableLevel::Forward)
         {
-        case kIROp_ForwardDifferentiate:
-        case kIROp_BackwardDifferentiate:
-            return true;
-        default:
-            break;
+            switch (func->getOp())
+            {
+            case kIROp_ForwardDifferentiate:
+            case kIROp_BackwardDifferentiate:
+                return true;
+            default:
+                break;
+            }
         }
 
-        func = getSpecializedVal(func);
+        func = getLeafFunc(func);
         if (!func)
             return false;
 
-        if (differentiableFunctions.Contains(func))
-            return true;
+        
+        if (auto existingLevel = differentiableFunctions.TryGetValue(func))
+            return *existingLevel >= level;
 
         if (func->findDecoration<IRTreatAsDifferentiableDecoration>())
             return true;
@@ -125,7 +138,10 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
                 {
                     if (entry->getOperand(0) == lookupInterfaceMethod->getRequirementKey())
                     {
-                        return true;
+                        if (as<IRBackwardDifferentiableMethodRequirementDictionaryItem>(child) && level == DifferentiableLevel::Backward)
+                            return true;
+                        if (as<IRForwardDifferentiableMethodRequirementDictionaryItem>(child) && level == DifferentiableLevel::Forward)
+                            return true;
                     }
                 }
             }
@@ -135,7 +151,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
         {
             if (as<IRGeneric>(func))
             {
-                return differentiableFunctions.Contains(func);
+                if (auto existingLevel = differentiableFunctions.TryGetValue(func))
+                {
+                    if (*existingLevel >= level)
+                        return true;
+                }
             }
         }
         return false;
@@ -235,14 +255,18 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
         if (differentiableInputs == 0)
             sink->diagnose(funcInst, Diagnostics::differentiableFuncMustHaveInput);
 
+        DifferentiableLevel requiredDiffLevel = DifferentiableLevel::Forward;
+        if (isBackwardDifferentiableFunc(funcInst))
+            requiredDiffLevel = DifferentiableLevel::Backward;
+
         auto isInstProducingDiff = [&](IRInst* inst) -> bool
         {
             switch (inst->getOp())
             {
             case kIROp_FloatLit:
                 return true;
             case kIROp_Call:
-                return inst->findDecoration<IRTreatAsDifferentiableDecoration>() || isDifferentiableFunc(as<IRCall>(inst)->getCallee());
+                return inst->findDecoration<IRTreatAsDifferentiableDecoration>() || isDifferentiableFunc(as<IRCall>(inst)->getCallee(), requiredDiffLevel);
             case kIROp_Load:
                 // We don't have more knowledge on whether diff is available at the destination address.
                 // Just assume it is producing diff.
@@ -310,7 +334,7 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
                     switch (inst->getOp())
                     {
                     case kIROp_Call:
-                        if (isDifferentiableFunc(as<IRCall>(inst)->getCallee()))
+                        if (isDifferentiableFunc(as<IRCall>(inst)->getCallee(), requiredDiffLevel))
                         {
                             addToExpectDiffWorkList(inst);
                         }
@@ -349,7 +373,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
             {
                 if (auto call = as<IRCall>(inst))
                 {
-                    sink->diagnose(inst, Diagnostics::lossOfDerivativeDueToCallOfNonDifferentiableFunction, getLeafFunc(call->getCallee()));
+                    sink->diagnose(
+                        inst,
+                        Diagnostics::lossOfDerivativeDueToCallOfNonDifferentiableFunction,
+                        getLeafFunc(call->getCallee()),
+                        requiredDiffLevel == DifferentiableLevel::Forward ? "forward" : "backward");
                 }
             }
             switch (inst->getOp())
@@ -395,22 +423,30 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
     void processModule()
     {
         // Collect set of differentiable functions.
-        HashSet<UnownedStringSlice> differentiableSymbolNames;
+        HashSet<UnownedStringSlice> fwdDifferentiableSymbolNames, bwdDifferentiableSymbolNames;
         for (auto inst : module->getGlobalInsts())
         {
-            if (_isDifferentiableFuncImpl(inst))
+            if (_isDifferentiableFuncImpl(inst, DifferentiableLevel::Backward))
+            {
+                if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())
+                    bwdDifferentiableSymbolNames.Add(linkageDecor->getMangledName());
+                differentiableFunctions.Add(inst, DifferentiableLevel::Backward);
+            }
+            else if (_isDifferentiableFuncImpl(inst, DifferentiableLevel::Forward))
             {
                 if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())
-                    differentiableSymbolNames.Add(linkageDecor->getMangledName());
-                differentiableFunctions.Add(inst);
+                    fwdDifferentiableSymbolNames.Add(linkageDecor->getMangledName());
+                differentiableFunctions.Add(inst, DifferentiableLevel::Forward);
             }
         }
         for (auto inst : module->getGlobalInsts())
         {
             if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())
             {
-                if (differentiableSymbolNames.Contains(linkageDecor->getMangledName()))
-                    differentiableFunctions.Add(inst);
+                if (bwdDifferentiableSymbolNames.Contains(linkageDecor->getMangledName()))
+                    differentiableFunctions[inst] = DifferentiableLevel::Backward;
+                else if (fwdDifferentiableSymbolNames.Contains(linkageDecor->getMangledName()))
+                    differentiableFunctions.AddIfNotExists(inst, DifferentiableLevel::Forward);
             }
         }
 
diff --git a/tests/diagnostics/autodiff-data-flow-2.slang b/tests/diagnostics/autodiff-data-flow-2.slang
@@ -0,0 +1,28 @@
+//DIAGNOSTIC_TEST:SIMPLE:
+
+float nonDiff(float x)
+{
+    return x;
+}
+
+[ForwardDifferentiable]
+float f(float x)
+{
+    return x * x;
+}
+
+[BackwardDifferentiable]
+float g(float x)
+{
+    float val = f(x + 1); // Error: f must also be backward-differentiable
+    return val;
+}
+
+[BackwardDifferentiable]
+float h(float x)
+{
+    float val = 0;
+    // no diagnostic by clarifying intention.
+    val = no_diff(f(x + 1));
+    return val;
+}
diff --git a/tests/diagnostics/autodiff-data-flow-2.slang.expected b/tests/diagnostics/autodiff-data-flow-2.slang.expected
@@ -0,0 +1,8 @@
+result code = -1
+standard error = {
+tests/diagnostics/autodiff-data-flow-2.slang(18): error 41020: derivative cannot be propagated through call to non-backward-differentiable function `f`, use 'no_diff' to clarify intention.
+    float val = f(x + 1); // Error: f must also be backward-differentiable
+                 ^
+}
+standard output = {
+}
diff --git a/tests/diagnostics/autodiff-data-flow.slang.expected b/tests/diagnostics/autodiff-data-flow.slang.expected
@@ -1,6 +1,6 @@
 result code = -1
 standard error = {
-tests/diagnostics/autodiff-data-flow.slang(15): error 41020: derivative cannot be propagated through call to non-differentiable function `nonDiff`, use 'no_diff' to clarify intention.
+tests/diagnostics/autodiff-data-flow.slang(15): error 41020: derivative cannot be propagated through call to non-forward-differentiable function `nonDiff`, use 'no_diff' to clarify intention.
         val = nonDiff(x * 2.0f);
                      ^
 tests/diagnostics/autodiff-data-flow.slang(22): error 41021: a differentiable function must have at least one differentiable output.

Original file line number	Diff line number	Diff line change
`@@ -1967,6 +1967,10 @@ namespace Slang`
`1967`	`1967`
`1968`	`1968`	`if (m_parentDifferentiableAttr)`
`1969`	`1969`	`{`
	`1970`	`+ FunctionDifferentiableLevel callerDiffLevel = FunctionDifferentiableLevel::None;`
	`1971`	`+ if (m_parentFunc)`
	`1972`	`+ callerDiffLevel = getShared()->getFuncDifferentiableLevel(m_parentFunc);`
	`1973`	`+`
`1970`	`1974`	`if (auto checkedInvokeExpr = as<InvokeExpr>(checkedExpr))`
`1971`	`1975`	`{`
`1972`	`1976`	`// Register types for final resolved invoke arguments again.`
`@@ -1978,7 +1982,8 @@ namespace Slang`
`1978`	`1982`	`{`
`1979`	`1983`	`if (auto calleeDecl = as<FunctionDeclBase>(calleeExpr->declRef.getDecl()))`
`1980`	`1984`	`{`
`1981`		`- if (getShared()->isDifferentiableFunc(calleeDecl))`
	`1985`	`+ auto calleeDiffLevel = getShared()->getFuncDifferentiableLevel(calleeDecl);`
	`1986`	`+ if (calleeDiffLevel >= callerDiffLevel)`
`1982`	`1987`	`{`
`1983`	`1988`	`if (!m_treatAsDifferentiableExpr)`
`1984`	`1989`	`{`
Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`12`	`12`	`DiagnosticSink* sink;`
`13`	`13`	`AutoDiffSharedContext sharedContext;`
`14`	`14`
`15`		`- HashSet<IRInst*> differentiableFunctions;`
	`15`	`+ enum DifferentiableLevel`
	`16`	`+ {`
	`17`	`+ Forward, Backward`
	`18`	`+ };`
	`19`	`+ Dictionary<IRInst*, DifferentiableLevel> differentiableFunctions;`
`16`	`20`
`17`	`21`	`CheckDifferentiabilityPassContext(IRModule* inModule, DiagnosticSink* inSink)`
`18`	`22`	`: InstPassBase(inModule), sink(inSink), sharedContext(inModule->getModuleInst())`
`@@ -59,7 +63,7 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`59`	`63`	`}`
`60`	`64`
`61`	`65`
`62`		`- bool _isDifferentiableFuncImpl(IRInst* func)`
	`66`	`+ bool _isDifferentiableFuncImpl(IRInst* func, DifferentiableLevel level)`
`63`	`67`	`{`
`64`	`68`	`func = getLeafFunc(func);`
`65`	`69`	`if (!func)`
`@@ -71,32 +75,41 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`71`	`75`	`{`
`72`	`76`	`case kIROp_ForwardDerivativeDecoration:`
`73`	`77`	`case kIROp_ForwardDifferentiableDecoration:`
	`78`	`+ if (level == DifferentiableLevel::Forward)`
	`79`	`+ return true;`
	`80`	`+ break;`
`74`	`81`	`case kIROp_UserDefinedBackwardDerivativeDecoration:`
`75`	`82`	`case kIROp_BackwardDerivativeDecoration:`
`76`	`83`	`case kIROp_BackwardDifferentiableDecoration:`
`77`	`84`	`return true;`
	`85`	`+ default:`
	`86`	`+ break;`
`78`	`87`	`}`
`79`	`88`	`}`
`80`	`89`	`return false;`
`81`	`90`	`}`
`82`	`91`
`83`		`- bool isDifferentiableFunc(IRInst* func)`
	`92`	`+ bool isDifferentiableFunc(IRInst* func, DifferentiableLevel level)`
`84`	`93`	`{`
`85`		`- switch (func->getOp())`
	`94`	`+ if (level == DifferentiableLevel::Forward)`
`86`	`95`	`{`
`87`		`- case kIROp_ForwardDifferentiate:`
`88`		`- case kIROp_BackwardDifferentiate:`
`89`		`- return true;`
`90`		`- default:`
`91`		`- break;`
	`96`	`+ switch (func->getOp())`
	`97`	`+ {`
	`98`	`+ case kIROp_ForwardDifferentiate:`
	`99`	`+ case kIROp_BackwardDifferentiate:`
	`100`	`+ return true;`
	`101`	`+ default:`
	`102`	`+ break;`
	`103`	`+ }`
`92`	`104`	`}`
`93`	`105`
`94`		`- func = getSpecializedVal(func);`
	`106`	`+ func = getLeafFunc(func);`
`95`	`107`	`if (!func)`
`96`	`108`	`return false;`
`97`	`109`
`98`		`- if (differentiableFunctions.Contains(func))`
`99`		`- return true;`
	`110`	`+`
	`111`	`+ if (auto existingLevel = differentiableFunctions.TryGetValue(func))`
	`112`	`+ return *existingLevel >= level;`
`100`	`113`
`101`	`114`	`if (func->findDecoration<IRTreatAsDifferentiableDecoration>())`
`102`	`115`	`return true;`
`@@ -125,7 +138,10 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`125`	`138`	`{`
`126`	`139`	`if (entry->getOperand(0) == lookupInterfaceMethod->getRequirementKey())`
`127`	`140`	`{`
`128`		`- return true;`
	`141`	`+ if (as<IRBackwardDifferentiableMethodRequirementDictionaryItem>(child) && level == DifferentiableLevel::Backward)`
	`142`	`+ return true;`
	`143`	`+ if (as<IRForwardDifferentiableMethodRequirementDictionaryItem>(child) && level == DifferentiableLevel::Forward)`
	`144`	`+ return true;`
`129`	`145`	`}`
`130`	`146`	`}`
`131`	`147`	`}`
`@@ -135,7 +151,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`135`	`151`	`{`
`136`	`152`	`if (as<IRGeneric>(func))`
`137`	`153`	`{`
`138`		`- return differentiableFunctions.Contains(func);`
	`154`	`+ if (auto existingLevel = differentiableFunctions.TryGetValue(func))`
	`155`	`+ {`
	`156`	`+ if (*existingLevel >= level)`
	`157`	`+ return true;`
	`158`	`+ }`
`139`	`159`	`}`
`140`	`160`	`}`
`141`	`161`	`return false;`
`@@ -235,14 +255,18 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`235`	`255`	`if (differentiableInputs == 0)`
`236`	`256`	`sink->diagnose(funcInst, Diagnostics::differentiableFuncMustHaveInput);`
`237`	`257`
	`258`	`+ DifferentiableLevel requiredDiffLevel = DifferentiableLevel::Forward;`
	`259`	`+ if (isBackwardDifferentiableFunc(funcInst))`
	`260`	`+ requiredDiffLevel = DifferentiableLevel::Backward;`
	`261`	`+`
`238`	`262`	`auto isInstProducingDiff = [&](IRInst* inst) -> bool`
`239`	`263`	`{`
`240`	`264`	`switch (inst->getOp())`
`241`	`265`	`{`
`242`	`266`	`case kIROp_FloatLit:`
`243`	`267`	`return true;`
`244`	`268`	`case kIROp_Call:`
`245`		`- return inst->findDecoration<IRTreatAsDifferentiableDecoration>() \|\| isDifferentiableFunc(as<IRCall>(inst)->getCallee());`
	`269`	`+ return inst->findDecoration<IRTreatAsDifferentiableDecoration>() \|\| isDifferentiableFunc(as<IRCall>(inst)->getCallee(), requiredDiffLevel);`
`246`	`270`	`case kIROp_Load:`
`247`	`271`	`// We don't have more knowledge on whether diff is available at the destination address.`
`248`	`272`	`// Just assume it is producing diff.`
`@@ -310,7 +334,7 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`310`	`334`	`switch (inst->getOp())`
`311`	`335`	`{`
`312`	`336`	`case kIROp_Call:`
`313`		`- if (isDifferentiableFunc(as<IRCall>(inst)->getCallee()))`
	`337`	`+ if (isDifferentiableFunc(as<IRCall>(inst)->getCallee(), requiredDiffLevel))`
`314`	`338`	`{`
`315`	`339`	`addToExpectDiffWorkList(inst);`
`316`	`340`	`}`
`@@ -349,7 +373,11 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`349`	`373`	`{`
`350`	`374`	`if (auto call = as<IRCall>(inst))`
`351`	`375`	`{`
`352`		`- sink->diagnose(inst, Diagnostics::lossOfDerivativeDueToCallOfNonDifferentiableFunction, getLeafFunc(call->getCallee()));`
	`376`	`+ sink->diagnose(`
	`377`	`+ inst,`
	`378`	`+ Diagnostics::lossOfDerivativeDueToCallOfNonDifferentiableFunction,`
	`379`	`+ getLeafFunc(call->getCallee()),`
	`380`	`+ requiredDiffLevel == DifferentiableLevel::Forward ? "forward" : "backward");`
`353`	`381`	`}`
`354`	`382`	`}`
`355`	`383`	`switch (inst->getOp())`
`@@ -395,22 +423,30 @@ struct CheckDifferentiabilityPassContext : public InstPassBase`
`395`	`423`	`void processModule()`
`396`	`424`	`{`
`397`	`425`	`// Collect set of differentiable functions.`
`398`		`- HashSet<UnownedStringSlice> differentiableSymbolNames;`
	`426`	`+ HashSet<UnownedStringSlice> fwdDifferentiableSymbolNames, bwdDifferentiableSymbolNames;`
`399`	`427`	`for (auto inst : module->getGlobalInsts())`
`400`	`428`	`{`
`401`		`- if (_isDifferentiableFuncImpl(inst))`
	`429`	`+ if (_isDifferentiableFuncImpl(inst, DifferentiableLevel::Backward))`
	`430`	`+ {`
	`431`	`+ if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())`
	`432`	`+ bwdDifferentiableSymbolNames.Add(linkageDecor->getMangledName());`
	`433`	`+ differentiableFunctions.Add(inst, DifferentiableLevel::Backward);`
	`434`	`+ }`
	`435`	`+ else if (_isDifferentiableFuncImpl(inst, DifferentiableLevel::Forward))`
`402`	`436`	`{`
`403`	`437`	`if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())`
`404`		`- differentiableSymbolNames.Add(linkageDecor->getMangledName());`
`405`		`- differentiableFunctions.Add(inst);`
	`438`	`+ fwdDifferentiableSymbolNames.Add(linkageDecor->getMangledName());`
	`439`	`+ differentiableFunctions.Add(inst, DifferentiableLevel::Forward);`
`406`	`440`	`}`
`407`	`441`	`}`
`408`	`442`	`for (auto inst : module->getGlobalInsts())`
`409`	`443`	`{`
`410`	`444`	`if (auto linkageDecor = inst->findDecoration<IRLinkageDecoration>())`
`411`	`445`	`{`
`412`		`- if (differentiableSymbolNames.Contains(linkageDecor->getMangledName()))`
`413`		`- differentiableFunctions.Add(inst);`
	`446`	`+ if (bwdDifferentiableSymbolNames.Contains(linkageDecor->getMangledName()))`
	`447`	`+ differentiableFunctions[inst] = DifferentiableLevel::Backward;`
	`448`	`+ else if (fwdDifferentiableSymbolNames.Contains(linkageDecor->getMangledName()))`
	`449`	`+ differentiableFunctions.AddIfNotExists(inst, DifferentiableLevel::Forward);`
`414`	`450`	`}`
`415`	`451`	`}`
`416`	`452`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`result code = -1`
`2`	`2`	`standard error = {`
`3`		-tests/diagnostics/autodiff-data-flow.slang(15): error 41020: derivative cannot be propagated through call to non-differentiable function `nonDiff`, use 'no_diff' to clarify intention.
	`3`	+tests/diagnostics/autodiff-data-flow.slang(15): error 41020: derivative cannot be propagated through call to non-forward-differentiable function `nonDiff`, use 'no_diff' to clarify intention.
`4`	`4`	`val = nonDiff(x * 2.0f);`
`5`	`5`	`^`
`6`	`6`	`tests/diagnostics/autodiff-data-flow.slang(22): error 41021: a differentiable function must have at least one differentiable output.`