Using a keyword argument leads to enormous allocations

General Usage Performance
1

I’ve been running into a bizarre problem with huge allocations, and I’ve just discovered that the problem can be avoided if I don’t use a keyword argument to a function inside my main computation function. The allocations aren’t happening in a line that actually uses the kwarg; the allocations come later. @code_warntype shows no problems, and specifically everything directly involved in the worst lines are just Int64 and Float64. Even for my MWE below, the worst line — which I wouldn’t expect to allocate at all — averages ~128 B of allocation per iteration. I don’t even see that many bytes in all the variables involved in the line! Even the ranges of the for loops are allocating lots of memory.

But all these allocations go away if I just don’t use the keyword argument in a function call inside my big and ugly function. The function with the kwarg is type stable either way. But looking more closely at the @code_warntype output, I see that that function is represented as a Core.kwfunc. I guess this screws things up??? Should I have known this somehow? Is this a bug?

My use case is a pretty big and ugly recurrence computation, but I’ve managed to simplify it as much as possible. Here, index is the function with the kwarg that I’d like to use, inplace! is the core computation (drastically simplified here), and compute_a just sets things up and measures the allocations.

using Profile

function index(n, mp, m; n_max=n)
    n + mp + m + n_max
end

function inplace!(a, n_max, index_func=index)
    i1 = index_func(1, 2, 3; n_max=n_max)  # This version allocates
    # i1 = index_func(1, 2, 3)             # This version doesn't allocate
    i2 = size(a, 1) - 2i1
    for i in 1:i2                   # Allocates 3182688 B if using kwarg above
        a[i + i1] = a[i + i1 - 1]   # Allocates 9573120 B if using kwarg above
    end
    for i in 3:i2-4                 # Allocates 3182576 B if using kwarg above
        a[i + i1] -= a[i + i1 - 2]  # Allocates 12771408 B if using kwarg above
    end
end


function compute_a(n_max::Int64)
    a = randn(Float64, 100_000)

    inplace!(a, n_max, index)

    Profile.clear_malloc_data()

    inplace!(a, n_max, index)
end


compute_a(10)

[Curiously, I need both for loops, or the allocations disappear.]

If I just remove the kwarg from the call to index_func, the allocations all competely disappear; there are no allocations inside inplace! in that case.

2

This is likely because Julia is not specialising on the type of the function index_func, see here.

To fix this, index_func should be typed:

function inplace!(a, n_max, index_func::F = index) where {F}
    # ...
end
3

Even if it is solved by annotating the type of the function, it is a bug to me.

1 Like
4

I have run the code (more than one time and cleaning the allocations between timings) and in Julia 1.5.3 calling inplace! from inside compute_a with and without the parameter leads to the same number of allocations. Interesting that not passing explicitly the argument to the outer function has no effect.

Not really how bugs work. It is documented behavior:

As a heuristic, Julia avoids automatically specializing on argument type parameters in three specific cases: Type, Function, and Vararg. Julia will always specialize when the argument is used within the method, but not if the argument is just passed through to another function. This usually has no performance impact at runtime and improves compiler performance. If you find it does have a performance impact at runtime in your case, you can trigger specialization by adding a type parameter to the method declaration.

1 Like
5

But the function is being used, isn’t It?

1 Like
6

I guess this is what confuses me. I did use the argument within the method. Now, looking very closely at the output of @code_warntype, I see that julia actually takes that argument, and passes it through to Core.kwfunc. So a lawyer could certainly argue that if I happen to actually use kwargs that are present in a function, the argument “is” passed through — just not by me. It’s asking a lot to say that users, even having read that part of the docs, also have to know what julia is doing behind the scenes.

This is doubly so when @code_warntype says everything is copacetic with all the nice blue colors. And then if I really want to dig in, I can read all the way to the bottom of that performance tip, and see that “@code_typed and friends” (which I guess includes @code_warntype) will always show you specialized code, so I should check

(@which inplace!(a, n_max, index)).specializations

instead, which I do and see that its output includes

MethodInstance for inplace!(::Vector{Float64}, ::Int64, ::typeof(index))

which makes me think that it has indeed specialized on the type of my function. Now, surely you could come back with some detailed explanation of why I’m misinterpreting that somehow, but that would be missing the point.

My point is that I do not expect to have to learn the internals to figure out why, if I happen to use a keyword argument, my code runs 100 times slower and allocates GiBs of memory when it wasn’t allocating any before. (Those were the numbers I saw in my real code, not this MWE.) While I’m no julia expert, I’m not terrible either. In my real code, the index function was located pretty far from the code that was actually allocating (loops like the ones in my MWE), so it took me a lot of luck and days of work to figure out the source of the problem. At the very least, this represents a documentation bug, in that keyword arguments should be mentioned in that performance tip. It would also be nice if there were some tool available to find problems like this. Maybe @code_warntype could show such functions as red or yellow. Of course, ideally, julia would be able to handle this scenario on its own…

2 Likes
7

Ah, yes, I got confused by the code. There are two keyword parameters there and none seem to fill the whole pre-requisite. The index_func is a keyword parameter that takes a Function but it is used inside the immediate receiver; n_max is a positional parameter in the same signature and it is passed through (as a keyword parameter) to index_func. I kinda mixed the two and saw a Function parameter passing throught.

Truly, there is every reason to be confused in this case. The non-specializing behaviour is a corner case, and this is a corner case within the corner case. The fact that keyword parameters are implemented this way probably should be regarded as a implementation detail, and not something you need to know to reason about code.

I think some time ago Julia also did not specialize on keyword parameters in general, so maybe this is a special case that was forgotten when this changed.

That said, were you able to confirm that forcing specialization using the trick described above really solves your problem? Or it may be yet that your problem is unrelated to that?

8

Sorry, I should have mentioned. Yes, it does. (Thanks @jipolanco.) Of course, I realized that my index function made more sense with a different signature anyway, so I’ve already changed it in my code to not use a kwarg at all. But it sure took me a long time to figure out that’s what I needed to do.

1 Like
9

I submitted this as an issue on julia. Jeff Bezanson labeled it a bug, and fixed it on master, and it has been labeled for backporting to julia 1.8. No code is perfect, but julia is constantly moving toward to the goal.

Thanks for the helpful discussion everyone!

4 Likes