From 259a015feb9d4ab65e8fbba32f6c777e92780cc7 Mon Sep 17 00:00:00 2001 From: Yong He Date: Wed, 22 Mar 2023 21:16:35 -0700 Subject: Type legalization and autodiff bug fixes. (#2722) * Bug fixes. * Fix. * Only perform autodiff for functions whose derivative is actually used. * Fix loop optimize bug. * Fix high order diff. * Fix trivial diff func generation. * Fixes. * Cleanup. --------- Co-authored-by: Yong He --- source/slang/slang-ir-autodiff.cpp | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) (limited to 'source/slang/slang-ir-autodiff.cpp') diff --git a/source/slang/slang-ir-autodiff.cpp b/source/slang/slang-ir-autodiff.cpp index 224cca9e0..f173aaa8b 100644 --- a/source/slang/slang-ir-autodiff.cpp +++ b/source/slang/slang-ir-autodiff.cpp @@ -1143,9 +1143,8 @@ struct AutoDiffPass : public InstPassBase { bool changed = false; List autoDiffWorkList; - // Collect all `ForwardDifferentiate`/`BackwardDifferentiate` insts from the module. - autoDiffWorkList.clear(); - processAllInsts([&](IRInst* inst) + // Collect all `ForwardDifferentiate`/`BackwardDifferentiate` insts from the call graph. + processAllReachableInsts([&](IRInst* inst) { switch (inst->getOp()) { @@ -1164,11 +1163,15 @@ struct AutoDiffPass : public InstPassBase { // Skip functions whose body still has a differentiate inst (higher order func). if (!isFullyDifferentiated(innerFunc)) + { + addToWorkList(inst->getOperand(0)); return; + } } autoDiffWorkList.add(inst); break; default: + autoDiffWorkList.add(inst->getOperand(0)); break; } break; @@ -1176,6 +1179,11 @@ struct AutoDiffPass : public InstPassBase // Explicit primal subst operator is not yet supported. SLANG_UNIMPLEMENTED_X("explicit primal_subst operator."); default: + for (UInt i = 0; i < inst->getOperandCount(); i++) + { + auto operand = inst->getOperand(i); + addToWorkList(operand); + } break; } }); @@ -1199,7 +1207,7 @@ struct AutoDiffPass : public InstPassBase } break; case kIROp_BackwardDifferentiatePrimal: - { + { auto baseFunc = differentiateInst->getOperand(0); diffFunc = backwardPrimalTranscriber.transcribe(&subBuilder, baseFunc); } @@ -1300,7 +1308,6 @@ struct AutoDiffPass : public InstPassBase hasChanges |= changed; } - return hasChanges; } -- cgit v1.2.3