Revert inlining change in #3217. (#3229)

Co-authored-by: Yong He <yhe@nvidia.com>

4 files changed, 31 insertions, 3 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index cd98e44c3..f2906dfc3 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -413,8 +413,6 @@ Result linkAndOptimizeIR(
         // since we may be missing out cases prevented by the functions that we just specialzied.
         performMandatoryEarlyInlining(irModule);
 
-        performPreAutoDiffForceInlining(irModule);
-
         // Unroll loops.
         if (codeGenContext->getSink()->getErrorCount() == 0)
         {
diff --git a/source/slang/slang-ir-autodiff-fwd.cpp b/source/slang/slang-ir-autodiff-fwd.cpp
index 10c8cdc51..560f3f64e 100644
--- a/source/slang/slang-ir-autodiff-fwd.cpp
+++ b/source/slang/slang-ir-autodiff-fwd.cpp
@@ -1691,6 +1691,8 @@ SlangResult ForwardDiffTranscriber::prepareFuncForForwardDiff(IRFunc* func)
     insertTempVarForMutableParams(autoDiffSharedContext->moduleInst->getModule(), func);
     removeLinkageDecorations(func);
     
+    performPreAutoDiffForceInlining(func);
+
     initializeLocalVariables(autoDiffSharedContext->moduleInst->getModule(), func);
 
     lowerSwizzledStores(autoDiffSharedContext->moduleInst->getModule(), func);
diff --git a/source/slang/slang-ir-autodiff-rev.cpp b/source/slang/slang-ir-autodiff-rev.cpp
index 7d5659425..335b6572e 100644
--- a/source/slang/slang-ir-autodiff-rev.cpp
+++ b/source/slang/slang-ir-autodiff-rev.cpp
@@ -533,6 +533,8 @@ namespace Slang
     {
         removeLinkageDecorations(func);
 
+        performPreAutoDiffForceInlining(func);
+
         DifferentiableTypeConformanceContext diffTypeContext(autoDiffSharedContext);
         diffTypeContext.setFunc(func);
 
diff --git a/source/slang/slang-ir-inline.cpp b/source/slang/slang-ir-inline.cpp
index 1d308c507..8412a1912 100644
--- a/source/slang/slang-ir-inline.cpp
+++ b/source/slang/slang-ir-inline.cpp
@@ -839,6 +839,8 @@ struct PreAutoDiffForceInliningPass : InliningPassBase
         : Super(module)
     {}
 
+    Dictionary<IRInst*, bool> m_funcCanInline;
+
     bool shouldInline(CallSiteInfo const& info)
     {
         if (info.callee->findDecoration<IRUnsafeForceInlineEarlyDecoration>() ||
@@ -862,7 +864,31 @@ struct PreAutoDiffForceInliningPass : InliningPassBase
                 break;
             }
         }
-        return (hasForceInline && !hasUserDefinedDerivative);
+        if (!hasForceInline || hasUserDefinedDerivative)
+        {
+            return false;
+        }
+        if (auto result = m_funcCanInline.tryGetValue(info.callee))
+            return *result;
+        bool canInline = true;
+        for (auto block : info.callee->getBlocks())
+        {
+            for (auto inst : block->getChildren())
+            {
+                switch (inst->getOp())
+                {
+                case kIROp_ForwardDifferentiate:
+                case kIROp_BackwardDifferentiate:
+                case kIROp_BackwardDifferentiatePrimal:
+                case kIROp_BackwardDifferentiatePropagate:
+                    canInline = false;
+                    goto end;
+                }
+            }
+        }
+    end:;
+        m_funcCanInline[info.callee] = canInline;
+        return canInline;
     }
 };