Change AD checkpointing policy to recompute more. (#2836)

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

1 files changed, 18 insertions, 0 deletions

diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp
index 03b74b36a..9348dfe8a 100644
--- a/source/slang/slang-ir-util.cpp
+++ b/source/slang/slang-ir-util.cpp
@@ -741,6 +741,24 @@ void moveParams(IRBlock* dest, IRBlock* src)
     }
 }
 
+void removePhiArgs(IRInst* phiParam)
+{
+    auto block = cast<IRBlock>(phiParam->getParent());
+    UInt paramIndex = 0;
+    for (auto p = block->getFirstParam(); p; p = p->getNextParam())
+    {
+        if (p == phiParam)
+            break;
+        paramIndex++;
+    }
+    for (auto predBlock : block->getPredecessors())
+    {
+        auto termInst = as<IRUnconditionalBranch>(predBlock->getTerminator());
+        SLANG_ASSERT(paramIndex < termInst->getArgCount());
+        termInst->removeArgument(paramIndex);
+    }
+}
+
 struct GenericChildrenMigrationContextImpl
 {
     IRCloneEnv cloneEnv;