From a00de3198d8ad9f9736da8b493b003b7dc95aed0 Mon Sep 17 00:00:00 2001 From: Maxime Chevalier-Boisvert Date: Wed, 2 Aug 2023 16:15:41 -0400 Subject: [PATCH] Refactor expandarray to use chain guards --- bootstraptest/test_yjit.rb | 11 +++++ yjit/src/codegen.rs | 91 +++++++++++++++++++++++++++----------- yjit/src/stats.rs | 2 +- 3 files changed, 78 insertions(+), 26 deletions(-) diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb index 80d0ef00d881c6..6cc509f212a959 100644 --- a/bootstraptest/test_yjit.rb +++ b/bootstraptest/test_yjit.rb @@ -2276,6 +2276,17 @@ def expandarray_rhs_too_small expandarray_rhs_too_small } +assert_equal '[nil, 2, nil]', %q{ + def foo(arr) + a, b, c = arr + end + + a, b, c1 = foo([0, 1]) + a, b, c2 = foo([0, 1, 2]) + a, b, c3 = foo([0, 1]) + [c1, c2, c3] +} + assert_equal '[1, [2]]', %q{ def expandarray_splat a, *b = [1, 2] diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 4e00f945bfb39a..fe933c5f5c36d9 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -1467,7 +1467,7 @@ fn guard_object_is_not_ruby2_keyword_hash( fn gen_expandarray( jit: &mut JITState, asm: &mut Assembler, - _ocb: &mut OutlinedCb, + ocb: &mut OutlinedCb, ) -> Option { // Both arguments are rb_num_t which is unsigned let num = jit.get_arg(0).as_usize(); @@ -1500,6 +1500,23 @@ fn gen_expandarray( return Some(KeepCompiling); } + // Defer compilation so we can specialize on a runtime `self` + if !jit.at_current_insn() { + defer_compilation(jit, asm, ocb); + return Some(EndBlock); + } + + let comptime_recv = jit.peek_at_stack(&asm.ctx, 0); + + // If the comptime receiver is not an array, bail + if comptime_recv.class_of() != unsafe { rb_cArray } { + return None; + } + + // Get the compile-time array length + let comptime_len = unsafe { rb_yjit_array_len(comptime_recv) as u32 }; + //println!("comptime_len={}", comptime_len); + // Move the array from the stack and check that it's an array. guard_object_is_array( asm, @@ -1507,41 +1524,65 @@ fn gen_expandarray( array_opnd.into(), Counter::expandarray_not_array, ); - let array_opnd = asm.stack_pop(1); // pop after using the type info // If we don't actually want any values, then just return. if num == 0 { + asm.stack_pop(1); // pop the array return Some(KeepCompiling); } + let array_opnd = asm.stack_opnd(0); let array_reg = asm.load(array_opnd); let array_len_opnd = get_array_len(asm, array_reg); - // Load the address of the embedded array into REG1. - // (struct RArray *)(obj)->as.ary - let array_reg = asm.load(array_opnd); - let ary_opnd = asm.lea(Opnd::mem(VALUE_BITS, array_reg, RUBY_OFFSET_RARRAY_AS_ARY)); - - // Conditionally load the address of the heap array into REG1. - // (struct RArray *)(obj)->as.heap.ptr - let flags_opnd = Opnd::mem(VALUE_BITS, array_reg, RUBY_OFFSET_RBASIC_FLAGS); - asm.test(flags_opnd, Opnd::UImm(RARRAY_EMBED_FLAG as u64)); - let heap_ptr_opnd = Opnd::mem( - usize::BITS as u8, - asm.load(array_opnd), - RUBY_OFFSET_RARRAY_AS_HEAP_PTR, + // Guard on the comptime/expected array length + asm.comment(&format!("guard array length == {}", comptime_len)); + asm.cmp(array_len_opnd, comptime_len.into()); + jit_chain_guard( + JCC_JNE, + jit, + asm, + ocb, + OPT_AREF_MAX_CHAIN_DEPTH, + Counter::expandarray_chain_max_depth, ); - let ary_opnd = asm.csel_nz(ary_opnd, heap_ptr_opnd); - // Loop backward through the array and push each element onto the stack. - for i in (0..num).rev() { - let top = asm.stack_push(Type::Unknown); - let offset = i32::try_from(i * SIZEOF_VALUE).unwrap(); + let array_opnd = asm.stack_pop(1); // pop after using the type info + + //println!("generating loads"); + + // If the array has length 0, then we don't even need the array pointer + if comptime_len == 0 { + // Loop backward through the array and push Qnils onto the stack. + for _ in 0..num { + let top = asm.stack_push(Type::Nil); + asm.mov(top, Qnil.into()); + } + } else { + let array_reg = asm.load(array_opnd); + let ary_opnd = asm.lea(Opnd::mem(VALUE_BITS, array_reg, RUBY_OFFSET_RARRAY_AS_ARY)); - // If the element index is less than the length of the array, load it - asm.cmp(array_len_opnd, i.into()); - let elem_opnd = asm.csel_g(Opnd::mem(64, ary_opnd, offset), Qnil.into()); - asm.mov(top, elem_opnd); + // Conditionally load the address of the heap array + // (struct RArray *)(obj)->as.heap.ptr + let flags_opnd = Opnd::mem(VALUE_BITS, array_reg, RUBY_OFFSET_RBASIC_FLAGS); + asm.test(flags_opnd, Opnd::UImm(RARRAY_EMBED_FLAG as u64)); + let heap_ptr_opnd = Opnd::mem( + usize::BITS as u8, + asm.load(array_opnd), + RUBY_OFFSET_RARRAY_AS_HEAP_PTR, + ); + let ary_opnd = asm.csel_nz(ary_opnd, heap_ptr_opnd); + + // Loop backward through the array and push each element onto the stack. + for i in (0..num).rev() { + let top = asm.stack_push(Type::Unknown); + let offset = i32::try_from(i * SIZEOF_VALUE).unwrap(); + + // Missing elements are Qnil + asm.comment(&format!("load array[{}]", i)); + let elem_opnd = if (i as u32) < comptime_len { Opnd::mem(64, ary_opnd, offset) } else { Qnil.into() }; + asm.mov(top, elem_opnd); + } } Some(KeepCompiling) @@ -7393,7 +7434,7 @@ fn gen_leave( ocb: &mut OutlinedCb, ) -> Option { // Only the return value should be on the stack - assert_eq!(1, asm.ctx.get_stack_size()); + assert_eq!(1, asm.ctx.get_stack_size(), "leave instruction expects stack size 1"); let ocb_asm = Assembler::new(); diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index 686516cb9b5834..bc95cc44d4f7e1 100644 --- a/yjit/src/stats.rs +++ b/yjit/src/stats.rs @@ -358,7 +358,7 @@ make_counters! { expandarray_splat, expandarray_postarg, expandarray_not_array, - expandarray_rhs_too_small, + expandarray_chain_max_depth, // getblockparam gbp_wb_required,