riscv64: Dynamically emit islands for return calls (bytecodealliance#…

…8868)

* riscv64: Increase max inst size

* riscv64: Emit islands in return call sequence

* riscv64: Update worst case size tests

Having duplicate registers was preventing
some moves from being generated

cranelift/codegen/src/isa/riscv64/abi.rs

@@ -855,7 +855,7 @@ fn compute_clobber_size(clobbers: &[Writable<RealReg>]) -> u32 {
     align_to(clobbered_size, 16)
 }
 
-pub(crate) const DEFAULT_CLOBBERS: PRegSet = PRegSet::empty()
+const DEFAULT_CLOBBERS: PRegSet = PRegSet::empty()
     .with(px_reg(1))
     .with(px_reg(5))
     .with(px_reg(6))

cranelift/codegen/src/isa/riscv64/inst/emit.rs

@@ -47,14 +47,14 @@ pub enum EmitVState {
 #[derive(Default, Clone, Debug)]
 pub struct EmitState {
     /// Safepoint stack map for upcoming instruction, as provided to `pre_safepoint()`.
-    pub(crate) stack_map: Option<StackMap>,
+    stack_map: Option<StackMap>,
     /// Only used during fuzz-testing. Otherwise, it is a zero-sized struct and
     /// optimized away at compiletime. See [cranelift_control].
-    pub(crate) ctrl_plane: ControlPlane,
+    ctrl_plane: ControlPlane,
     /// Vector State
     /// Controls the current state of the vector unit at the emission point.
-    pub(crate) vstate: EmitVState,
-    pub(crate) frame_layout: FrameLayout,
+    vstate: EmitVState,
+    frame_layout: FrameLayout,
 }
 
 impl EmitState {
@@ -242,14 +242,21 @@ impl MachInstEmit for Inst {
             self.emit_uncompressed(sink, emit_info, state, &mut start_off);
         }
 
-        let end_off = sink.cur_offset();
-        assert!(
-            (end_off - start_off) <= Inst::worst_case_size(),
-            "Inst:{:?} length:{} worst_case_size:{}",
+        // We exclude br_table and return call from these checks since they emit
+        // their own islands, and thus are allowed to exceed the worst case size.
+        if !matches!(
             self,
-            end_off - start_off,
-            Inst::worst_case_size()
-        );
+            Inst::BrTable { .. } | Inst::ReturnCall { .. } | Inst::ReturnCallInd { .. }
+        ) {
+            let end_off = sink.cur_offset();
+            assert!(
+                (end_off - start_off) <= Inst::worst_case_size(),
+                "Inst:{:?} length:{} worst_case_size:{}",
+                self,
+                end_off - start_off,
+                Inst::worst_case_size()
+            );
+        }
     }
 
     fn pretty_print_inst(&self, state: &mut Self::State) -> String {
@@ -2608,6 +2615,40 @@ fn emit_return_call_common_sequence(
     emit_info: &EmitInfo,
     state: &mut EmitState,
     info: &ReturnCallInfo,
+) {
+    // The return call sequence can potentially emit a lot of instructions (up to 634 bytes!)
+    // So lets emit an island here if we need it.
+    //
+    // It is difficult to calculate exactly how many instructions are going to be emitted, so
+    // we calculate it by emitting it into a disposable buffer, and then checking how many instructions
+    // were actually emitted.
+    let mut buffer = MachBuffer::new();
+    let mut fake_emit_state = state.clone();
+
+    return_call_emit_impl(&mut buffer, emit_info, &mut fake_emit_state, info);
+
+    // Finalize the buffer and get the number of bytes emitted.
+    let buffer = buffer.finish(&Default::default(), &mut Default::default());
+    let length = buffer.data().len() as u32;
+
+    // And now emit the island inline with this instruction.
+    if sink.island_needed(length) {
+        let jump_around_label = sink.get_label();
+        Inst::gen_jump(jump_around_label).emit(sink, emit_info, state);
+        sink.emit_island(length + 4, &mut state.ctrl_plane);
+        sink.bind_label(jump_around_label, &mut state.ctrl_plane);
+    }
+
+    // Now that we're done, emit the *actual* return sequence.
+    return_call_emit_impl(sink, emit_info, state, info);
+}
+
+/// This should not be called directly, Instead prefer to call [emit_return_call_common_sequence].
+fn return_call_emit_impl(
+    sink: &mut MachBuffer<Inst>,
+    emit_info: &EmitInfo,
+    state: &mut EmitState,
+    info: &ReturnCallInfo,
 ) {
     let sp_to_fp_offset = {
         let frame_layout = state.frame_layout();

cranelift/codegen/src/isa/riscv64/inst/emit_tests.rs

@@ -1,6 +1,6 @@
 #[allow(unused)]
 use crate::ir::LibCall;
-use crate::isa::riscv64::{abi::DEFAULT_CLOBBERS, inst::*};
+use crate::isa::riscv64::inst::*;
 use std::borrow::Cow;
 
 fn fa7() -> Reg {
@@ -2190,47 +2190,21 @@ fn riscv64_worst_case_instruction_size() {
             .map(|op| Inst::AtomicRmwLoop {
                 op: *op,
                 offset: a0(),
-                dst: writable_a0(),
+                dst: writable_a1(),
                 ty: I16,
-                p: a0(),
-                x: a0(),
+                p: a1(),
+                x: a2(),
                 t0: writable_a0(),
             }),
     );
 
-    candidates.push(Inst::ReturnCallInd {
-        callee: a0(),
-        info: Box::new(ReturnCallInfo {
-            opcode: Opcode::ReturnCallIndirect,
-            new_stack_arg_size: 64,
-            uses: DEFAULT_CLOBBERS
-                .into_iter()
-                .map(|reg| CallArgPair {
-                    vreg: reg.into(),
-                    preg: reg.into(),
-                })
-                .collect(),
-        }),
-    });
+    // Return Call Indirect and BrTable are the largest instructions possible. However they
+    // emit their own island, so we don't account them here.
 
     let mut max: (u32, MInst) = (0, Inst::Nop0);
     for i in candidates {
         let mut buffer = MachBuffer::new();
-        let mut emit_state = EmitState {
-            // This frame layout is important to ensure that the ReturnCallIndirect
-            // instruction in this test, becomes as large as practically possible.
-            frame_layout: FrameLayout {
-                tail_args_size: 64,
-                setup_area_size: 8192,
-                clobbered_callee_saves: DEFAULT_CLOBBERS
-                    .into_iter()
-                    .filter(|r| r.class() != RegClass::Vector)
-                    .map(|r| Writable::from_reg(r.into()))
-                    .collect(),
-                ..Default::default()
-            },
-            ..Default::default()
-        };
+        let mut emit_state = Default::default();
         i.emit(&mut buffer, &emit_info, &mut emit_state);
         let buffer = buffer.finish(&Default::default(), &mut Default::default());
         let length = buffer.data().len() as u32;

cranelift/codegen/src/isa/riscv64/inst/mod.rs

@@ -845,7 +845,7 @@ impl MachInst for Inst {
 
     fn worst_case_size() -> CodeOffset {
         // Our worst case size is determined by the riscv64_worst_case_instruction_size test
-        168
+        84
     }
 
     fn ref_type_regclass(_settings: &settings::Flags) -> RegClass {