diff options
| author | Yong He <yonghe@outlook.com> | 2023-04-24 19:44:23 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-04-24 19:44:23 -0700 |
| commit | 284cee1f246c072f190c87c8fb60c1d2181e458f (patch) | |
| tree | 6f8b4ff3d619ad518e733000464daae233890962 /source/slang/slang-ir-autodiff-unzip.cpp | |
| parent | fbe37ea6d90f7bfe18506b042657c6e533eaf9b2 (diff) | |
Change AD checkpointing policy to recompute more. (#2836)
Co-authored-by: Yong He <yhe@nvidia.com>
Diffstat (limited to 'source/slang/slang-ir-autodiff-unzip.cpp')
| -rw-r--r-- | source/slang/slang-ir-autodiff-unzip.cpp | 19 |
1 files changed, 12 insertions, 7 deletions
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(); } |