From 284cee1f246c072f190c87c8fb60c1d2181e458f Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 24 Apr 2023 19:44:23 -0700
Subject: Change AD checkpointing policy to recompute more. (#2836)

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-ir-autodiff-unzip.cpp | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

(limited to 'source/slang/slang-ir-autodiff-unzip.cpp')

diff --git a/source/slang/slang-ir-autodiff-unzip.cpp b/source/slang/slang-ir-autodiff-unzip.cpp
index 44e981404..a864a74b2 100644
--- a/source/slang/slang-ir-autodiff-unzip.cpp
+++ b/source/slang/slang-ir-autodiff-unzip.cpp
@@ -216,12 +216,15 @@ struct ExtractPrimalFuncContext
                 {
                     if (as<IRVar>(inst))
                     {
-                        auto field = addIntermediateContextField(cast<IRPtrTypeBase>(inst->getDataType())->getValueType(), outIntermediary);
-                        builder.setInsertBefore(inst);
-                        auto fieldAddr = builder.emitFieldAddress(
-                            inst->getFullType(), outIntermediary, field->getKey());
-                        inst->replaceUsesWith(fieldAddr);
-                        builder.addPrimalValueStructKeyDecoration(inst, field->getKey());
+                        if (inst->hasUses())
+                        {
+                            auto field = addIntermediateContextField(cast<IRPtrTypeBase>(inst->getDataType())->getValueType(), outIntermediary);
+                            builder.setInsertBefore(inst);
+                            auto fieldAddr = builder.emitFieldAddress(
+                                inst->getFullType(), outIntermediary, field->getKey());
+                            inst->replaceUsesWith(fieldAddr);
+                            builder.addPrimalValueStructKeyDecoration(inst, field->getKey());
+                        }
                     }
                     else
                     {
@@ -359,7 +362,7 @@ IRFunc* DiffUnzipPass::extractPrimalFunc(
     List<IRInst*> instsToRemove;
     for (auto block : func->getBlocks())
     {
-        for (auto inst : block->getOrdinaryInsts())
+        for (auto inst : block->getChildren())
         {
             if (auto structKeyDecor = inst->findDecoration<IRPrimalValueStructKeyDecoration>())
             {
@@ -420,6 +423,8 @@ IRFunc* DiffUnzipPass::extractPrimalFunc(
 
     for (auto inst : instsToRemove)
     {
+        if (as<IRParam>(inst))
+            removePhiArgs(inst);
         inst->removeAndDeallocate();
     }
 
-- 
cgit v1.2.3