I need to remove the runtime allocations for a project that is targeting embedding Julia code in a satellite computer. One big problem I have is to create static matrices programmatically. For example, in C++, we can use Eigen and static matrices (which usually are allocated in the stack) and we can change its coefficients to whatever we want given that the type is conserved.
I am trying to replicate this behavior in Julia using StaticArrays.jl and Accessors.jl. However, the code leads to a huge number of allocations and it is considerable slower than just allocating the matrix initially.
Here is one example of an algorithm that I need to use static matrices:
function fully_normalized_legendre(
ϕ::T,
n_max::Integer = -1,
m_max::Integer = -1;
ph_term::Bool = false
) where T<:Number
n_max < 0 && throw(ArgumentError("n_max must be positive."))
if (m_max < 0) || (m_max > n_max)
m_max = n_max
end
P = zeros(float(T), n_max + 1, m_max + 1)
# Auxiliary variables to improve code performance.
s, c = sincos(T(ϕ))
# The sine must be always positive. In fact, `s` was previously computed using
# `√(1 - c^2)`. However, we had numerical problems for very small angles that lead to
# `cos(ϕ) = 1`.
if s < 0
s = -s
end
s_fact = !ph_term ? +s : -s
# Get the first indices in `P` to take into account offset arrays.
i₀, j₀ = first.(axes(P))
sq3 = √T(3)
@inbounds for n in 0:n_max
# Starting values.
if n == 0
P[i₀, j₀] = 1
continue
elseif n == 1
P[i₀+1, j₀] = +sq3 * c
if m_max > 0
P[i₀+1, j₀+1] = +sq3 * s_fact
end
continue
end
aux_an = T(2n - 1) * T(2n + 1)
aux_bn = T(2n + 1) / T(2n - 3)
for m in 0:n
if n == m
P[i₀+n, j₀+n] = s_fact * √(T(2n + 1) / T(2n)) * P[i₀+n-1, j₀+n-1]
else
aux_nm = T(n - m) * T(n + m)
a_nm = √(aux_an / aux_nm) * c
b_nm = √(T(n + m - 1) * T(n - m - 1) * aux_bn / aux_nm)
# We assume that the matrix is not initialized. Hence, we must not access
# elements on the upper triangle.
if m != n - 1
P[i₀+n, j₀+m] = a_nm * P[i₀+n-1, j₀+m] - b_nm * P[i₀+n-2, j₀+m]
else
P[i₀+n, j₀+m] = a_nm * P[i₀+n-1, j₀+m]
end
end
# Check if the maximum desired order has been reached.
m == m_max && break
end
end
return P
end
The point of StaticArrays is not that they are “on the stack” (which is not necessarily true), it is that the sizes are known statically (i.e. at compile time), so that loops can be unrolled and the code can be specialized on the size. This fundamentally isn’t going to work if the matrix size is a runtime variable like your n_max and m_max.
If you try to use StaticArrays with runtime/dynamic sizes, it will be type unstable and (usually) slow.
I see, thanks! However, in my application, I know the sizes during compilation time. In C++, it is a template. The problem is that if I use static arrays, I can allocate them, but not build them using a function unless I unroll the loops manually. This function will be called in different algorithms to build matrices with different sizes, but all known in compile time.
You still would have to provide the dimensions at compile time, with Val parameters, for example. But it seems that your P could be constructed as an MMatrix, and if you return it after conversion to a SMatrix it may not allocate.
using StaticArrays
function fully_normalized_legendre(
ϕ::T,
::Val{n_max},
::Val{m_max};
ph_term::Bool = false
) where {T<:Number, n_max, m_max}
P = @MMatrix zeros(float(T), n_max + 1, m_max + 1)
# Auxiliary variables to improve code performance.
s, c = sincos(T(ϕ))
# The sine must be always positive. In fact, `s` was previously computed using
# `√(1 - c^2)`. However, we had numerical problems for very small angles that lead to
# `cos(ϕ) = 1`.
if s < 0
s = -s
end
s_fact = !ph_term ? +s : -s
# Get the first indices in `P` to take into account offset arrays.
i₀, j₀ = first.(axes(P))
sq3 = √T(3)
@inbounds for n in 0:n_max
# Starting values.
if n == 0
P[i₀, j₀] = 1
continue
elseif n == 1
P[i₀+1, j₀] = +sq3 * c
if m_max > 0
P[i₀+1, j₀+1] = +sq3 * s_fact
end
continue
end
aux_an = T(2n - 1) * T(2n + 1)
aux_bn = T(2n + 1) / T(2n - 3)
for m in 0:n
if n == m
P[i₀+n, j₀+n] = s_fact * √(T(2n + 1) / T(2n)) * P[i₀+n-1, j₀+n-1]
else
aux_nm = T(n - m) * T(n + m)
a_nm = √(aux_an / aux_nm) * c
b_nm = √(T(n + m - 1) * T(n - m - 1) * aux_bn / aux_nm)
# We assume that the matrix is not initialized. Hence, we must not access
# elements on the upper triangle.
if m != n - 1
P[i₀+n, j₀+m] = a_nm * P[i₀+n-1, j₀+m] - b_nm * P[i₀+n-2, j₀+m]
else
P[i₀+n, j₀+m] = a_nm * P[i₀+n-1, j₀+m]
end
end
# Check if the maximum desired order has been reached.
m == m_max && break
end
end
return SMatrix(P)
end