[AutoDiff] Support curry thunks differentiation in fragile funcs #77615
Conversation
Inside fragile functions, we expect function derivatives to be public, which could be achieved by either explicitly marking the functions as differentiable or having a public explicit derivative defined for them. This is obviously not possible for single and double curry thunks which are a special case of `AutoClosureExpr`. Instead of looking at the thunk itself, we unwrap it and look at the function being wrapped. While the thunk itself and its differentiability witness will not have public visibility, it's not an issue for the case where the function being wrapped (and its witness) have public visibility. Fixes swiftlang#54819 Fixes swiftlang#75776
|
Tagging @asl |
kovdan01
requested review from
eeckstein,
hborla,
slavapestov and
xedin
as code owners
|
@swift-ci please test |
lib/AST/Expr.cpp
Outdated
| @@ -2211,6 +2211,20 @@ Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const { | |||
| return nullptr; | |||
| } | |||
|
|
|||
| Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const { | |||
There was a problem hiding this comment.
Do we still need getUnwrappedCurryThunkExpr()? It's not clear to me under what circumstances ae->getFn() will return something other than ae->getCalledValue(/*skipFunctionConversions=*/true);
There was a problem hiding this comment.
@slavapestov Do you mean that getUnwrappedCurryThunkCalledValue() might be used instead of getUnwrappedCurryThunkExpr()? I don't think so: I intentionally moved the actual implementation into a separate getUnwrappedCurryThunkImpl() function, so we can either get an Expr or a ValueDecl depending of what we need.
getUnwrappedCurryThunkExpr() seems to have a number of uses over the code base, and that uses seem to expect Expr and not ValueDecl, so it's probably still needed. Anyway, even if some refactoring might be done (while I do not see a room for it in this particular case), probably it would be better to keep it separate from functional changes.
If I got your question wrong and this does not answer it, it would be great if you can clarify your point a bit.
|
Tagging @rxwei |
|
Tagging @JaapWijnen |
|
Would be glad to see feedback from everyone interested |
It's already included in test/AutoDiff/SILOptimizer/fragile_curry_thunk.swift
|
@swift-ci please test |
|
Would be glad to see feedback from everyone interested |
|
Great to have this fixed thanks @kovdan01. Looking forward to have this reviewed and merged! :) |
|
Would be glad to see feedback from everyone interested |
3 similar comments
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
4 similar comments
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
|
Would be glad to see feedback from everyone interested |
| auto *abstractCE = originalFn->getDeclRef().getAbstractClosureExpr(); | ||
| if (abstractCE == nullptr) | ||
| return nullptr; | ||
| auto *autoCE = dyn_cast<AutoClosureExpr>(abstractCE); |
There was a problem hiding this comment.
nit: could simplify a bit with dyn_cast_or_null
There was a problem hiding this comment.
Applied the suggestion, thanks, see c674f80
| // for that due to the crash | ||
| // https://github.com/swiftlang/swift/issues/77613 | ||
| if (currentAFD->hasCurriedSelf()) | ||
| afdToSILFn.insert({currentAFD, ¤tFunc}); |
There was a problem hiding this comment.
try_emplace + assert on result?
There was a problem hiding this comment.
Applied the suggestion, thanks, see c674f80
|
I've limited scope of the changes only to AutoDiff code - see 2cadfa0 |
|
@swift-ci please test |
| /*skipFunctionConversions=*/true)); | ||
| break; | ||
| } | ||
| assert(afd); |
There was a problem hiding this comment.
I would prefer some message attached to assert. I would probably also change switch to ordinary if, this was you can error out in else case and this assert will be unnecessary at all, as afd will always be a result of cast.
There was a problem hiding this comment.
I've changed the logic a bit in 4d74d34. It actually looks like that other thunk types (not single and double curry thunks) could go here, and this is valid, we just need to return nullptr (just as we do right above if we detect not AutoClosureExpr).
Since getUnwrappedCurryThunkExpr already has checks against single and double curry thunks and returns nullptr for other thunk types, using it in conjunction with dyn_cast_or_null with subsequent check against nullptr should be enough.
| if (currentAFD->hasCurriedSelf()) { | ||
| auto [_, wasEmplace] = | ||
| afdToSILFn.try_emplace(currentAFD, ¤tFunc); | ||
| assert(wasEmplace); |
There was a problem hiding this comment.
Addressed in 4d74d34, thanks
|
|
||
| auto silFnIt = afdToSILFn.find(afd); | ||
| if (silFnIt == afdToSILFn.end()) { | ||
| assert(afdToSILFn.empty()); |
There was a problem hiding this comment.
Addressed in 4d74d34, thanks
|
@swift-ci please test |
Inside fragile functions, we expect function derivatives to be public, which could be achieved by either explicitly marking the functions as differentiable or having a public explicit derivative defined for them. This is obviously not possible for single and double curry thunks which are a special case of
AutoClosureExpr.Instead of looking at the thunk itself, we unwrap it and look at the function being wrapped. While the thunk itself and its differentiability witness will not have public visibility, it's not an issue for the case where the function being wrapped (and its witness) have public visibility.
Fixes #54819
Fixes #75776