add compile-time optimization based on parameter values

[baggepinnen]

Status: Waiting for JSCompiler release

In multibody models, many parameters are sparse, e.g., arrays like [1, 0, 0] or diagonal inertia matrices. Allowing JSCompiler to make use of this information when simplifying the system and generating code can lead to a 2x performance boost for multibody simulations. The downside is of course that those parameters cannot be changed without recompiling. It's thus only worthwhile for simulations where runtime is longer than compile time.

To see why this makes such a big difference, consider the expression R'*I*R where the matrices are both 3x3. Without any structure, this expands to the computation

julia> collect(R'I*R)
3×3 Matrix{Num}:
 (I[1, 1]*R[1, 1] + I[2, 1]*R[2, 1] + I[3, 1]*R[3, 1])*R[1, 1] + (I[1, 2]*R[1, 1] + I[2, 2]*R[2, 1] + I[3, 2]*R[3, 1])*R[2, 1] + (I[1, 3]*R[1, 1] + I[2, 3]*R[2, 1] + I[3, 3]*R[3, 1])*R[3, 1]  …  (I[1, 1]*R[1, 1] + I[2, 1]*R[2, 1] + I[3, 1]*R[3, 1])*R[1, 3] + (I[1, 2]*R[1, 1] + I[2, 2]*R[2, 1] + I[3, 2]*R[3, 1])*R[2, 3] + (I[1, 3]*R[1, 1] + I[2, 3]*R[2, 1] + I[3, 3]*R[3, 1])*R[3, 3]
 (I[1, 1]*R[1, 2] + I[2, 1]*R[2, 2] + I[3, 1]*R[3, 2])*R[1, 1] + (I[1, 2]*R[1, 2] + I[2, 2]*R[2, 2] + I[3, 2]*R[3, 2])*R[2, 1] + (I[1, 3]*R[1, 2] + I[2, 3]*R[2, 2] + I[3, 3]*R[3, 2])*R[3, 1]     (I[1, 1]*R[1, 2] + I[2, 1]*R[2, 2] + I[3, 1]*R[3, 2])*R[1, 3] + (I[1, 2]*R[1, 2] + I[2, 2]*R[2, 2] + I[3, 2]*R[3, 2])*R[2, 3] + (I[1, 3]*R[1, 2] + I[2, 3]*R[2, 2] + I[3, 3]*R[3, 2])*R[3, 3]
 (I[1, 1]*R[1, 3] + I[2, 1]*R[2, 3] + I[3, 1]*R[3, 3])*R[1, 1] + (I[1, 2]*R[1, 3] + I[2, 2]*R[2, 3] + I[3, 2]*R[3, 3])*R[2, 1] + (I[1, 3]*R[1, 3] + I[2, 3]*R[2, 3] + I[3, 3]*R[3, 3])*R[3, 1]     (I[1, 1]*R[1, 3] + I[2, 1]*R[2, 3] + I[3, 1]*R[3, 3])*R[1, 3] + (I[1, 2]*R[1, 3] + I[2, 2]*R[2, 3] + I[3, 2]*R[3, 3])*R[2, 3] + (I[1, 3]*R[1, 3] + I[2, 3]*R[2, 3] + I[3, 3]*R[3, 3])*R[3, 3]

If we know that R is diagonal (very common case), we get

julia> collect(Rr'I*Rr)
3×3 Matrix{Num}:
  I[1, 1]*(r[1]^2)  I[1, 2]*r[1]*r[2]  I[1, 3]*r[1]*r[3]
 I[2, 1]*r[1]*r[2]   I[2, 2]*(r[2]^2)  I[2, 3]*r[2]*r[3]
 I[3, 1]*r[1]*r[3]  I[3, 2]*r[2]*r[3]   I[3, 3]*(r[3]^2)

If we know that I is diagonal we get

julia> R'Ii*R
3×3 Matrix{Num}:
             (R[1, 1]^2)*i[1] + (R[2, 1]^2)*i[2] + (R[3, 1]^2)*i[3]  R[1, 1]*R[1, 2]*i[1] + R[2, 1]*R[2, 2]*i[2] + R[3, 1]*R[3, 2]*i[3]  R[1, 1]*R[1, 3]*i[1] + R[2, 1]*R[2, 3]*i[2] + R[3, 1]*R[3, 3]*i[3]
 R[1, 1]*R[1, 2]*i[1] + R[2, 1]*R[2, 2]*i[2] + R[3, 1]*R[3, 2]*i[3]              (R[1, 2]^2)*i[1] + (R[2, 2]^2)*i[2] + (R[3, 2]^2)*i[3]  R[1, 2]*R[1, 3]*i[1] + R[2, 2]*R[2, 3]*i[2] + R[3, 2]*R[3, 3]*i[3]
 R[1, 1]*R[1, 3]*i[1] + R[2, 1]*R[2, 3]*i[2] + R[3, 1]*R[3, 3]*i[3]  R[1, 2]*R[1, 3]*i[1] + R[2, 2]*R[2, 3]*i[2] + R[3, 2]*R[3, 3]*i[3]              (R[1, 3]^2)*i[1] + (R[2, 3]^2)*i[2] + (R[3, 3]^2)*i[3]

and if both are diagonal, we get

julia> Rr'Ii*Rr
3×3 Diagonal{Num, Vector{Num}}:
 i[1]*(r[1]^2)              ⋅              ⋅
             ⋅  i[2]*(r[2]^2)              ⋅
             ⋅              ⋅  i[3]*(r[3]^2)

The PR #123 adds an explicit option to handle structural zeros in inertia matrices only, but this PR adds more of a nuclear option using JuliaSimCompiler.freeze_parameters, replacing all parameters that have the value 0 or 1 with their value. For a simple model of a car, this improved both simulation and RHS time by 2x.

Code for the above

@parameters R[1:3, 1:3] I[1:3, 1:3];
R,I = collect.((R,I));
collect(R'I*R)

# # With known zeros
@parameters r[1:3];
r = collect(r);
Rr = Diagonal(collect(r));
collect(Rr'I*Rr)

@parameters i[1:3];
i = collect(i);
Ii = Diagonal(collect(i));
R'Ii*R