Skip to content

Commit 4391f42

Browse files
authored
[AutoDiff] Fix Optional differentiation crash. (#33386)
Fix `Optional` differentiation crash for non-resilient `Wrapped` reference type. Add `NonresilientTracked` type to `DifferentiationUnittest` for testing. Resolves SR-13377.
1 parent b9cf249 commit 4391f42

File tree

4 files changed

+387
-129
lines changed

4 files changed

+387
-129
lines changed

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2056,8 +2056,8 @@ void PullbackCloner::Implementation::accumulateAdjointForOptional(
20562056
// Find `Optional<T.TangentVector>.some` EnumElementDecl.
20572057
auto someEltDecl = builder.getASTContext().getOptionalSomeDecl();
20582058

2059-
// Initialize a `Optional<T.TangentVector>` buffer from `wrappedAdjoint`as the
2060-
// input for `Optional<T>.TangentVector.init`.
2059+
// Initialize an `Optional<T.TangentVector>` buffer from `wrappedAdjoint` as
2060+
// the input for `Optional<T>.TangentVector.init`.
20612061
auto *optArgBuf = builder.createAllocStack(pbLoc, optionalOfWrappedTanType);
20622062
if (optionalOfWrappedTanType.isLoadableOrOpaque(builder.getFunction())) {
20632063
// %enum = enum $Optional<T.TangentVector>, #Optional.some!enumelt,
@@ -2066,7 +2066,7 @@ void PullbackCloner::Implementation::accumulateAdjointForOptional(
20662066
optionalOfWrappedTanType);
20672067
// store %enum to %optArgBuf
20682068
builder.emitStoreValueOperation(pbLoc, enumInst, optArgBuf,
2069-
StoreOwnershipQualifier::Trivial);
2069+
StoreOwnershipQualifier::Init);
20702070
} else {
20712071
// %enumAddr = init_enum_data_addr %optArgBuf $Optional<T.TangentVector>,
20722072
// #Optional.some!enumelt
@@ -2279,14 +2279,15 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
22792279
for (auto pair : incomingValues) {
22802280
auto *predBB = std::get<0>(pair);
22812281
auto incomingValue = std::get<1>(pair);
2282-
blockTemporaries[getPullbackBlock(predBB)].insert(concreteBBArgAdjCopy);
22832282
// Handle `switch_enum` on `Optional`.
22842283
auto termInst = bbArg->getSingleTerminator();
2285-
if (isSwitchEnumInstOnOptional(termInst))
2284+
if (isSwitchEnumInstOnOptional(termInst)) {
22862285
accumulateAdjointForOptional(bb, incomingValue, concreteBBArgAdjCopy);
2287-
else
2286+
} else {
2287+
blockTemporaries[getPullbackBlock(predBB)].insert(concreteBBArgAdjCopy);
22882288
setAdjointValue(predBB, incomingValue,
22892289
makeConcreteAdjointValue(concreteBBArgAdjCopy));
2290+
}
22902291
}
22912292
break;
22922293
}

stdlib/private/DifferentiationUnittest/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
add_swift_target_library(swiftDifferentiationUnittest ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS_STDLIB
22
# This file should be listed first. Module name is inferred from the filename.
3-
DifferentiationUnittest.swift
3+
GYB_SOURCES DifferentiationUnittest.swift.gyb
44

55
SWIFT_MODULE_DEPENDS _Differentiation StdlibUnittest
66
INSTALL_IN_COMPONENT stdlib-experimental

0 commit comments

Comments
 (0)