Strange memory allocations with structs containing functions

I’m working on a piece of code that takes a struct containing a Function and then calls that function. However, I’m fighting allocation issues. Ideally, I would think that the ability to specify the function signature, rather than just func::Function would help… but I’m sure if that’s possible and/or how to do it.

Also, in the course of investigating this, I found some odd behavior… see the MWE below. I’m trying to understand why one of the these function invocations allocates memory, while the other doesn’t. Anyone have any insight?

using Printf

struct Foo
    func::Function
end

function f1(λ::F, m::Foo) where F
    x = acos(√(rand()))
    y = 2π*rand()
    z::F = m.func(λ)
    z
end

function f2(λ::F, m::Foo) where F
    z::F = m.func(λ)
    z
end

const foo = Foo(x -> x^2)

@show f1(3.14, foo)
@time f1(3.14, foo)

@show f2(3.14, foo)
@time f2(3.14, foo)

Output is:

f1(3.14, foo) = 9.8596
  0.000003 seconds (2 allocations: 32 bytes)
f2(3.14, foo) = 9.8596
  0.000000 seconds
9.8596

Function is an abstract type. If you don’t want allocations, all the fields of your struct should have a concrete type. A quick fix for this problem is:

struct Foo{F}
    func::F
end
2 Likes

Thanks! That’s what I was looking for… but was unsure of how to specify a concrete type T <: Function. Just curious, is the inconsistent behavior to be expected? Or is that potentially a bug?

It’s expected. See Performance Tips · The Julia Language

1 Like

Follow on question… is there a way to place a constraint on the Foo.func? For example, maybe it has to have the signature func(x::Float64)::Float64?

Using the suggestion above is helpful, but

struct Foo{F}
    func::F
end

means (I think) that the instances Foo(x -> x^2) and Foo(x -> x^3) are different types.

So if I have a Vector{Foo}, it cannot be specialized for the case when all funcs have the same signature, i.e.

const foos = [Foo(x -> x^2), Foo(x -> x^3)]
typeof(foos) <: Vector{typeof(first(foos))}

is false. I believe this is the reason the my example below allocates when running run_foos

const foos = [Foo(x -> x^2), Foo(x -> x^3)]
typeof(foos) <: Vector{typeof(first(foos))}

function run_foos(foos)
    total = 0.0
    for foo in foos
        total += foo.func(1.0)
    end
    total
end

run_foos(foos)
@time run_foos(foos) # result:   0.000005 seconds (4 allocations: 64 bytes)

This is correct, as you observed. In this very particular example, you can use either of

foos = [Foo(Base.Fix2(^,2)), Foo(Base.Fix2(^,3))]
# Vector{Foo{Base.Fix2{typeof(^), Int64}}}

foos = [Foo(x->x^pow) for pow in 2:3] # or use `Base.Fix2` again
# Vector{Foo{var"#6#8"{Int64}}} # anonymous function closing over an `Int64`

which do have uniform types. As I showed here, sometimes this can be done using functors (such as Base.Fix2 or one you make yourself) or closures. Note the comprehension over x->x^pow did produce a consistent function type because it was the same declaration that resulted in both, just capturing different values of the pow::Int64 in a closure.

Other times you might have to be more creative (for example, making [cos,sqrt] stable doesn’t really work with this pattern). If they’re a fixed set of functions that doesn’t ever change, a tuple of functions (cos,sqrt) can sometimes work. Other cases may be more challenging.

At the fully-flexible extreme, you can also look into FunctionWrappers.jl for performantly wrapping different functions with matching type signatures. But I don’t have a lot of experience with that package so can’t help more.


You can typeassert the result of a function call

x = foos[2].func(3.14)::Float64
# you can additionally typeassert the `3.14::Float64`,
# but that doesn't change you anything if the compiler
# already knows the type of the input

This doesn’t fully resolve instability if foos has differently-typed functions, but it means that the output is known to be Float64 (or else it will error) so any instability will not propagate.

Note that functions in Julia don’t have input/output types, so you can’t look at a function to see its signature. You can look at a function’s list of methods to see which (if any) corresponds to a particular set of input types (and also attempt to look at the resulting output type), but this isn’t usually something that’s done in Julia.

FunctionWrappers, however, requires concrete input/output type declarations and can be used for this purpose. But again, I’ll recommend you try to do this without FunctionWrappers if possible.

3 Likes

Sorry for the delay. Thanks for the reply. This is generally useful to know, but as you pointed out it doesn’t work for my specific use-case. Fortunately, even without this optimization my actual code performs well enough.