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
This commit is contained in:
Aiden Fox Ivey 2025-10-10 13:22:15 -04:00 committed by GitHub
parent 17a5a5e2ef
commit 50cd34c4e8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
Notes: git 2025-10-10 17:22:46 +00:00
Merged-By: k0kubun <takashikkbn@gmail.com>
5 changed files with 72 additions and 0 deletions

View File

@ -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

View File

@ -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")

View File

@ -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,

View File

@ -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;

View File

@ -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@<compiled>: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
");
}
}