ZJIT: Call megamorphic ISEQ targets without leaving JIT code

The send class table made the *lookup* at a megamorphic site cheap, but
the call itself still left JIT code: `rb_zjit_send_cached_without_block`,
`vm_call_iseq_setup`'s generic argument setup, an interpreter frame push,
the `setjmp` in `vm_exec`, and a trip through `jit_exec` and the entry
trampoline -- all to arrive at code ZJIT had already compiled. On lobsters
that is 4.4M calls a run, and instrumenting what they resolve to shows 45%
of them are the same small shape: an ISEQ method with fixed arity, no
locals beyond its parameters, and a compiled entry point sitting ready.

    ISEQ  2,488,687   of which simple, no extra locals, compiled:  2,065,410
    CFUNC   743,217
    IVAR    639,155
    MISSING 478,689
    OPTIMIZED 269,758

So probe the table inline and, for that shape, push the frame and call the
callee from JIT code, the way SendDirect does for a statically known
target. The entry point it calls is the *interpreter* one,
`ISEQ_BODY(iseq)->jit_entry`: a JIT-to-JIT entry takes its arguments in C
registers and its address only exists inside the compiler's patching
machinery, while the interpreter entry reads its parameters back out of
the frame -- which is exactly where a simple callee's arguments already
are, because its local table is its parameter list and the caller's
operand stack holds them in order. The frame push is then the whole
calling convention, and the entry address is a field on the ISEQ.

Correctness rests on re-deriving everything per call rather than caching
it. The callcache validates itself against `CLASS_OF(recv)` and
`METHOD_ENTRY_INVALIDATED`, the two checks `vm_cc_hit_p` makes, which is
what covers redefinition, undef, aliasing, visibility changes,
include/prepend, refinements and object death. `body->jit_entry` is loaded
fresh, so an ISEQ whose code was reset -- by a TracePoint, a patch point,
a recompile -- reads as null and takes the slow path. And because the path
skips `vm_call_method`, the fill side only marks a target directly
callable when that function would have let the call straight through:
public, or private from an FCALL site, never protected.

The slot grows a second word for that decision, holding the method entry
rather than the ISEQ so JIT code can check it against the one the
callcache produced. That makes the pair self-validating: a reader that
catches another ractor's fill half-done sees a mismatch and searches,
so neither word needs a barrier or an atomic.

lobsters, 15 iterations:

    ccall rb_zjit_send_cached_without_block  4,410,274 -> 2,423,746
    send_megamorphic_direct                          0 -> 1,986,528
    send_megamorphic                         4,187,398 -> 4,187,397

`--zjit-disable-megamorphic-direct` turns the new path off for A/B runs.
The wider slot doubles the tables' memory, 6.7MB to 13.3MB on lobsters;
they are 98% empty at 512 slots, so `--zjit-send-cache-entries=256` gets
that back with no measurable change in direct dispatches.

[zjit/min port note] Uses asm.new_block() directly: the FrameCaches-inheriting
JITState::new_block() wrapper the original called comes from an excluded commit
(66ff4fd) and master has no per-LIR-block frame caches. state.rs picks up only
gen_jit_entry_call_trampoline from the original's import list; the rest of that list
belongs to excluded commits.

Author: k0kubun

vm_insnhelper.c

@@ -6312,6 +6312,70 @@ zjit_send_cache_cacheable_p(const struct rb_callcache *cc)
         && !vm_cc_refinement_p(cc);
 }
 
+/* The callee's method entry when a hit on `cc` is one that JIT code may enter
+ * with a direct call, or NULL.
+ *
+ * The inline dispatch path in gen_send_megamorphic_direct() reproduces exactly
+ * what vm_call_iseq_setup() does for the simplest possible callee, and nothing
+ * else, so the callee has to be one this list of conditions describes:
+ *
+ *  - an ordinary ISEQ method, not a bmethod, alias, refinement or anything else
+ *    with its own frame shape;
+ *  - rb_simple_iseq_p(), i.e. no optional, rest, post, keyword or block
+ *    parameters and not `...`-forwardable, so vm_callee_setup_arg() would take
+ *    its no-op fast path and the frame is no taller than the local table;
+ *  - arity exactly the site's argc, since the inline path cannot raise
+ *    ArgumentError;
+ *  - no locals beyond the parameters, so the frame's local area is exactly the
+ *    arguments already sitting on the caller's stack: nothing to nil-fill, and
+ *    the frame size is a compile-time constant at the call site;
+ *  - a `stack_max` within ZJIT_MEGA_DIRECT_MAX_STACK, the bound the call site's
+ *    stack overflow check is compiled against.
+ *
+ * The *method entry* rather than the ISEQ is what goes in the slot, and JIT code
+ * checks it against the one it read out of the callcache. That is what makes the
+ * two-word slot safe to fill with two plain stores: a reader that sees a fresh
+ * `cc` beside a previous class's `direct_cme` (or the reverse) fails the compare
+ * and searches, exactly as it would for an empty slot. No store ordering, and so
+ * no atomics or barriers, are needed on either side.
+ *
+ * Nothing about the compiled code is cached: the entry point is re-read from
+ * ISEQ_BODY(iseq)->jit_entry on every call, which rb_iseq_reset_jit_func()
+ * clears whenever that code stops being valid, and the ISEQ is only reached
+ * through a method entry that has already validated. */
+static const rb_callable_method_entry_t *
+zjit_send_cache_direct_cme(const struct rb_zjit_send_cache *cache, const struct rb_callcache *cc)
+{
+    if (!cache->direct_ok) return NULL;
+
+    const rb_callable_method_entry_t *cme = vm_cc_cme(cc);
+    if (cme->def->type != VM_METHOD_TYPE_ISEQ) return NULL;
+
+    /* The direct path skips vm_call_method(), so it must only take over calls
+     * that function would have let straight through to vm_call_method_each_type().
+     * A protected method's check reads the *caller's* self, which the table
+     * cannot be keyed on, and a private method is only callable without an
+     * explicit receiver. Everything else raises NoMethodError through
+     * vm_call_method_missing(), which is very much not a frame push. */
+    switch (METHOD_ENTRY_VISI(cme)) {
+      case METHOD_VISI_PUBLIC:
+        break;
+      case METHOD_VISI_PRIVATE:
+        if (!(cache->direct_flags & VM_CALL_FCALL)) return NULL;
+        break;
+      default:
+        return NULL;
+    }
+
+    const rb_iseq_t *iseq = def_iseq_ptr(cme->def);
+    const struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
+    if (!rb_simple_iseq_p(iseq)) return NULL;
+    if (body->param.lead_num != (int)cache->direct_argc) return NULL;
+    if (body->local_table_size != (unsigned int)body->param.lead_num) return NULL;
+    if (body->stack_max > ZJIT_MEGA_DIRECT_MAX_STACK) return NULL;
+    return cme;
+}
+
 /* Resolve the callcache for `recv` at this site, out of `cache` when possible.
  *
  * The hit condition is exactly vm_cc_hit_p()'s: the cache's own class matches
@@ -6326,7 +6390,7 @@ zjit_send_cache_search(rb_control_frame_t *reg_cfp, struct rb_call_data *cd,
 {
     VALUE klass = CLASS_OF(recv);
     unsigned int slot = zjit_send_cache_slot(cache, klass);
-    const struct rb_callcache *cached = cache->slots[slot];
+    const struct rb_callcache *cached = cache->slots[slot].cc;
 
     if (LIKELY(cached != NULL && cached->klass == klass &&
                !METHOD_ENTRY_INVALIDATED(vm_cc_cme(cached)))) {
@@ -6340,11 +6404,17 @@ zjit_send_cache_search(rb_control_frame_t *reg_cfp, struct rb_call_data *cd,
     const struct rb_callcache *cc = vm_search_method_fastpath(reg_cfp, cd, klass);
     bool cacheable = zjit_send_cache_cacheable_p(cc);
     if (LIKELY(cacheable)) {
-        /* One naturally-aligned pointer store, published the same way
+        /* Two naturally-aligned stores, published the same way
          * vm_search_method_slowpath0() publishes cd->cc. No write barrier: the
          * table is a GC root, scanned on every collection, not a heap object
-         * that could be missed by a minor GC. */
-        cache->slots[slot] = cc;
+         * that could be missed by a minor GC.
+         *
+         * The two words need no ordering between them either: JIT code checks
+         * `direct_cme` against the method entry it read out of `cc`, so a reader
+         * that catches this fill half-done sees a mismatched pair and searches.
+         * See zjit_send_cache_direct_cme(). */
+        cache->slots[slot].direct_cme = zjit_send_cache_direct_cme(cache, cc);
+        cache->slots[slot].cc = cc;
     }
 
     if (UNLIKELY(cache->hit_counter != NULL)) {
@@ -6407,6 +6477,20 @@ rb_zjit_send_cached(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
     VM_EXEC(ec, val);
     return val;
 }
+
+/* Field offsets and flag masks that the inline send-cache probe bakes into JIT
+ * code. See the declarations in zjit.h. */
+size_t rb_zjit_cc_klass_offset(void) { return offsetof(struct rb_callcache, klass); }
+size_t rb_zjit_cc_cme_offset(void) { return offsetof(struct rb_callcache, cme_); }
+size_t rb_zjit_iseq_body_offset(void) { return offsetof(struct rb_iseq_struct, body); }
+size_t rb_zjit_iseq_body_jit_entry_offset(void) { return offsetof(struct rb_iseq_constant_body, jit_entry); }
+size_t rb_zjit_send_cache_entry_size(void) { return sizeof(struct rb_zjit_send_cache_entry); }
+size_t rb_zjit_send_cache_entry_direct_cme_offset(void) { return offsetof(struct rb_zjit_send_cache_entry, direct_cme); }
+size_t rb_zjit_cme_def_offset(void) { return offsetof(rb_callable_method_entry_t, def); }
+size_t rb_zjit_def_iseqptr_offset(void) { return offsetof(rb_method_definition_t, body.iseq.iseqptr); }
+VALUE rb_zjit_method_entry_invalidated_flag(void) { return IMEMO_FL_USER5; }
+size_t rb_zjit_mega_direct_max_stack(void) { return ZJIT_MEGA_DIRECT_MAX_STACK; }
+
 #endif // USE_ZJIT
 
 VALUE

zjit.h

@@ -108,20 +108,48 @@ ZJIT_STACK_MAP_BASE_PTR_STACK_SIZE(VALUE entry)
 
 // Class -> callcache table for a send site that dispatches over too many
 // classes for ZJIT's inline class-guard chain. Allocated and owned by Rust; the
-// layout of these four fields must stay in step with `struct SendCache` in
+// layout of these fields must stay in step with `struct SendCache` in
 // zjit/src/send_cache.rs, which documents what the table caches and why a stale
 // entry cannot be wrong.
+//
+// One slot of the table. Two words so that a hit can answer both questions a
+// megamorphic send asks: which callcache dispatches this class (`cc`), and, when
+// the answer is an ISEQ method JIT code can enter without the interpreter,
+// which method that is (`direct_cme`). See `gen_send_megamorphic_direct()` in
+// zjit/src/codegen.rs for the inline probe that reads both.
+struct rb_zjit_send_cache_entry {
+    // The cached callcache, or NULL when the slot is empty. Validates itself:
+    // see zjit_send_cache_search().
+    const struct rb_callcache *cc;
+    // `vm_cc_cme(cc)` when that method is one JIT code may enter with a direct
+    // call (see zjit_send_cache_direct_cme()), NULL otherwise. JIT code checks
+    // it against the method entry it read out of `cc`, so a slot caught
+    // half-written by another ractor reads as a miss rather than as a call to
+    // the wrong method -- which is what lets both words be plain stores.
+    const rb_callable_method_entry_t *direct_cme;
+};
+
 struct rb_zjit_send_cache {
     // Number of slots. A power of two.
     uint32_t len;
     // 64 - log2(len): the shift that turns the hash product into a slot index.
     uint32_t shift;
-    // Slot 0. One callcache pointer per slot, NULL when the slot is empty.
-    const struct rb_callcache **slots;
+    // Slot 0.
+    struct rb_zjit_send_cache_entry *slots;
     // The ZJIT hit counter under --zjit-stats, NULL otherwise. Doubles as the
     // flag for whether to report misses to rb_zjit_send_cache_record_miss(), so
     // that a build without stats pays a never-taken branch rather than a call.
     uint64_t *hit_counter;
+    // `argc` of the call shape this table serves, and whether that shape lets a
+    // hit be dispatched with a direct JIT-to-JIT call at all. Both are constant
+    // for the life of the table; Rust sets them when it allocates it.
+    uint32_t direct_argc;
+    // Zero when the call shape rules direct dispatch out (a splat, a block
+    // argument, a tail call, ...), so `direct_cme` stays NULL in every slot.
+    uint32_t direct_ok;
+    // The call shape's `vm_ci_flag()`, for the visibility test the fill path
+    // runs (a private method is directly callable only from an FCALL site).
+    uint32_t direct_flags;
 };
 
 // Why a probe of a `struct rb_zjit_send_cache` did not produce a callcache.
@@ -137,6 +165,25 @@ struct rb_zjit_send_cache {
 
 void rb_zjit_send_cache_record_miss(int kind);
 
+// Field offsets and flag masks the inline send-cache probe in JIT code needs.
+// They are functions rather than bindgen constants because the structs they
+// reach into (rb_callcache, rb_iseq_constant_body) are opaque to Rust.
+size_t rb_zjit_cc_klass_offset(void);
+size_t rb_zjit_cc_cme_offset(void);
+size_t rb_zjit_iseq_body_offset(void);
+size_t rb_zjit_iseq_body_jit_entry_offset(void);
+size_t rb_zjit_send_cache_entry_size(void);
+size_t rb_zjit_send_cache_entry_direct_cme_offset(void);
+size_t rb_zjit_cme_def_offset(void);
+size_t rb_zjit_def_iseqptr_offset(void);
+VALUE rb_zjit_method_entry_invalidated_flag(void);
+size_t rb_zjit_mega_direct_max_stack(void);
+
+// Largest `stack_max` a callee may have and still be entered by the inline
+// megamorphic dispatch path, which checks for stack overflow against this bound
+// instead of the callee's own (unknown at compile time) requirement.
+#define ZJIT_MEGA_DIRECT_MAX_STACK 64
+
 extern void *rb_zjit_entry;
 extern bool rb_zjit_compiling_p;
 extern const zjit_jit_frame_t rb_zjit_c_frame;