Skip to content

[AutoDiff] Allow closure differentiation when used as default arg val #78373

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 1 commit into from
Feb 20, 2025
Merged

[AutoDiff] Allow closure differentiation when used as default arg val #78373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
31 changes: 22 additions & 9 deletions lib/SILOptimizer/Mandatory/Differentiation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,15 +681,28 @@ DifferentiationTransformer::emitDerivativeFunctionReference(
}
assert(minimalWitness);
if (original->getFunction()->isSerialized()) {
// When dealing with curry thunk, look at the function being wrapped
// inside implicit closure. If it has public visibility, the corresponding
// differentiability witness also has public visibility. It should be OK
// for implicit wrapper closure and its witness to have private linkage.
SILFunction *unwrappedFn = getUnwrappedCurryThunkFunction(originalFn);
bool isWitnessPublic =
unwrappedFn == nullptr
? hasPublicVisibility(minimalWitness->getLinkage())
: hasPublicVisibility(unwrappedFn->getLinkage());
bool isWitnessPublic;
if (SILFunction *unwrappedFn =
getUnwrappedCurryThunkFunction(originalFn)) {
// When dealing with curry thunk, look at the function being wrapped
// inside implicit closure. If it has public visibility, the
// corresponding differentiability witness also has public visibility.
// It should be OK for implicit wrapper closure and its witness to have
// private linkage.
isWitnessPublic = hasPublicVisibility(unwrappedFn->getLinkage());
} else if (originalFn->getDeclRef().getAbstractClosureExpr() &&
originalFRI->getFunction()
->getDeclRef()
.isDefaultArgGenerator()) {
// If we reference a closure from inside default argument generator,
// check against generator's visibility. If the function having this
// default argument has public visibility, it's OK to have a closure
// (which always has private visibility) as its default value.
isWitnessPublic =
hasPublicVisibility(originalFRI->getFunction()->getLinkage());
} else {
isWitnessPublic = hasPublicVisibility(minimalWitness->getLinkage());
}
if (!isWitnessPublic) {
enum { Inlinable = 0, DefaultArgument = 1 };
unsigned fragileKind = Inlinable;
Expand Down
4 changes: 1 addition & 3 deletions test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -751,9 +751,7 @@ public func hasImplicitlyDifferentiatedTopLevelDefaultArgument(
_ f: @differentiable(reverse) (Float) -> Float = implicitlyDifferentiableFromFragile
) {}

// TODO(TF-1030): This will eventually not be an error.
// expected-error @+2 {{function is not differentiable}}
// 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}}
// No error expected
public func hasImplicitlyDifferentiatedClosureDefaultArgument(_ f: @differentiable(reverse) (Float) -> Float = { $0 }) {}

@inlinable
Expand Down