Mark Differentiable related array methods with inlinable for big performance boost #75778
Conversation
…ed specialization possibilities
|
tagging @asl |
|
also tagging @rxwei |
|
@swift-ci please test |
|
preset=buildbot,tools=RA,stdlib=RA |
|
preset=buildbot,tools=RA,stdlib=DA |
1 similar comment
|
preset=buildbot,tools=RA,stdlib=DA |
|
How much is the performance boost? Could you quantify that in the PR description? |
|
@rxwei updated my original post with some benchmark numbers |
wow! |
|
Anyone specific I could tag for the other required review? @rxwei |
|
@swiftlang/standard-librarians |
|
@rxwei @swiftlang/standard-librarians What is the policy for such stdlib changes? Should it be additionally reviewed / approved by someone? |
|
@rxwei I've been going through all the methods in stdlib/public/Differentiation and there's also some candidates for adding @inlinable in OptionalDifferentiation.swift and FloatingPointDifferentiation.swift.gyb |
There was a problem hiding this comment.
This looks plausible to me as a stdlib engineer; the implementations exposed look relatively obvious and unlikely to need to change, and making them inlinable will allow client modules to specialize them.
As noted though, this change is only safe to make as long as Differentiation is not expected to be ABI stable in any context.
| @@ -21,27 +21,29 @@ extension Array where Element: Differentiable { | |||
| /// multiplied with itself `count` times. | |||
| @frozen | |||
| public struct DifferentiableView { | |||
| @usableFromInline | |||
There was a problem hiding this comment.
Beware, adding @usableFromInline introduces a newly exported symbol without availability. It would not be safe to make this change in a module that's expected to have stable ABI. Newly built code after this change may not be binary compatible with Differentiation in earlier Swift releases.
(However, as I understand it, while Differentation is being build with library evolution enabled, it is not distributed as such on any platform where Swift is ABI stable. It is also unclear if the problematic stored property accessor exports are ever actually called in practice for a @frozen structure like this.)
There was a problem hiding this comment.
Thanks for the review @lorentey !
Just to double check, this is considered unsafe because the struct was already decorated with @frozen? Or is using @usableFromInline without public availability generally ABI unstable? I might not fully grasp the subtlety here but would love to understand better to keep these in mind for future changes.
But indeed Differentiation is not yet distributed to a platform that is ABI stable, so it's not an issue yet!
There was a problem hiding this comment.
@rxwei in your opinion, are we good to go here? As @lorentey points out we break ABI here but that's currently not an issue yet right? I'll add some comments to the main PR message regarding ABI stability.
As an additional question, just for my understanding. Is the following change binary incompatible? And if so why? It seems to me that we're only adding information to the interface not changing the binary layout of the struct. But I also don't have a lot of experience here so would like to understand the details if possible!
@frozen
struct Thing {
var storage: Float
}
->
@frozen
struct Thing {
@usableFromInline
var storage: Float
}
Also tagging @asl
|
@asl Can we merge this? Seems to be fully approved! :) |
This PR marks several methods in ArrayDifferentiation.swift with @inlinable adding much more opportunity for specialisation. This leads to huge performance increases and much lower memory usage in a few Differentiable example programs using Array and Array.DifferentiableView.
Main increase is due to being able to specialise Array.DifferentiableView's + operator from it's conformance to AdditiveArithmetic and therefore further being able to inline function calls and not having to go through the protocol witness table.
Technically we do restrict ourselves wrt future changes to the
DifferentiableViewimplementation since the internals are now marked@usableFromInline. We currently think this is acceptable since Differentiation currently is not shipped to an ABI stable platform.Some performance numbers:
The benchmark I'm using is a reimplementation of the original SwiftForTensorflow example found here: https://github.com/tensorflow/swift-models/tree/2fb0b92e1291b730fd1a5cd8a3b107c8e75c7d7a/Examples/Shallow-Water-PDE
This was implemented on top of the Tensor type that S4TF introduced.
My example uses an Array for storage.
For the benchmark I've set the amount of iterations to 1 so I'm benchmarking the wall clock time and malloc of running a forward pass and pullback once through an array of
resolution * resolutionanddurationamount of time stepsEvery tilmestep the benchmark applies the Laplace operator to every cell of the array.
I ran my benchmark for values of 10 and 20 for both of these variables. You can find the results below. As you can see the difference in performance is quite enormous ranging from at least 20x in both categories to 60-70x