Poor codegen with `setindex` on a `NamedTuple`

[tecosaur]

I'm seeing a lot more generated LLVM (and much worse performance) with Base.setindex(nt::NamedTuple, value, key::Symbol) on 1.12 and nightly.

This may depend somewhat on the system, since @MasonProtter tried this and wasn't able to reproduce the poor codegen on his machine. For reference:

Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 20 × 13th Gen Intel(R) Core(TM) i5-13600K
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, alderlake)
Threads: 18 default, 0 interactive, 9 GC (on 20 virtual cores)
Environment:
  JULIA_NUM_THREADS = 18
  JULIA_PROJECT = @.
  JULIA_EDITOR = emacs -nw
  JULIA_PKG_PRESERVE_TIERED_INSTALLED = true

Now that that's out of the way, this is what I'm seeing on 1.10 and 1.11:

julia> @code_llvm ((a, b) -> Base.setindex(a, b, :next))((next = zero(UInt32), prev = zero(UInt32)), 2)
; Function Signature: var"#1"(NamedTuple{(:next, :prev), Tuple{UInt32, UInt32}}, Int64)
;  @ REPL[1]:1 within `#1`
define nonnull ptr @"julia_#1_773"(ptr nocapture noundef nonnull readonly align 4 dereferenceable(8) %"a::NamedTuple", i64 signext %"b::Int64") #0 {
top:
  %0 = call nonnull ptr @j_setindex_776(ptr nocapture nonnull readonly %"a::NamedTuple", i64 signext %"b::Int64", ptr nonnull @"jl_sym#next#777.jit")
  ret ptr %0
}

and on 1.12/nightly

Lots of LLVM

julia> @code_llvm ((a, b) -> Base.setindex(a, b, :next))((next = zero(UInt32), prev = zero(UInt32)), 2)
; Function Signature: var"#2"(NamedTuple{(:next, :prev), Tuple{UInt32, UInt32}}, Int64)
;  @ REPL[1]:1 within `#2`
define nonnull ptr @"julia_#2_1195"(ptr nocapture noundef nonnull readonly align 4 dereferenceable(8) %"a::NamedTuple", i64 signext %"b::Int64") #0 {
top:
  %jlcallframe1 = alloca [5 x ptr], align 8
  %gcframe2 = alloca [5 x ptr], align 16
  call void @llvm.memset.p0.i64(ptr align 16 %gcframe2, i8 0, i64 40, i1 true)
  %thread_ptr = call ptr asm "movq %fs:0, $0", "=r"() #12
  %tls_ppgcstack = getelementptr inbounds i8, ptr %thread_ptr, i64 -8
  %tls_pgcstack = load ptr, ptr %tls_ppgcstack, align 8
  store i64 12, ptr %gcframe2, align 8
  %frame.prev = getelementptr inbounds ptr, ptr %gcframe2, i64 1
  %task.gcstack = load ptr, ptr %tls_pgcstack, align 8
  store ptr %task.gcstack, ptr %frame.prev, align 8
  store ptr %gcframe2, ptr %tls_pgcstack, align 8
; ┌ @ namedtuple.jl:484 within `setindex`
; │┌ @ boot.jl:792 within `NamedTuple`
    %0 = call ptr @jl_get_builtin_fptr(ptr nonnull @"+Core.#_compute_sparams#1197.jit")
    %ptls_field = getelementptr inbounds i8, ptr %tls_pgcstack, i64 16
    %ptls_load = load ptr, ptr %ptls_field, align 8
    %"box::NamedTuple" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_small_alloc(ptr %ptls_load, i32 360, i32 16, i64 139744856668560) #8
    %"box::NamedTuple.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple", i64 -1
    store atomic i64 139744856668560, ptr %"box::NamedTuple.tag_addr" unordered, align 8
    %1 = load i64, ptr %"a::NamedTuple", align 4
    store i64 %1, ptr %"box::NamedTuple", align 8
    %gc_slot_addr_1 = getelementptr inbounds ptr, ptr %gcframe2, i64 3
    store ptr %"box::NamedTuple", ptr %gc_slot_addr_1, align 8
    %ptls_load35 = load ptr, ptr %ptls_field, align 8
    %"box::NamedTuple3" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_small_alloc(ptr %ptls_load35, i32 360, i32 16, i64 139744871045968) #8
    %"box::NamedTuple3.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple3", i64 -1
    store atomic i64 139744871045968, ptr %"box::NamedTuple3.tag_addr" unordered, align 8
    store i64 %"b::Int64", ptr %"box::NamedTuple3", align 8
    %gc_slot_addr_0 = getelementptr inbounds ptr, ptr %gcframe2, i64 2
    store ptr %"box::NamedTuple3", ptr %gc_slot_addr_0, align 8
    store ptr @"-Main.Base.merge#1199.jit", ptr %jlcallframe1, align 8
    %2 = getelementptr inbounds ptr, ptr %jlcallframe1, i64 1
    store ptr @"jl_global#1200.jit", ptr %2, align 8
    %3 = getelementptr inbounds ptr, ptr %jlcallframe1, i64 2
    store ptr %"box::NamedTuple", ptr %3, align 8
    %4 = getelementptr inbounds ptr, ptr %jlcallframe1, i64 3
    store ptr %"box::NamedTuple3", ptr %4, align 8
    %Builtin_ret = call nonnull ptr %0(ptr nonnull @"jl_global#1198.jit", ptr nonnull %jlcallframe1, i32 4)
    store ptr null, ptr %gc_slot_addr_1, align 8
    store ptr %Builtin_ret, ptr %gc_slot_addr_0, align 8
; │└
; │┌ @ namedtuple.jl:338 within `merge`
    store ptr %Builtin_ret, ptr %jlcallframe1, align 8
    store ptr @"jl_global#1203.jit", ptr %2, align 8
    %jl_f__svec_ref_ret = call nonnull ptr @jl_f__svec_ref(ptr null, ptr nonnull %jlcallframe1, i32 2)
    %jl_f__svec_ref_ret.tag_addr = getelementptr inbounds i64, ptr %jl_f__svec_ref_ret, i64 -1
    %jl_f__svec_ref_ret.tag = load atomic i64, ptr %jl_f__svec_ref_ret.tag_addr unordered, align 8
    %5 = and i64 %jl_f__svec_ref_ret.tag, -16
    %6 = inttoptr i64 %5 to ptr
    %7 = icmp ult ptr %6, inttoptr (i64 1024 to ptr)
    br i1 %7, label %guard_pass, label %guard_exit

L6:                                               ; preds = %guard_exit
    store ptr %"box::NamedTuple16", ptr %jlcallframe1, align 8
    store ptr %"box::NamedTuple20", ptr %2, align 8
    store ptr @"jl_global#1207.jit", ptr %3, align 8
    store ptr %jl_f__svec_ref_ret, ptr %4, align 8
    %8 = call nonnull ptr @j1_merge_fallback_1205(ptr nonnull @"jl_global#1206.jit", ptr nonnull %jlcallframe1, i32 4)
    %frame.prev49 = load ptr, ptr %frame.prev, align 8
    store ptr %frame.prev49, ptr %tls_pgcstack, align 8
    ret ptr %8

L9:                                               ; preds = %guard_exit
    store ptr @"jl_global#1206.jit", ptr %jlcallframe1, align 8
    store ptr %"box::NamedTuple16", ptr %2, align 8
    store ptr %"box::NamedTuple20", ptr %3, align 8
    store ptr @"jl_global#1207.jit", ptr %4, align 8
    %9 = getelementptr inbounds ptr, ptr %jlcallframe1, i64 4
    store ptr %jl_f__svec_ref_ret, ptr %9, align 8
    %jl_f_throw_methoderror_ret = call nonnull ptr @jl_f_throw_methoderror(ptr null, ptr nonnull %jlcallframe1, i32 5)
    call void @llvm.trap()
    unreachable

guard_pass:                                       ; preds = %top
    %10 = getelementptr inbounds i8, ptr @jl_small_typeof, i64 %5
    %11 = load ptr, ptr %10, align 8
    br label %guard_exit

guard_exit:                                       ; preds = %guard_pass, %top
    %typeof = phi ptr [ %6, %top ], [ %11, %guard_pass ]
    store ptr null, ptr %gc_slot_addr_0, align 8
    %gc_slot_addr_2 = getelementptr inbounds ptr, ptr %gcframe2, i64 4
    store ptr %jl_f__svec_ref_ret, ptr %gc_slot_addr_2, align 8
    %12 = call i32 @ijl_subtype(ptr nonnull %typeof, ptr nonnull @"+Core.Tuple#1204.jit")
    %.not = icmp eq i32 %12, 0
    %ptls_load45 = load ptr, ptr %ptls_field, align 8
    %"box::NamedTuple16" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_small_alloc(ptr %ptls_load45, i32 360, i32 16, i64 139744856668560) #8
    %"box::NamedTuple16.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple16", i64 -1
    store atomic i64 139744856668560, ptr %"box::NamedTuple16.tag_addr" unordered, align 8
    %13 = load i64, ptr %"a::NamedTuple", align 4
    store i64 %13, ptr %"box::NamedTuple16", align 8
    store ptr %"box::NamedTuple16", ptr %gc_slot_addr_1, align 8
    %ptls_load48 = load ptr, ptr %ptls_field, align 8
    %"box::NamedTuple20" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_small_alloc(ptr %ptls_load48, i32 360, i32 16, i64 139744871045968) #8
    %"box::NamedTuple20.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple20", i64 -1
    store atomic i64 139744871045968, ptr %"box::NamedTuple20.tag_addr" unordered, align 8
    store i64 %"b::Int64", ptr %"box::NamedTuple20", align 8
    store ptr %"box::NamedTuple20", ptr %gc_slot_addr_0, align 8
    br i1 %.not, label %L9, label %L6
; └└
}

[MasonProtter]

Okay so I double checked and the REPL I thought was v1.12 was v1.11, so I can actually reproduce this.

The relevant method implementation for Base.setindex here is

julia/base/namedtuple.jl

Lines 483 to 485 in d82ccc1

	function setindex(nt::NamedTuple, v, idx::Symbol)
	merge(nt, (; idx => v))
	end

which then hits

julia/base/namedtuple.jl

Lines 331 to 340 in d82ccc1

	function merge(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn}
	if @generated
	names = merge_names(an, bn)
	types = merge_types(names, a, b)
	vals = Any[ :(getfield($(sym_in(names[n], bn) ? :b : :a), $(QuoteNode(names[n])))) for n in 1:length(names) ]
	:( _new_NamedTuple(NamedTuple{$names,$types}, ($(vals...),)) )
	else
	merge_fallback(a, b, an, bn)
	end
	end

and on v1.12 it seems to be hitting the else branch, resulting in the bad performance. My understanding was that these else branches should only be hit when we're interpreting or if the type level expressions are so complicated that the compiler should give up. If this is happening in relatively pedestrian code like this now, that's going to be really bad for all the optionally generated functions out there...

[giordano]

Now that that's out of the way, this is what I'm seeing on 1.10 and 1.11:

That's simply not inlining setindex, unlike v1.12?

[KristofferC]

As @giordano, the "clean" LLVM on pre-1.12 is just a function call to something that got inlined on 1.12.

Regarding the

on v1.12 it seems to be hitting the else branch, resulting in the bad performance.

I think that can be confirmed with:

@eval Base begin function merge(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn} 
    if @generated 
        Core.println("merging generated")
        names = merge_names(an, bn) 
        types = merge_types(names, a, b) 
        vals = Any[ :(getfield($(sym_in(names[n], bn) ? :b : :a), $(QuoteNode(names[n])))) for n in 1:length(names) ] 
        :( _new_NamedTuple(NamedTuple{$names,$types}, ($(vals...),)) ) 
    else 
        Core.println("merging non-generated")
        merge_fallback(a, b, an, bn) 
     end 
end
end

nt = (next = zero(UInt32), prev = zero(UInt32))
Base.setindex(nt, 2, :next)

pasting that gives on 1.11:

julia> Base.setindex(nt, 2, :next)
merging generated
(next = 2, prev = 0x00000000)

on 1.12:

julia> Base.setindex(nt, 2, :next)
merging non-generated
(next = 2, prev = 0x00000000)

I'll try a bisect to see when that happens. (edit: Maybe won't get to that today).

[tecosaur]

That's simply not inlining setindex, unlike v1.12?

It's not just a matter of inlining, you can see in the 1.12 LLVM a whole lot of code around the "within setindex" block, which you see if you add @inline as well (probably would have been a good idea to do that in the initial post).

It seems like Kristoffer's been able to start digging into this (thanks!).

1.10/1.11 LLVM with `@inline`

julia> @code_llvm ((a, b) -> @inline Base.setindex(a, b, :next))((next = zero(UInt32), prev = zero(UInt32)), 2)
; Function Signature: var"#7"(NamedTuple{(:next, :prev), Tuple{UInt32, UInt32}}, Int64)
;  @ REPL[4]:1 within `#7`
define void @"julia_#7_1228"(ptr noalias nocapture noundef nonnull sret({ i64, i32 }) align 8 dereferenceable(16) %sret_return, ptr nocapture noundef nonnull readonly align 4 dereferenceable(8) %"a::NamedTuple", i64 signext %"b::Int64") #0 {
top:
  %jlcallframe1 = alloca [2 x ptr], align 8
  %gcframe2 = alloca [4 x ptr], align 16
  call void @llvm.memset.p0.i64(ptr align 16 %gcframe2, i8 0, i64 32, i1 true)
  %thread_ptr = call ptr asm "movq %fs:0, $0", "=r"() #10
  %tls_ppgcstack = getelementptr i8, ptr %thread_ptr, i64 -8
  %tls_pgcstack = load ptr, ptr %tls_ppgcstack, align 8
  store i64 8, ptr %gcframe2, align 16
  %frame.prev = getelementptr inbounds ptr, ptr %gcframe2, i64 1
  %task.gcstack = load ptr, ptr %tls_pgcstack, align 8
  store ptr %task.gcstack, ptr %frame.prev, align 8
  store ptr %gcframe2, ptr %tls_pgcstack, align 8
; ┌ @ namedtuple.jl:469 within `setindex`
   %ptls_field = getelementptr inbounds ptr, ptr %tls_pgcstack, i64 2
   %ptls_load = load ptr, ptr %ptls_field, align 8
   %"box::NamedTuple" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_pool_alloc_instrumented(ptr %ptls_load, i32 800, i32 16, i64 140030798228432) #6
   %"box::NamedTuple.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple", i64 -1
   store atomic i64 140030798228432, ptr %"box::NamedTuple.tag_addr" unordered, align 8
   %0 = load i64, ptr %"a::NamedTuple", align 4
   store i64 %0, ptr %"box::NamedTuple", align 8
   %gc_slot_addr_1 = getelementptr inbounds ptr, ptr %gcframe2, i64 3
   store ptr %"box::NamedTuple", ptr %gc_slot_addr_1, align 8
   %ptls_load9 = load ptr, ptr %ptls_field, align 8
   %"box::NamedTuple4" = call noalias nonnull align 8 dereferenceable(16) ptr @ijl_gc_pool_alloc_instrumented(ptr %ptls_load9, i32 800, i32 16, i64 140030794299472) #6
   %"box::NamedTuple4.tag_addr" = getelementptr inbounds i64, ptr %"box::NamedTuple4", i64 -1
   store atomic i64 140030794299472, ptr %"box::NamedTuple4.tag_addr" unordered, align 8
   store i64 %"b::Int64", ptr %"box::NamedTuple4", align 8
   %gc_slot_addr_0 = getelementptr inbounds ptr, ptr %gcframe2, i64 2
   store ptr %"box::NamedTuple4", ptr %gc_slot_addr_0, align 16
   store ptr %"box::NamedTuple", ptr %jlcallframe1, align 8
   %1 = getelementptr inbounds ptr, ptr %jlcallframe1, i64 1
   store ptr %"box::NamedTuple4", ptr %1, align 8
   %2 = call nonnull ptr @ijl_apply_generic(ptr nonnull @"jl_global#1232.jit", ptr nonnull %jlcallframe1, i32 2)
   call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 8 dereferenceable(16) %sret_return, ptr noundef nonnull align 8 dereferenceable(16) %2, i64 16, i1 false)
   %frame.prev11 = load ptr, ptr %frame.prev, align 8
   store ptr %frame.prev11, ptr %tls_pgcstack, align 8
   ret void
; └
}

[MasonProtter]

FWIW @tecosaur Accessors handles this well on all julia versions without using internals:

julia> using Accessors, BenchmarkTools

julia> let nt = Ref((next = zero(UInt32), prev = zero(UInt32))), val = Ref(2)
           @btime @set(nt.next = val) setup=(nt=$nt[]; val=$val[])
       end
  1.953 ns (0 allocations: 0 bytes)
(next = 2, prev = 0x00000000)

[KristofferC]

It's not just a matter of inlining

No, but the exmaple given of the LLVM code on 1.11 just says that it calls setindex which is just due to no inlining and it is not informative.

[JeffBezanson]

Amazing. In 1.9 we generated the same good code as Accessors. In 1.10 and 1.11 the code is still pretty good except merge is not inlined. One change in 1.10 was the introduction of _new_NamedTuple (which should not have any effect on this). In 1.12 it falls apart and the code is almost totally unspecialized.

[KristofferC]

The change to what @generated branch that runs was bisected to #54972.

Running

@eval Base begin function merge(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn} 
    if @generated 
        names = merge_names(an, bn) 
        types = merge_types(names, a, b) 
        vals = Any[ :(getfield($(sym_in(names[n], bn) ? :b : :a), $(QuoteNode(names[n])))) for n in 1:length(names) ] 
        :( _new_NamedTuple(NamedTuple{$names,$types}, ($(vals...),)) ) 
    else 
        Core.println("GIT-BISECT-RESULT: BAD")
        merge_fallback(a, b, an, bn) 
     end 
end
end

nt = (next = zero(UInt32), prev = zero(UInt32))
Base.setindex(nt, 2, :next)

on 61c044c gives

GIT-BISECT-RESULT: BAD

While on the previous commit, it doesn't print anything.

I haven't checked if this is the cause of the full regression.

cc @topolarity @gbaraldi

[KristofferC]

The following also regressed from 61c044c:

julia> nt = (next = zero(UInt32), prev = zero(UInt32))
(next = 0x00000000, prev = 0x00000000)

julia> f(nt) = @inline Base.setindex(nt, 2, :next)
f (generic function with 1 method)

julia> @code_warntype optimize=true f(nt)

On f808606:

julia> @code_warntype optimize=true f(nt)
MethodInstance for f(::@NamedTuple{next::UInt32, prev::UInt32})
  from f(nt) @ Main REPL[2]:1
Arguments
  #self#::Core.Const(Main.f)
  nt::@NamedTuple{next::UInt32, prev::UInt32}
Body::@NamedTuple{next::Int64, prev::UInt32}
1 ─ %1 = Base.merge(nt, $(QuoteNode((next = 2,))))::@NamedTuple{next::Int64, prev::UInt32}
└──      return %1

while on 61c044c:

julia> @code_warntype optimize=true f(nt)
MethodInstance for f(::@NamedTuple{next::UInt32, prev::UInt32})
  from f(nt) @ Main REPL[5]:1
Arguments
  #self#::Core.Const(Main.f)
  nt::@NamedTuple{next::UInt32, prev::UInt32}
Body::NamedTuple
1 ─ %1 = Base.merge::typeof(merge)
│   %2 = (Core._compute_sparams)(merge(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn} @ Base namedtuple.jl:334, %1, nt, $(QuoteNode((next = 2,))))::Core.SimpleVector
│   %3 = (Core._svec_ref)(%2, 2)::Any
│   %4 = Base.merge_fallback::typeof(Base.merge_fallback)
│   %5 = (isa)(%3, Tuple{Vararg{Symbol}})::Bool
└──      goto #3 if not %5
2 ─ %7 = π (%3, Tuple{Vararg{Symbol}})
│   %8 = invoke %4(nt::NamedTuple, $(QuoteNode((next = 2,)))::NamedTuple, (:next, :prev)::Tuple{Vararg{Symbol}}, %7::Tuple{Vararg{Symbol}})::NamedTuple
└──      goto #4
3 ─      Core.throw_methoderror(%4, nt, $(QuoteNode((next = 2,))), (:next, :prev), %3)
└──      unreachable
4 ─      goto #5
5 ─      goto #6
6 ─      return %8