[Auto-diff] Overhaul auto-diff type tracking + Overhaul dynamic dispatch for differentiable functions (#5866)

* Overhauled the auto-diff system for dynamic dispatch

* More fixes

* remove intermediate dumps

* Update slang-ast-type.h

* More fixes + add a workaround for existential no-diff

* Update reverse-control-flow-3.slang

* remove dumps

* remove more dumps

* Delete working-reverse-control-flow-3.hlsl

* Cleanup comments + unused variables

* More comment cleanup

* Add support for lowering `DiffPairType(TypePack)` & `MakePair(MakeValuePack, MakeValuePack)`

* Fix array of issues in Falcor tests.

* Update slang-ir-autodiff-pairs.cpp

* More fixes for Falcor image tests

* Small fixups.

---------

Co-authored-by: Yong He <yonghe@outlook.com>

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);