1 files changed, 80 insertions, 2 deletions
diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index dd7819e1a..a3cf28a68 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -1215,8 +1215,65 @@ bool isSideEffectFreeFunctionalCall(IRCall* call, SideEffectAnalysisOptions opti
     return false;
 }
 
+// Enumerate any associated functions of 'func'
+// that might be used by a pass (e.g. auto-diff)
+//
+template<typename TFunc>
+void forEachAssociatedFunction(IRInst* func, TFunc callback)
+{
+    // Resolve the function to get all its decorations
+    auto resolvedFunc = getResolvedInstForDecorations(func);
+    if (!resolvedFunc)
+        return;
+
+    // We'll scan for appropriate decorations and return
+    // the function references.
+    //
+    // TODO: In the future, as we get more function transformation
+    // passes, we might want to create a parent class for such
+    // decorations that associate functions with each other.
+    //
+    for (auto decor : resolvedFunc->getDecorations())
+    {
+        switch (decor->getOp())
+        {
+        case kIROp_UserDefinedBackwardDerivativeDecoration:
+            if (as<IRUserDefinedBackwardDerivativeDecoration>(decor))
+            {
+                auto associatedCallee = as<IRUserDefinedBackwardDerivativeDecoration>(decor)
+                                            ->getBackwardDerivativeFunc();
+                callback(associatedCallee);
+            }
+            break;
+
+        case kIROp_ForwardDerivativeDecoration:
+            if (as<IRForwardDerivativeDecoration>(decor))
+            {
+                auto associatedCallee =
+                    as<IRForwardDerivativeDecoration>(decor)->getForwardDerivativeFunc();
+                callback(associatedCallee);
+            }
+            break;
+
+        case kIROp_PrimalSubstituteDecoration:
+            if (as<IRPrimalSubstituteDecoration>(decor))
+            {
+                auto associatedCallee =
+                    as<IRPrimalSubstituteDecoration>(decor)->getPrimalSubstituteFunc();
+                callback(associatedCallee);
+            }
+            break;
+
+        default:
+            break;
+        }
+    }
+}
+
 bool doesCalleeHaveSideEffect(IRInst* callee)
 {
+    bool sideEffect = true;
+
     for (auto decor : getResolvedInstForDecorations(callee)->getDecorations())
     {
         switch (decor->getOp())
@@ -1224,10 +1281,31 @@ bool doesCalleeHaveSideEffect(IRInst* callee)
         case kIROp_NoSideEffectDecoration:
         case kIROp_ReadNoneDecoration:
         case kIROp_IgnoreSideEffectsDecoration:
-            return false;
+            sideEffect = false;
+            break;
+        default:
+            break;
         }
     }
-    return true;
+
+    // If the callee has no side effect, check if any of its associated functions have side effect.
+    // If so, we want to keep the callee around.
+    //
+    // Typically, once the relevant pass has completed, the association is removed,
+    // and at that point we can remove the function.
+    //
+    if (!sideEffect)
+    {
+        forEachAssociatedFunction(
+            callee,
+            [&](IRInst* associatedCallee)
+            {
+                sideEffect |= doesCalleeHaveSideEffect(associatedCallee);
+                return;
+            });
+    }
+
+    return sideEffect;
 }
 
 IRInst* findInterfaceRequirement(IRInterfaceType* type, IRInst* key)