From 68f80f0793cff96b1fc784059a76605888593cd8 Mon Sep 17 00:00:00 2001 From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com> Date: Fri, 14 Jul 2023 15:47:06 -0400 Subject: Robustness fixes around reverse-mode differentiation of variables & inout parameters (#2985) * Add a new test case for checking loop in the reverse mode * Create duplicate var for primal value to avoid inconsistent inputs to backward call Fixes an issue with inout parameters where the backprop call may use the post-call value of the var instead of the pre-call value. * `IRStore`s transpose to `IRLoad` and an `IRStore(0)` to clear differential. Fixes some subtle issues around transposing * Simplify test * Delete out-edited.hlsl --------- Co-authored-by: Lifan Wu Co-authored-by: Yong He --- source/slang/slang-ir-autodiff-transpose.h | 8 ++++++++ source/slang/slang-ir-autodiff-unzip.h | 16 +++++++++++++++- 2 files changed, 23 insertions(+), 1 deletion(-) (limited to 'source') diff --git a/source/slang/slang-ir-autodiff-transpose.h b/source/slang/slang-ir-autodiff-transpose.h index 843612df9..739ce6553 100644 --- a/source/slang/slang-ir-autodiff-transpose.h +++ b/source/slang/slang-ir-autodiff-transpose.h @@ -1560,6 +1560,14 @@ struct DiffTransposePass TranspositionResult transposeStore(IRBuilder* builder, IRStore* fwdStore, IRInst*) { IRInst* revVal = builder->emitLoad(fwdStore->getPtr()); + + auto primalType = tryGetPrimalTypeFromDiffInst(fwdStore->getVal()); + SLANG_ASSERT(primalType); + + // Clear the value at the differential address, by setting to 0. + IRInst* emptyVal = emitDZeroOfDiffInstType(builder, primalType); + builder->emitStore(fwdStore->getPtr(), emptyVal); + if (auto diffPairType = as(revVal->getDataType())) { revVal = builder->emitDifferentialPairGetDifferential( diff --git a/source/slang/slang-ir-autodiff-unzip.h b/source/slang/slang-ir-autodiff-unzip.h index 1d05a8081..2a5d73a01 100644 --- a/source/slang/slang-ir-autodiff-unzip.h +++ b/source/slang/slang-ir-autodiff-unzip.h @@ -340,7 +340,21 @@ struct DiffUnzipPass // the remerged diff pair. auto diffPairType = as(as(arg->getDataType())->getValueType()); auto primalValueType = diffPairType->getValueType(); - auto diffPairRef = diffBuilder->emitReverseGradientDiffPairRef(arg->getDataType(), primalArg, diffArg); + + // We can't simply reuse primalArg for an inout parameter since this will represent the value + // after the primal call which can potentially alter primalArg. Therefore, we will find the + // first store into primalArg, and create a temp var holding that value (i.e. value prior to primal call) + // + auto storeUse = findUniqueStoredVal(cast(primalArg)); + auto storeInst = cast(storeUse->getUser()); + + auto storedVal = storeInst->getVal(); + + // Emit the temp var into the primal blocks since it's holding a primal value. + auto tempPrimalVar = primalBuilder->emitVar(primalValueType); + primalBuilder->emitStore(tempPrimalVar, storedVal); + + auto diffPairRef = diffBuilder->emitReverseGradientDiffPairRef(arg->getDataType(), tempPrimalVar, diffArg); diffBuilder->markInstAsDifferential(diffPairRef, primalValueType); diffArgs.add(diffPairRef); } -- cgit v1.2.3