Implement JIT-to-JIT calls (Shopify/zjit#109)

* Implement JIT-to-JIT calls

* Use a closer dummy address for Arm64

* Revert an obsoleted change

* Revert a few more obsoleted changes

* Fix outdated comments

* Explain PosMarkers for CCall

* s/JIT code/machine code/

* Get rid of ParallelMov

Author: k0kubun

test/ruby/test_zjit.rb

@@ -389,6 +389,21 @@ def test = callee + callee
     }
   end
 
+  def test_method_call
+    assert_compiles '12', %q{
+      def callee(a, b)
+        a - b
+      end
+
+      def test
+        callee(4, 2) + 10
+      end
+
+      test # profile test
+      test
+    }, call_threshold: 2
+  end
+
   def test_recursive_fact
     assert_compiles '[1, 6, 720]', %q{
       def fact(n)
@@ -401,6 +416,19 @@ def fact(n)
     }
   end
 
+  def test_profiled_fact
+    assert_compiles '[1, 6, 720]', %q{
+      def fact(n)
+        if n == 0
+          return 1
+        end
+        return n * fact(n-1)
+      end
+      fact(1) # profile fact
+      [fact(0), fact(3), fact(6)]
+    }, call_threshold: 3, num_profiles: 2
+  end
+
   def test_recursive_fib
     assert_compiles '[0, 2, 3]', %q{
       def fib(n)
@@ -413,11 +441,24 @@ def fib(n)
     }
   end
 
+  def test_profiled_fib
+    assert_compiles '[0, 2, 3]', %q{
+      def fib(n)
+        if n < 2
+          return n
+        end
+        return fib(n-1) + fib(n-2)
+      end
+      fib(3) # profile fib
+      [fib(0), fib(3), fib(4)]
+    }, call_threshold: 5, num_profiles: 3
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT
   # at a given call_threshold
-  def assert_compiles(expected, test_script, call_threshold: 1)
+  def assert_compiles(expected, test_script, **opts)
     pipe_fd = 3
 
     script = <<~RUBY
@@ -429,7 +470,7 @@ def assert_compiles(expected, test_script, call_threshold: 1)
       IO.open(#{pipe_fd}).write(result.inspect)
     RUBY
 
-    status, out, err, actual = eval_with_jit(script, call_threshold:, pipe_fd:)
+    status, out, err, actual = eval_with_jit(script, pipe_fd:, **opts)
 
     message = "exited with status #{status.to_i}"
     message << "\nstdout:\n```\n#{out}```\n" unless out.empty?
@@ -440,10 +481,11 @@ def assert_compiles(expected, test_script, call_threshold: 1)
   end
 
   # Run a Ruby process with ZJIT options and a pipe for writing test results
-  def eval_with_jit(script, call_threshold: 1, timeout: 1000, pipe_fd:, debug: true)
+  def eval_with_jit(script, call_threshold: 1, num_profiles: 1, timeout: 1000, pipe_fd:, debug: true)
     args = [
       "--disable-gems",
       "--zjit-call-threshold=#{call_threshold}",
+      "--zjit-num-profiles=#{num_profiles}",
     ]
     args << "--zjit-debug" if debug
     args << "-e" << script_shell_encode(script)

zjit/src/asm/mod.rs

@@ -111,6 +111,28 @@ impl CodeBlock {
         self.get_ptr(self.write_pos)
     }
 
+    /// Set the current write position from a pointer
+    fn set_write_ptr(&mut self, code_ptr: CodePtr) {
+        let pos = code_ptr.as_offset() - self.mem_block.borrow().start_ptr().as_offset();
+        self.write_pos = pos.try_into().unwrap();
+    }
+
+    /// Invoke a callback with write_ptr temporarily adjusted to a given address
+    pub fn with_write_ptr(&mut self, code_ptr: CodePtr, callback: impl Fn(&mut CodeBlock)) {
+        // Temporarily update the write_pos. Ignore the dropped_bytes flag at the old address.
+        let old_write_pos = self.write_pos;
+        let old_dropped_bytes = self.dropped_bytes;
+        self.set_write_ptr(code_ptr);
+        self.dropped_bytes = false;
+
+        // Invoke the callback
+        callback(self);
+
+        // Restore the original write_pos and dropped_bytes flag.
+        self.dropped_bytes = old_dropped_bytes;
+        self.write_pos = old_write_pos;
+    }
+
     /// Get a (possibly dangling) direct pointer into the executable memory block
     pub fn get_ptr(&self, offset: usize) -> CodePtr {
         self.mem_block.borrow().start_ptr().add_bytes(offset)

zjit/src/backend/arm64/mod.rs

@@ -491,19 +491,21 @@ impl Assembler
                     // register.
                     // Note: the iteration order is reversed to avoid corrupting x0,
                     // which is both the return value and first argument register
-                    let mut args: Vec<(Reg, Opnd)> = vec![];
-                    for (idx, opnd) in opnds.into_iter().enumerate().rev() {
-                        // If the value that we're sending is 0, then we can use
-                        // the zero register, so in this case we'll just send
-                        // a UImm of 0 along as the argument to the move.
-                        let value = match opnd {
-                            Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
-                            Opnd::Mem(_) => split_memory_address(asm, *opnd),
-                            _ => *opnd
-                        };
-                        args.push((C_ARG_OPNDS[idx].unwrap_reg(), value));
+                    if !opnds.is_empty() {
+                        let mut args: Vec<(Reg, Opnd)> = vec![];
+                        for (idx, opnd) in opnds.into_iter().enumerate().rev() {
+                            // If the value that we're sending is 0, then we can use
+                            // the zero register, so in this case we'll just send
+                            // a UImm of 0 along as the argument to the move.
+                            let value = match opnd {
+                                Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
+                                Opnd::Mem(_) => split_memory_address(asm, *opnd),
+                                _ => *opnd
+                            };
+                            args.push((C_ARG_OPNDS[idx].unwrap_reg(), value));
+                        }
+                        asm.parallel_mov(args);
                     }
-                    asm.parallel_mov(args);
 
                     // Now we push the CCall without any arguments so that it
                     // just performs the call.

zjit/src/backend/lir.rs

@@ -345,7 +345,19 @@ pub enum Insn {
     CPushAll,
 
     // C function call with N arguments (variadic)
-    CCall { opnds: Vec<Opnd>, fptr: *const u8, out: Opnd },
+    CCall {
+        opnds: Vec<Opnd>,
+        fptr: *const u8,
+        /// Optional PosMarker to remember the start address of the C call.
+        /// It's embedded here to insert the PosMarker after push instructions
+        /// that are split from this CCall on alloc_regs().
+        start_marker: Option<PosMarkerFn>,
+        /// Optional PosMarker to remember the end address of the C call.
+        /// It's embedded here to insert the PosMarker before pop instructions
+        /// that are split from this CCall on alloc_regs().
+        end_marker: Option<PosMarkerFn>,
+        out: Opnd,
+    },
 
     // C function return
     CRet(Opnd),
@@ -1455,6 +1467,23 @@ impl Assembler
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), live_ranges.len());
 
         while let Some((index, mut insn)) = iterator.next() {
+            let before_ccall = match (&insn, iterator.peek().map(|(_, insn)| insn)) {
+                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
+                (Insn::CCall { .. }, _) if !pool.is_empty() => {
+                    // If C_RET_REG is in use, move it to another register.
+                    // This must happen before last-use registers are deallocated.
+                    if let Some(vreg_idx) = pool.vreg_for(&C_RET_REG) {
+                        let new_reg = pool.alloc_reg(vreg_idx).unwrap(); // TODO: support spill
+                        asm.mov(Opnd::Reg(new_reg), C_RET_OPND);
+                        pool.dealloc_reg(&C_RET_REG);
+                        reg_mapping[vreg_idx] = Some(new_reg);
+                    }
+
+                    true
+                },
+                _ => false,
+            };
+
             // Check if this is the last instruction that uses an operand that
             // spans more than one instruction. In that case, return the
             // allocated register to the pool.
@@ -1477,32 +1506,20 @@ impl Assembler
                 }
             }
 
-            // If we're about to make a C call, save caller-saved registers
-            match (&insn, iterator.peek().map(|(_, insn)| insn)) {
-                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
-                (Insn::CCall { .. }, _) if !pool.is_empty() => {
-                    // If C_RET_REG is in use, move it to another register
-                    if let Some(vreg_idx) = pool.vreg_for(&C_RET_REG) {
-                        let new_reg = pool.alloc_reg(vreg_idx).unwrap(); // TODO: support spill
-                        asm.mov(Opnd::Reg(new_reg), C_RET_OPND);
-                        pool.dealloc_reg(&C_RET_REG);
-                        reg_mapping[vreg_idx] = Some(new_reg);
-                    }
-
-                    // Find all live registers
-                    saved_regs = pool.live_regs();
+            // Save caller-saved registers on a C call.
+            if before_ccall {
+                // Find all live registers
+                saved_regs = pool.live_regs();
 
-                    // Save live registers
-                    for &(reg, _) in saved_regs.iter() {
-                        asm.cpush(Opnd::Reg(reg));
-                        pool.dealloc_reg(&reg);
-                    }
-                    // On x86_64, maintain 16-byte stack alignment
-                    if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
-                        asm.cpush(Opnd::Reg(saved_regs.last().unwrap().0));
-                    }
+                // Save live registers
+                for &(reg, _) in saved_regs.iter() {
+                    asm.cpush(Opnd::Reg(reg));
+                    pool.dealloc_reg(&reg);
+                }
+                // On x86_64, maintain 16-byte stack alignment
+                if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
+                    asm.cpush(Opnd::Reg(saved_regs.last().unwrap().0));
                 }
-                _ => {},
             }
 
             // If the output VReg of this instruction is used by another instruction,
@@ -1590,13 +1607,24 @@ impl Assembler
 
             // Push instruction(s)
             let is_ccall = matches!(insn, Insn::CCall { .. });
-            if let Insn::ParallelMov { moves } = insn {
-                // Now that register allocation is done, it's ready to resolve parallel moves.
-                for (reg, opnd) in Self::resolve_parallel_moves(&moves) {
-                    asm.load_into(Opnd::Reg(reg), opnd);
+            match insn {
+                Insn::ParallelMov { moves } => {
+                    // Now that register allocation is done, it's ready to resolve parallel moves.
+                    for (reg, opnd) in Self::resolve_parallel_moves(&moves) {
+                        asm.load_into(Opnd::Reg(reg), opnd);
+                    }
                 }
-            } else {
-                asm.push_insn(insn);
+                Insn::CCall { opnds, fptr, start_marker, end_marker, out } => {
+                    // Split start_marker and end_marker here to avoid inserting push/pop between them.
+                    if let Some(start_marker) = start_marker {
+                        asm.push_insn(Insn::PosMarker(start_marker));
+                    }
+                    asm.push_insn(Insn::CCall { opnds, fptr, start_marker: None, end_marker: None, out });
+                    if let Some(end_marker) = end_marker {
+                        asm.push_insn(Insn::PosMarker(end_marker));
+                    }
+                }
+                _ => asm.push_insn(insn),
             }
 
             // After a C call, restore caller-saved registers
@@ -1720,38 +1748,30 @@ impl Assembler {
         self.push_insn(Insn::Breakpoint);
     }
 
+    /// Call a C function without PosMarkers
     pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
-        /*
-        // Let vm_check_canary() assert this ccall's leafness if leaf_ccall is set
-        let canary_opnd = self.set_stack_canary(&opnds);
-
-        let old_temps = self.ctx.get_reg_mapping(); // with registers
-        // Spill stack temp registers since they are caller-saved registers.
-        // Note that this doesn't spill stack temps that are already popped
-        // but may still be used in the C arguments.
-        self.spill_regs();
-        let new_temps = self.ctx.get_reg_mapping(); // all spilled
-
-        // Temporarily manipulate RegMappings so that we can use registers
-        // to pass stack operands that are already spilled above.
-        self.ctx.set_reg_mapping(old_temps);
-        */
-
-        // Call a C function
         let out = self.new_vreg(Opnd::match_num_bits(&opnds));
-        self.push_insn(Insn::CCall { fptr, opnds, out });
-
-        /*
-        // Registers in old_temps may be clobbered by the above C call,
-        // so rollback the manipulated RegMappings to a spilled version.
-        self.ctx.set_reg_mapping(new_temps);
-
-        // Clear the canary after use
-        if let Some(canary_opnd) = canary_opnd {
-            self.mov(canary_opnd, 0.into());
-        }
-        */
+        self.push_insn(Insn::CCall { fptr, opnds, start_marker: None, end_marker: None, out });
+        out
+    }
 
+    /// Call a C function with PosMarkers. This is used for recording the start and end
+    /// addresses of the C call and rewriting it with a different function address later.
+    pub fn ccall_with_pos_markers(
+        &mut self,
+        fptr: *const u8,
+        opnds: Vec<Opnd>,
+        start_marker: impl Fn(CodePtr, &CodeBlock) + 'static,
+        end_marker: impl Fn(CodePtr, &CodeBlock) + 'static,
+    ) -> Opnd {
+        let out = self.new_vreg(Opnd::match_num_bits(&opnds));
+        self.push_insn(Insn::CCall {
+            fptr,
+            opnds,
+            start_marker: Some(Box::new(start_marker)),
+            end_marker: Some(Box::new(end_marker)),
+            out,
+        });
         out
     }
 

zjit/src/backend/x86_64/mod.rs

@@ -82,8 +82,8 @@ impl From<&Opnd> for X86Opnd {
 /// This has the same number of registers for x86_64 and arm64.
 /// SCRATCH_REG is excluded.
 pub const ALLOC_REGS: &'static [Reg] = &[
-    RSI_REG,
     RDI_REG,
+    RSI_REG,
     RDX_REG,
     RCX_REG,
     R8_REG,
@@ -338,11 +338,13 @@ impl Assembler
                     assert!(opnds.len() <= C_ARG_OPNDS.len());
 
                     // Load each operand into the corresponding argument register.
-                    let mut args: Vec<(Reg, Opnd)> = vec![];
-                    for (idx, opnd) in opnds.into_iter().enumerate() {
-                        args.push((C_ARG_OPNDS[idx].unwrap_reg(), *opnd));
+                    if !opnds.is_empty() {
+                        let mut args: Vec<(Reg, Opnd)> = vec![];
+                        for (idx, opnd) in opnds.into_iter().enumerate() {
+                            args.push((C_ARG_OPNDS[idx].unwrap_reg(), *opnd));
+                        }
+                        asm.parallel_mov(args);
                     }
-                    asm.parallel_mov(args);
 
                     // Now we push the CCall without any arguments so that it
                     // just performs the call.