Skip to content

[AutoDiff] Correctly propagate optional adjoint through switch_enum #74985

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 12, 2024
Merged

[AutoDiff] Correctly propagate optional adjoint through switch_enum #74985

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9db059c
Correctly propagate optional adjoint through `switch_enum`
asl Jul 5, 2024
0c5ae93
Accumulate adjoints for normal values as well
asl Jul 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 12 additions & 10 deletions lib/SILOptimizer/Differentiation/PullbackCloner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -765,14 +765,15 @@ class PullbackCloner::Implementation final
SILValue wrappedAdjoint,
SILType optionalTy);

/// Accumulate optional buffer from `wrappedAdjoint`.
/// Accumulate adjoint of `wrappedAdjoint` into optionalBuffer.
void accumulateAdjointForOptionalBuffer(SILBasicBlock *bb,
SILValue optionalBuffer,
SILValue wrappedAdjoint);

/// Set optional value from `wrappedAdjoint`.
void setAdjointValueForOptional(SILBasicBlock *bb, SILValue optionalValue,
SILValue wrappedAdjoint);
/// Accumulate adjoint of `wrappedAdjoint` into optionalValue.
void accumulateAdjointValueForOptional(SILBasicBlock *bb,
SILValue optionalValue,
SILValue wrappedAdjoint);

//--------------------------------------------------------------------------//
// Array literal initialization differentiation
Expand Down Expand Up @@ -2733,8 +2734,8 @@ void PullbackCloner::Implementation::accumulateAdjointForOptionalBuffer(
builder.createDeallocStack(pbLoc, optTanAdjBuf);
}

// Set the adjoint value for the incoming `Optional` value.
void PullbackCloner::Implementation::setAdjointValueForOptional(
// Accumulate adjoint for the incoming `Optional` value.
void PullbackCloner::Implementation::accumulateAdjointValueForOptional(
SILBasicBlock *bb, SILValue optionalValue, SILValue wrappedAdjoint) {
assert(getTangentValueCategory(optionalValue) == SILValueCategory::Object);
auto pbLoc = getPullback().getLocation();
Expand All @@ -2746,10 +2747,11 @@ void PullbackCloner::Implementation::setAdjointValueForOptional(

auto optTanAdjVal = builder.emitLoadValueOperation(
pbLoc, optTanAdjBuf, LoadOwnershipQualifier::Take);

recordTemporary(optTanAdjVal);
builder.createDeallocStack(pbLoc, optTanAdjBuf);

setAdjointValue(bb, optionalValue, makeConcreteAdjointValue(optTanAdjVal));
addAdjointValue(bb, optionalValue, makeConcreteAdjointValue(optTanAdjVal), pbLoc);
}

SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
Expand Down Expand Up @@ -2960,12 +2962,12 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
// Handle `switch_enum` on `Optional`.
auto termInst = bbArg->getSingleTerminator();
if (isSwitchEnumInstOnOptional(termInst)) {
setAdjointValueForOptional(bb, incomingValue, concreteBBArgAdjCopy);
accumulateAdjointValueForOptional(bb, incomingValue, concreteBBArgAdjCopy);
} else {
blockTemporaries[getPullbackBlock(predBB)].insert(
concreteBBArgAdjCopy);
setAdjointValue(predBB, incomingValue,
makeConcreteAdjointValue(concreteBBArgAdjCopy));
addAdjointValue(predBB, incomingValue,
makeConcreteAdjointValue(concreteBBArgAdjCopy), pbLoc);
}
}
break;
Expand Down
18 changes: 18 additions & 0 deletions test/AutoDiff/validation-test/optional.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,30 @@ func optional_nil_coalescing(_ maybeX: Float?) -> Float {
*/

OptionalTests.test("Active") {
@differentiable(reverse)
func id(y: Float) -> Float? {
return y
}

@differentiable(reverse)
func id2(y: Float?) -> Float {
return y!
}

@differentiable(reverse)
func square(y: Float) -> Float? {
return y * y
}

@differentiable(reverse)
func square2(y: Float?) -> Float {
return y! * y!
}

expectEqual(gradient(at: 10, of: {y in id(y:y)!}), .init(1.0))
expectEqual(gradient(at: 10, of: {y in id2(y:y)}), .init(1.0))
expectEqual(gradient(at: 10, of: {y in square(y:y)!}), .init(20.0))
expectEqual(gradient(at: 10, of: {y in square2(y:y)}), .init(20.0))
}

OptionalTests.test("Let") {
Expand Down