ZJIT: Back ivar shape-chain misses with a per-name shape table

ZJIT specializes an instance variable access by profiling the receiver's
shapes and emitting a chain of shape_id guards, one arm per profiled
shape. That covers a site that sees a handful of shapes. It does not
cover code like Shopify's Storefront Renderer, where one class has
hundreds of live shapes because instances stop at different points in a
long chain of conditionally-assigned ivars: the chain cannot be widened
far enough (each arm is code, and MAX_IVAR_RESPECIALIZATIONS bounds the
recompiles on purpose), so those receivers fall off the end of the chain
into a generic rb_ivar_get / rb_vm_getinstancevariable call. On SFR that
is 10.8M getivar and 2.4M setivar generic calls per ~350 requests, each
one walking the shape tree for an answer that only depends on the
receiver's shape id.

Put a lookup table between the guard chain and that call. Each ivar
*name* accessed by such a site gets one direct-mapped table of 8-byte
entries keyed by the raw shape id, holding the byte offset of the ivar
within the receiver. Reads probe it inline -- the shape id is already in
a register for the guard chain, so it is a multiply, a shift, a mask, two
32-bit loads, two compares and a load -- and produce the value with no
call at all, including when the answer is nil because the shape does not
have the ivar. Writes have no inline probe (a store also needs a frozen
check and a write barrier) but resolve out of the same table.

Keying by name, not by site: shape_id -> location of @name does not
depend on the site, so sites sharing an ivar warm each other's entries,
and the table size multiplies by the number of ivar names accessed
polymorphically rather than by the number of compiled sites (and their
recompiles).

Nothing needs invalidating. Shapes are immutable and the shape tree is
append-only, so shape_id -> index is a pure function; and every mutation
that could move an object's ivars -- adding one, remove_instance_variable,
freeze, object_id, going too-complex, compaction changing embedded
capacity or layout -- changes a bit of the object's shape_id, which is
stored unmasked as the key. The table holds no VALUEs, so the GC has
nothing to mark. Entries are single naturally-aligned words published
with one relaxed atomic store, so a concurrent ractor sees an old entry
or a new one, never a mix.

Measured with benchmark/zjit_ivar_megashape.rb (110-deep ivar chain,
220 shapes per name), instructions retired per access relative to
--zjit-disable-ivar-cache, for a cyclic access pattern (the worst case
for a direct-mapped table) and a Zipf-skewed one (what applications look
like):

                cyclic  skewed
  attr_reader     0.37    0.30
  plain read      0.47    0.36
  absent read     0.33    0.41
  attr_writer     0.43    0.48
  plain write     0.35    0.40
  monomorphic     1.00    1.00

Wall clock on the same benchmark moves 220 -> 26 ns/op for an absent
read and 55 -> 21 ns/op for an attr_reader. The monomorphic and low-
polymorphism paths are untouched by construction: only the codegen of
Insn::GetIvar and Insn::SetIvar changed, which a shape-specialized site
never emits. HIR and fast-path disassembly for a monomorphic site are
identical to before.

Memory is 8 bytes times --zjit-ivar-cache-entries (default 512, i.e.
4KiB) per ivar name, reported as mem_ivar_cache_bytes and
mem_ivar_cache_count in the mem_* breakdown. Undersizing the table is
not a graceful degradation -- a direct-mapped table smaller than the
shape working set misses on nearly every access -- so the default is
chosen from the measured curve; see DEFAULT_CACHE_ENTRIES.

New counters: getivar_cache_hit (served inline, no call),
getivar_cache_helper_hit, _fill, _evict, _uncacheable, _immediate, the
setivar_cache_* equivalents, and ivar_cache_alloc_count.
--zjit-disable-ivar-cache restores the previous code exactly, for A/B.

[zjit/min port note] Ported without mem_stats.rs (the memdiet accounting), keeping
master's plain-HashMap profile layout. The two emit_ivar_reprofile() call sites are
re-added by hand: upstream gates them on ShapeMiss::CallFallback, an enum that comes
from excluded work, so here they sit on the two paths that call the generic fallback
instead (no-profile under a no-exit policy, and a shape-chain miss), which is the
same set of sites for this tree.

[reorder port note] Ported ahead of the megamorphic send class table and
the ivar-reprofile machinery: send_cache options/counters/state and the
IvarReprofile/gen_ivar_reprofile pieces were dropped from conflict hunks;
they return with their own commits.

Author: k0kubun

benchmark/zjit_ivar_megashape.rb

@@ -0,0 +1,123 @@
+# Shape-polymorphic instance variable access.
+#
+# Models the Shopify Storefront Renderer pattern that motivates ZJIT's ivar
+# shape table (see zjit/src/ivar_cache.rs): a single class whose instances stop
+# at many different points in one long ivar chain, so every ivar site sees
+# hundreds of distinct shape ids and no bounded inline guard chain can cover the
+# traffic.
+#
+# Two access distributions are measured, because they stress the table in
+# opposite ways:
+#
+#   cyclic  every shape is touched exactly once per pass, in a stride order. The
+#           worst case for a direct-mapped table: any slot holding two live
+#           shapes misses on every access.
+#   skewed  shape rank r is picked with weight 1/(r+1), which is what a real
+#           application's shape distribution at one site looks like.
+#
+# The population is DEPTHS instances at distinct chain depths, doubled by
+# freezing half of them (the frozen bit is part of shape_id, so it doubles the
+# shape count without adding shape *variations*, which past SHAPE_MAX_VARIATIONS
+# would demote the class to hash-backed complex shapes).
+#
+#   ruby --zjit benchmark/zjit_ivar_megashape.rb
+#   ruby --zjit --zjit-stats benchmark/zjit_ivar_megashape.rb   # table counters
+#   ruby --zjit --zjit-disable-ivar-cache benchmark/zjit_ivar_megashape.rb
+#
+# Wall and CPU time both swing by 2x on a loaded machine; for A/B work prefer
+# `perf stat -e instructions` over two iteration counts and compare the slope.
+#
+# Env knobs: ITERS, REPS, DEPTHS.
+
+DEPTHS = Integer(ENV.fetch("DEPTHS", 110))
+ITERS  = Integer(ENV.fetch("ITERS", 2000))
+REPS   = Integer(ENV.fetch("REPS", 5))
+
+class Mega
+  body = +"def initialize(depth)\n"
+  DEPTHS.times do |i|
+    body << "  @i#{i} = #{i}\n"
+    body << "  return if depth == #{i + 1}\n"
+  end
+  body << "end\n"
+  eval(body) # rubocop:disable Security/Eval
+
+  attr_reader :i0
+  attr_writer :i0
+
+  def read_first = @i0
+  def read_absent = @never_assigned
+  def write_first(v) = (@i0 = v)
+end
+
+class Mono
+  def initialize
+    @a = 1
+    @b = 2
+  end
+  attr_reader :a
+  attr_writer :a
+  def read_a = @a
+end
+
+Mega.new(DEPTHS) # prime RCLASS_MAX_IV_COUNT so every instance is embedded
+
+READ_POP = ((1..DEPTHS).map { |d| Mega.new(d) } +
+            (1..DEPTHS).map { |d| Mega.new(d).freeze })
+WRITE_POP = (1..DEPTHS).map { |d| Mega.new(d) }
+
+def cyclic(pop, n) = Array.new(n) { |i| pop[(i * 97) % pop.size] }
+
+def skewed(pop, n)
+  cdf = []
+  acc = 0.0
+  total = (0...pop.size).sum { |r| 1.0 / (r + 1) }
+  pop.each_index { |r| acc += 1.0 / (r + 1); cdf << acc / total }
+  rng = Random.new(1234)
+  Array.new(n) { x = rng.rand; pop[cdf.index { |c| c >= x } || pop.size - 1] }
+end
+
+POPS = {
+  "cyclic" => [cyclic(READ_POP, 2 * DEPTHS).freeze, cyclic(WRITE_POP, 2 * DEPTHS).freeze],
+  "skewed" => [skewed(READ_POP, 2 * DEPTHS).freeze, skewed(WRITE_POP, 2 * DEPTHS).freeze],
+}
+MONO = Array.new(2 * DEPTHS) { Mono.new }.freeze
+
+def bench_attr_reader(pop) = pop.each { |o| o.i0 }
+def bench_plain_read(pop) = pop.each { |o| o.read_first }
+def bench_absent_read(pop) = pop.each { |o| o.read_absent }
+def bench_attr_writer(pop) = pop.each { |o| o.i0 = 7 }
+def bench_plain_write(pop) = pop.each { |o| o.write_first(9) }
+def bench_mono_reader(pop) = pop.each { |o| o.a }
+def bench_mono_plain(pop) = pop.each { |o| o.read_a }
+def bench_mono_writer(pop) = pop.each { |o| o.a = 3 }
+
+benches = []
+POPS.each do |kind, (read_pop, write_pop)|
+  benches << ["attr_reader/#{kind}", method(:bench_attr_reader), read_pop]
+  benches << ["plain read/#{kind}", method(:bench_plain_read), read_pop]
+  benches << ["absent read/#{kind}", method(:bench_absent_read), read_pop]
+  benches << ["attr_writer/#{kind}", method(:bench_attr_writer), write_pop]
+  benches << ["plain write/#{kind}", method(:bench_plain_write), write_pop]
+end
+benches << ["attr_reader/mono", method(:bench_mono_reader), MONO]
+benches << ["plain read/mono", method(:bench_mono_plain), MONO]
+benches << ["attr_writer/mono", method(:bench_mono_writer), MONO]
+
+# Warm every site up so ZJIT compiles it (and finishes respecializing) first.
+benches.each { |_, m, pop| 60.times { m.call(pop) } }
+
+width = benches.map { |name, _, _| name.length }.max
+totals = Hash.new(0.0)
+benches.each do |name, m, pop|
+  best = Float::INFINITY
+  REPS.times do
+    t0 = Process.clock_gettime(Process::CLOCK_PROCESS_CPUTIME_ID)
+    ITERS.times { m.call(pop) }
+    t1 = Process.clock_gettime(Process::CLOCK_PROCESS_CPUTIME_ID)
+    best = [best, t1 - t0].min
+  end
+  totals[name.split("/").last] += best
+  puts format("%-#{width}s  %8.4f s  %7.2f ns/op", name, best, best / (ITERS * pop.size) * 1e9)
+end
+totals.each { |kind, secs| puts format("%-#{width}s  %8.4f s", "TOTAL #{kind}", secs) }

zjit/src/codegen.rs

@@ -1223,7 +1223,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
             let CCallVariadicData { cfunc, recv, name, args, cme, state, block, .. } = &**insn;
             gen_ccall_variadic(jit, asm, function, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, &function.frame_state(*state))
         }
-        Insn::GetIvar { self_val, id, ic, state } => gen_getivar(asm, opnd!(self_val), *id, *ic, &function.frame_state(*state)),
+        &Insn::GetIvar { self_val, id, ic, state } => gen_getivar(jit, asm, opnd!(self_val), function.type_of(self_val), id, ic, &function.frame_state(state)),
         Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, function, *id, opnd!(val), &function.frame_state(*state))),
         Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, function, *id, &function.frame_state(*state)),
         &Insn::IsBlockParamModified { flags } => gen_is_block_param_modified(asm, opnd!(flags)),
@@ -1668,8 +1668,52 @@ fn gen_ccall_variadic(
     result
 }
 
-/// Emit an uncached instance variable lookup
-fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry, state: &FrameState) -> Opnd {
+/// Emit an instance variable lookup with no compile-time shape information.
+///
+/// The shape guard chain [`crate::hir::Function::dispatch_ivar`] emits covers
+/// the shapes the site's profile named; this is what runs for every other
+/// receiver. A generic call is fine when that is rare, but on shape-polymorphic
+/// code (hundreds of live shapes for one class) it *is* the hot path, so we put a
+/// per-ivar-name shape table in front of it: see [`gen_ivar_cache_probe`] and
+/// [`crate::ivar_cache`].
+fn gen_getivar(jit: &mut JITState, asm: &mut Assembler, recv: Opnd, recv_type: Type, id: ID, ic: *const iseq_inline_iv_cache_entry, state: &FrameState) -> Opnd {
+    if get_option!(disable_ivar_cache) {
+        return gen_getivar_generic(asm, recv, id, ic, state);
+    }
+
+    let cache = crate::ivar_cache::ivar_cache_for(id);
+    let hir_block_id = asm.current_block().hir_block_id;
+    let rpo_idx = asm.current_block().rpo_index;
+    let slow_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let result_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let slow_edge = Target::Block(Box::new(lir::BranchEdge { target: slow_block, args: vec![] }));
+    let result_edge = |val: Opnd| Target::Block(Box::new(lir::BranchEdge { target: result_block, args: vec![val] }));
+
+    let val = gen_ivar_cache_probe(jit, asm, recv, recv_type, id, cache, slow_edge);
+    gen_incr_counter(asm, Counter::getivar_cache_hit);
+    asm.jmp(result_edge(val));
+
+    asm.set_current_block(slow_block);
+    let label = jit.get_label(asm, slow_block, hir_block_id);
+    asm.write_label(label);
+    // rb_zjit_getivar_cached only ever raises what rb_ivar_get raises, which is
+    // what this site called before, so it needs no more frame setup than the
+    // generic call it replaces.
+    gen_trace_fallback(asm, "getivar");
+    use crate::ivar_cache::rb_zjit_getivar_cached;
+    let val = asm_ccall!(asm, rb_zjit_getivar_cached, recv, Opnd::const_ptr(cache as *const u8));
+    asm.jmp(result_edge(val));
+
+    asm.set_current_block(result_block);
+    let label = jit.get_label(asm, result_block, hir_block_id);
+    asm.write_label(label);
+    let param = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(param);
+    param
+}
+
+/// The generic instance variable lookup, with no shape table in front of it.
+fn gen_getivar_generic(asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry, state: &FrameState) -> Opnd {
     gen_trace_fallback(asm, "getivar");
     if ic.is_null() {
         asm_ccall!(asm, rb_ivar_get, recv, id.0.into())
@@ -1679,17 +1723,105 @@ fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_i
     }
 }
 
-/// Emit an uncached instance variable store
+/// Emit the inline half of an ivar shape table lookup: hash the receiver's shape
+/// id into the table, load one entry, and on a hit produce the ivar's value with
+/// no call. Jumps to `miss` for anything else -- an immediate receiver, a shape
+/// the table does not hold, or an entry whose kind is neither
+/// [`crate::ivar_cache::EntryKind::Direct`] nor [`crate::ivar_cache::EntryKind::Nil`].
+///
+/// The encoding is documented on [`crate::ivar_cache::Entry`]; every constant
+/// this reads comes from that module so the two halves cannot drift apart.
+fn gen_ivar_cache_probe(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    recv: Opnd,
+    recv_type: Type,
+    id: ID,
+    cache: *const crate::ivar_cache::IvarCache,
+    miss: Target,
+) -> Opnd {
+    use crate::ivar_cache::{IVAR_CACHE_HASH_MULT, IVAR_CACHE_INFO_OFFSET, IVAR_CACHE_KEY_OFFSET, IVAR_CACHE_NIL_BIT, IVAR_CACHE_NIL_SLOT, IVAR_CACHE_NOT_INLINE_MASK, cache_byte_mask, cache_hash_shift};
+
+    asm_comment!(asm, "ivar shape table probe for :{}", id.contents_lossy());
+    let recv = asm.load_mem(recv);
+    if !recv_type.is_subtype(types::HeapBasicObject) {
+        // Immediates and false have no shape id word to load.
+        asm.cmp(recv, Opnd::Value(Qfalse));
+        asm.je(jit, miss.clone());
+        asm.test(recv, Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
+        asm.jnz(jit, miss.clone());
+    }
+
+    let shape_offset = unsafe { rb_shape_id_offset() };
+    // Two loads of the same word, because `Assembler::mul` is destructive: the
+    // x86 lowering emits `imul left, right` and copies the result out afterwards,
+    // so whatever register held the multiplicand is clobbered. (Same for lshift,
+    // rshift, urshift and not; add/and/or/xor copy into the output first and are
+    // safe.) The second load is what the key comparison needs, and it is a
+    // guaranteed L1 hit on a line the first load just brought in.
+    let shape_to_hash = asm.load(Opnd::mem(SHAPE_ID_NUM_BITS as u8, recv, shape_offset));
+    let shape = asm.load(Opnd::mem(SHAPE_ID_NUM_BITS as u8, recv, shape_offset));
+
+    // Fibonacci hash, taking the byte offset out of the product's high bits;
+    // this has to agree with `crate::ivar_cache::slot_of` or the helper will fill
+    // slots the probe does not read. Widening the loaded shape id to 64 bits
+    // relies on a 32-bit load zero-extending its destination register, which both
+    // backends do; if one ever did not, the only consequence would be probing a
+    // slot the helper does not fill -- a miss, not a wrong answer.
+    let hash = asm.mul(shape_to_hash.with_num_bits(64), Opnd::UImm(IVAR_CACHE_HASH_MULT));
+    let shifted = asm.urshift(hash, Opnd::UImm(cache_hash_shift()));
+    let slot = asm.and(shifted, Opnd::UImm(cache_byte_mask()));
+    let table = unsafe { (*cache).table_ptr() };
+    let slot_ptr = asm.add(slot, Opnd::const_ptr(table));
+
+    asm_comment!(asm, "check the entry's shape id");
+    let key = asm.load(Opnd::mem(32, slot_ptr, IVAR_CACHE_KEY_OFFSET));
+    asm.cmp(key, shape);
+    asm.jne(jit, miss.clone());
+
+    // The entry's high half is the byte offset plus the kind. Direct and Nil are
+    // servable here; everything else goes to the helper.
+    let info = asm.load(Opnd::mem(32, slot_ptr, IVAR_CACHE_INFO_OFFSET));
+    asm.test(info, Opnd::UImm(IVAR_CACHE_NOT_INLINE_MASK));
+    asm.jnz(jit, miss);
+
+    asm_comment!(asm, "load the ivar at the cached offset, or nil if absent");
+    let offset = asm.and(info.with_num_bits(64), Opnd::UImm(u16::MAX as u64));
+    let ivar_ptr = asm.add(offset, recv);
+    // A Nil entry means the shape has no such ivar. Rather than branch, read a
+    // static word that holds Qnil: this keeps the absent case -- the most
+    // expensive one to resolve generically, because the search has to fail --
+    // on the inline path for two extra instructions.
+    asm.test(info, Opnd::UImm(IVAR_CACHE_NIL_BIT));
+    let nil_slot = Opnd::const_ptr(std::ptr::addr_of!(IVAR_CACHE_NIL_SLOT) as *const u8);
+    let load_from = asm.csel_nz(nil_slot, ivar_ptr);
+    asm.load(Opnd::mem(VALUE_BITS, load_from, 0))
+}
+
+/// Emit an instance variable store with no compile-time shape information.
+///
+/// The counterpart of [`gen_getivar`]: the shape guard chain covers what the
+/// profile named, and this runs for everything else. There is no inline probe on
+/// the write side -- a store also needs a frozen check and a write barrier, which
+/// is a lot of code to inline for a quarter of the traffic reads see -- but the
+/// call goes to a helper that resolves the location out of the same per-name
+/// shape table instead of walking the shape tree. See [`crate::ivar_cache`].
 fn gen_setivar(jit: &mut JITState, asm: &mut Assembler, function: &Function, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry, val: Opnd, state: &FrameState) {
     gen_trace_fallback(asm, "setivar");
     // Setting an ivar can raise FrozenError, so we need proper frame state for exception handling.
     gen_prepare_non_leaf_call(jit, asm, function, state);
-    if ic.is_null() {
-        asm_ccall!(asm, rb_ivar_set, recv, id.0.into(), val);
-    } else {
-        let iseq = Opnd::Value(state.iseq.into());
-        asm_ccall!(asm, rb_vm_setinstancevariable, iseq, recv, id.0.into(), val, Opnd::const_ptr(ic));
+    if get_option!(disable_ivar_cache) {
+        if ic.is_null() {
+            asm_ccall!(asm, rb_ivar_set, recv, id.0.into(), val);
+        } else {
+            let iseq = Opnd::Value(state.iseq.into());
+            asm_ccall!(asm, rb_vm_setinstancevariable, iseq, recv, id.0.into(), val, Opnd::const_ptr(ic));
+        }
+        return;
     }
+    let cache = crate::ivar_cache::ivar_cache_for(id);
+    use crate::ivar_cache::rb_zjit_setivar_cached;
+    asm_ccall!(asm, rb_zjit_setivar_cached, recv, val, Opnd::const_ptr(cache as *const u8));
 }
 
 fn gen_getclassvar(jit: &mut JITState, asm: &mut Assembler, function: &Function, id: ID, ic: *const iseq_inline_cvar_cache_entry, state: &FrameState) -> Opnd {