Matrix exponential compiles slowly on Julia 1.0

[sethaxen]

Here's a simple example:

julia> using Pkg, LinearAlgebra, StaticArrays;

julia> v = @MMatrix zeros(Float64, 5, 5);

julia> v[2, 3] = -π;

julia> v[3, 2] = π;

julia> v
5×5 MArray{Tuple{5,5},Float64,2,25}:
0.0  0.0       0.0      0.0  0.0
0.0  0.0      -3.14159  0.0  0.0
0.0  3.14159   0.0      0.0  0.0
0.0  0.0       0.0      0.0  0.0
0.0  0.0       0.0      0.0  0.0

On Julia 1.1, this compiles and runs quickly

julia> @time exp(v)
8.741098 seconds (23.58 M allocations: 983.001 MiB, 8.48% gc time)
5×5 MArray{Tuple{5,5},Float64,2,25}:
1.0   0.0           0.0           0.0   0.0
-0.0  -1.0          -2.35127e-16  -0.0  -0.0
0.0   2.35127e-16  -1.0           0.0   0.0
0.0   0.0           0.0           1.0   0.0
0.0   0.0           0.0           0.0   1.0

julia> versioninfo()
Julia Version 1.1.0
Commit 80516ca202 (2019-01-21 21:24 UTC)
Platform Info:
OS: macOS (x86_64-apple-darwin14.5.0)
CPU: Intel(R) Core(TM) i7-4870HQ CPU @ 2.50GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-6.0.1 (ORCJIT, haswell)

julia> Pkg.status("StaticArrays")
    Status `~/.julia/environments/v1.1/Project.toml`
  [90137ffa] StaticArrays v0.11.0 #master (https://github.com/JuliaArrays/StaticArrays.jl)

On Julia 1.0, it takes nearly 50x longer:

julia> @time exp(v)
406.894653 seconds (24.31 M allocations: 1.512 GiB, 0.31% gc time)
5×5 MArray{Tuple{5,5},Float64,2,25}:
1.0   0.0           0.0           0.0   0.0
-0.0  -1.0          -2.35127e-16  -0.0  -0.0
0.0   2.35127e-16  -1.0           0.0   0.0
0.0   0.0           0.0           1.0   0.0
0.0   0.0           0.0           0.0   1.0

julia> versioninfo()
Julia Version 1.0.4
Commit 38e9fb7f80 (2019-05-16 03:38 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin15.6.0)
  CPU: Intel(R) Core(TM) i7-4870HQ CPU @ 2.50GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, haswell)

julia> Pkg.installed()["StaticArrays"]
v"0.11.0"

The culprit appears to be these lines

StaticArrays.jl/src/expm.jl

Lines 105 to 112 in e77e3ef

	U = A6*(S(1)*A6 + S(16380)*A4 + S(40840800)*A2) +
	(S(33522128640)*A6 + S(10559470521600)*A4 + S(1187353796428800)*A2) +
	S(32382376266240000)*I
	U = A*U
	V = A6*(S(182)*A6 + S(960960)*A4 + S(1323241920)*A2) +
	(S(670442572800)*A6 + S(129060195264000)*A4 + S(7771770303897600)*A2) +
	S(64764752532480000)*I
	expA = (V - U) \ (V + U)

If I remove them and everything afterward, the function compiles and returns quickly. Also, if I interrupt the stalled command (exp(v)) with Ctrl+C, and then re-run the command, it runs instantly.

[c42f]

I assume the problem here is in type inference. I tried a few rearrangements but nothing I tried worked very well. For example, it's possible to get compile time down to 90 s (from ~300 s on my machine) by moving the last part of the function into a separate function.

Maybe we should just fall back to Base for this operation on julia 1.0.

[schmrlng]

I'm not sure about MMatrix but it used to be the case that the LU-based solve (\) took ages to compile because lu was implemented in a somewhat unorthodox way (#368) to account for SMatrix immutability. Basically inference needs to unroll and reason about a huge recursion encompassing the entire LU decomposition algorithm which gets worse and worse as matrix size increases. It's possible that the more standard for loop algorithm (i.e., Base's generic LU) would work and compile a lot faster for MMatrix. There was some investigation into making this approach work for SMatrix in #424 (comment) but I'm not sure whatever happened to that. Perhaps @martinholters would recall more?

I think there was also something funny with inference in evaluating the polynomial expansions for U and V which is why the middle terms (e.g., (S(33522128640)*A6 + S(10559470521600)*A4 + S(1187353796428800)*A2) +) have a seemingly unnecessary set of parentheses surrounding them to group the computation into smaller chunks. I vaguely recall inference not even succeeding (let alone taking a long time, as seen here) when this code was developed on 0.7 if those parentheses are removed.

[tkoolen]

It's possible that the more standard for loop algorithm (i.e., Base's generic LU) would work and compile a lot faster for MMatrix.

Given that stuff like

StaticArrays.jl/src/solve.jl

Lines 34 to 38 in d82be15

	c = similar(b, T)
	for col = 1:Sb[2]
	@inbounds c[:, col] = _solve(Size($Sa), Size($Sa[1],), a, b[:, col])
	end
	return similar_type(b, T)(c)

now doesn't allocate for SMatrix inputs, maybe that hypothetical for-loop-based re-implementation could also just be used for SMatrix.

Edit: ah yes, that's exactly what Martin did.

[c42f]

Excellent point; MMatrix temporaries + some loops are much kinder on the compiler. Ideally we'd use those a lot more.

It would be great if we could eventually rely on LinearAlgebra for all medium size matrix algorithms and only fall back to specialized / unrolled code for really tiny cases.

[andyferris]

Yes, agreed - I actually kind-of think this package has succeeded when we have just the Size specialized methods (for 2x2, 3x3 etc) and dispatch to Base otherwise.

On a side note... are we actually supporting Julia 1.0?

[c42f]

I think we should try to support 1.0 as long as it's not too much of a drain on forward progress. For example fixing the current issue by going to an MArray based implementation would probably be a net improvement with 1.2 as well as 1.0.

Side note - according to here we have some 500 reverse dependencies (!) Yikes.

[andyferris]

For context, there are currently just under 3000 registered packages total:

ferris@ferris-laptop:~/.julia/registries/General (master)$ tree | grep Package.toml | wc -l
2930

I suppose we shouldn't intentially break one sixth of the ecosystem...