-
Notifications
You must be signed in to change notification settings - Fork 10.3k
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.
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
[Autodiff] Adds logic to rewrite call-sites using functions specialized by the closure-spec optimization #73688
Conversation
@swift-ci please test |
if !function.isAutodiffVJP { | ||
return | ||
|
||
if function.isOptimizable, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is not needed because the pass manager doesn't invoke function passes for functions which shouldn't be optimized
!function.isAvailableExternally, | ||
function.isAutodiffVJP, | ||
function.hasSingleBlock | ||
{ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
coding style: It's better to turn this into an early-exit if (or guard):
guard ... else
{
return
}
var (specializedFunction, alreadyExists, invalidatedStackNestingWhileCreatingSpecializedFunc) = | ||
getOrCreateSpecializedFunction(basedOn: callSite, context) | ||
|
||
invalidatedStackNesting = invalidatedStackNesting || invalidatedStackNestingWhileCreatingSpecializedFunc |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You can use context.notifyInvalidatedStackNesting
/context.needFixStackNesting
instead. Then you don't need to pass around the flag
.forEach { deadClosures.insert($0) } | ||
} | ||
|
||
for deadClosure in deadClosures { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please don't iterate over a Set
. The order is not deterministic.
You can use an InstructionWorklist
instead
function.fixStackNesting(context) | ||
} | ||
|
||
context.notifyFunctionBodyChanged() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You don't need this. Notifications are sent automatically if instructions are changed/inserted or deleted
@@ -39,6 +39,8 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash | |||
|
|||
public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() } | |||
|
|||
public var isOptimizable: Bool { bridged.isOptimizable() } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Beside that you don't need this, I would call it shouldOptimize
and add a comment that this is false if the function has an @_optimize(none)
attribute
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not needed after your suggested changes.
@@ -64,6 +66,10 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash | |||
public var blocks : BasicBlockList { BasicBlockList(first: bridged.getFirstBlock().block) } | |||
|
|||
public var entryBlock: BasicBlock { blocks.first! } | |||
|
|||
public var hasSingleBlock: Bool { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You don't need this. Just use function.blocks.singleElement
@@ -1634,6 +1634,9 @@ void BridgedChangeNotificationHandler::notifyChanges(Kind changeKind) const { | |||
case Kind::effectsChanged: | |||
invocation->notifyChanges(SILAnalysis::InvalidationKind::Effects); | |||
break; | |||
case Kind::functionBodyChanged: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As I said in my other comment, you don't need this
@@ -735,6 +735,17 @@ static void addMidLevelFunctionPipeline(SILPassPipelinePlan &P) { | |||
P.addLoopUnroll(); | |||
} | |||
|
|||
static void addAutodiffClosureSpecializationPassPipeline(SILPassPipelinePlan &P) { | |||
P.startPipeline("AutodiffClosureSpecialization"); | |||
P.addDeadStoreElimination(); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Adding such a number of new passes in the pipeline is not completely unproblematic. We have to be careful to not increase compile time.
Is it really needed to add those passes? If yes, can you check the impact on compile time, e.g. by measuring the compile time of the stdlib core module.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was modelling the pipeline after the existing closure-specialization pipeline and kept only the basic optimization passes that I knew could affect autodiff.
That said, those other passes are not strictly needed. I've removed them.
|
||
P.addAutodiffClosureSpecialization(); | ||
|
||
P.addDeadFunctionAndGlobalElimination(); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
dead-function-elimination is especially problematic because it's a module pass. Module passes split the function pass pipeline and you can get a completely different optimization behavior. (Unfortunately this is not really visible and documented). Isn't it good enough that the later function-elimination pass cleans up dead functions?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Addressed in a previous comment. Removed extra optimization passes.
@swift-ci please test |
@swift-ci please test |
@swift-ci please test macos |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
lgtm!
@@ -38,7 +38,7 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash | |||
public var isTrapNoReturn: Bool { bridged.isTrapNoReturn() } | |||
|
|||
public var isAutodiffVJP: Bool { bridged.isAutodiffVJP() } | |||
|
|||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
accidental change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks like it. Fixed now.
@@ -215,7 +215,6 @@ class SILPerformanceInliner { | |||
bool isAutoDiffLinearMapWithControlFlow(FullApplySite AI); | |||
|
|||
bool isTupleWithAllocsOrPartialApplies(SILValue retVal); | |||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
accidental change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks like it. Fixed now.
@swift-ci please test macos |
…ed by the closure-spec optimization
…he closure-spec optimization pass The OSSA elimination pass has not yet been moved below all high level function passes. Until that work has been completed the Autodiff closure-spec optimization pass cannot solely support OSSA instructions.
@swift-ci please test |
The inlining benefit for VJPs and pullbacks seems to be a bit different on macos and linux. This difference seems to be arising due to certain function calls such as, `sin`, `cos` etc. always being inlined on macos but not on linux. Until I figure out a better way, I'm modifying these tests to avoid matching the exact value of the inlining benefit, which causes them to fail on Linux.
@swift-ci please test |
@swift-ci please test macos |
@swift-ci please test linux |
1 similar comment
@swift-ci please test linux |
@swift-ci please test |
…d" test failures on Linux builds
@swift-ci please test |
@swift-ci please test macos |
@shahmishal The macos CI builds seem to have been a bit flaky recently. I've had some rare successful builds but mostly failures. The logs seem to be scattered with messages like these:
Do you know what's going on here? |
This PR adds the final part of the initial autodiff specific closure-specialization optimization implementation.
The previous parts of the optimization have been implemented in: