From 50cd34c4e837466ce041adf114ea474e6627bb91 Mon Sep 17 00:00:00 2001 From: Aiden Fox Ivey Date: Fri, 10 Oct 2025 13:22:15 -0400 Subject: [PATCH] ZJIT: Add Insn:: ArrayArefFixnum to accelerate Array#[] (#14717) * ZJIT: Add Insn:: ArrayArefFixnum to accelerate Array#[] * ZJIT: Use result from GuardType in ArrayArefFixnum * ZJIT: Unbox index for aref_fixnum * ZJIT: Change condition and add ArrayArefFixnum test * ZJIT: Fix ArrayArefFixnum display for InsnPrinter * ZJIT: Change insta test --- test/ruby/test_zjit.rb | 8 ++++++ zjit/bindgen/src/main.rs | 1 + zjit/src/codegen.rs | 11 ++++++++ zjit/src/cruby_bindings.inc.rs | 1 + zjit/src/hir.rs | 51 ++++++++++++++++++++++++++++++++++ 5 files changed, 72 insertions(+) diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb index c6dbc01dc2..faf717096a 100644 --- a/test/ruby/test_zjit.rb +++ b/test/ruby/test_zjit.rb @@ -1153,6 +1153,14 @@ class TestZJIT < Test::Unit::TestCase } end + def test_array_fixnum_aref + assert_compiles '3', %q{ + def test(x) = [1,2,3][x] + test(2) + test(2) + }, call_threshold: 2, insns: [:opt_aref] + end + def test_new_range_inclusive assert_compiles '1..5', %q{ def test(a, b) = a..b diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs index 64b235b838..b40986c0f6 100644 --- a/zjit/bindgen/src/main.rs +++ b/zjit/bindgen/src/main.rs @@ -125,6 +125,7 @@ fn main() { .allowlist_function("rb_ary_unshift_m") .allowlist_function("rb_ec_ary_new_from_values") .allowlist_function("rb_ary_tmp_new_from_values") + .allowlist_function("rb_ary_entry") .allowlist_function("rb_class_attached_object") .allowlist_function("rb_singleton_class") .allowlist_function("rb_define_class") diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 4b9331e05b..35791bc0d7 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -355,6 +355,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::NewRange { low, high, flag, state } => gen_new_range(jit, asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)), Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)), Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)), + Insn::ArrayArefFixnum { array, index, .. } => gen_aref_fixnum(asm, opnd!(array), opnd!(index)), Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)), &Insn::ObjectAllocClass { class, state } => gen_object_alloc_class(asm, class, &function.frame_state(state)), Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)), @@ -1241,6 +1242,16 @@ fn gen_new_array( new_array } +/// Compile array access (array[index]) +fn gen_aref_fixnum( + asm: &mut Assembler, + array: Opnd, + index: Opnd, +) -> lir::Opnd { + let unboxed_idx = asm.rshift(index, Opnd::UImm(1)); + asm_ccall!(asm, rb_ary_entry, array, unboxed_idx) +} + /// Compile a new hash instruction fn gen_new_hash( jit: &mut JITState, diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs index ab442841ff..ea1bf68acc 100644 --- a/zjit/src/cruby_bindings.inc.rs +++ b/zjit/src/cruby_bindings.inc.rs @@ -786,6 +786,7 @@ unsafe extern "C" { pub fn rb_ary_resurrect(ary: VALUE) -> VALUE; pub fn rb_ary_cat(ary: VALUE, train: *const VALUE, len: ::std::os::raw::c_long) -> VALUE; pub fn rb_ary_push(ary: VALUE, elem: VALUE) -> VALUE; + pub fn rb_ary_entry(ary: VALUE, off: ::std::os::raw::c_long) -> VALUE; pub fn rb_ary_clear(ary: VALUE) -> VALUE; pub fn rb_ary_concat(lhs: VALUE, rhs: VALUE) -> VALUE; pub fn rb_hash_new() -> VALUE; diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 2fe8eb7970..1ab6032832 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -577,6 +577,7 @@ pub enum Insn { ArrayExtend { left: InsnId, right: InsnId, state: InsnId }, /// Push `val` onto `array`, where `array` is already `Array`. ArrayPush { array: InsnId, val: InsnId, state: InsnId }, + ArrayArefFixnum { array: InsnId, index: InsnId }, HashDup { val: InsnId, state: InsnId }, @@ -888,6 +889,10 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> { } Ok(()) } + Insn::ArrayArefFixnum { array, index, .. } => { + write!(f, "ArrayArefFixnum {array}, {index}")?; + Ok(()) + } Insn::NewHash { elements, .. } => { write!(f, "NewHash")?; let mut prefix = " "; @@ -1579,6 +1584,7 @@ impl Function { &NewHash { ref elements, state } => NewHash { elements: find_vec!(elements), state: find!(state) }, &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) }, &NewRangeFixnum { low, high, flag, state } => NewRangeFixnum { low: find!(low), high: find!(high), flag, state: find!(state) }, + &ArrayArefFixnum { array, index } => ArrayArefFixnum { array: find!(array), index: find!(index) }, &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) }, &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state }, &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state }, @@ -1664,6 +1670,7 @@ impl Function { Insn::ToRegexp { .. } => types::RegexpExact, Insn::NewArray { .. } => types::ArrayExact, Insn::ArrayDup { .. } => types::ArrayExact, + Insn::ArrayArefFixnum { .. } => types::BasicObject, Insn::NewHash { .. } => types::HashExact, Insn::HashDup { .. } => types::HashExact, Insn::NewRange { .. } => types::RangeExact, @@ -1969,6 +1976,16 @@ impl Function { } } } + if self.type_of(idx_val).is_subtype(types::Fixnum) { + self.push_insn(block, Insn::PatchPoint { invariant: Invariant::BOPRedefined { klass: ARRAY_REDEFINED_OP_FLAG, bop: BOP_AREF }, state }); + let fixnum_idx = self.push_insn(block, Insn::GuardType { val: idx_val, guard_type: types::Fixnum, state }); + let result = self.push_insn(block, Insn::ArrayArefFixnum { + array: self_val, + index: fixnum_idx, + }); + self.make_equal_to(orig_insn_id, result); + return; + } } self.push_insn_id(block, orig_insn_id); } @@ -2687,6 +2704,10 @@ impl Function { worklist.push_back(val); worklist.push_back(state); } + &Insn::ArrayArefFixnum { array, index } => { + worklist.push_back(array); + worklist.push_back(index); + } &Insn::Send { recv, ref args, state, .. } | &Insn::SendForward { recv, ref args, state, .. } | &Insn::SendWithoutBlock { recv, ref args, state, .. } @@ -12564,4 +12585,34 @@ mod opt_tests { Return v23 "); } + + #[test] + fn test_array_aref_fixnum() { + eval(" + def test + arr = [1, 2, 3] + arr[0] + end + "); + assert_snapshot!(hir_string("test"), @r" + fn test@:3: + bb0(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + v2:NilClass = Const Value(nil) + Jump bb2(v1, v2) + bb1(v5:BasicObject): + EntryPoint JIT(0) + v6:NilClass = Const Value(nil) + Jump bb2(v5, v6) + bb2(v8:BasicObject, v9:NilClass): + v13:ArrayExact[VALUE(0x1000)] = Const Value(VALUE(0x1000)) + v15:ArrayExact = ArrayDup v13 + v18:Fixnum[0] = Const Value(0) + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF) + v30:BasicObject = ArrayArefFixnum v15, v18 + CheckInterrupts + Return v30 + "); + } }