Skip to content

Commit 3286556

Browse files
tekknolagitekknolagi
committed
ZJIT: Optimize Integer#[]
This is used a lot in optcarrot.
1 parent ee1aa78 commit 3286556

4 files changed

Lines changed: 136 additions & 0 deletions

File tree

‎zjit/src/codegen.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,6 +439,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
439439
gen_fixnum_rshift(asm, opnd!(left), shift_amount)
440440
}
441441
&Insn::FixnumMod { left, right, state } => gen_fixnum_mod(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
442+
&Insn::FixnumAref { recv, index } => gen_fixnum_aref(asm, opnd!(recv), opnd!(index)),
442443
Insn::IsNil { val } => gen_isnil(asm, opnd!(val)),
443444
&Insn::IsMethodCfunc { val, cd, cfunc, state: _ } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc),
444445
&Insn::IsBitEqual { left, right } => gen_is_bit_equal(asm, opnd!(left), opnd!(right)),
@@ -1911,6 +1912,10 @@ fn gen_fixnum_mod(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, righ
19111912
asm_ccall!(asm, rb_fix_mod_fix, left, right)
19121913
}
19131914

1915+
fn gen_fixnum_aref(asm: &mut Assembler, recv: lir::Opnd, index: lir::Opnd) -> lir::Opnd {
1916+
asm_ccall!(asm, rb_fix_aref, recv, index)
1917+
}
1918+
19141919
// Compile val == nil
19151920
fn gen_isnil(asm: &mut Assembler, val: lir::Opnd) -> lir::Opnd {
19161921
asm.cmp(val, Qnil.into());

‎zjit/src/cruby_methods.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,7 @@ pub fn init() -> Annotations {
255255
annotate!(rb_cInteger, "<=", inline_integer_le);
256256
annotate!(rb_cInteger, "<<", inline_integer_lshift);
257257
annotate!(rb_cInteger, ">>", inline_integer_rshift);
258+
annotate!(rb_cInteger, "[]", inline_integer_aref);
258259
annotate!(rb_cInteger, "to_s", types::StringExact);
259260
annotate!(rb_cString, "to_s", inline_string_to_s, types::StringExact);
260261
let thread_singleton = unsafe { rb_singleton_class(rb_cThread) };
@@ -679,6 +680,17 @@ fn inline_integer_rshift(fun: &mut hir::Function, block: hir::BlockId, recv: hir
679680
try_inline_fixnum_op(fun, block, &|left, right| hir::Insn::FixnumRShift { left, right }, BOP_GTGT, recv, other, state)
680681
}
681682

683+
fn inline_integer_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
684+
let &[index] = args else { return None; };
685+
if fun.likely_a(recv, types::Fixnum, state) && fun.likely_a(index, types::Fixnum, state) {
686+
let recv = fun.coerce_to(block, recv, types::Fixnum, state);
687+
let index = fun.coerce_to(block, index, types::Fixnum, state);
688+
let result = fun.push_insn(block, hir::Insn::FixnumAref { recv, index });
689+
return Some(result);
690+
}
691+
None
692+
}
693+
682694
fn inline_basic_object_eq(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
683695
let &[other] = args else { return None; };
684696
let c_result = fun.push_insn(block, hir::Insn::IsBitEqual { left: recv, right: other });

‎zjit/src/hir.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -765,6 +765,7 @@ pub enum Insn {
765765
/// Convert a C `long` to a Ruby `Fixnum`. Side exit on overflow.
766766
BoxFixnum { val: InsnId, state: InsnId },
767767
UnboxFixnum { val: InsnId },
768+
FixnumAref { recv: InsnId, index: InsnId },
768769
// TODO(max): In iseq body types that are not ISEQ_TYPE_METHOD, rewrite to Constant false.
769770
Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, state: InsnId },
770771
GetConstantPath { ic: *const iseq_inline_constant_cache, state: InsnId },
@@ -1050,6 +1051,7 @@ impl Insn {
10501051
Insn::FixnumXor { .. } => false,
10511052
Insn::FixnumLShift { .. } => false,
10521053
Insn::FixnumRShift { .. } => false,
1054+
Insn::FixnumAref { .. } => false,
10531055
Insn::GetLocal { .. } => false,
10541056
Insn::IsNil { .. } => false,
10551057
Insn::LoadPC => false,
@@ -1254,6 +1256,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
12541256
Insn::BoxBool { val } => write!(f, "BoxBool {val}"),
12551257
Insn::BoxFixnum { val, .. } => write!(f, "BoxFixnum {val}"),
12561258
Insn::UnboxFixnum { val } => write!(f, "UnboxFixnum {val}"),
1259+
Insn::FixnumAref { recv, index } => write!(f, "FixnumAref {recv}, {index}"),
12571260
Insn::Jump(target) => { write!(f, "Jump {target}") }
12581261
Insn::IfTrue { val, target } => { write!(f, "IfTrue {val}, {target}") }
12591262
Insn::IfFalse { val, target } => { write!(f, "IfFalse {val}, {target}") }
@@ -1971,6 +1974,7 @@ impl Function {
19711974
&BoxBool { val } => BoxBool { val: find!(val) },
19721975
&BoxFixnum { val, state } => BoxFixnum { val: find!(val), state: find!(state) },
19731976
&UnboxFixnum { val } => UnboxFixnum { val: find!(val) },
1977+
&FixnumAref { recv, index } => FixnumAref { recv: find!(recv), index: find!(index) },
19741978
Jump(target) => Jump(find_branch_edge!(target)),
19751979
&IfTrue { val, ref target } => IfTrue { val: find!(val), target: find_branch_edge!(target) },
19761980
&IfFalse { val, ref target } => IfFalse { val: find!(val), target: find_branch_edge!(target) },
@@ -2184,6 +2188,7 @@ impl Function {
21842188
Insn::BoxBool { .. } => types::BoolExact,
21852189
Insn::BoxFixnum { .. } => types::Fixnum,
21862190
Insn::UnboxFixnum { .. } => types::CInt64,
2191+
Insn::FixnumAref { .. } => types::Fixnum,
21872192
Insn::StringCopy { .. } => types::StringExact,
21882193
Insn::StringIntern { .. } => types::Symbol,
21892194
Insn::StringConcat { .. } => types::StringExact,
@@ -4150,6 +4155,10 @@ impl Function {
41504155
&Insn::ObjectAllocClass { state, .. } |
41514156
&Insn::SideExit { state, .. } => worklist.push_back(state),
41524157
&Insn::UnboxFixnum { val } => worklist.push_back(val),
4158+
&Insn::FixnumAref { recv, index } => {
4159+
worklist.push_back(recv);
4160+
worklist.push_back(index);
4161+
}
41534162
&Insn::IsA { val, class } => {
41544163
worklist.push_back(val);
41554164
worklist.push_back(class);
@@ -4817,6 +4826,10 @@ impl Function {
48174826
Insn::UnboxFixnum { val } => {
48184827
self.assert_subtype(insn_id, val, types::Fixnum)
48194828
}
4829+
Insn::FixnumAref { recv, index } => {
4830+
self.assert_subtype(insn_id, recv, types::Fixnum)?;
4831+
self.assert_subtype(insn_id, index, types::Fixnum)
4832+
}
48204833
Insn::FixnumAdd { left, right, .. }
48214834
| Insn::FixnumSub { left, right, .. }
48224835
| Insn::FixnumMult { left, right, .. }

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1201,6 +1201,112 @@ mod hir_opt_tests {
12011201
");
12021202
}
12031203

1204+
#[test]
1205+
fn integer_aref_with_fixnum_emits_fixnum_aref() {
1206+
eval("
1207+
def test(a, b) = a[b]
1208+
test(3, 4)
1209+
");
1210+
assert_snapshot!(hir_string("test"), @r"
1211+
fn test@<compiled>:2:
1212+
bb0():
1213+
EntryPoint interpreter
1214+
v1:BasicObject = LoadSelf
1215+
v2:BasicObject = GetLocal :a, l0, SP@5
1216+
v3:BasicObject = GetLocal :b, l0, SP@4
1217+
Jump bb2(v1, v2, v3)
1218+
bb1(v6:BasicObject, v7:BasicObject, v8:BasicObject):
1219+
EntryPoint JIT(0)
1220+
Jump bb2(v6, v7, v8)
1221+
bb2(v10:BasicObject, v11:BasicObject, v12:BasicObject):
1222+
PatchPoint MethodRedefined(Integer@0x1000, []@0x1008, cme:0x1010)
1223+
v26:Fixnum = GuardType v11, Fixnum
1224+
v27:Fixnum = GuardType v12, Fixnum
1225+
v28:Fixnum = FixnumAref v26, v27
1226+
IncrCounter inline_cfunc_optimized_send_count
1227+
CheckInterrupts
1228+
Return v28
1229+
");
1230+
}
1231+
1232+
#[test]
1233+
fn elide_fixnum_aref() {
1234+
eval("
1235+
def test
1236+
1[2]
1237+
5
1238+
end
1239+
");
1240+
assert_snapshot!(hir_string("test"), @r"
1241+
fn test@<compiled>:3:
1242+
bb0():
1243+
EntryPoint interpreter
1244+
v1:BasicObject = LoadSelf
1245+
Jump bb2(v1)
1246+
bb1(v4:BasicObject):
1247+
EntryPoint JIT(0)
1248+
Jump bb2(v4)
1249+
bb2(v6:BasicObject):
1250+
v10:Fixnum[1] = Const Value(1)
1251+
v12:Fixnum[2] = Const Value(2)
1252+
PatchPoint MethodRedefined(Integer@0x1000, []@0x1008, cme:0x1010)
1253+
IncrCounter inline_cfunc_optimized_send_count
1254+
v19:Fixnum[5] = Const Value(5)
1255+
CheckInterrupts
1256+
Return v19
1257+
");
1258+
}
1259+
1260+
#[test]
1261+
fn do_not_optimize_integer_aref_with_too_many_args() {
1262+
eval("
1263+
def test = 1[2, 3]
1264+
");
1265+
assert_snapshot!(hir_string("test"), @r"
1266+
fn test@<compiled>:2:
1267+
bb0():
1268+
EntryPoint interpreter
1269+
v1:BasicObject = LoadSelf
1270+
Jump bb2(v1)
1271+
bb1(v4:BasicObject):
1272+
EntryPoint JIT(0)
1273+
Jump bb2(v4)
1274+
bb2(v6:BasicObject):
1275+
v10:Fixnum[1] = Const Value(1)
1276+
v12:Fixnum[2] = Const Value(2)
1277+
v14:Fixnum[3] = Const Value(3)
1278+
PatchPoint MethodRedefined(Integer@0x1000, []@0x1008, cme:0x1010)
1279+
v23:BasicObject = CCallVariadic v10, :Integer#[]@0x1038, v12, v14
1280+
CheckInterrupts
1281+
Return v23
1282+
");
1283+
}
1284+
1285+
#[test]
1286+
fn do_not_optimize_integer_aref_with_non_fixnum() {
1287+
eval(r#"
1288+
def test = 1["x"]
1289+
"#);
1290+
assert_snapshot!(hir_string("test"), @r"
1291+
fn test@<compiled>:2:
1292+
bb0():
1293+
EntryPoint interpreter
1294+
v1:BasicObject = LoadSelf
1295+
Jump bb2(v1)
1296+
bb1(v4:BasicObject):
1297+
EntryPoint JIT(0)
1298+
Jump bb2(v4)
1299+
bb2(v6:BasicObject):
1300+
v10:Fixnum[1] = Const Value(1)
1301+
v12:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
1302+
v13:StringExact = StringCopy v12
1303+
PatchPoint MethodRedefined(Integer@0x1008, []@0x1010, cme:0x1018)
1304+
v23:BasicObject = CCallVariadic v10, :Integer#[]@0x1040, v13
1305+
CheckInterrupts
1306+
Return v23
1307+
");
1308+
}
1309+
12041310
#[test]
12051311
fn test_optimize_send_into_fixnum_lt_both_profiled() {
12061312
eval("

0 commit comments

Comments
 (0)