1 files changed, 16 insertions, 0 deletions
diff --git a/source/slang/slang-ir-autodiff-unzip.cpp b/source/slang/slang-ir-autodiff-unzip.cpp
index 9ee2cb4d2..49c1d9ff7 100644
--- a/source/slang/slang-ir-autodiff-unzip.cpp
+++ b/source/slang/slang-ir-autodiff-unzip.cpp
@@ -93,6 +93,22 @@ struct ExtractPrimalFuncContext
             as<IRGeneric>(findOuterGeneric(destFunc)),
             destFunc);
 
+        if (auto origGeneric = as<IRGeneric>(findOuterGeneric(originalFunc)))
+        {
+            // Clone in everything else except the return value.
+            IRBuilder subBuilder(destFunc);
+            builder.setInsertAfter(findOuterGeneric(destFunc)->getFirstBlock()->getLastParam());
+
+            // Clone in any hoistable insts.
+            for (auto child = origGeneric->getFirstBlock()->getFirstOrdinaryInst(); child;
+                 child = child->getNextInst())
+            {
+                if ((child != originalFunc) && !as<IRReturn>(child) &&
+                    !as<IRGlobalValueWithCode>(child))
+                    migrationContext.cloneInst(&subBuilder, child);
+            }
+        }
+
         originalFuncType = as<IRFuncType>(originalFunc->getDataType());
 
         SLANG_RELEASE_ASSERT(originalFuncType);