How to understand reasons for some unexpected heap allocations

What would be good is an analysis tool that can give a reason for each allocation. E.g.:

  • Mutable object with unknown lifetime
  • Immutable isbits function return value that is too large to be stack based
  • Return value of dynamically dispatched function

Etc.

Incredibly irritating but hopefully enlightening update. I hadn’t really looked through the codebase and it occurred to me that “callback” might be an input function, and a longstanding problem with @code_ is that it assumes those arguments are specialized during compilation when they’re not. That flaw can propagate to more sophisticated static analysis packages. Another factor is that type inference (and presumably the resulting static dispatches) can still occur even when specialization is suppressed (see docstrings of @nospecialize and the internal Base.@nospecializeinfer) and affects compilation of callers, so the fix would not simply infer from abstract input types.

With the edit callback::F=nullcallback) where F to optimize!, the check_allocs report increases to 27 entries, one of which is a dynamic dispatch to optimizeinternal!, which also has an unannotated callback argument. The allocation of Int64 is also gone because it was processed into something else before the optimizeinternal! call.

With the edit callback::F) where F to optimizeinternal!, the check_allocs report increases to 62 entries. The allocation of NLLSOptions is gone.

These newer entries include the allocations shown in the original runtime profile graph. Some of them are conditional allocations that likely did not occur during the run.

I still don’t know much at this point. I had assumed that the non-specialization heuristic for unparameterized and uncalled Function/Type/Vararg arguments didn’t cause dynamic dispatches, so I’m not certain if check_allocs is entirely accurate here despite being consistent with the boxing of NLLSOptions and Int64. But @code_llvm and Cthulhu.@descend are visibly mis-specializing callback, and JET.@report_opt likely does as well.

Thanks for this. If I remove the call to the callback function (the default, nullcallback, does nothing), then all the allocations to the immutable values disappear, leaving only the mutable ones. So that function is the cause of the most unexpected allocations.

I’ve tried annotating the return type of the callback at the call site (callback(...)::Tuple{Float64, Int}), but rather than remove any allocations, this just creates one extra, for the tuple.

I also tried annotating all the instances of the callback input argument, thus: callback::F) where F and this similarly removed all the allocations of immutable values. Great! But I find it odd that adding a type parameter to a function signature improves things, because my understanding was that Julia will specialize on the input type when it can anyway.

There are exceptions to input concrete type specialization because it can easily result in compilation bloat of forwarding higher-order functions for little gain (here we only save 2 small allocations). Method parameters or an unbroken chain of inlining are how to opt back in.

And the issue about @code_ failing to consider unspecialized arguments:
Can we support a way to make `@code_typed` stop lying to us all? :slight_smile: · Issue #32834 · JuliaLang/julia

Just a short comment on the @warn and other infrequently called code in critical code sections. A @warn inserts quite a bit of code, and can prevent inlining and other transformations of the code, which may affect optimization of the code in various ways. Though, I don’t know if that’s the issue here.

It might be a good idea to isolate infrequently called code in separate functions, either explicitly or at the use site. Done the wrong way it may incur boxing, so one should be a bit careful.

# direct use risks code bloat:
function sqr1(i)
    if i < zero(i)
        @warn "i is negative"
    end
    return i^2
end

# isolate the @warn
function sqr2(i)
    if i < zero(i)
        @noinline (() -> @warn "i is negative")()
    end
    return i^2
end

I have a little @outline macro that I like to use for these situations, i.e.

macro outline(args...)
    captures = esc.(args[1:end-1])
    body = esc(args[end])
    quote
        @noinline thunk($(captures...),) = $body
        thunk($(captures...),)
    end
end

function sqr2(i)
    if i < zero(i)
        @outline(i, @warn "i is negative" i)
    end
    return i^2
end