ZJIT: Stop optimizing toplevel locals (ruby#14458)

Author: k0kubun

.github/workflows/zjit-macos.yml

@@ -47,6 +47,10 @@ jobs:
             configure: '--enable-zjit=dev'
             zjit_opts: '--zjit-call-threshold=1'
 
+          - test_task: 'test-spec'
+            configure: '--enable-zjit=dev'
+            specopts: '-T --zjit-call-threshold=1'
+
     env:
       GITPULLOPTIONS: --no-tags origin ${{ github.ref }}
       RUN_OPTS: ${{ matrix.zjit_opts }}

.github/workflows/zjit-ubuntu.yml

@@ -68,6 +68,10 @@ jobs:
             configure: '--enable-zjit=dev'
             zjit_opts: '--zjit-call-threshold=1'
 
+          - test_task: 'test-spec'
+            configure: '--enable-zjit=dev'
+            specopts: '-T --zjit-call-threshold=1'
+
     env:
       GITPULLOPTIONS: --no-tags origin ${{ github.ref }}
       RUN_OPTS: ${{ matrix.zjit_opts }}

test/ruby/test_zjit.rb

@@ -196,6 +196,32 @@ def test
     }, insns: [:setglobal]
   end
 
+  def test_toplevel_binding
+    # Not using assert_compiles, which doesn't use the toplevel frame for `test_script`.
+    out, err, status = eval_with_jit(%q{
+      a = 1
+      b = 2
+      TOPLEVEL_BINDING.local_variable_set(:b, 3)
+      c = 4
+      print [a, b, c]
+    })
+    assert_success(out, err, status)
+    assert_equal "[1, 3, 4]", out
+  end
+
+  def test_toplevel_local_after_eval
+    # Not using assert_compiles, which doesn't use the toplevel frame for `test_script`.
+    out, err, status = eval_with_jit(%q{
+      a = 1
+      b = 2
+      eval('b = 3')
+      c = 4
+      print [a, b, c]
+    })
+    assert_success(out, err, status)
+    assert_equal "[1, 3, 4]", out
+  end
+
   def test_getlocal_after_eval
     assert_compiles '2', %q{
       def test
@@ -2428,12 +2454,8 @@ def assert_runs(expected, test_script, insns: [], assert_compiles: false, **opts
       IO.open(#{pipe_fd}).write(Marshal.dump(result))
     RUBY
 
-    status, out, err, result = eval_with_jit(script, pipe_fd:, **opts)
-
-    message = "exited with status #{status.to_i}"
-    message << "\nstdout:\n```\n#{out}```\n" unless out.empty?
-    message << "\nstderr:\n```\n#{err}```\n" unless err.empty?
-    assert status.success?, message
+    out, err, status, result = eval_with_jit(script, pipe_fd:, **opts)
+    assert_success(out, err, status)
 
     result = Marshal.load(result)
     assert_equal(expected, result.fetch(:ret_val).inspect)
@@ -2462,7 +2484,7 @@ def eval_with_jit(
     debug: true,
     allowed_iseqs: nil,
     timeout: 1000,
-    pipe_fd:
+    pipe_fd: nil
   )
     args = ["--disable-gems"]
     if zjit
@@ -2478,18 +2500,25 @@ def eval_with_jit(
       end
     end
     args << "-e" << script_shell_encode(script)
-    pipe_r, pipe_w = IO.pipe
-    # Separate thread so we don't deadlock when
-    # the child ruby blocks writing the output to pipe_fd
-    pipe_out = nil
-    pipe_reader = Thread.new do
-      pipe_out = pipe_r.read
-      pipe_r.close
+    ios = {}
+    if pipe_fd
+      pipe_r, pipe_w = IO.pipe
+      # Separate thread so we don't deadlock when
+      # the child ruby blocks writing the output to pipe_fd
+      pipe_out = nil
+      pipe_reader = Thread.new do
+        pipe_out = pipe_r.read
+        pipe_r.close
+      end
+      ios[pipe_fd] = pipe_w
     end
-    out, err, status = EnvUtil.invoke_ruby(args, '', true, true, rubybin: RbConfig.ruby, timeout: timeout, ios: { pipe_fd => pipe_w })
-    pipe_w.close
-    pipe_reader.join(timeout)
-    [status, out, err, pipe_out]
+    result = EnvUtil.invoke_ruby(args, '', true, true, rubybin: RbConfig.ruby, timeout: timeout, ios:)
+    if pipe_fd
+      pipe_w.close
+      pipe_reader.join(timeout)
+      result << pipe_out
+    end
+    result
   ensure
     pipe_reader&.kill
     pipe_reader&.join(timeout)
@@ -2498,6 +2527,13 @@ def eval_with_jit(
     jitlist&.unlink
   end
 
+  def assert_success(out, err, status)
+    message = "exited with status #{status.to_i}"
+    message << "\nstdout:\n```\n#{out}```\n" unless out.empty?
+    message << "\nstderr:\n```\n#{err}```\n" unless err.empty?
+    assert status.success?, message
+  end
+
   def script_shell_encode(s)
     # We can't pass utf-8-encoded characters directly in a shell arg. But we can use Ruby \u constants.
     s.chars.map { |c| c.ascii_only? ? c : "\\u%x" % c.codepoints[0] }.join

zjit/src/cruby.rs

@@ -294,6 +294,17 @@ pub fn iseq_opcode_at_idx(iseq: IseqPtr, insn_idx: u32) -> u32 {
     unsafe { rb_iseq_opcode_at_pc(iseq, pc) as u32 }
 }
 
+/// Return true if the ISEQ always uses a frame with escaped EP.
+pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
+    match unsafe { get_iseq_body_type(iseq) } {
+        // The EP of the <main> frame points to TOPLEVEL_BINDING
+        ISEQ_TYPE_MAIN |
+        // eval frames point to the EP of another frame or scope
+        ISEQ_TYPE_EVAL => true,
+        _ => false,
+    }
+}
+
 /// Iterate over all existing ISEQs
 pub fn for_each_iseq<F: FnMut(IseqPtr)>(mut callback: F) {
     unsafe extern "C" fn callback_wrapper(iseq: IseqPtr, data: *mut c_void) {

zjit/src/hir.rs

@@ -3019,6 +3019,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     let payload = get_or_create_iseq_payload(iseq);
     let mut profiles = ProfileOracle::new(payload);
     let mut fun = Function::new(iseq);
+
     // Compute a map of PC->Block by finding jump targets
     let BytecodeInfo { jump_targets, has_blockiseq } = compute_bytecode_info(iseq);
     let mut insn_idx_to_block = HashMap::new();
@@ -3029,6 +3030,10 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
         insn_idx_to_block.insert(insn_idx, fun.new_block(insn_idx));
     }
 
+    // Check if the EP is escaped for the ISEQ from the beginning. We give up
+    // optimizing locals in that case because they're shared with other frames.
+    let ep_escaped = iseq_escapes_ep(iseq);
+
     // Iteratively fill out basic blocks using a queue
     // TODO(max): Basic block arguments at edges
     let mut queue = std::collections::VecDeque::new();
@@ -3066,7 +3071,6 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     let mut visited = HashSet::new();
 
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
-    let iseq_type = unsafe { get_iseq_body_type(iseq) };
     while let Some((incoming_state, block, mut insn_idx, mut local_inval)) = queue.pop_front() {
         if visited.contains(&block) { continue; }
         visited.insert(block);
@@ -3361,8 +3365,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_getlocal_WC_0 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    if iseq_type == ISEQ_TYPE_EVAL || has_blockiseq {
-                        // On eval, the locals are always on the heap, so read the local using EP.
+                    if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here
+                        // Read the local using EP
                         let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
                         state.setlocal(ep_offset, val); // remember the result to spill on side-exits
                         state.stack_push(val);
@@ -3374,15 +3378,16 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                             fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id });
                             local_inval = false;
                         }
+                        // Read the local from FrameState
                         let val = state.getlocal(ep_offset);
                         state.stack_push(val);
                     }
                 }
                 YARVINSN_setlocal_WC_0 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let val = state.stack_pop()?;
-                    if iseq_type == ISEQ_TYPE_EVAL || has_blockiseq {
-                        // On eval, the locals are always on the heap, so write the local using EP.
+                    if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here
+                        // Write the local using EP
                         fun.push_insn(block, Insn::SetLocal { val, ep_offset, level: 0 });
                     } else if local_inval {
                         // If there has been any non-leaf call since JIT entry or the last patch point,
@@ -3391,6 +3396,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id });
                         local_inval = false;
                     }
+                    // Write the local into FrameState
                     state.setlocal(ep_offset, val);
                 }
                 YARVINSN_getlocal_WC_1 => {