diff options
| author | Yong He <yonghe@outlook.com> | 2023-03-13 10:57:28 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-03-13 10:57:28 -0700 |
| commit | a911ca6e06ce41e403b80fe6054162393491c8ac (patch) | |
| tree | 6c8d56a3060b1887e7fd3126fe54a1241160eddd /source/slang/slang-ir-autodiff-transcriber-base.cpp | |
| parent | 3fea56ef77a33273bf5af6f432163b30c0a0e1dc (diff) | |
Support high order diff pattern: `bwd_diff(fwd_diff(f))`. (#2695)
* Support high order diff pattern: `bwd_diff(fwd_diff(f))`.
* Fix.
---------
Co-authored-by: Yong He <yhe@nvidia.com>
Diffstat (limited to 'source/slang/slang-ir-autodiff-transcriber-base.cpp')
| -rw-r--r-- | source/slang/slang-ir-autodiff-transcriber-base.cpp | 21 |
1 files changed, 19 insertions, 2 deletions
diff --git a/source/slang/slang-ir-autodiff-transcriber-base.cpp b/source/slang/slang-ir-autodiff-transcriber-base.cpp index ed122c862..091e7f1ab 100644 --- a/source/slang/slang-ir-autodiff-transcriber-base.cpp +++ b/source/slang/slang-ir-autodiff-transcriber-base.cpp @@ -304,8 +304,16 @@ IRType* AutoDiffTranscriberBase::_differentiateTypeImpl(IRBuilder* builder, IRTy auto primalPairType = as<IRDifferentialPairType>(primalType); return getOrCreateDiffPairType( builder, - pairBuilder->getDiffTypeFromPairType(builder, primalPairType), - pairBuilder->getDiffTypeWitnessFromPairType(builder, primalPairType)); + differentiableTypeConformanceContext.getDiffTypeFromPairType(builder, primalPairType), + differentiableTypeConformanceContext.getDiffTypeWitnessFromPairType(builder, primalPairType)); + } + + case kIROp_DifferentialPairUserCodeType: + { + auto primalPairType = as<IRDifferentialPairUserCodeType>(primalType); + return builder->getDifferentialPairUserCodeType( + (IRType*)differentiableTypeConformanceContext.getDiffTypeFromPairType(builder, primalPairType), + differentiableTypeConformanceContext.getDiffTypeWitnessFromPairType(builder, primalPairType)); } case kIROp_FuncType: @@ -634,6 +642,15 @@ IRInst* AutoDiffTranscriberBase::getDifferentialZeroOfType(IRBuilder* builder, I builder->markInstAsDifferential(makeDiffPair, as<IRDifferentialPairType>(diffType)->getValueType()); return makeDiffPair; } + case kIROp_DifferentialPairUserCodeType: + { + auto makeDiffPair = builder->emitMakeDifferentialPairUserCode( + diffType, + getDifferentialZeroOfType(builder, as<IRDifferentialPairUserCodeType>(diffType)->getValueType()), + getDifferentialZeroOfType(builder, as<IRDifferentialPairUserCodeType>(diffType)->getValueType())); + builder->markInstAsDifferential(makeDiffPair, as<IRDifferentialPairUserCodeType>(diffType)->getValueType()); + return makeDiffPair; + } } if (auto arrayType = as<IRArrayType>(primalType)) |