ZJIT: Use SP for reloading locals (ruby#16746)

* ZJIT: Use SP for reloading locals

* ZJIT: Deduplicate post-send NoEPEscape patch point generation

Author: nozomemein

zjit/src/codegen_tests.rs

@@ -652,6 +652,24 @@ fn test_send_with_local_written_by_blockiseq() {
     "), @"[1, 2]");
 }
 
+#[test]
+fn test_no_ep_escape_patch_point_after_send_does_not_repeat_send() {
+    eval(r#"
+        $send_count = 0
+
+        def test
+          captured = nil
+          tap do |_|
+            $send_count += 1
+            -> { captured } if $send_count == 2
+          end
+          $send_count
+        end
+    "#);
+    assert_contains_opcode("test", YARVINSN_send);
+    assert_snapshot!(assert_compiles_allowing_exits("[test, test, test]"), @"[1, 2, 3]");
+}
+
 #[test]
 fn test_send_without_block() {
     assert_snapshot!(inspect("

zjit/src/hir.rs

@@ -4720,6 +4720,17 @@ impl Function {
         self.push_insn(block, Insn::PatchPoint { invariant: Invariant::MethodRedefined { klass: recv_class, method: method_id, cme }, state });
     }
 
+    /// Side exit back to the state after a block-backed send.
+    /// Using the pre-send snapshot would re-execute the send in the interpreter.
+    fn gen_post_send_no_ep_escape_patch_point(&mut self, block: BlockId, state: &FrameState, insn_idx: u32) {
+        let iseq = state.iseq;
+        let mut reload_state = state.clone();
+        reload_state.insn_idx = insn_idx as usize;
+        reload_state.pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) };
+        let reload_exit_id = self.push_insn(block, Insn::Snapshot { state: reload_state.without_locals() });
+        self.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: reload_exit_id });
+    }
+
     fn count_not_inlined_cfunc(&mut self, block: BlockId, cme: *const rb_callable_method_entry_t) {
         let owner = unsafe { (*cme).owner };
         let called_id = unsafe { (*cme).called_id };
@@ -8004,13 +8015,23 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         // Reload locals that may have been modified by the blockiseq.
                         // TODO: Avoid reloading locals that are not referenced by the blockiseq
                         // or not used after this. Max thinks we could eventually DCE them.
-                        let ep = fun.push_insn(block, Insn::GetEP { level: 0 });
+                        if !ep_escaped && !state.locals.is_empty() {
+                            fun.gen_post_send_no_ep_escape_patch_point(block, &state, insn_idx);
+                        }
+                        let mut base: Option<InsnId> = None;
                         for local_idx in 0..state.locals.len() {
                             let ep_offset = local_idx_to_ep_offset(iseq, local_idx);
                             let ep_offset_u32 = u32::try_from(ep_offset)
                                 .unwrap_or_else(|_| panic!("Could not convert ep_offset {ep_offset} to u32"));
-                            // TODO: We could use `use_sp: true` with PatchPoint.
-                            let val = fun.get_local_from_ep(block, ep, ep_offset_u32, 0, types::BasicObject);
+                            let recv = *base.get_or_insert_with(|| {
+                                let base_insn = if !ep_escaped { Insn::LoadSP } else { Insn::GetEP { level: 0 } };
+                                fun.push_insn(block, base_insn)
+                            });
+                            let val = if !ep_escaped {
+                                fun.get_local_from_sp(block, recv, ep_offset_u32, types::BasicObject)
+                            } else {
+                                fun.get_local_from_ep(block, recv, ep_offset_u32, 0, types::BasicObject)
+                            };
                             state.setlocal(ep_offset_u32, val);
                         }
                     }
@@ -8040,13 +8061,23 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     if !blockiseq.is_null() {
                         // Reload locals that may have been modified by the blockiseq.
-                        let ep = fun.push_insn(block, Insn::GetEP { level: 0 });
+                        if !ep_escaped && !state.locals.is_empty() {
+                            fun.gen_post_send_no_ep_escape_patch_point(block, &state, insn_idx);
+                        }
+                        let mut base: Option<InsnId> = None;
                         for local_idx in 0..state.locals.len() {
                             let ep_offset = local_idx_to_ep_offset(iseq, local_idx);
                             let ep_offset_u32 = u32::try_from(ep_offset)
                                 .unwrap_or_else(|_| panic!("Could not convert ep_offset {ep_offset} to u32"));
-                            // TODO: We could use `use_sp: true` with PatchPoint.
-                            let val = fun.get_local_from_ep(block, ep, ep_offset_u32, 0, types::BasicObject);
+                            let recv = *base.get_or_insert_with(|| {
+                                let base_insn = if !ep_escaped { Insn::LoadSP } else { Insn::GetEP { level: 0 } };
+                                fun.push_insn(block, base_insn)
+                            });
+                            let val = if !ep_escaped {
+                                fun.get_local_from_sp(block, recv, ep_offset_u32, types::BasicObject)
+                            } else {
+                                fun.get_local_from_ep(block, recv, ep_offset_u32, 0, types::BasicObject)
+                            };
                             state.setlocal(ep_offset_u32, val);
                         }
                     }
@@ -8077,13 +8108,23 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         // Reload locals that may have been modified by the blockiseq.
                         // TODO: Avoid reloading locals that are not referenced by the blockiseq
                         // or not used after this. Max thinks we could eventually DCE them.
-                        let ep = fun.push_insn(block, Insn::GetEP { level: 0 });
+                        if !ep_escaped && !state.locals.is_empty() {
+                            fun.gen_post_send_no_ep_escape_patch_point(block, &state, insn_idx);
+                        }
+                        let mut base: Option<InsnId> = None;
                         for local_idx in 0..state.locals.len() {
                             let ep_offset = local_idx_to_ep_offset(iseq, local_idx);
                             let ep_offset_u32 = u32::try_from(ep_offset)
                                 .unwrap_or_else(|_| panic!("Could not convert ep_offset {ep_offset} to u32"));
-                            // TODO: We could use `use_sp: true` with PatchPoint.
-                            let val = fun.get_local_from_ep(block, ep, ep_offset_u32, 0, types::BasicObject);
+                            let recv = *base.get_or_insert_with(|| {
+                                let base_insn = if !ep_escaped { Insn::LoadSP } else { Insn::GetEP { level: 0 } };
+                                fun.push_insn(block, base_insn)
+                            });
+                            let val = if !ep_escaped {
+                                fun.get_local_from_sp(block, recv, ep_offset_u32, types::BasicObject)
+                            } else {
+                                fun.get_local_from_ep(block, recv, ep_offset_u32, 0, types::BasicObject)
+                            };
                             state.setlocal(ep_offset_u32, val);
                         }
                     }
@@ -8114,13 +8155,23 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         // Reload locals that may have been modified by the blockiseq.
                         // TODO: Avoid reloading locals that are not referenced by the blockiseq
                         // or not used after this. Max thinks we could eventually DCE them.
-                        let ep = fun.push_insn(block, Insn::GetEP { level: 0 });
+                        if !ep_escaped && !state.locals.is_empty() {
+                            fun.gen_post_send_no_ep_escape_patch_point(block, &state, insn_idx);
+                        }
+                        let mut base: Option<InsnId> = None;
                         for local_idx in 0..state.locals.len() {
                             let ep_offset = local_idx_to_ep_offset(iseq, local_idx);
                             let ep_offset_u32 = u32::try_from(ep_offset)
                                 .unwrap_or_else(|_| panic!("Could not convert ep_offset {ep_offset} to u32"));
-                            // TODO: We could use `use_sp: true` with PatchPoint.
-                            let val = fun.get_local_from_ep(block, ep, ep_offset_u32, 0, types::BasicObject);
+                            let recv = *base.get_or_insert_with(|| {
+                                let base_insn = if !ep_escaped { Insn::LoadSP } else { Insn::GetEP { level: 0 } };
+                                fun.push_insn(block, base_insn)
+                            });
+                            let val = if !ep_escaped {
+                                fun.get_local_from_sp(block, recv, ep_offset_u32, types::BasicObject)
+                            } else {
+                                fun.get_local_from_ep(block, recv, ep_offset_u32, 0, types::BasicObject)
+                            };
                             state.setlocal(ep_offset_u32, val);
                         }
                     }

zjit/src/hir/opt_tests.rs

@@ -1410,12 +1410,13 @@ mod hir_opt_tests {
         bb3(v9:BasicObject, v10:BasicObject):
           PatchPoint NoSingletonClass(C@0x1008)
           PatchPoint MethodRedefined(C@0x1008, fun_new_map@0x1010, cme:0x1018)
-          v25:ArraySubclass[class_exact:C] = GuardType v10, ArraySubclass[class_exact:C]
-          v26:BasicObject = SendDirect v25, 0x1040, :fun_new_map (0x1050)
-          v16:CPtr = GetEP 0
-          v17:BasicObject = LoadField v16, :o@0x1058
+          v27:ArraySubclass[class_exact:C] = GuardType v10, ArraySubclass[class_exact:C]
+          v28:BasicObject = SendDirect v27, 0x1040, :fun_new_map (0x1050)
+          PatchPoint NoEPEscape(test)
+          v18:CPtr = LoadSP
+          v19:BasicObject = LoadField v18, :o@0x1000
           CheckInterrupts
-          Return v26
+          Return v28
         ");
     }
 
@@ -1446,12 +1447,13 @@ mod hir_opt_tests {
         bb3(v9:BasicObject, v10:BasicObject):
           PatchPoint NoSingletonClass(C@0x1008)
           PatchPoint MethodRedefined(C@0x1008, bar@0x1010, cme:0x1018)
-          v26:ObjectSubclass[class_exact:C] = GuardType v10, ObjectSubclass[class_exact:C]
-          v27:BasicObject = CCallWithFrame v26, :Enumerable#bar@0x1040, block=0x1048
-          v16:CPtr = GetEP 0
-          v17:BasicObject = LoadField v16, :o@0x1050
+          v28:ObjectSubclass[class_exact:C] = GuardType v10, ObjectSubclass[class_exact:C]
+          v29:BasicObject = CCallWithFrame v28, :Enumerable#bar@0x1040, block=0x1048
+          PatchPoint NoEPEscape(test)
+          v18:CPtr = LoadSP
+          v19:BasicObject = LoadField v18, :o@0x1000
           CheckInterrupts
-          Return v27
+          Return v29
         ");
     }
 
@@ -3641,10 +3643,10 @@ mod hir_opt_tests {
           v11:Fixnum[1] = Const Value(1)
           v13:Fixnum[2] = Const Value(2)
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
-          v23:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
-          v24:BasicObject = SendDirect v23, 0x1038, :foo (0x1048), v11, v13
+          v22:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v23:BasicObject = SendDirect v22, 0x1038, :foo (0x1048), v11, v13
           CheckInterrupts
-          Return v24
+          Return v23
         ");
     }
 
@@ -3675,12 +3677,54 @@ mod hir_opt_tests {
         bb3(v8:BasicObject, v9:NilClass):
           v13:Fixnum[1] = Const Value(1)
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
-          v32:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v8, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
-          v19:CPtr = GetEP 0
-          v20:BasicObject = LoadField v19, :a@0x1038
+          v34:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v8, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v36:Fixnum[1] = Const Value(1)
           PatchPoint NoEPEscape(test)
+          v21:CPtr = LoadSP
+          v22:BasicObject = LoadField v21, :a@0x1038
           CheckInterrupts
-          Return v20
+          Return v22
+        ");
+    }
+
+    #[test]
+    fn reload_local_across_send_after_ep_escape() {
+        eval("
+            def foo(&block) = 1
+            def test
+              a = 1
+              lambda { a }
+              foo {|| }
+              a
+            end
+            test
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:4:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          Jump bb3(v1, v2)
+        bb2():
+          EntryPoint JIT(0)
+          v5:BasicObject = LoadArg :self@0
+          v6:NilClass = Const Value(nil)
+          Jump bb3(v5, v6)
+        bb3(v8:BasicObject, v9:NilClass):
+          v13:Fixnum[1] = Const Value(1)
+          SetLocal :a, l0, EP@3, v13
+          PatchPoint MethodRedefined(Object@0x1000, lambda@0x1008, cme:0x1010)
+          v45:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v8, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v46:BasicObject = CCallWithFrame v45, :Kernel#lambda@0x1038, block=0x1040
+          v20:CPtr = GetEP 0
+          v21:BasicObject = LoadField v20, :a@0x1048
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1049, cme:0x1050)
+          v32:CPtr = GetEP 0
+          v33:BasicObject = LoadField v32, :a@0x1048
+          CheckInterrupts
+          Return v33
         ");
     }
 
@@ -4248,12 +4292,12 @@ mod hir_opt_tests {
           v17:ArrayExact = NewArray
           v22:TrueClass = Const Value(true)
           v24:BasicObject = Send v12, 0x1008, :each_line, v22 # SendFallbackReason: Complex argument passing
-          v25:CPtr = GetEP 0
-          v26:BasicObject = LoadField v25, :s@0x1030
-          v27:BasicObject = LoadField v25, :a@0x1031
           PatchPoint NoEPEscape(test)
+          v27:CPtr = LoadSP
+          v28:BasicObject = LoadField v27, :s@0x1000
+          v29:BasicObject = LoadField v27, :a@0x1030
           CheckInterrupts
-          Return v27
+          Return v29
         ");
     }
 
@@ -8150,9 +8194,9 @@ mod hir_opt_tests {
           v11:ArrayExact = ArrayDup v10
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, map@0x1010, cme:0x1018)
-          v22:BasicObject = SendDirect v11, 0x1040, :map (0x1050)
+          v21:BasicObject = SendDirect v11, 0x1040, :map (0x1050)
           CheckInterrupts
-          Return v22
+          Return v21
         ");
     }
 
@@ -8186,17 +8230,17 @@ mod hir_opt_tests {
           v13:ArrayExact = NewArray
           PatchPoint SingleRactorMode
           PatchPoint StableConstantNames(0x1000, A)
-          v36:ArrayExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+          v38:ArrayExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint StableConstantNames(0x1010, B)
-          v39:ArrayExact[VALUE(0x1018)] = Const Value(VALUE(0x1018))
+          v41:ArrayExact[VALUE(0x1018)] = Const Value(VALUE(0x1018))
           PatchPoint NoSingletonClass(Array@0x1020)
           PatchPoint MethodRedefined(Array@0x1020, zip@0x1028, cme:0x1030)
-          v43:BasicObject = CCallVariadic v36, :Array#zip@0x1058, v39
-          v22:CPtr = GetEP 0
-          v23:BasicObject = LoadField v22, :result@0x1060
+          v45:BasicObject = CCallVariadic v38, :Array#zip@0x1058, v41
           PatchPoint NoEPEscape(test)
+          v24:CPtr = LoadSP
+          v25:BasicObject = LoadField v24, :result@0x1060
           CheckInterrupts
-          Return v23
+          Return v25
         ");
     }
 
@@ -8341,10 +8385,10 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
-          v19:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
-          v20:BasicObject = SendDirect v19, 0x1038, :foo (0x1048)
+          v18:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)]
+          v19:BasicObject = SendDirect v18, 0x1038, :foo (0x1048)
           CheckInterrupts
-          Return v20
+          Return v19
         ");
     }
 

zjit/src/hir/tests.rs

@@ -1983,8 +1983,9 @@ pub(crate) mod hir_build_tests {
           Jump bb3(v6, v7)
         bb3(v9:BasicObject, v10:BasicObject):
           v15:BasicObject = Send v10, 0x1008, :each # SendFallbackReason: Uncategorized(send)
-          v16:CPtr = GetEP 0
-          v17:BasicObject = LoadField v16, :a@0x1030
+          PatchPoint NoEPEscape(test)
+          v18:CPtr = LoadSP
+          v19:BasicObject = LoadField v18, :a@0x1000
           CheckInterrupts
           Return v15
         ");
@@ -2266,8 +2267,9 @@ pub(crate) mod hir_build_tests {
           Jump bb3(v6, v7)
         bb3(v9:BasicObject, v10:BasicObject):
           v16:BasicObject = InvokeSuperForward v9, 0x1008, v10 # SendFallbackReason: InvokeSuperForward: not yet specialized
-          v17:CPtr = GetEP 0
-          v18:BasicObject = LoadField v17, :...@0x1010
+          PatchPoint NoEPEscape(test)
+          v19:CPtr = LoadSP
+          v20:BasicObject = LoadField v19, :...@0x1000
           CheckInterrupts
           Return v16
         ");