Type instability (allocations) regression in Julia v1.10-alpha1

[ranocha]

We have some code in Trixi.jl that checks

if :i_backward in secondary_indices

where, secondary_indices is a tuple of Symbols, see https://github.com/trixi-framework/Trixi.jl/blob/07981c3e50943a9bdb1a845f918a4e8831714338/src/solvers/dgsem_p4est/dg_2d.jl#L148

We observed a significantly increased amount of allocations when checking CI with Julia v1.10.0-alpha1 in trixi-framework/Trixi.jl#1562. Locally, I checked that the allocations are gone when modifying the code above to

if any(==(:i_backward), secondary_indices)

[aviatesk]

Is secondary_indices compile-time constant?

[ranocha]

No, it is only known at runtime

[aviatesk]

Can you please give more details about the code? I can't confirm any allocation difference between the code snippets you shared above:

julia> global secondary_indices::Tuple{Symbol,Symbol,Symbol} = (:a, :b, :c)
(:a, :b, :c)

julia> f1() = :i_backward in secondary_indices
f1 (generic function with 1 method)

julia> f2() = any(==(:i_backward), secondary_indices)
f2 (generic function with 1 method)

julia> @benchmark f2()
[ Info: Loading BenchmarkTools...
BenchmarkTools.Trial: 10000 samples with 1000 evaluations.
 Range (min … max):  1.708 ns … 11.125 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     1.833 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.842 ns ±  0.213 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

                         █                      ▃             
  ▂▁▁▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▆▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▄ ▂
  1.71 ns        Histogram: frequency by time        1.92 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

julia> @benchmark f1()
BenchmarkTools.Trial: 10000 samples with 1000 evaluations.
 Range (min … max):  4.084 ns … 14.917 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     4.250 ns              ┊ GC (median):    0.00%
 Time  (mean ± σ):   4.261 ns ±  0.324 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

                     ▂                  █         ▆           
  ▂▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▃ ▂
  4.08 ns        Histogram: frequency by time        4.33 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

[ranocha]

An MWE seems to be

julia> v = Vector{NTuple{2, Symbol}}(undef, 100);

julia> for i in eachindex(v)
           values = (:i, :i_backward, :j, :j_backward)
           v[i] = (rand(values), rand(values))
       end

julia> function f1(v)
           res = 0
           for secondary_indices in v
               if :i_backward in secondary_indices
                   res += 1
               else
                   res += 2
               end
           end
           return res
       end
f1 (generic function with 1 method)

julia> function f2(v)
           res = 0
           for secondary_indices in v
               if any(==(:i_backward), secondary_indices)
                   res += 1
               else
                   res += 2
               end
           end
           return res
       end
f2 (generic function with 1 method)

julia> using BenchmarkTools

julia> f1(v)
162

julia> f2(v)
162

julia> @benchmark f1(v)
BenchmarkTools.Trial: 10000 samples with 152 evaluations.
 Range (min … max):  672.553 ns …   6.856 μs  ┊ GC (min … max): 0.00% … 83.65%
 Time  (median):     736.796 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   794.691 ns ± 429.613 ns  ┊ GC (mean ± σ):  5.35% ±  8.51%

  █▇▂▁▂                                                         ▂
  ████████▅▅▄▁▃▄▆█▃▃▄▁▁▃▁▁▁▁▃▁▁▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▅ █
  673 ns        Histogram: log(frequency) by time       4.64 μs <

 Memory estimate: 3.12 KiB, allocs estimate: 100.

julia> @benchmark f2(v)
BenchmarkTools.Trial: 10000 samples with 965 evaluations.
 Range (min … max):  81.077 ns … 273.202 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     81.771 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   83.106 ns ±   7.061 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▅█▄                                                        ▁ ▁
  ███▄▆▇▃▃▁▁▁▁▁▃▁▁▃▁▁▁▁▃▁▁▃▅█▇▆▄▇▇▇▆▆▄▃▁▃▃▃▃▃▁▁▁▁▁▁▁▁▃▄▃▄▇▇▅▇█ █
  81.1 ns       Histogram: log(frequency) by time       113 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

That's with Julia v1.10.0-alpha1. With the current stable release Julia v1.9.2, I get

julia> @benchmark f1(v)
BenchmarkTools.Trial: 10000 samples with 961 evaluations.
 Range (min … max):  78.634 ns … 204.733 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     86.562 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   87.574 ns ±   6.437 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▁          ▇█  ▁                                             ▁
  █▇▁▁▁▁▁▃▁▁▁██▇██▅▁▅▃▁▄▃▄▁▁▃▃▁▄▃▅█▇▅▇█▆▅▆▄▃▃▄▃▁▁▃▃▁▁▁▄▆▄▄▇▇▆█ █
  78.6 ns       Histogram: log(frequency) by time       118 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

julia> @benchmark f2(v)
BenchmarkTools.Trial: 10000 samples with 961 evaluations.
 Range (min … max):  80.238 ns … 267.342 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     86.915 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   88.485 ns ±   7.884 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

         ▂█ ▂                                  ▁               ▁
  ▇▁▃▁▁▁▁██▅█▅▄▁▁▁▃▃▁▁▃▁▁▄█▆▆▇▆▇▅▁▁▄▄▁▁▃▄▁▅▆▄▇▇█▇▆▅▄▃▄▁▁▃▄▃▄▄▄ █
  80.2 ns       Histogram: log(frequency) by time       129 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

[aviatesk]

Yes, thanks. I could reproduce it with this MRE also:

julia> f1(rt) = Base.sym_in(:i_backward, rt[]);

julia> f2(rt) = @inline Base.sym_in(:i_backward, rt[]);

julia> rt = Ref((:b_back, :foofakgka, :i_backw));

julia> f1(rt); @allocated f1(rt)
32

julia> f2(rt); @allocated f2(rt)
0

So this is caused by the @noinline annotation added on Base.sym_in.

julia/base/tuple.jl

Lines 609 to 610 in d215d91

	function sym_in(x::Symbol, @nospecialize itr::Tuple{Vararg{Symbol}})
	@noinline

As the MRE above shows, we can eliminate the allocation if we remove it. Alternatively, we can achieve the same result if we remove the @nospecialize annotation.

The @noinline annotation was added in #48066, which I am not sure about its rationale. @Keno Could you tell me why we want to have @noinline on Base.sym_in? Do you think we can remove either of the annotations?

[aviatesk]

FWIW, the sym_in method was added in #36596 by @vtjnash along with @nospecialize annotation on itr argument.

[Keno]

@Keno Could you tell me why we want to have @noinline on Base.sym_in? Do you think we can remove either of the annotations?

The problem is that inlining these functions destroys the effect assumptions, so they can't be constproped, even if irinterp later finds the constants. What object is being allocated here? Can we improve the codegen to avoid that happening?

[aviatesk]

It appears that the allocation is triggered by the construction of a Vararg tuple, which is used for dispatching to Base.sym_in:

julia> @code_typed f1(rt)
CodeInfo(
1 ─ %1 = Base.getfield(rt, :x)::Tuple{Symbol, Symbol, Symbol}
│   %2 = invoke Base.sym_in(:i_backward::Symbol, %1::Tuple{Vararg{Symbol}})::Bool
└──      return %2
) => Bool

julia> @code_llvm f1(rt)
;  @ REPL[1]:1 within `f1`
define i8 @julia_f1_1641({}* noundef nonnull align 8 dereferenceable(24) %"rt::RefValue") #0 {
top:
  %gcframe13 = alloca [5 x {}*], align 16
  %gcframe13.sub = getelementptr inbounds [5 x {}*], [5 x {}*]* %gcframe13, i64 0, i64 0
  %0 = bitcast [5 x {}*]* %gcframe13 to i8*
  call void @llvm.memset.p0i8.i64(i8* align 16 %0, i8 0, i64 40, i1 true)
  %1 = call {}*** inttoptr (i64 6855994896 to {}*** (i64)*)(i64 261) #9
; ┌ @ refvalue.jl:59 within `getindex`
; │┌ @ Base.jl:37 within `getproperty`
    %2 = bitcast [5 x {}*]* %gcframe13 to i64*
    store i64 12, i64* %2, align 16
    %3 = load {}**, {}*** %1, align 8
    %4 = getelementptr inbounds [5 x {}*], [5 x {}*]* %gcframe13, i64 0, i64 1
    %5 = bitcast {}** %4 to {}***
    store {}** %3, {}*** %5, align 8
    %6 = bitcast {}*** %1 to {}***
    store {}** %gcframe13.sub, {}*** %6, align 8
    %7 = bitcast {}* %"rt::RefValue" to {}**
    %.unpack = load {}*, {}** %7, align 8
    %.not = icmp eq {}* %.unpack, null
    br i1 %.not, label %fail, label %pass

fail:                                             ; preds = %top
    call void @ijl_throw({}* inttoptr (i64 5353759904 to {}*))
    unreachable

pass:                                             ; preds = %top
    %8 = bitcast {}* %"rt::RefValue" to [3 x {}*]*
    %.elt2 = getelementptr inbounds [3 x {}*], [3 x {}*]* %8, i64 0, i64 1
    %.unpack3 = load {}*, {}** %.elt2, align 8
    %.elt4 = getelementptr inbounds [3 x {}*], [3 x {}*]* %8, i64 0, i64 2
    %.unpack5 = load {}*, {}** %.elt4, align 8
    %9 = getelementptr inbounds [5 x {}*], [5 x {}*]* %gcframe13, i64 0, i64 4
    store {}* %.unpack, {}** %9, align 16
    %10 = getelementptr inbounds [5 x {}*], [5 x {}*]* %gcframe13, i64 0, i64 3
    store {}* %.unpack5, {}** %10, align 8
    %11 = getelementptr inbounds [5 x {}*], [5 x {}*]* %gcframe13, i64 0, i64 2
    store {}* %.unpack3, {}** %11, align 16
; └└
  %ptls_field14 = getelementptr inbounds {}**, {}*** %1, i64 2
  %12 = bitcast {}*** %ptls_field14 to i8**
  %ptls_load1516 = load i8*, i8** %12, align 8
  %box = call noalias nonnull dereferenceable(32) {}* @ijl_gc_pool_alloc(i8* %ptls_load1516, i32 1184, i32 32) #7
  %13 = bitcast {}* %box to i64*
  %14 = getelementptr inbounds i64, i64* %13, i64 -1
  store atomic i64 5280480960, i64* %14 unordered, align 8
  %15 = bitcast {}* %box to [3 x {}*]*
  %.repack = bitcast {}* %box to {}**
  store {}* %.unpack, {}** %.repack, align 8
  %.repack8 = getelementptr inbounds [3 x {}*], [3 x {}*]* %15, i64 0, i64 1
  store {}* %.unpack3, {}** %.repack8, align 8
  %.repack10 = getelementptr inbounds [3 x {}*], [3 x {}*]* %15, i64 0, i64 2
  store {}* %.unpack5, {}** %.repack10, align 8
  store {}* %box, {}** %11, align 16
  %16 = call i8 @j_sym_in_1643({}* inttoptr (i64 4336660832 to {}*), {}* nonnull readonly %box)
  %17 = load {}*, {}** %4, align 8
  %18 = bitcast {}*** %1 to {}**
  store {}* %17, {}** %18, align 8
  ret i8 %16
}

In particular, this seems to be the source of allocation %box = call noalias nonnull dereferenceable(32) {}* @ijl_gc_pool_alloc(i8* %ptls_load1516, i32 1184, i32 32) #7.

Alternatively we can remove the @nospecialize annotation, although it forces us to compile Base.sym_in multiple times for each tuples with different length.

[Keno]

We could potentially consider adjusting the specsig codegen ABI for Tuple{Vararg{T}} to separate pass the pointer and the length to avoid having to do this allocation.

[vtjnash]

That seems like what the ABI was supposed to be for this. Why is it getting a specsig for the Vararg, when non-specsig should have "done the right thing" instead.

[Keno]

We have special handling for a vararg signature, but not for tuples of unknown size.

[ranocha]

Are there any plans to fix this in Julia v1.10.0 or do we need to introduce workarounds in our packages?

[oscardssmith]

A 10x regression seems like it should be on the milestone

[ranocha]

Thanks for adding it to the milestone!

[JeffBezanson]

For 1.10 we can remove the noinline, but the correct fix is to add ABI support for passing a tuple like this without boxing (or boxing on the stack).

[aviatesk]

This issue was fixed. Closing.