StaticArrays.jl
StaticArrays.jl copied to clipboard
Improve latency of matrix-`exp`
As discussed in https://github.com/JuliaArrays/StaticArrays.jl/pull/985#issuecomment-1044720641, exp is very slow to compile for larger StaticArrays. This tries to reduce the latency a bit by just pulling out code into smaller functions, aiming to reduce the amount of inlined code.
There's little difference in the latency at small matrix sizes - and a slight reduction in performance for 3×3 matrices. For all other sizes, this PR actually produces a clear a performance improvement however. The latency improvements at larger sizes are also pronounced. Still, 3×3 matrices are definitely more important than bigger matrices, so it's not a no-brainer I guess - but I wanted to put the PR up for discussion anyway.
The summary of timings before/after this PR are (the performance timings depend on which branch is taken; also, this is on v1.7.2):
| Size | @time (1st) |
@btime (branch 1) |
@btime (branch 2) |
Latency win | Performance win |
|---|---|---|---|---|---|
| 3×3 | 1.2 s/1.3 s | 62 ns/77 ns | 68 ns/73 ns | ÷ | ÷ |
| 4×4 | 2.9 s/3.2 s | 208 ns/201 ns | 194 ns/161 ns | ÷ | ✓ |
| 10×10 | 46 s/21 s | 1.77 μs/1.62 μs | 8.43 μs/5.80 μs | ✓ | ✓ |
| 15×15 | 47 min/4.5 min | ?/11.9 μs | ?/37-50 μs | ✓ | ? |
I wonder how much compilation time reduction is just due to not duplicating U = A*U in each if branch? Maybe we don't have to decrease performance for the 3x3 case. Also, we could use the old variant for size 3x3 and the new one for larger matrices.
I tested various permutations of this quite extensively yesterday (pretty sure I also tested the A*U shuffle).
Puzzlingly, I found that almost all the compilation time comes from the presence of the if and for statements that rescale A in the second branch: without those, compilation time is drastically reduced. I don't understand why they make such a difference.
Alternatively, we could also use the new in-function @inline in 1.8 to force inlining in the 3x3 case (to avoid code duplication). Just having the code duplicated would indeed also be a solution.
almost all the compilation time comes from the presence of the two loops that rescale
Ain the second branch: without those, compilation time is drastically reduced
Interesting, though confusing! Did you try moving just that part out into a separate function?