diff options
| author | Yong He <yonghe@outlook.com> | 2023-03-22 21:16:35 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-03-22 21:16:35 -0700 |
| commit | 259a015feb9d4ab65e8fbba32f6c777e92780cc7 (patch) | |
| tree | 45bd4cb9217325c67f5a27d8562b0e7e6b79bb77 /source/slang/slang-ir-autodiff.cpp | |
| parent | d4f99c8bac8b28f18c864a717d8833db6a1c872d (diff) | |
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 <yhe@nvidia.com>
Diffstat (limited to 'source/slang/slang-ir-autodiff.cpp')
| -rw-r--r-- | source/slang/slang-ir-autodiff.cpp | 17 |
1 files changed, 12 insertions, 5 deletions
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<IRInst*> 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; } |