Using Symbolics.jl to compile sparse matrix operations

main.jl

using LinearAlgebra, BenchmarkTools, SparseArrays, SymbolicUtils, Symbolics

## Problem setup
const z = zeros;

wx = Float64[ 0 -1  0;
              1  0  0;
              0  0  0;];

Ak_1 = [ wx           -I   zeros(3,12);
         zeros(3,3) .01*I  zeros(3,12);
                zeros(12,18)      ;]; 

Fk_1 = exp(Ak_1); sFk_1 = sparse(Fk_1);


# Compiling a Symbolics expression is just a couple of lines of code!
function compile(f, m, n)
    @variables M[1:m, 1:n]
    eval.(build_function(f(M), M))[2]
end

aux1 = zeros(18,18); aux2 = zeros(18,18);

f = compile(M -> sFk_1 * M, 18, 18)

@btime mul!(aux1, Fk_1, Pu);
# 617.824 ns (0 allocations: 0 bytes)

@btime f(aux2, Pu);
# 217.575 ns (0 allocations: 0 bytes)
# Already a ~3x speed-up!

# And we see even more benefits on more complicated functions
g = compile(M -> sFk_1 * M * sFk_1', 18, 18)

@btime begin
    mul!(aux1, Fk_1, Pu)
    mul!(aux2, aux1, Fk_1', 1.0, 0.0)
end;
#  1.347 μs (2 allocations: 48 bytes)

@btime g(aux1, Pu);
#  227.832 ns (0 allocations: 0 bytes)
# Here, compiling the chained operation brings a ~6x speed-up!