How to ensure function is type-stable and allocation-free

General Usage Performance
1

I am writing a function and I’m trying to ensure that it is type-stable and allocation-free. I am using the @code_warntype to ensure type-stability and @btime to determine memory allocation.

Here is my code:

# we assume u is a 2 x num_particles matrix
function rhs!(du, u, parameters, t) 
    C = parameters
    v_x(x, y, t) = *(-sin(pi * y),cos(pi * t / C)) 
    v_y(x, y, t) = *(sin(pi * x), cos(pi * t / C))

    for i in 1:size(u, 2)
        du[1, i] = v_x(u[1, i], u[2, i], t)
        du[2, i] = v_y(u[1, i], u[2, i], t)
    end
    
end
u = 0.1 .+ .45 * rand(2, 100)

parameters = 0

du = similar(u)
t = 0.0
using BenchmarkTools; @btime rhs!(du,u,parameters,t)
@code_warntype rhs!(du,u,parameters,t)

How do I determine if my function is type-stable? Here is the output of @code_warntype:

MethodInstance for rhs!(::Matrix{Float64}, ::Matrix{Float64}, ::Int64, ::Float64)
  from rhs!(du, u, parameters, t) in Main at c:\Users\William\Downloads\part_2.jl:2
Arguments
  #self#::Core.Const(rhs!)
  du::Matrix{Float64}
  u::Matrix{Float64}
  parameters::Int64
  t::Float64
Locals
  @_6::Union{Nothing, Tuple{Int64, Int64}}
  v_y::var"#v_y#33"{Int64}
  v_x::var"#v_x#32"{Int64}
  C::Int64
  i::Int64
Body::Nothing
1 ─       (C = parameters)
│   %2  = Main.:(var"#v_x#32")::Core.Const(var"#v_x#32")
│   %3  = Core.typeof(C)::Core.Const(Int64)
│   %4  = Core.apply_type(%2, %3)::Core.Const(var"#v_x#32"{Int64})
│         (v_x = %new(%4, C))
│   %6  = Main.:(var"#v_y#33")::Core.Const(var"#v_y#33")
│   %7  = Core.typeof(C)::Core.Const(Int64)
│   %8  = Core.apply_type(%6, %7)::Core.Const(var"#v_y#33"{Int64})
│         (v_y = %new(%8, C))
│   %10 = Main.size(u, 2)::Int64
│   %11 = (1:%10)::Core.PartialStruct(UnitRange{Int64}, Any[Core.Const(1), Int64])
│         (@_6 = Base.iterate(%11))
│   %13 = (@_6 === nothing)::Bool
│   %14 = Base.not_int(%13)::Bool
└──       goto #4 if not %14
2 ┄ %16 = @_6::Tuple{Int64, Int64}
│         (i = Core.getfield(%16, 1))
│   %18 = Core.getfield(%16, 2)::Int64
│   %19 = Base.getindex(u, 1, i)::Float64
│   %20 = Base.getindex(u, 2, i)::Float64
│   %21 = (v_x)(%19, %20, t)::Float64
│         Base.setindex!(du, %21, 1, i)
│   %23 = Base.getindex(u, 1, i)::Float64
│   %24 = Base.getindex(u, 2, i)::Float64
│   %25 = (v_y)(%23, %24, t)::Float64
│         Base.setindex!(du, %25, 2, i)
│         (@_6 = Base.iterate(%11, %18))
│   %28 = (@_6 === nothing)::Bool
│   %29 = Base.not_int(%28)::Bool
└──       goto #4 if not %29
3 ─       goto #2
4 ┄       return nothing

And how do I modify the function to make it allocation free?

2

Change this to:

@btime rhs!($du,$u,$parameters,$t)

For accurate allocations.

(This might not make a difference but…) I would also refactor the function definitions v_x and v_y as you are creating a closure with them both (you should add C to the function definitions too). I am not sure if this does allocate, but it would be worth seeing if it makes a difference.

My version would be:

function rhs!(du, u, parameters, t) 
    C = parameters
    phase = cos(pi* t / C)
    for i in 1:size(u, 2)
        du[1, i] = -sin(pi*u[2, i]) * phase
        du[2, i] = sin(pi*u[1, i]) * phase
    end
    nothing
end
2 Likes
3

Also consider using cospi and sinpi here rather than cos and sin, and using axes(u,2) instead of 1:size(u,2).

4 Likes