Skip to content

Commit 6cdff4f

Browse files
committed
[AutoDiff] Allow closure differentiation when used as default arg val
If the default argument generator (and, consequently, the function taking this default argument) has public visibility, it's OK to have a closure (which always has private visibility) as the default value of the argument.
1 parent 93518a5 commit 6cdff4f

File tree

2 files changed

+23
-12
lines changed

2 files changed

+23
-12
lines changed

lib/SILOptimizer/Mandatory/Differentiation.cpp

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -673,15 +673,28 @@ DifferentiationTransformer::emitDerivativeFunctionReference(
673673
}
674674
assert(minimalWitness);
675675
if (original->getFunction()->isSerialized()) {
676-
// When dealing with curry thunk, look at the function being wrapped
677-
// inside implicit closure. If it has public visibility, the corresponding
678-
// differentiability witness also has public visibility. It should be OK
679-
// for implicit wrapper closure and its witness to have private linkage.
680-
SILFunction *unwrappedFn = getUnwrappedCurryThunkFunction(originalFn);
681-
bool isWitnessPublic =
682-
unwrappedFn == nullptr
683-
? hasPublicVisibility(minimalWitness->getLinkage())
684-
: hasPublicVisibility(unwrappedFn->getLinkage());
676+
bool isWitnessPublic;
677+
if (SILFunction *unwrappedFn =
678+
getUnwrappedCurryThunkFunction(originalFn)) {
679+
// When dealing with curry thunk, look at the function being wrapped
680+
// inside implicit closure. If it has public visibility, the
681+
// corresponding differentiability witness also has public visibility.
682+
// It should be OK for implicit wrapper closure and its witness to have
683+
// private linkage.
684+
isWitnessPublic = hasPublicVisibility(unwrappedFn->getLinkage());
685+
} else if (originalFn->getDeclRef().getAbstractClosureExpr() &&
686+
originalFRI->getFunction()
687+
->getDeclRef()
688+
.isDefaultArgGenerator()) {
689+
// If we reference a closure from inside default argument generator,
690+
// check against generator's visibility. If the function having this
691+
// default argument has public visibility, it's OK to have a closure
692+
// (which always has private visibility) as its default value.
693+
isWitnessPublic =
694+
hasPublicVisibility(originalFRI->getFunction()->getLinkage());
695+
} else {
696+
isWitnessPublic = hasPublicVisibility(minimalWitness->getLinkage());
697+
}
685698
if (!isWitnessPublic) {
686699
enum { Inlinable = 0, DefaultArgument = 1 };
687700
unsigned fragileKind = Inlinable;

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -751,9 +751,7 @@ public func hasImplicitlyDifferentiatedTopLevelDefaultArgument(
751751
_ f: @differentiable(reverse) (Float) -> Float = implicitlyDifferentiableFromFragile
752752
) {}
753753

754-
// TODO(TF-1030): This will eventually not be an error.
755-
// expected-error @+2 {{function is not differentiable}}
756-
// expected-note @+1 {{differentiated functions in default arguments must be marked '@differentiable' or have a public '@derivative'; this is not possible with a closure, make a top-level function instead}}
754+
// No error expected
757755
public func hasImplicitlyDifferentiatedClosureDefaultArgument(_ f: @differentiable(reverse) (Float) -> Float = { $0 }) {}
758756

759757
@inlinable

0 commit comments

Comments
 (0)