fix allocation counts for primals (plus _new_ frule)

[penelopeysm]

This PR fixes failing allocation tests in #714 for a number of primal functions, such as sum(abs2, ::Vector) on 1.12, and rand(::AbstractRNG, Float64) on 1.10 / 1.11. It also fixes the allocs for forward-mode _new_.

YES, it's ugly.

Specifically, this avoids the reporting of spurious allocations by:

1. using @allocations on 1.11, which Just Works;

2. creating a closure on 1.12, which effectively does the same as interpolating the arguments with a package like BenchmarkTools, which is needed for sum(abs2, xs).

3. manually splatting the varargs on all versions, which removes allocations which I believe are caused by something like incomplete specialisation. See for example https://discourse.julialang.org/t/specialization-on-vararg-of-types/108251 and also the following example:

using Random

vararg(f::F, x::Vararg{Any,N}) where {F,N} = nothing
vararg(randn, Xoshiro(468), Float64) # generate the specialisation
Base.specializations(@which vararg(randn, Xoshiro(468), Float64))
# -> Base.MethodSpecializations(MethodInstance for vararg(::typeof(randn), ::Xoshiro, ::Type))

explicit(f::F, x1::X1, x2::X2) where {F,X1,X2} = nothing
explicit(randn, Xoshiro(468), Float32) # generate the specialisation
Base.specializations(@which explicit(randn, Xoshiro(468), Float32))
# -> Base.MethodSpecializations(MethodInstance for explicit(::typeof(randn), ::Xoshiro, ::Type{Float32}))

Because Base.allocations is also a varargs function (https://github.com/JuliaLang/julia/blob/ba1e628ee49351af0b704afd2b2903d253bd3564/base/timing.jl#L485-L491), going via Base.allocations sometimes causes extra allocations (grr). So this PR also just inlines its definition into count_allocs.

Effectively, count_allocs is now a version of Base.allocations that uses a closure + explicit numbered arguments.

---

Note that I tested various combinations of BenchmarkTools and Chairmarks and I couldn't find a way to correctly get interpolations all the way down to 0, so I think this PR is legitimately the only way to correctly detect zero-allocation function calls. For example:

using BenchmarkTools
using Mooncake
function allocs()
    local f = Mooncake._new_
    local x = (Mooncake.TestResources.TypeStableMutableStruct{Float64}, 5.0, 4.0)
    display(@btime $f($x...))
    local closure = () -> f(x...)
    display(@btime $closure())
end
allocs() # 3, 3

[github-actions]

Mooncake.jl documentation for PR #843 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR843/

[github-actions]

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 100.0 ns │      1.8 │         1.9 │     1.1 │        5.61 │   8.21 │
│                  _sum_1000 │ 941.0 ns │     6.64 │        1.01 │  1500.0 │        33.7 │   1.09 │
│               sum_sin_1000 │  6.55 μs │      2.5 │         1.4 │    1.69 │        10.3 │   2.21 │
│              _sum_sin_1000 │  5.22 μs │     3.09 │        2.19 │   274.0 │        13.2 │    2.5 │
│                   kron_sum │ 317.0 μs │     41.5 │        2.73 │    5.22 │       225.0 │   8.35 │
│              kron_view_sum │ 331.0 μs │     42.0 │        3.37 │    10.7 │       210.0 │   8.63 │
│      naive_map_sin_cos_exp │  2.13 μs │      2.4 │        1.41 │ missing │        7.17 │   2.37 │
│            map_sin_cos_exp │  2.17 μs │     2.66 │        1.42 │    1.55 │         6.0 │   2.32 │
│      broadcast_sin_cos_exp │  2.26 μs │     2.38 │        1.39 │    2.33 │        1.46 │   2.22 │
│                 simple_mlp │ 160.0 μs │     7.16 │        3.04 │    1.82 │        14.3 │    3.4 │
│                     gp_lml │ 189.0 μs │     10.7 │        2.43 │    4.68 │     missing │   7.63 │
│ turing_broadcast_benchmark │  1.87 ms │     4.53 │        3.25 │ missing │        27.8 │   2.49 │
│         large_single_block │ 380.0 ns │     4.53 │        2.03 │  4420.0 │        32.9 │   2.24 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[penelopeysm]

The TLDR of where I'm stuck is 1.12. Benchmarking with a closure, i.e.

clos = () -> f(x1, x2, x3)
clos()
Base.allocations(clos)

gives 0 allocations for sum(abs2, randn(10)), BUT gives 1 allocation for Mooncake._new_(Mooncake.TestResources.TypeStableMutableStruct{Float64}, 5.0, 4.0).

So, OK, remove the closure:

f(x1, x2, x3)
Base.allocations(f, x1, x2, x3)

and that gives 0 allocations for _new_ but 1 allocation for sum.

Obviously, we could fix this quickly by special casing _new_ but that feels wrong and I'd like to know when the closure is needed and when it isn't needed.

On top of that, a lot of this is surprisingly impossible to reproduce locally in a simple script. I can at least get it to fail reliably with act.

[Technici4n]

Did you look at the generated IR? That might provide some insights

[yebai]

Just a quick question — is there any particular reason we can’t continue using Base.@allocations, given that it hasn’t been removed?

[penelopeysm]

Did you look at the generated IR?

Yup, but it wasn't very helpful and also it's the same on 1.11 and 1.12 for _new_, so I think it's certainly just down to changes in the way allocations works.

is there any particular reason we can’t continue using Base.@allocations

I mostly took my cue from #714 (comment). Just tried it anyway, and no luck on 1.12:

@test (@allocations f(x...)) == 0
# Fails for both `_new_` and `sum`

pclos = () -> f(x...)
@test (@allocations pclos()) == 0
# Fails for `_new_` and `randn(Xoshiro(3), Float64)`

[penelopeysm]

Base.allocations itself is stuck behind a Vararg argument, so now I've inlined its definition, and it is looking good when I run this locally with act. Will see what CI says.

https://github.com/JuliaLang/julia/blob/ba1e628ee49351af0b704afd2b2903d253bd3564/base/timing.jl#L485-L491

[penelopeysm]

This should fix all the allocations now. Remaining issues are lacpy, dtrsv, and JET test.

I'll try to find an MWE for the Base.allocations issue with Varargs and file it upstream, although goodness knows I've struggled to reproduce this locally, so I'm not optimistic about the chances.

[yebai]

Thanks, @penelopeysm!

[penelopeysm]

I kind of think that this count_allocs should be in a standalone package 😄

[codecov]

Codecov Report

❌ Patch coverage is 58.06452% with 13 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/test_utils.jl	58.06%	13 Missing ⚠️

📢 Thoughts on this report? Let us know!