# Need help implementing non-allocating normal distribution with StaticArrays

I’m trying to implement a zero-allocating version of a normal distribution `pdf`, but get some unexpected allocations. Below is my code for a script `./test_allocations.jl`, which could be ran to test allocations quickly.

``````using StaticArrays
using Profile

module T

using StaticArrays
using LinearAlgebra

CovMat = SMatrix{3, 3, Float64}
MeanVec = SVector{3, Float64}

struct MvNormalF
μ::MeanVec;
Σ::CovMat;
pdf_divider::Float64;
Σ_inv::CovMat;

MvNormalF(μ::MeanVec, Σ::CovMat) = new(μ, Σ, 0.5 * log((2π)^3 * det(Σ)), inv(Σ))
end

@inbounds function vXvt(a::Float64, b::Float64, c::Float64, m::CovMat)
return a*a*m[1, 1] + b*b*m[2, 2] + c*c*m[3, 3] + 2*a*b*m[1, 2] + 2*a*c*m[1, 3] + 2*b*c*m[2, 3]
end

@inbounds log_pdf(d::MvNormalF, x::SVector{3, Float64})::Float64 =
vXvt(x[1] - d.μ[1], x[2] - d.μ[2], x[3] - d.μ[3], d.Σ_inv)

function test_pdf_perf(comp::MvNormalF, x::SVector{3, Float64}; n::Int)
for i in 1:n
log_pdf(comp, x)
end
end

end

x = SVector{3}(0., 1., 1.);
comp = T.MvNormalF(T.MeanVec(0., 0., 0.), T.CovMat(1., 0., 0., 0., 1., 0., 0., 0., 1.))

T.test_pdf_perf(comp, x; n=10);
Profile.clear_malloc_data()

T.test_pdf_perf(comp, x; n=10);
``````

If I run it with `julia --track-allocation=all ./test_allocations.jl`, it shows that the only allocation is 640 bytes for the line `@inbounds log_pdf(...)`. What is the reason of these allocations or how can I find its source?

1 Like

My first step when unexpected allocations occur is to look at the output of `@code_warntype`, because often they occur due to type inference issues.

``````julia> @code_warntype T.log_pdf(comp, x)
MethodInstance for Main.T.log_pdf(::Main.T.MvNormalF, ::SVector{3, Float64})
from log_pdf(d::Main.T.MvNormalF, x::SVector{3, Float64}) in Main.T at REPL[88]:22
Arguments
#self#::Core.Const(Main.T.log_pdf)
d::Main.T.MvNormalF
x::SVector{3, Float64}
Body::Float64
1 ─ %1  = Main.T.Float64::Core.Const(Float64)
│   %2  = Base.getindex(x, 1)::Float64
│   %3  = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %4  = Base.getindex(%3, 1)::Float64
│   %5  = (%2 - %4)::Float64
│   %6  = Base.getindex(x, 2)::Float64
│   %7  = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %8  = Base.getindex(%7, 2)::Float64
│   %9  = (%6 - %8)::Float64
│   %10 = Base.getindex(x, 3)::Float64
│   %11 = Base.getproperty(d, :μ)::SVector{3, Float64}
│   %12 = Base.getindex(%11, 3)::Float64
│   %13 = (%10 - %12)::Float64
│   %14 = Base.getproperty(d, :Σ_inv)::SMatrix{3, 3, Float64} # <-- marked in red
│   %15 = Main.T.vXvt(%5, %9, %13, %14)::Float64
│   %16 = Base.convert(%1, %15)::Float64
│   %17 = Core.typeassert(%16, %1)::Float64
└──       return %17
``````

It doesn’t show it in the above code, but the line
`%14 = Base.getproperty(d, :Σ_inv)::SMatrix{3, 3, Float64} ` shows an issue.

That is because `SMatrix{3, 3, Float64}` is not a concrete type.

``````julia> isconcretetype(SMatrix{3,3,Float64})
false
``````

In fact, `SMatrix` has a fourth type parameter equal to the total number of elements (see the help entry for more info)

Solution: Change the definition to `CovMat = SMatrix{3, 3, Float64, 9}` and your code should be type stable and allocation free.

Note that you can construct an `MvNormal` from `Distributions` with static arrays too.

Also, it seems there is a bug in your `log_pdf`

``````julia> comp = T.MvNormalF(T.MeanVec(0., 0., 0.), T.CovMat(1., 0., 0., 0., 1., 0., 0., 0., 1.))
Main.T.MvNormalF([0.0, 0.0, 0.0], [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0], 2.756815599614018, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0])

julia> D = MvNormal(comp.μ, comp.Σ)
MvNormal{Float64, PDMats.PDMat{Float64, SMatrix{3, 3, Float64, 9}}, SVector{3, Float64}}(
dim: 3
μ: [0.0, 0.0, 0.0]
Σ: [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]
)

julia> logpdf(D, x) - T.log_pdf(comp, x)
-5.7568155996140185
``````
1 Like

That solves it, thank you so much! As for the bug, I deliberately removed some parts of the formula to narrow down the problem, so it’s fine.