No allocs after restart

First off let me say that I realize it’s a long shot that someone can help me because I’m not providing minimal working example. However I’m hoping that this will sound familiar to someone.

I’ve make all my code entirely type stable and entirely allocation free. I have checked that there are no allocations by both running @time and with PProf. However, when I edit and re-run it, it starts allocating. If afterwards I restart it, it no longer allocates.

Schematically, I have a top level function

function model(...)
    (...)
    y = f1(x)
    (...)
end

If I restart Julia the above doesn’t allocate (the above isn’t called by me, it’s called by other code). In particular f1 is type stable.

I’m iterating on some research so some times I want to replace f1 with f2 and all of a sudden it allocates. This is surprising because f2 is also type stable, and takes the same types and returns the same types as f1. However, if I restart Julia again, all of a sudden this

function model(...)
    (...)
    y = f2(x)
    (...)
end

no longer allocates. If I replace f2 with f1 again it allocates again.

How is this possible?

I’m hoping that someone here has seen this before. If not I can put in the work to come up with a minimal example. It’s just that it might be significant work because the code is complex-ish, involving a lot of callbacks and closures and my own defined types and it’s hard to know a priori which of these aspects are causing the problem.

How are measuring allocation? You might be measuring allocation during compilation.

No, in everything that I’ve described I always run it once to get it compiled and ignore the results of the first run. I’m talking afterwards.

So to be clear the scenarios are:

Scenario 1

  • edit model to use function f1
  • restart julia
  • run once
  • run second time with @time, verify that there are no allocations

Scenario 2

  • edit model to use function f2
  • restart julia
  • run once
  • run second time with @time, verify that there are no allocations

Scenario 3

  • edit model to use one of the functions
  • restart julia
  • run once
  • run second time with @time, verify that there are no allocations
  • then edit model to use the other function (don’t restart)
  • run once
  • run second time with @time, now it’s allocating on every iteration of the inner loop!

It’s known that type inference can sometimes exhibit hysteresis, i.e., it may succeed or fail in inferring every type (and thus statically resolving all method dispatch) depending on the order in which methods are compiled. So you may be observing allocations from a dynamic dispatch due to incomplete inference in scenario 3, while in scenarios 1 and 2 methods are compiled in a different order such that inference succeeds.

This is most commonly observed in the context of recursion. Does your code have any of that?

The information about the issue I’m referring to seems a bit scattered, but looks like maybe this is the best place to start reading if you’re curious: problem with the unreliable approximation of `Core.Compiler.return_type` · Issue #35800 · JuliaLang/julia · GitHub. Anyway, I’d experiment with minor refactoring and/or strategically placed type assertions (i.e., return value::Type) anywhere you have a recursive loop, and see if that fixes the problem.

Exactly, my suspicion is that the allocations in scenario 3 are coming from some dynamic dispatch that doesn’t happen under some unspecified ideal conditions but breaks the moment you touch anything.

However my code doesn’t have any recursion whatsoever, no. Must be something else :expressionless:

A brittle type inference optimization that comes to mind occurs when a function has 3(?) or fewer methods, so the compiler could branch to and infer them all for an argument-type-unstable call. But I don’t have enough details about how you’re editing the models to rule that out. If you’re evaluating all the code (just not ran and compiled yet) right after each restart, then I don’t think this would be relevant; if you’re evaluating the code for each function on demand, then it could be. A way to remove this factor completely is to have f1 and f2 run the same code, just duplicating their method body; worth a try to see if you still replicate this issue, but it’s just as likely such a large change just fails to reproduce anything for no important reason.

1 Like

Regarding code evaluation, the first time I evaluate the code after restart I activate the environment, then shift-enter a block of code (on vscode) that runs all the includes which in particular define the model function. After that, on I literally edit the function model to call f1 or f2 and then press shift-enter.

I didn’t really understand the experiment you proposed. Do you want me to have two functions f1 and f2 contain the same code, for example that of f1? I can do the experiment, but what does this show?

I’ve just made the following experiment: edit f2, copy past the code of f1 into it. Same phenomenon happens.

Another experiment that I’ve made: Try the same scenarios from the julia repl directly (as opposed to vscode), as same phenomenon happens.

Yet another experiment: In scenario 3 first do the experiment of replacing f1 with f2 and then replace it with f1 again. It stays allocating. Facepalm.

Yeah basically on restart you have all your code evaluated, and f1 and f2 have the exact same code, just differing in function name (and thus type). Then if you do your experiment and replicate this errant allocation, then you have ruled out any source code differences as the cause. Sounds like you have done something similar enough (and I would try it the other way, pasting f2’s original code into f1, just to see if it’s not particular to f1’s original code), but at this point it’s worth trying to make a MWE to properly demonstrate it and allow people to fiddle with it themselves. It could very well be something that you naturally wouldn’t even think to describe.

It’s also worth saying that in most contexts, an extra small allocation is negligible over a decent period of time, so without more information it’s very likely this side project won’t get in the way of your primary work.

This vaguely reminds me of the following.

Another possibility is there is a bug in @time if that is the sole way you are checking allocations.

Have you tried AllocCheck.jl?

Haha believe it or not I had also considered that. I’ve used PProf and it shows the allocs. In Scenario 3 I get 10M allocations (consistent with the number from @time) and they’re all inside FunctionWrapper which basically means it’s above my paygrade. FunctionWrapper calls something called do_ccall which calls something called macro expansion, which leads to an allocation named just Profile.Allocs.UnkownType, 10 M of these. To be clear, these are not present in scenarios 1 or 2.

I’m currently debugging some heisenallocations myself which seem to be related to how Julia avoids specializing on function arguments unless they’re given a type variable with where, see Performance Tips · The Julia Language. I’ve confirmed with both JET.jl and DispatchDoctor.jl that my code is stable and inferred end-to-end, so that’s not the issue. In fact, I can completely eliminate the allocations by decorating a particular subset of functions with @stable from DispatchDoctor—but I get allocations both when I remove all occurences of @stable and when I add it to every method (in all cases, the code works and returns the correct result, no errors). Truly mysterious!

I think this is just showing it’s a generated function. Can’t really guess where the 10M allocations are from (over how many iterations?) without details. I do recall that it’s the reinit_wrapper call that gets flagged as a runtime dispatch by JET.jl, but that is only supposed to be run automatically at precompile-time, not at runtime.

Even less certain about this causing allocations, but worth saying that reinit_wrapper mutation is necessary to make even newly instantiated FunctionWrappers respond to dynamic method changes. This isn’t documented or public but I spot no alternative in the source.

julia> using FunctionWrappers: FunctionWrappers as FWs, FunctionWrapper as FW

julia> foo() = 1; fwoo = FW{Int, Tuple{}}(foo);

julia> foo(), fwoo() # works as expected
(1, 1)

julia> foo() = 2; foo(), fwoo() # previous wrapper is obsolete
(2, 1)

julia> fwoo2 = FW{Int, Tuple{}}(foo); foo(), fwoo(), fwoo2() # new wrapper is obsolete
(2, 1, 1)

julia> FWs.reinit_wrapper(fwoo); foo(), fwoo(), fwoo2() # reinit_wrapper updates a wrapper
(2, 2, 1)

julia> foo() = 3; fwoo3 = FW{Int, Tuple{}}(foo); foo(), fwoo(), fwoo2(), fwoo3() # new wrapper is obsolete to 1st version
(3, 2, 1, 1)