Optimisation Opportunities

[willtebbutt]

This issue will remain open on an on-going basis. Items should be added as they are discovered.

Examples which allocate that should not, but which currently do:

• [x] (false, :allocs, nothing, TestResources.kron!, randn(400, 400), Diagonal(randn(20)), randn(20, 20)) (to do with increment!! allocating for non isbits types contained in PossiblyUninitTangent)
• [x] (false, :stability, nothing, lsetfield!, TestResources.FullyInitMutableStruct(5.0, [1.0, 2.0]), Val(:y), [1.0, 3.0, 4.0]) (to do with increment implementation for PossiblyUninitTangent.
• [ ] dynamic dispatch in general -- the current implementation is not performant.
• [x] small-unions -- the performance of TestResources._sum is okay, but not fantastic. Ought to be improvable.
• [ ] for register / arguments which are provably single-use (arguments always are), modify the way that tangent stacks are interacted with in order to avoid needing to perform bounds checks.
• [x] don't log tangent stack refs when it can be proven that the top of a specific tangent stack should always be used.
• [x] don't pop the block stack when there is only a single predecessor to the current block -- just move the stack pointer without retrieving anything, and shove an IDGoToNode at the end of the way back.
• [x] do not push to the block stack in blocks where all successors have this block as their unique predecessor. This happens e.g. in functions which branch, and then immediately return.
• [x] do not push / pop the block stack in exit blocks which are the only exit block
• [x] don't insert an additional block into the reverse-pass to handle exiting if the entry block to the primal has no predecessors -- just insert a ReturnNode. Also don't push / pop the block stack for the entry block in this case -- this will completely remove usage of the block stack in single-block code.
• [x] completely eliminate shared data for :calls / :invokes whose pullback is provablyNoPullback.
• [x] non-differentiable isbits types which appear as literals / QuoteNodes in PhiNodes should have their AugmentedRegisters appear inline as QuoteNodes, rather than living in shared storage. Currently they tend to take up quite a bit of storage.
• [x] Verify that the block_stack is never used in single-block functions.
• [x] Continue to hunt down cases in which tangents ought to be NoTangent, but aren't.

[willtebbutt]

Note: at the time of writing, there are roughly 80 elements in shared data for the _kron! function, and I believe that there should only really be 3.

[yebai]

Let’s consider introducing a benchmark/test for each optimisation candidate above. It is much easier to work with the codebase if the performance implications of new PRs can be clearly understood and pinned down to a specific use case.