Need help implementing non-allocating normal distribution with StaticArrays

General Usage Performance
1

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
2

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
3

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.